Add tests for heading angle

movement/kinematics.py

@@ -4,8 +4,8 @@
 from typing import Literal
 
 import numpy as np
-import numpy.typing as npt
 import xarray as xr
+from numpy.typing import ArrayLike
 from scipy.spatial.distance import cdist
 
 from movement.utils.logging import log_error, log_warning
@@ -297,13 +297,17 @@ def compute_forward_vector(
     upward_vector = xr.DataArray(
         np.tile(upward_vector.reshape(1, -1), [len(data.time), 1]),
         dims=["time", "space"],
+        coords={
+            "space": ["x", "y", "z"],
+        },
     )
     # Compute forward direction as the cross product
     # (right-to-left) cross (forward) = up
     forward_vector = xr.cross(
         right_to_left_vector, upward_vector, dim="space"
-    )[:, :, :-1]  # keep only the first 2 dimensions of the result
-
+    ).drop_sel(
+        space="z"
+    )  # keep only the first 2 spatal dimensions of the result
     # Return unit vector
     return convert_to_unit(forward_vector)
 
@@ -351,24 +355,25 @@ def compute_head_direction_vector(
     )
 
 
-def compute_heading_angle(
+def compute_forward_vector_angle(
     data: xr.DataArray,
     left_keypoint: str,
     right_keypoint: str,
-    reference_vector: npt.NDArray | list | tuple = (1, 0),
+    reference_vector: xr.DataArray | ArrayLike = (1, 0),
     camera_view: Literal["top_down", "bottom_up"] = "top_down",
-    in_radians=False,
+    in_radians: bool = False,
+    angle_rotates: Literal[
+        "ref to forward", "forward to ref"
+    ] = "ref to forward",
 ) -> xr.DataArray:
-    r"""Compute the 2D heading angle given two left-right symmetric keypoints.
+    r"""Compute the signed angle between a forward and reference vector.
 
-    Heading angle is defined as the :func:`signed angle\
-    <movement.utils.vector.signed_angle_between_2d_vectors>`
-    between the animal's :func:`forward vector\
-    <movement.kinematics.compute_forward_vector>`)
-    and a reference vector. By default, the reference vector
-    corresponds to the direction of the positive x-axis.
-    The returned angles are in degrees, unless ``in_radians=True``,
-    and span the range (-180, 180] (or :math:`(-\pi, \pi]` in radians).
+    Forward vector angle is the :func:`signed angle\
+    <movement.utils.vector.compute_signed_angle_2d>`
+    between the reference vector and the animal's :func:`forward vector\
+    <movement.kinematics.compute_forward_vector>`).
+    The returned angles are in degrees, spanning the range :math:`(-180, 180]`,
+    unless ``in_radians`` is set to ``True``.
 
     Parameters
     ----------
@@ -377,14 +382,16 @@ def compute_heading_angle(
         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"
+        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"
-    reference_vector : ndt.NDArray | list | tuple, optional
-        The reference vector against which the ```forward_vector`` is
+        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
+        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
@@ -396,41 +403,60 @@ def compute_heading_angle(
     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 heading angles,
+        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.signed_angle_between_2d_vectors : The underlying
+    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.
 
     """
-    # Convert reference vector to np.array if list or tuple
-    if isinstance(reference_vector, (list | tuple)):
-        reference_vector = np.array(reference_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}'")
 
-    # Validate that reference vector has correct dimensionality
-    if reference_vector.shape != (2,):
-        raise log_error(
-            ValueError,
-            f"Reference vector must be two-dimensional (with"
-            f" shape ``(2,)``), but got {reference_vector.shape}.",
-        )
+    # 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)
+    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: