Compute Forward Vector (#276)

* Basic implementation of `compute_head_direction_vector()`

* Minor fixes docstring

* Added unit test for `compute_head_direction_vector()`

* Bug fixes for `test_compute_head_direction_vector()`

* Added validator (and test) to ensure input is 2D

* Refactored `navigation.py` and implemented PR review feedback

* Extended testing and added `front_keypoint` argument to `compute_head_direction_vector()`

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Implemented PR feedback and bugfixes for `compute_2d_head_direction_vector()`

* Removed uppercase letters from function names

* Tweaked `compute_polar_coordinates.py` to use new function

* Implemented feedback from Zulip discussion and created `compute_head_direction_vector()` alias function

* Fixed typo in docstring

* Tweaked `compute_forward_vector()` to use new validator

* More fixes for `test_kinematics.py`

* Bugfix for `compute_forward_vector()` and expanded `test_compute_forward_vector` to cover both `camera_view` options

* Added test coverage for `compute_head_direction_vector()` alias

* Reversed changes to `compute_polar_coordinates.py` and implemented final feedback

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

movement/analysis/kinematics.py

@@ -1,8 +1,12 @@
 """Compute kinematic variables like velocity and acceleration."""
 
+from typing import Literal
+
+import numpy as np
 import xarray as xr
 
 from movement.utils.logging import log_error
+from movement.utils.vector import compute_norm
 from movement.validators.arrays import validate_dims_coords
 
 
@@ -165,3 +169,177 @@ def compute_time_derivative(data: xr.DataArray, order: int) -> xr.DataArray:
     for _ in range(order):
         result = result.differentiate("time")
     return result
+
+
+def compute_forward_vector(
+    data: xr.DataArray,
+    left_keypoint: str,
+    right_keypoint: str,
+    camera_view: Literal["top_down", "bottom_up"] = "top_down",
+):
+    """Compute a 2D forward vector given two left-right symmetric keypoints.
+
+    The forward vector is computed as a vector perpendicular to the
+    line connecting two symmetrical keypoints on either side of the body
+    (i.e., symmetrical relative to the mid-sagittal plane), and pointing
+    forwards (in the rostral direction). A top-down or bottom-up view of the
+    animal is assumed (see Notes).
+
+    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"
+    right_keypoint : str
+        Name of the right keypoint, e.g., "right_ear"
+    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"``.
+
+    Returns
+    -------
+    xarray.DataArray
+        An xarray DataArray representing the forward vector, with
+        dimensions matching the input data array, but without the
+        ``keypoints`` dimension.
+
+    Notes
+    -----
+    To determine the forward direction of the animal, we need to specify
+    (1) the right-to-left direction of the animal and (2) its upward direction.
+    We determine the right-to-left direction via the input left and right
+    keypoints. The upwards direction, in turn, can be determined by passing the
+    ``camera_view`` argument with either ``"top_down"`` or ``"bottom_up"``. If
+    the camera view is specified as being ``"top_down"``, or if no additional
+    information is provided, we assume that the upwards direction matches that
+    of the vector ``[0, 0, -1]``. If the camera view is ``"bottom_up"``, the
+    upwards direction is assumed to be given by ``[0, 0, 1]``. For both cases,
+    we assume that position values are expressed in the image coordinate
+    system (where the positive X-axis is oriented to the right, the positive
+    Y-axis faces downwards, and positive Z-axis faces away from the person
+    viewing the screen).
+
+    If one of the required pieces of information is missing for a frame (e.g.,
+    the left keypoint is not visible), then the computed head direction vector
+    is set to NaN.
+
+    """
+    # Validate input data
+    _validate_type_data_array(data)
+    validate_dims_coords(
+        data,
+        {
+            "time": [],
+            "keypoints": [left_keypoint, right_keypoint],
+            "space": [],
+        },
+    )
+    if len(data.space) != 2:
+        raise log_error(
+            ValueError,
+            "Input data must have exactly 2 spatial dimensions, but "
+            f"currently has {len(data.space)}.",
+        )
+
+    # Validate input keypoints
+    if left_keypoint == right_keypoint:
+        raise log_error(
+            ValueError, "The left and right keypoints may not be identical."
+        )
+
+    # Define right-to-left vector
+    right_to_left_vector = data.sel(
+        keypoints=left_keypoint, drop=True
+    ) - data.sel(keypoints=right_keypoint, drop=True)
+
+    # Define upward vector
+    # default: negative z direction in the image coordinate system
+    if camera_view == "top_down":
+        upward_vector = np.array([0, 0, -1])
+    else:
+        upward_vector = np.array([0, 0, 1])
+
+    upward_vector = xr.DataArray(
+        np.tile(upward_vector.reshape(1, -1), [len(data.time), 1]),
+        dims=["time", "space"],
+    )
+
+    # 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
+
+    # Return unit vector
+
+    return forward_vector / compute_norm(forward_vector)
+
+
+def compute_head_direction_vector(
+    data: xr.DataArray,
+    left_keypoint: str,
+    right_keypoint: str,
+    camera_view: Literal["top_down", "bottom_up"] = "top_down",
+):
+    """Compute the 2D head direction vector given two keypoints on the head.
+
+    This function is an alias for :func:`compute_forward_vector()\
+    <movement.analysis.kinematics.compute_forward_vector>`. For more
+    detailed information on how the head direction vector is computed,
+    please refer to the documentation for that function.
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+        The input data representing position. This must contain
+        the two chosen keypoints corresponding to the left and
+        right of the head.
+    left_keypoint : str
+        Name of the left keypoint, e.g., "left_ear"
+    right_keypoint : str
+        Name of the right keypoint, e.g., "right_ear"
+    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"``.
+
+    Returns
+    -------
+    xarray.DataArray
+        An xarray DataArray representing the head direction vector, with
+        dimensions matching the input data array, but without the
+        ``keypoints`` dimension.
+
+    """
+    return compute_forward_vector(
+        data, left_keypoint, right_keypoint, camera_view=camera_view
+    )
+
+
+def _validate_type_data_array(data: xr.DataArray) -> None:
+    """Validate the input data is an xarray DataArray.
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+        The input data to validate.
+
+    Raises
+    ------
+    ValueError
+        If the input data is not an xarray DataArray.
+
+    """
+    if not isinstance(data, xr.DataArray):
+        raise log_error(
+            TypeError,
+            f"Input data must be an xarray.DataArray, but got {type(data)}.",
+        )

tests/test_unit/test_kinematics.py

@@ -1,3 +1,5 @@
+import re
+
 import numpy as np
 import pytest
 import xarray as xr
@@ -182,3 +184,168 @@ def test_approximate_derivative_with_invalid_order(order):
     expected_exception = ValueError if isinstance(order, int) else TypeError
     with pytest.raises(expected_exception):
         kinematics.compute_time_derivative(data, order=order)
+
+
+@pytest.fixture
+def valid_data_array_for_forward_vector():
+    """Return a position data array for an individual with 3 keypoints
+    (left ear, right ear and nose), tracked for 4 frames, in x-y space.
+    """
+    time = [0, 1, 2, 3]
+    individuals = ["individual_0"]
+    keypoints = ["left_ear", "right_ear", "nose"]
+    space = ["x", "y"]
+
+    ds = xr.DataArray(
+        [
+            [[[1, 0], [-1, 0], [0, -1]]],  # time 0
+            [[[0, 1], [0, -1], [1, 0]]],  # time 1
+            [[[-1, 0], [1, 0], [0, 1]]],  # time 2
+            [[[0, -1], [0, 1], [-1, 0]]],  # time 3
+        ],
+        dims=["time", "individuals", "keypoints", "space"],
+        coords={
+            "time": time,
+            "individuals": individuals,
+            "keypoints": keypoints,
+            "space": space,
+        },
+    )
+    return ds
+
+
+@pytest.fixture
+def invalid_input_type_for_forward_vector(valid_data_array_for_forward_vector):
+    """Return a numpy array of position values by individual, per keypoint,
+    over time.
+    """
+    return valid_data_array_for_forward_vector.values
+
+
+@pytest.fixture
+def invalid_dimensions_for_forward_vector(valid_data_array_for_forward_vector):
+    """Return a position DataArray in which the ``keypoints`` dimension has
+    been dropped.
+    """
+    return valid_data_array_for_forward_vector.sel(keypoints="nose", drop=True)
+
+
+@pytest.fixture
+def invalid_spatial_dimensions_for_forward_vector(
+    valid_data_array_for_forward_vector,
+):
+    """Return a position DataArray containing three spatial dimensions."""
+    dataarray_3d = valid_data_array_for_forward_vector.pad(
+        space=(0, 1), constant_values=0
+    )
+    return dataarray_3d.assign_coords(space=["x", "y", "z"])
+
+
+@pytest.fixture
+def valid_data_array_for_forward_vector_with_nans(
+    valid_data_array_for_forward_vector,
+):
+    """Return a position DataArray where position values are NaN for the
+    ``left_ear`` keypoint at time ``1``.
+    """
+    nan_dataarray = valid_data_array_for_forward_vector.where(
+        (valid_data_array_for_forward_vector.time != 1)
+        | (valid_data_array_for_forward_vector.keypoints != "left_ear")
+    )
+    return nan_dataarray
+
+
+def test_compute_forward_vector(valid_data_array_for_forward_vector):
+    """Test that the correct output forward direction vectors
+    are computed from a valid mock dataset.
+    """
+    forward_vector = kinematics.compute_forward_vector(
+        valid_data_array_for_forward_vector,
+        "left_ear",
+        "right_ear",
+        camera_view="bottom_up",
+    )
+    forward_vector_flipped = kinematics.compute_forward_vector(
+        valid_data_array_for_forward_vector,
+        "left_ear",
+        "right_ear",
+        camera_view="top_down",
+    )
+    head_vector = kinematics.compute_head_direction_vector(
+        valid_data_array_for_forward_vector,
+        "left_ear",
+        "right_ear",
+        camera_view="bottom_up",
+    )
+    known_vectors = np.array([[[0, -1]], [[1, 0]], [[0, 1]], [[-1, 0]]])
+
+    assert (
+        isinstance(forward_vector, xr.DataArray)
+        and ("space" in forward_vector.dims)
+        and ("keypoints" not in forward_vector.dims)
+    )
+    assert np.equal(forward_vector.values, known_vectors).all()
+    assert np.equal(forward_vector_flipped.values, known_vectors * -1).all()
+    assert head_vector.equals(forward_vector)
+
+
+@pytest.mark.parametrize(
+    "input_data, expected_error, expected_match_str, keypoints",
+    [
+        (
+            "invalid_input_type_for_forward_vector",
+            TypeError,
+            "must be an xarray.DataArray",
+            ["left_ear", "right_ear"],
+        ),
+        (
+            "invalid_dimensions_for_forward_vector",
+            ValueError,
+            "Input data must contain ['keypoints']",
+            ["left_ear", "right_ear"],
+        ),
+        (
+            "invalid_spatial_dimensions_for_forward_vector",
+            ValueError,
+            "must have exactly 2 spatial dimensions",
+            ["left_ear", "right_ear"],
+        ),
+        (
+            "valid_data_array_for_forward_vector",
+            ValueError,
+            "keypoints may not be identical",
+            ["left_ear", "left_ear"],
+        ),
+    ],
+)
+def test_compute_forward_vector_with_invalid_input(
+    input_data, keypoints, expected_error, expected_match_str, request
+):
+    """Test that ``compute_forward_vector`` catches errors
+    correctly when passed invalid inputs.
+    """
+    # Get fixture
+    input_data = request.getfixturevalue(input_data)
+
+    # Catch error
+    with pytest.raises(expected_error, match=re.escape(expected_match_str)):
+        kinematics.compute_forward_vector(
+            input_data, keypoints[0], keypoints[1]
+        )
+
+
+def test_nan_behavior_forward_vector(
+    valid_data_array_for_forward_vector_with_nans,
+):
+    """Test that ``compute_forward_vector()`` generates the
+    expected output for a valid input DataArray containing ``NaN``
+    position values at a single time (``1``) and keypoint
+    (``left_ear``).
+    """
+    forward_vector = kinematics.compute_forward_vector(
+        valid_data_array_for_forward_vector_with_nans, "left_ear", "right_ear"
+    )
+    assert (
+        np.isnan(forward_vector.values[1, 0, :]).all()
+        and not np.isnan(forward_vector.values[[0, 2, 3], 0, :]).any()
+    )