Compute Head Direction Vector

examples/compute_polar_coordinates.py

@@ -15,6 +15,7 @@
 from matplotlib import pyplot as plt
 
 from movement import sample_data
+from movement.analysis.kinematics import compute_head_direction_vector
 from movement.io import load_poses
 from movement.utils.vector import cart2pol, pol2cart
 
@@ -48,20 +49,14 @@
 # To demonstrate how polar coordinates can be useful in behavioural analyses,
 # we will compute the head vector of the mouse.
 #
-# We define it as the vector from the midpoint between the ears to the snout.
+# In ``movement``, head vector is defined as the vector perpendicular to the
+# line connecting two symmetrical keypoints on either side of the head (usually
+# the ears), pointing forwards. (See :func:`here\
+# <movement.analysis.kinematics.compute_forward_vector>` for a more
+# detailed explanation).
 
-# compute the midpoint between the ears
-midpoint_ears = position.sel(keypoints=["left_ear", "right_ear"]).mean(
-    dim="keypoints"
-)
+head_vector = compute_head_direction_vector(position, "left_ear", "right_ear")
 
-# compute the head vector
-head_vector = position.sel(keypoints="snout") - midpoint_ears
-
-# drop the keypoints dimension
-# (otherwise the `head_vector` data array retains a `snout` keypoint from the
-# operation above)
-head_vector = head_vector.drop_vars("keypoints")
 
 # %%
 # Visualise the head trajectory
@@ -72,9 +67,12 @@
 # We can start by plotting the trajectory of the midpoint between the ears. We
 # will refer to this as the head trajectory.
 
-fig, ax = plt.subplots(1, 1)
+midpoint_ears = position.sel(keypoints=["left_ear", "right_ear"]).mean(
+    dim="keypoints"
+)
 mouse_name = ds.individuals.values[0]
 
+fig, ax = plt.subplots(1, 1)
 sc = ax.scatter(
     midpoint_ears.sel(individuals=mouse_name, space="x"),
     midpoint_ears.sel(individuals=mouse_name, space="y"),

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
 
 
 def compute_displacement(data: xr.DataArray) -> xr.DataArray:
@@ -161,6 +165,150 @@ def _compute_approximate_time_derivative(
     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 symmetrical 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.
+
+    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.
+
+    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.
+
+    """
+    # Validate input data
+    _validate_type_data_array(data)
+    _validate_time_keypoints_space_dimensions(data)
+    if len(data.space) != 2:
+        raise log_error(
+            ValueError,
+            "Input data must have 2 (and only 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 this 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_time_dimension(data: xr.DataArray) -> None:
     """Validate the input data contains a ``time`` dimension.
 
@@ -179,3 +327,46 @@ def _validate_time_dimension(data: xr.DataArray) -> None:
         raise log_error(
             ValueError, "Input data must contain 'time' as a dimension."
         )
+
+
+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)}.",
+        )
+
+
+def _validate_time_keypoints_space_dimensions(data: xr.DataArray) -> None:
+    """Validate if input data contains ``time``, ``keypoints`` and ``space``.
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+        The input data to validate.
+
+    Raises
+    ------
+    ValueError
+        If the input data is not an xarray DataArray.
+
+    """
+    if not all(coord in data.dims for coord in ["time", "keypoints", "space"]):
+        raise log_error(
+            AttributeError,
+            "Input data must contain 'time', 'space', and 'keypoints' as "
+            "dimensions.",
+        )