diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index b2bbbf9bc..b1bce6ac4 100644 --- a/movement/analysis/kinematics.py +++ b/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()\ + `. 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)}.", + ) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index a1b933e0b..a54d199ca 100644 --- a/tests/test_unit/test_kinematics.py +++ b/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() + )