diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml index 4b7b4c11..cd295632 100644 --- a/.github/workflows/test_and_deploy.yml +++ b/.github/workflows/test_and_deploy.yml @@ -71,8 +71,14 @@ jobs: upload_all: name: Publish build distributions needs: [build_sdist_wheels] + if: github.event_name == 'push' && github.ref_type == 'tag' runs-on: ubuntu-latest steps: - - uses: neuroinformatics-unit/actions/upload_pypi@v2 + - uses: actions/download-artifact@v4 + with: + name: artifact + path: dist + - uses: pypa/gh-action-pypi-publish@v1.12.3 with: - secret-pypi-key: ${{ secrets.TWINE_API_KEY }} + user: __token__ + password: ${{ secrets.TWINE_API_KEY }} diff --git a/movement/kinematics.py b/movement/kinematics.py index c8f1635a..4abe9e99 100644 --- a/movement/kinematics.py +++ b/movement/kinematics.py @@ -5,11 +5,16 @@ import numpy as np import xarray as xr +from numpy.typing import ArrayLike from scipy.spatial.distance import cdist from movement.utils.logging import log_error, log_warning from movement.utils.reports import report_nan_values -from movement.utils.vector import compute_norm +from movement.utils.vector import ( + compute_norm, + compute_signed_angle_2d, + convert_to_unit, +) from movement.validators.arrays import validate_dims_coords @@ -203,7 +208,7 @@ def compute_forward_vector( left_keypoint: str, right_keypoint: str, camera_view: Literal["top_down", "bottom_up"] = "top_down", -): +) -> xr.DataArray: """Compute a 2D forward vector given two left-right symmetric keypoints. The forward vector is computed as a vector perpendicular to the @@ -304,7 +309,7 @@ def compute_forward_vector( space="z" ) # keep only the first 2 spatal dimensions of the result # Return unit vector - return forward_vector / compute_norm(forward_vector) + return convert_to_unit(forward_vector) def compute_head_direction_vector( @@ -350,6 +355,116 @@ def compute_head_direction_vector( ) +def compute_forward_vector_angle( + data: xr.DataArray, + left_keypoint: str, + right_keypoint: str, + reference_vector: xr.DataArray | ArrayLike = (1, 0), + camera_view: Literal["top_down", "bottom_up"] = "top_down", + in_radians: bool = False, + angle_rotates: Literal[ + "ref to forward", "forward to ref" + ] = "ref to forward", +) -> xr.DataArray: + r"""Compute the signed angle between a forward and reference vector. + + Forward vector angle is the :func:`signed angle\ + ` + between the reference vector and the animal's :func:`forward vector\ + `). + The returned angles are in degrees, spanning the range :math:`(-180, 180]`, + unless ``in_radians`` is set to ``True``. + + Parameters + ---------- + data : xarray.DataArray + The input data representing position. This must contain + the two symmetrical keypoints located on the left and + right sides of the body, respectively. + left_keypoint : str + Name of the left keypoint, e.g., "left_ear", used to compute the + forward vector. + right_keypoint : str + Name of the right keypoint, e.g., "right_ear", used to compute the + forward vector. + reference_vector : xr.DataArray | ArrayLike, optional + The reference vector against which the ``forward_vector`` is + compared to compute 2D heading. Must be a two-dimensional vector, + in the form [x,y] - where ``reference_vector[0]`` corresponds to the + x-coordinate and ``reference_vector[1]`` corresponds to the + y-coordinate. If left unspecified, the vector [1, 0] is used by + default. + camera_view : Literal["top_down", "bottom_up"], optional + The camera viewing angle, used to determine the upwards + direction of the animal. Can be either ``"top_down"`` (where the + upwards direction is [0, 0, -1]), or ``"bottom_up"`` (where the + upwards direction is [0, 0, 1]). If left unspecified, the camera + view is assumed to be ``"top_down"``. + in_radians : bool, optional + If true, the returned heading array is given in radians. + If false, the array is given in degrees. False by default. + angle_rotates : Literal["ref to forward", "forward to ref"], optional + Determines the sign convention for the returned signed angle. + (See Notes). + + Returns + ------- + xarray.DataArray + An xarray DataArray containing the computed forward vector angles, + with dimensions matching the input data array, + but without the ``keypoints`` and ``space`` dimensions. + + Notes + ----- + There are different conventions for the sign of the forward vector angle. + The ``angle_rotates`` argument can be used to select which behaviour the + function should use. + + By default, ``angle_rotates`` is set to ``"ref to forward"``, which + results in the signed angle between the reference vector and the forward + vector being returned. That is, the angle which the reference vector would + need to be rotated by, to align with the forward vector. + Setting ``angle_rotates`` to ``"forward to ref"`` reverses this convention, + returning the signed angle between the forward vector and the reference + vector; that is, the rotation that would need to be applied to the forward + vector to return the reference vector. + + See Also + -------- + movement.utils.vector.compute_signed_angle_2d : The underlying + function used to compute the signed angle between two 2D vectors. + movement.kinematics.compute_forward_vector : The function used + to compute the forward vector. + + """ + if angle_rotates.lower() == "ref to forward": + ref_is_left_operand = True + elif angle_rotates.lower() == "forward to ref": + ref_is_left_operand = False + else: + raise ValueError(f"Unknown angle convention: '{angle_rotates}'") + + # Convert reference vector to np.array if not already a valid array + if not isinstance(reference_vector, np.ndarray | xr.DataArray): + reference_vector = np.array(reference_vector) + + # Compute forward vector + forward_vector = compute_forward_vector( + data, left_keypoint, right_keypoint, camera_view=camera_view + ) + + # Compute signed angle between forward vector and reference vector + heading_array = compute_signed_angle_2d( + forward_vector, reference_vector, v_as_left_operand=ref_is_left_operand + ) + + # Convert to degrees + if not in_radians: + heading_array = np.rad2deg(heading_array) + + return heading_array + + def _cdist( a: xr.DataArray, b: xr.DataArray, @@ -389,6 +504,7 @@ def _cdist( Additional keyword arguments to pass to :func:`scipy.spatial.distance.cdist`. + Returns ------- xarray.DataArray diff --git a/movement/utils/vector.py b/movement/utils/vector.py index c91e43ec..965d2696 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -4,7 +4,10 @@ import xarray as xr from movement.utils.logging import log_error -from movement.validators.arrays import validate_dims_coords +from movement.validators.arrays import ( + validate_dims_coords, + validate_reference_vector, +) def compute_norm(data: xr.DataArray) -> xr.DataArray: @@ -165,6 +168,96 @@ def pol2cart(data: xr.DataArray) -> xr.DataArray: ).transpose(*dims) +def compute_signed_angle_2d( + u: xr.DataArray, + v: xr.DataArray | np.ndarray, + v_as_left_operand: bool = False, +) -> xr.DataArray: + r"""Compute the signed angle from the vector ``u`` to the vector ``v``. + + The signed angle between ``u`` and ``v`` is the rotation that needs to be + applied to ``u`` to have it point in the same direction as ``v`` (see + Notes). Angles are returned in radians, spanning :math:`(-\pi, \pi]` + according to the ``arctan2`` convention. + + Parameters + ---------- + u : xarray.DataArray + An array of position vectors containing the ``space`` + dimension with only ``"x"`` and ``"y"`` coordinates. + v : xarray.DataArray | numpy.ndarray + A 2D vector (or array of 2D vectors) against which to + compare ``u``. May either be an xarray + DataArray containing the ``"space"`` dimension or a numpy + array containing one or more 2D vectors. (See Notes) + v_as_left_operand : bool, default False + If True, the signed angle between ``v`` and ``u`` is returned, instead + of the signed angle between ``u`` and ``v``. This is a convenience + wrapper for when one of the two vectors to be used does not have time + points, and the other does. + + Returns + ------- + xarray.DataArray : + An xarray DataArray containing signed angle between + ``u`` and ``v`` for every time point. Matches the dimensions of + ``u``, but without the ``space`` dimension. + + Notes + ----- + Given two vectors :math:`u = (u_x, u_y)` and :math:`v = (v_x, v_y)`, + the signed angle :math:`\alpha` between ``u`` and ``v`` is computed as + + .. math:: + \alpha &= \mathrm{arctan2}(u \times v, u\cdot v) \\ + &= \mathrm{arctan2}(u_x v_y - u_y v_x, u_x v_x + u_y v_y), + + which corresponds to the rotation that needs to be applied to ``u`` for it + to point in the direction of ``v``. + + If ``v`` is passed as an ``xarray.DataArray``, ``v`` must have the spatial + coordinates ``"x"`` and ``"y"`` only, and must have a ``time`` dimension + matching that of ``u``. + + If passed as a numpy array, the ``v`` must have one of three shapes: + + - ``(2,)``: where dimension ``0`` contains spatial + coordinates (x,y), and no time dimension is specified. + - ``(1,2)``:, where dimension ``0`` corresponds to a + single time-point and dimension ``1`` contains spatial + coordinates (x,y). + - ``(n,2)``: where dimension ``0`` corresponds to + time and dimension ``1`` contains spatial coordinates + (x,y), and where ``n == len(u.time)``. + + Vectors given as ``v`` that contain more dimensions, or have shapes + otherwise different from those defined above are considered invalid. + + """ + validate_dims_coords(u, {"space": ["x", "y"]}, exact_coords=True) + # ensure v can be broadcast over u + v_with_matching_dims = validate_reference_vector(v, u) + + u_unit = convert_to_unit(u) + u_x = u_unit.sel(space="x") + u_y = u_unit.sel(space="y") + + v_unit = convert_to_unit(v_with_matching_dims) + v_x = v_unit.sel(space="x") + v_y = v_unit.sel(space="y") + + cross = u_x * v_y - u_y * v_x + if v_as_left_operand: + cross *= -1.0 + dot = u_x * v_x + u_y * v_y + + 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 + return angles + + def _raise_error_for_missing_spatial_dim() -> None: raise log_error( ValueError, diff --git a/movement/validators/arrays.py b/movement/validators/arrays.py index 76847571..24417c50 100644 --- a/movement/validators/arrays.py +++ b/movement/validators/arrays.py @@ -1,26 +1,37 @@ """Validators for data arrays.""" +import numpy as np import xarray as xr from movement.utils.logging import log_error def validate_dims_coords( - data: xr.DataArray, required_dim_coords: dict + data: xr.DataArray, + required_dim_coords: dict[str, list[str]], + exact_coords: bool = False, ) -> None: """Validate dimensions and coordinates in a data array. This function raises a ValueError if the specified dimensions and - coordinates are not present in the input data array. + coordinates are not present in the input data array. By default, + each dimension must contain *at least* the specified coordinates. + Pass ``exact_coords=True`` to require that each dimension contains + *exactly* the specified coordinates (and no others). Parameters ---------- data : xarray.DataArray The input data array to validate. - required_dim_coords : dict - A dictionary of required dimensions and their corresponding - coordinate values. If you don't need to specify any - coordinate values, you can pass an empty list. + required_dim_coords : dict of {str: list of str} + A dictionary mapping required dimensions to a list of required + coordinate values along each dimension. + exact_coords : bool, optional + If False (default), checks only that the listed coordinates + exist in each dimension. If True, checks that each dimension + has exactly the specified coordinates and no more. + The exactness check is completely skipped for dimensions with + no required coordinates. Examples -------- @@ -34,6 +45,10 @@ def validate_dims_coords( >>> validate_dims_coords(data, {"time": [], "space": ["x", "y"]}) + Enforce that 'space' has *only* 'x' and 'y', and no other coordinates: + + >>> validate_dims_coords(data, {"space": ["x", "y"]}, exact_coords=True) + Raises ------ ValueError @@ -41,21 +56,136 @@ def validate_dims_coords( and/or the required coordinate(s). """ + # 1. Check that all required dimensions are present missing_dims = [dim for dim in required_dim_coords if dim not in data.dims] error_message = "" if missing_dims: error_message += ( f"Input data must contain {missing_dims} as dimensions.\n" ) - missing_coords = [] + + # 2. For each dimension, check the presence of required coords for dim, coords in required_dim_coords.items(): - missing_coords = [ - coord for coord in coords if coord not in data.coords.get(dim, []) - ] + dim_coords_in_data = data.coords.get(dim, []) + missing_coords = [c for c in coords if c not in dim_coords_in_data] if missing_coords: error_message += ( - "Input data must contain " - f"{missing_coords} in the '{dim}' coordinates." + f"Input data must contain {missing_coords} " + f"in the '{dim}' coordinates.\n" ) + + # 3. If exact_coords is True, verify no extra coords exist + if exact_coords and coords: + extra_coords = [c for c in dim_coords_in_data if c not in coords] + if extra_coords: + error_message += ( + f"Dimension '{dim}' must only contain " + f"{coords} as coordinates, " + f"but it also has {list(extra_coords)}.\n" + ) + if error_message: raise log_error(ValueError, error_message) + + +def validate_reference_vector( + reference_vector: xr.DataArray | np.ndarray, + test_vector: xr.DataArray, +) -> xr.DataArray: + """Validate a reference vector being used in a calculation. + + Reference vectors must contain the same number of time points as the + ``test_vector``, in order to be used in computations with them. + Vector-like objects that are passed in are converted into ``xr.DataArray``s + and inherit the ``test_vector``s ``"time"`` axis. + + Parameters + ---------- + reference_vector : xarray.DataArray | numpy.ndarray + The reference vector to validate, and convert if necessary. + test_vector : xarray.DataArray + The test vector to compare against. The reference vector must have + the same number of time points as the test vector. + + Returns + ------- + xarray.DataArray + A validated xarray DataArray. + + Raises + ------ + ValueError + If shape or dimensions do not match the expected form. + TypeError + If reference_vector is neither a NumPy array nor an xarray DataArray. + + """ + if isinstance(reference_vector, np.ndarray): + # Check shape: must be 1D or 2D + if reference_vector.ndim > 2: + raise log_error( + ValueError, + "Reference vector must be 1D or 2D, but got " + f"{reference_vector.ndim}D array.", + ) + # Reshape 1D -> (1, -1) so axis 0 can be 'time' + if reference_vector.ndim == 1: + reference_vector = reference_vector.reshape(1, -1) + + # If multiple time points, must match test_vector time length + if ( + reference_vector.shape[0] > 1 + and reference_vector.shape[0] != test_vector.sizes["time"] + ): + raise log_error( + ValueError, + "Reference vector must have the same number of time " + "points as the test vector.", + ) + + # Decide whether we have (time, space) or just (space) + if reference_vector.shape[0] == 1: + coords = {"space": ["x", "y"]} + reference_vector = reference_vector.squeeze() + else: + coords = { + "time": test_vector["time"], + "space": ["x", "y"], + } + + return xr.DataArray( + reference_vector, dims=list(coords.keys()), coords=coords + ) + + elif isinstance(reference_vector, xr.DataArray): + # Must contain exactly 'x' and 'y' in the space dimension + validate_dims_coords( + reference_vector, {"space": ["x", "y"]}, exact_coords=True + ) + + # If it has a time dimension, time size must match + if ( + "time" in reference_vector.dims + and reference_vector.sizes["time"] != test_vector.sizes["time"] + ): + raise log_error( + ValueError, + "Reference vector must have the same number of time " + "points as the test vector.", + ) + + # Only 'time' and 'space' are allowed + if any(d not in {"time", "space"} for d in reference_vector.dims): + raise log_error( + ValueError, + "Only 'time' and 'space' dimensions " + "are allowed in reference_vector.", + ) + return reference_vector + + # If it's neither a DataArray nor a NumPy array + raise log_error( + TypeError, + "Reference vector must be an xarray.DataArray or np.ndarray, " + f"but got {type(reference_vector)}.", + ) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 0735e22b..ff3ffc22 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -1,5 +1,6 @@ import re from contextlib import nullcontext as does_not_raise +from typing import Any, Literal import numpy as np import pytest @@ -761,3 +762,293 @@ def test_compute_pairwise_distances_with_invalid_input( kinematics.compute_pairwise_distances( request.getfixturevalue(ds).position, dim, pairs ) + + +class TestForwardVectorAngle: + """Test the compute_forward_vector_angle function. + + These tests are grouped together into a class to distinguish them from the + other methods that are tested in the Kinematics module. + + Note that since this method is a combination of calls to two lower-level + methods, we run limited input/output checks in this collection. + Correctness of the results is delegated to the tests of the dependent + methods, as appropriate. + """ + + x_axis = np.array([1.0, 0.0]) + y_axis = np.array([0.0, 1.0]) + sqrt_2 = np.sqrt(2.0) + + @staticmethod + def push_into_range( + numeric_values: xr.DataArray | np.ndarray, + lower: float = -180.0, + upper: float = 180.0, + ) -> xr.DataArray | np.ndarray: + """Translate values into the range (lower, upper]. + + The interval width is the value ``upper - lower``. + Each ``v`` in ``values`` that starts less than or equal to the + ``lower`` bound has multiples of the interval width added to it, + until the result lies in the desirable interval. + + Each ``v`` in ``values`` that starts greater than the ``upper`` + bound has multiples of the interval width subtracted from it, + until the result lies in the desired interval. + """ + translated_values = ( + numeric_values.values.copy() + if isinstance(numeric_values, xr.DataArray) + else numeric_values.copy() + ) + + interval_width = upper - lower + if interval_width <= 0: + raise ValueError( + f"Upper bound ({upper}) must be strictly " + f"greater than lower bound ({lower})" + ) + + while np.any( + (translated_values <= lower) | (translated_values > upper) + ): + translated_values[translated_values <= lower] += interval_width + translated_values[translated_values > upper] -= interval_width + + if isinstance(numeric_values, xr.DataArray): + translated_values = numeric_values.copy( + deep=True, data=translated_values + ) + return translated_values + + @pytest.fixture + def spinning_on_the_spot(self) -> xr.DataArray: + """Simulate data for an individual's head spinning on the spot. + + The left / right keypoints move in a circular motion counter-clockwise + around the unit circle centred on the origin, always opposite each + other. + The left keypoint starts on the negative x-axis, and the motion is + split into 8 time points of uniform rotation angles. + """ + data = np.zeros(shape=(8, 2, 2), dtype=float) + data[:, :, 0] = np.array( + [ + -self.x_axis, + (-self.x_axis - self.y_axis) / self.sqrt_2, + -self.y_axis, + (self.x_axis - self.y_axis) / self.sqrt_2, + self.x_axis, + (self.x_axis + self.y_axis) / self.sqrt_2, + self.y_axis, + (-self.x_axis + self.y_axis) / self.sqrt_2, + ] + ) + data[:, :, 1] = -data[:, :, 0] + return xr.DataArray( + data=data, + dims=["time", "space", "keypoints"], + coords={"space": ["x", "y"], "keypoints": ["left", "right"]}, + ) + + @pytest.mark.parametrize( + ["swap_left_right", "swap_camera_view", "swap_angle_rotation"], + [ + pytest.param(True, True, True, id="(TTT) LR, Camera, Angle"), + pytest.param(True, True, False, id="(TTF) LR, Camera"), + pytest.param(True, False, True, id="(TFT) LR, Angle"), + pytest.param(True, False, False, id="(TFF) LR"), + pytest.param(False, True, True, id="(FTT) Camera, Angle"), + pytest.param(False, True, False, id="(FTF) Camera"), + pytest.param(False, False, True, id="(FFT) Angle"), + pytest.param(False, False, False, id="(FFF)"), + ], + ) + def test_antisymmetry_properties( + self, + spinning_on_the_spot: xr.DataArray, + swap_left_right: bool, + swap_camera_view: bool, + swap_angle_rotation: bool, + ) -> None: + r"""Test antisymmetry arises where expected. + + Reversing the right and left keypoints, or the camera position, has the + effect of mapping angles to the "opposite side" of the unit circle. + Explicitly; + - :math:`\theta <= 0` is mapped to :math:`\theta + 180`, + - :math:`\theta > 0` is mapped to :math:`\theta - 180`. + + In theory, the antisymmetry of ``angle_rotates`` should be covered by + the underlying tests for ``compute_signed_angle_2d``, however we + include this case here for additional checks in conjunction with other + behaviour. + """ + reference_vector = self.x_axis + left_keypoint = "left" + right_keypoint = "right" + + args_to_function = {} + if swap_left_right: + args_to_function["left_keypoint"] = right_keypoint + args_to_function["right_keypoint"] = left_keypoint + else: + args_to_function["left_keypoint"] = left_keypoint + args_to_function["right_keypoint"] = right_keypoint + if swap_camera_view: + args_to_function["camera_view"] = "bottom_up" + if swap_angle_rotation: + args_to_function["angle_rotates"] = "forward to ref" + + # mypy call here is angry, https://github.com/python/mypy/issues/1969 + with_orientations_swapped = kinematics.compute_forward_vector_angle( + data=spinning_on_the_spot, + reference_vector=reference_vector, + **args_to_function, # type: ignore[arg-type] + ) + without_orientations_swapped = kinematics.compute_forward_vector_angle( + data=spinning_on_the_spot, + left_keypoint=left_keypoint, + right_keypoint=right_keypoint, + reference_vector=reference_vector, + ) + + expected_orientations = without_orientations_swapped.copy(deep=True) + if swap_left_right: + expected_orientations = self.push_into_range( + expected_orientations + 180.0 + ) + if swap_camera_view: + expected_orientations = self.push_into_range( + expected_orientations + 180.0 + ) + if swap_angle_rotation: + expected_orientations *= -1.0 + expected_orientations = self.push_into_range(expected_orientations) + + xr.testing.assert_allclose( + with_orientations_swapped, expected_orientations + ) + + @pytest.mark.parametrize( + ["transformation"], + [pytest.param("scale"), pytest.param("translation")], + ) + def test_transformation_invariance( + self, + spinning_on_the_spot: xr.DataArray, + transformation: Literal["scale", "translation"], + ) -> None: + """Test that certain transforms of the data have no effect on + the relative angle computed. + + - Translations applied to both keypoints (even if the translation + changes with time) should not affect the result, so long as both + keypoints receive the same translation (at each timepoint). + - Scaling the right to left keypoint vector should not produce a + different angle. + """ + left_keypoint = "left" + right_keypoint = "right" + reference_vector = self.x_axis + + translated_data = spinning_on_the_spot.values.copy() + n_time_pts = translated_data.shape[0] + + if transformation == "translation": + # Effectively, the data is being translated (1,1)/time-point, + # but its keypoints are staying in the same relative positions. + translated_data += np.arange(n_time_pts).reshape(n_time_pts, 1, 1) + elif transformation == "scale": + # The left keypoint position is "stretched" further away from the + # origin over time; for the time-point at index t, + # a scale factor of (t+1) is applied to the left keypoint. + # The right keypoint remains unscaled, but remains in the same + # direction away from the left keypoint. + translated_data[:, :, 0] *= np.arange(1, n_time_pts + 1).reshape( + n_time_pts, 1 + ) + else: + raise ValueError(f"Did not recognise case: {transformation}") + translated_data = spinning_on_the_spot.copy( + deep=True, data=translated_data + ) + + untranslated_output = kinematics.compute_forward_vector_angle( + spinning_on_the_spot, + left_keypoint=left_keypoint, + right_keypoint=right_keypoint, + reference_vector=reference_vector, + ) + translated_output = kinematics.compute_forward_vector_angle( + spinning_on_the_spot, + left_keypoint=left_keypoint, + right_keypoint=right_keypoint, + reference_vector=reference_vector, + ) + + xr.testing.assert_allclose(untranslated_output, translated_output) + + @pytest.mark.parametrize( + ["args_to_fn", "expected_error"], + [ + pytest.param( + { + "reference_vector": x_axis, + "angle_rotates": "elephants first", + }, + ValueError("Unknown angle convention: 'elephants first'"), + id="Unknown angle orientation", + ), + pytest.param( + { + "reference_vector": np.array([1.0, 0.0, 0.0]), + }, + ValueError( + "conflicting sizes for dimension 'space': length 3 on the " + "data but length 2 on coordinate 'space'" + ), + id="Reference vector not of shape (2,)", + ), + ], + ) + def test_error_cases( + self, + spinning_on_the_spot: xr.DataArray, + args_to_fn: dict[str, Any], + expected_error: Exception, + ) -> None: + """Test that the angle orientation provided has to be recognised.""" + with pytest.raises( + type(expected_error), + match=re.escape(str(expected_error)), + ): + kinematics.compute_forward_vector_angle( + spinning_on_the_spot, + "left", + "right", + **args_to_fn, + ) + + def test_casts_from_tuple( + self, spinning_on_the_spot: xr.DataArray + ) -> None: + """Test that tuples and lists are cast to numpy arrays, + when given as the reference vector. + """ + x_axis_as_tuple = (1.0, 0.0) + x_axis_as_list = [1.0, 0.0] + + pass_numpy = kinematics.compute_forward_vector_angle( + spinning_on_the_spot, "left", "right", self.x_axis + ) + pass_tuple = kinematics.compute_forward_vector_angle( + spinning_on_the_spot, "left", "right", x_axis_as_tuple + ) + pass_list = kinematics.compute_forward_vector_angle( + spinning_on_the_spot, "left", "right", x_axis_as_list + ) + + xr.testing.assert_allclose(pass_numpy, pass_tuple) + xr.testing.assert_allclose(pass_numpy, pass_list) diff --git a/tests/test_unit/test_validators/test_array_validators.py b/tests/test_unit/test_validators/test_array_validators.py index a1a4412c..674d7b80 100644 --- a/tests/test_unit/test_validators/test_array_validators.py +++ b/tests/test_unit/test_validators/test_array_validators.py @@ -11,46 +11,62 @@ def expect_value_error_with_message(error_msg): return pytest.raises(ValueError, match=re.escape(error_msg)) +# dims_coords: dict, exact_coords: bool, expected_exception valid_cases = [ - ({"time": []}, does_not_raise()), - ({"time": [0, 1]}, does_not_raise()), - ({"space": ["x", "y"]}, does_not_raise()), - ({"time": [], "space": []}, does_not_raise()), - ({"time": [], "space": ["x", "y"]}, does_not_raise()), + ({"time": []}, False, does_not_raise()), + ({"time": []}, True, does_not_raise()), + ({"time": [0, 1]}, False, does_not_raise()), + ({"space": ["x", "y"]}, False, does_not_raise()), + ({"space": ["x", "y"]}, True, does_not_raise()), + ({"time": [], "space": []}, False, does_not_raise()), + ({"time": [], "space": ["x", "y"]}, False, does_not_raise()), ] # Valid cases (no error) invalid_cases = [ ( {"spacetime": []}, + False, expect_value_error_with_message( "Input data must contain ['spacetime'] as dimensions." ), ), ( {"time": [0, 100], "space": ["x", "y"]}, + False, expect_value_error_with_message( "Input data must contain [100] in the 'time' coordinates." ), ), ( {"space": ["x", "y", "z"]}, + False, expect_value_error_with_message( "Input data must contain ['z'] in the 'space' coordinates." ), ), + ( + {"space": ["x"]}, + True, + expect_value_error_with_message( + "Dimension 'space' must only contain ['x'] as coordinates, " + ), + ), ] # Invalid cases (raise ValueError) @pytest.mark.parametrize( - "required_dims_coords, expected_exception", + "required_dims_coords, exact_coords, expected_exception", valid_cases + invalid_cases, ) def test_validate_dims_coords( valid_poses_dataset_uniform_linear_motion, # fixture from conftest.py required_dims_coords, + exact_coords, expected_exception, ): """Test validate_dims_coords for both valid and invalid inputs.""" position_array = valid_poses_dataset_uniform_linear_motion["position"] with expected_exception: - validate_dims_coords(position_array, required_dims_coords) + validate_dims_coords( + position_array, required_dims_coords, exact_coords=exact_coords + ) diff --git a/tests/test_unit/test_validators/test_ref_vector_validator.py b/tests/test_unit/test_validators/test_ref_vector_validator.py new file mode 100644 index 00000000..ca0ecdb6 --- /dev/null +++ b/tests/test_unit/test_validators/test_ref_vector_validator.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import re + +import numpy as np +import pytest +import xarray as xr + +from movement.validators.arrays import validate_reference_vector + + +@pytest.fixture +def x_axis() -> xr.DataArray: + return xr.DataArray( + data=np.array([1.0, 0.0]), + coords={"space": ["x", "y"]}, + ) + + +def x_axis_over_time(n_time_pts: int = 10) -> xr.DataArray: + """Return the x-axis as a constant DataArray, + with as many time points as requested. + """ + return xr.DataArray( + data=np.tile(np.array([1.0, 0.0]).reshape(1, -1), [n_time_pts, 1]), + dims=["time", "space"], + coords={"time": np.arange(n_time_pts), "space": ["x", "y"]}, + ) + + +@pytest.mark.parametrize( + ["ref_vector", "test_vector", "expected_result"], + [ + pytest.param( + np.array([1.0, 0.0]), + x_axis_over_time(), + "x_axis", + id="(2,) numpy array vs (n, 2) DataArray", + ), + pytest.param( + np.array([1.0, 0.0]), + "x_axis", + "x_axis", + id="(2,) numpy array vs (1, 2) DataArray", + ), + pytest.param( + np.array([1.0, 0.0]).reshape(1, -1), + x_axis_over_time(), + "x_axis", + id="(1, 2) numpy array vs (n, 2) DataArray", + ), + pytest.param( + "x_axis", + x_axis_over_time(), + "x_axis", + id="(1, 2) DataArray vs (n, 2) DataArray", + ), + pytest.param( + x_axis_over_time(), + x_axis_over_time(), + x_axis_over_time(), + id="(n, 2) DataArray vs (n, 2) DataArray", + ), + pytest.param( + np.tile(np.array([1.0, 0.0]).reshape(1, -1), [10, 1]), + x_axis_over_time(), + x_axis_over_time(), + id="(n, 2) numpy array vs (n, 2) DataArray", + ), + pytest.param( + x_axis_over_time(5), + x_axis_over_time(), + ValueError( + "Reference vector must have the same number " + "of time points as the test vector." + ), + id="Too few time points (numpy)", + ), + pytest.param( + np.tile(np.array([1.0, 0.0]).reshape(1, -1), [5, 1]), + x_axis_over_time(), + ValueError( + "Reference vector must have the same number " + "of time points as the test vector." + ), + id="Too few time points (DataArray)", + ), + pytest.param( + np.ones(shape=(2, 2, 2)), + x_axis_over_time(), + ValueError("Reference vector must be 1D or 2D, but got 3D array."), + id="Too many dimensions", + ), + pytest.param( + xr.DataArray( + data=np.ones(shape=(10, 2, 1)), + dims=["time", "space", "elephants"], + coords={ + "time": np.arange(10), + "space": ["x", "y"], + "elephants": ["e"], + }, + ), + x_axis_over_time(10), + ValueError( + "Only 'time' and 'space' dimensions " + "are allowed in reference_vector.", + ), + id="Extra dimension in reference vector", + ), + pytest.param( + np.pi, + "x_axis", + TypeError( + "Reference vector must be an xarray.DataArray or np.ndarray, " + "but got ." + ), + id="Wrong input type", + ), + ], +) +def test_validate_reference_vector( + ref_vector: xr.DataArray | np.ndarray, + test_vector: xr.DataArray, + expected_result: xr.DataArray | Exception, + request, +) -> None: + """Test that reference vectors or objects that can be cast to reference + vectors, are cast correctly. + """ + if isinstance(ref_vector, str): + ref_vector = request.getfixturevalue(ref_vector) + if isinstance(test_vector, str): + test_vector = request.getfixturevalue(test_vector) + if isinstance(expected_result, str): + expected_result = request.getfixturevalue(expected_result) + + if isinstance(expected_result, Exception): + with pytest.raises( + type(expected_result), match=re.escape(str(expected_result)) + ): + validate_reference_vector( + reference_vector=ref_vector, test_vector=test_vector + ) + else: + returned_vector = validate_reference_vector( + reference_vector=ref_vector, test_vector=test_vector + ) + + xr.testing.assert_allclose(returned_vector, expected_result) diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index 8787a468..db43b030 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -188,3 +188,112 @@ def test_convert_to_unit(self, ds, expected_exception, request): expected_unit_pol.loc[{"space_pol": "rho"}] = 1 expected_unit_pol = expected_unit_pol.where(~expected_nan_idxs) xr.testing.assert_allclose(unit_pol, expected_unit_pol) + + +class TestComputeSignedAngle: + """Tests for the compute_signed_angle_2d method.""" + + x_axis = np.array([1.0, 0.0]) + y_axis = np.array([0.0, 1.0]) + coord_axes_array = np.array([x_axis, y_axis, -x_axis, -y_axis]) + + @pytest.mark.parametrize( + ["left_vector", "right_vector", "expected_angles"], + [ + pytest.param( + x_axis.reshape(1, 2), + x_axis, + [0.0], + id="x-axis to x-axis", + ), + pytest.param( + coord_axes_array, + x_axis, + [0.0, -np.pi / 2.0, np.pi, np.pi / 2.0], + id="+/- axes to x-axis", + ), + pytest.param( + coord_axes_array, + -x_axis, + [np.pi, np.pi / 2.0, 0.0, -np.pi / 2.0], + id="+/- axes to -ve x-axis", + ), + pytest.param( + np.array([[1.0, 1.0], [-1.0, 1.0], [-1.0, -1.0], [1.0, -1.0]]), + coord_axes_array, + [-np.pi / 4.0] * 4, + id="-pi/4 trailing axes", + ), + pytest.param( + xr.DataArray( + data=coord_axes_array, + dims=["time", "space"], + coords={ + "time": np.arange(coord_axes_array.shape[0]), + "space": ["x", "y"], + }, + ), + xr.DataArray( + data=-1.0 * coord_axes_array, + dims=["time", "space"], + coords={ + "time": np.arange(coord_axes_array.shape[0]), + "space": ["x", "y"], + }, + ), + [np.pi] * 4, + id="Two DataArrays given", + ), + pytest.param( + np.array([-x_axis, x_axis]), + np.array([x_axis, -x_axis]), + [np.pi, np.pi], + id="Rotation by '-pi' should map to pi.", + ), + ], + ) + def test_compute_signed_angle_2d( + self, + left_vector: xr.DataArray | np.ndarray, + right_vector: xr.DataArray | np.ndarray, + expected_angles: xr.DataArray, + ) -> None: + """Test computed angles are what we expect. + + This test also checks the antisymmetry of the function in question. + Swapping the ``u`` and ``v`` arguments should produce an array of + angles with the same magnitude but opposite signs + (except for pi -> -pi). + """ + if not isinstance(left_vector, xr.DataArray): + left_vector = xr.DataArray( + data=left_vector, + dims=["time", "space"], + coords={ + "time": np.arange(left_vector.shape[0]), + "space": ["x", "y"], + }, + ) + if not isinstance(expected_angles, xr.DataArray): + expected_angles = xr.DataArray( + data=np.array(expected_angles), + dims=["time"], + coords={"time": left_vector["time"]} + if "time" in left_vector.dims + else None, + ) + # pi and -pi should map to the same angle, regardless! + expected_angles_reversed = expected_angles.copy(deep=True) + expected_angles_reversed[expected_angles < np.pi] *= -1.0 + + computed_angles = vector.compute_signed_angle_2d( + left_vector, right_vector + ) + computed_angles_reversed = vector.compute_signed_angle_2d( + left_vector, right_vector, v_as_left_operand=True + ) + + xr.testing.assert_allclose(computed_angles, expected_angles) + xr.testing.assert_allclose( + computed_angles_reversed, expected_angles_reversed + )