Merge branch 'sp/366-transforms-scale' of https://github.com/stellapr…

…ins/movement into sp/366-transforms-scale

.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/[email protected]
       with:
-        secret-pypi-key: ${{ secrets.TWINE_API_KEY }}
+        user: __token__
+        password: ${{ secrets.TWINE_API_KEY }}

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\
+    <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
+    ----------
+    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

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,