Repository: neuroinformatics-unit/movement

Issue:
compute_forward_vector_angle throws IndexError
**Describe the bug**
Calling the `compute_forward_vector_angle` on some "normal" movement datasets (e.g. our sample data) leads to an `IndexError`.

<details>
<summary>Expand to see the Traceback</summary>

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
File /Users/nsirmpilatze/Code/NIU/movement/examples/diagnose_error.py:7
      1 # %%
      2 # Load data
      3 # ---------
      4 # Load the sample data
      5 ds = sample_data.fetch_dataset("SLEAP_single-mouse_EPM.analysis.h5")
----> 7 angles = compute_forward_vector_angle(
      8     ds.position,
      9     left_keypoint="left_ear",
     10     right_keypoint="right_ear",
     11 )

File ~/Code/NIU/movement/movement/kinematics.py:457, in compute_forward_vector_angle(data, left_keypoint, right_keypoint, reference_vector, camera_view, in_radians, angle_rotates)
    452 forward_vector = compute_forward_vector(
    453     data, left_keypoint, right_keypoint, camera_view=camera_view
    454 )
    456 # Compute signed angle between forward vector and reference vector
--> 457 heading_array = compute_signed_angle_2d(
    458     forward_vector, reference_vector, v_as_left_operand=ref_is_left_operand
    459 )
    461 # Convert to degrees
    462 if not in_radians:

File ~/Code/NIU/movement/movement/utils/vector.py:257, in compute_signed_angle_2d(u, v, v_as_left_operand)
    254 angles = np.arctan2(cross, dot)
    255 # arctan2 returns values in [-pi, pi].
    256 # We need to map -pi angles to pi, to stay in the (-pi, pi] range
--> 257 angles[angles <= -np.pi] = np.pi
    258 return angles

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/dataarray.py:913, in DataArray.__setitem__(self, key, value)
    908     self.coords[key] = value
    909 else:
    910     # Coordinates in key, value and self[key] should be consistent.
    911     # TODO Coordinate consistency in key is checked here, but it
    912     # causes unnecessary indexing. It should be optimized.
--> 913     obj = self[key]
    914     if isinstance(value, DataArray):
    915         assert_coordinate_consistent(value, obj.coords.variables)

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/dataarray.py:904, in DataArray.__getitem__(self, key)
    901     return self._getitem_coord(key)
    902 else:
    903     # xarray-style array indexing
--> 904     return self.isel(indexers=self._item_key_to_dict(key))

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/dataarray.py:1540, in DataArray.isel(self, indexers, drop, missing_dims, **indexers_kwargs)
   1537 indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "isel")
   1539 if any(is_fancy_indexer(idx) for idx in indexers.values()):
-> 1540     ds = self._to_temp_dataset()._isel_fancy(
   1541         indexers, drop=drop, missing_dims=missing_dims
   1542     )
   1543     return self._from_temp_dataset(ds)
   1545 # Much faster algorithm for when all indexers are ints, slices, one-dimensional
   1546 # lists, or zero or one-dimensional np.ndarray's

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/dataset.py:3152, in Dataset._isel_fancy(self, indexers, drop, missing_dims)
   3148 var_indexers = {
   3149     k: v for k, v in valid_indexers.items() if k in var.dims
   3150 }
   3151 if var_indexers:
-> 3152     new_var = var.isel(indexers=var_indexers)
   3153     # drop scalar coordinates
   3154     # https://github.com/pydata/xarray/issues/6554
   3155     if name in self.coords and drop and new_var.ndim == 0:

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/variable.py:1079, in Variable.isel(self, indexers, missing_dims, **indexers_kwargs)
   1076 indexers = drop_dims_from_indexers(indexers, self.dims, missing_dims)
   1078 key = tuple(indexers.get(dim, slice(None)) for dim in self.dims)
-> 1079 return self[key]

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/variable.py:825, in Variable.__getitem__(self, key)
    812 def __getitem__(self, key) -> Self:
    813     """Return a new Variable object whose contents are consistent with
    814     getting the provided key from the underlying data.
    815 
   (...)
    823     array `x.values` directly.
    824     """
--> 825     dims, indexer, new_order = self._broadcast_indexes(key)
    826     indexable = as_indexable(self._data)
    828     data = indexing.apply_indexer(indexable, indexer)

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/variable.py:663, in Variable._broadcast_indexes(self, key)
    660 if all(isinstance(k, BASIC_INDEXING_TYPES) for k in key):
    661     return self._broadcast_indexes_basic(key)
--> 663 self._validate_indexers(key)
    664 # Detect it can be mapped as an outer indexer
    665 # If all key is unlabeled, or
    666 # key can be mapped as an OuterIndexer.
    667 if all(not isinstance(k, Variable) for k in key):

File ~/.miniconda3/envs/movement-dev-py312/lib/python3.12/site-packages/xarray/core/variable.py:712, in Variable._validate_indexers(self, key)
    707     raise IndexError(
    708         f"Boolean array size {len(k):d} is used to index array "
    709         f"with shape {self.shape}."
    710     )
    711 if k.ndim > 1:
--> 712     raise IndexError(
    713         f"{k.ndim}-dimensional boolean indexing is "
    714         "not supported. "
    715     )
    716 if is_duck_dask_array(k.data):
    717     raise KeyError(
    718         "Indexing with a boolean dask array is not allowed. "
    719         "This will result in a dask array of unknown shape. "
    720         "Such arrays are unsupported by Xarray."
    721         "Please compute the indexer first using .compute()"
    722     )

IndexError: 2-dimensional boolean indexing is not supported.
```

</details>

**To Reproduce**

```python
from movement import sample_data
from movement.kinematics import compute_forward_vector_angle

ds = sample_data.fetch_dataset("SLEAP_single-mouse_EPM.analysis.h5")

# Raises IndexError
angles = compute_forward_vector_angle(
    ds.position,
    left_keypoint="left_ear",
    right_keypoint="right_ear",
)

# Works fine if we squeeze out the singleton "individuals" dimension
angles_squeezed = compute_forward_vector_angle(
    ds.position.squeeze(),
    left_keypoint="left_ear",
    right_keypoint="right_ear",
)
```

**Expected behaviour**
There should be no error, and the returned array of angles should maintain the `individuals` dimension as it was.



Discussion hints:
sonarqubecloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=432) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=432&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=432&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=neuroinformatics-unit_movement&pullRequest=432&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=432&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=432&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarQube Cloud](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=432)


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/neuroinformatics-unit/movement/pull/432?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.84%. Comparing base [(`470cf6a`)](https://app.codecov.io/gh/neuroinformatics-unit/movement/commit/470cf6aadc301ca8c47020fef570b0cebe0f5c03?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit) to head [(`2bcbcb6`)](https://app.codecov.io/gh/neuroinformatics-unit/movement/commit/2bcbcb6903d62e9b93006238dfdc7ee067697f79?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main     #432   +/-   ##
=======================================
  Coverage   99.84%   99.84%           
=======================================
  Files          23       23           
  Lines        1250     1250           
=======================================
  Hits         1248     1248           
  Misses          2        2           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/neuroinformatics-unit/movement/pull/432?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit).


Solution patch:
diff --git a/movement/utils/vector.py b/movement/utils/vector.py
index 965d269..16beb38 100644
--- a/movement/utils/vector.py
+++ b/movement/utils/vector.py
@@ -254,7 +254,7 @@ def compute_signed_angle_2d(
     angles = np.arctan2(cross, dot)
     # arctan2 returns values in [-pi, pi].
     # We need to map -pi angles to pi, to stay in the (-pi, pi] range
-    angles[angles <= -np.pi] = np.pi
+    angles.values[angles <= -np.pi] = np.pi
     return angles
 
 


Test patch:
diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py
index db43b03..57cbcc3 100644
--- a/tests/test_unit/test_vector.py
+++ b/tests/test_unit/test_vector.py
@@ -1,3 +1,4 @@
+from collections.abc import Iterable
 from contextlib import nullcontext as does_not_raise
 
 import numpy as np
@@ -297,3 +298,68 @@ class TestComputeSignedAngle:
         xr.testing.assert_allclose(
             computed_angles_reversed, expected_angles_reversed
         )
+
+    @pytest.mark.parametrize(
+        ["extra_dim_sizes"],
+        [
+            pytest.param((1,), id="[1D] Trailing singleton dimension"),
+            pytest.param((2, 3), id="[2D] Individuals-keypoints esque"),
+            pytest.param(
+                (1, 2),
+                id="[2D] As if .sel had been used on an additional axis",
+            ),
+        ],
+    )
+    def test_multidimensional_input(
+        self, extra_dim_sizes: tuple[int, ...]
+    ) -> None:
+        """Test that non-spacetime dimensions are respected.
+
+        The user may pass in data that has more than just the time and space
+        dimensions, in which case the method should be able to cope with
+        such an input, and preserve the additional dimensions.
+
+        We simulate this case in this test by adding "junk" dimensions onto
+        an existing input. We should see the result also spit out these "junk"
+        dimensions, preserving the dimension names and coordinates too.
+        """
+        v_left = self.coord_axes_array.copy()
+        v_left_original_ndim = v_left.ndim
+        v_right = self.x_axis
+        expected_angles = np.array([0.0, -np.pi / 2.0, np.pi, np.pi / 2.0])
+        expected_angles_original_ndim = expected_angles.ndim
+
+        # Append additional "junk" dimensions
+        left_vector_coords: dict[str, Iterable[str] | Iterable[int]] = {
+            "time": [f"t{i}" for i in range(v_left.shape[0])],
+            "space": ["x", "y"],
+        }
+        for extra_dim_index, extra_dim_size in enumerate(extra_dim_sizes):
+            v_left = np.repeat(
+                v_left[..., np.newaxis],
+                extra_dim_size,
+                axis=extra_dim_index + v_left_original_ndim,
+            )
+            left_vector_coords[f"edim{extra_dim_index}"] = range(
+                extra_dim_size
+            )
+            expected_angles = np.repeat(
+                expected_angles[..., np.newaxis],
+                extra_dim_size,
+                axis=extra_dim_index + expected_angles_original_ndim,
+            )
+        v_left = xr.DataArray(
+            data=v_left,
+            dims=left_vector_coords.keys(),
+            coords=left_vector_coords,
+        )
+        expected_angles_coords = dict(left_vector_coords)
+        expected_angles_coords.pop("space")
+        expected_angles = xr.DataArray(
+            data=expected_angles,
+            dims=expected_angles_coords.keys(),
+            coords=expected_angles_coords,
+        )
+
+        computed_angles = vector.compute_signed_angle_2d(v_left, v_right)
+        xr.testing.assert_allclose(computed_angles, expected_angles)