Validate input data in kinematic functions

Author: lochhh

movement/analysis/kinematics.py

@@ -1,36 +1,37 @@
 import numpy as np
 import xarray as xr
 
+from movement.logging import log_error
+
 
 def compute_displacement(data: xr.DataArray) -> xr.DataArray:
-    """Compute the displacement between consecutive x, y
-    locations of each keypoint of each individual.
+    """Compute the displacement between consecutive locations
+    of each keypoint of each individual across time.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data, assumed to be of shape (..., 2), where
-        the last dimension contains the x and y coordinates.
+        The input data containing `time` as a dimension.
 
     Returns
     -------
     xarray.DataArray
         An xarray DataArray containing the computed displacement.
     """
-    displacement_xy = data.diff(dim="time")
-    displacement_xy = displacement_xy.reindex(data.coords, fill_value=0)
-    return displacement_xy
+    _validate_time_dimension(data)
+    result = data.diff(dim="time")
+    result = result.reindex(data.coords, fill_value=0)
+    return result
 
 
 def compute_velocity(data: xr.DataArray) -> xr.DataArray:
-    """Compute the velocity between consecutive x, y locations
-    of each keypoint of each individual.
+    """Compute the velocity between consecutive locations
+    of each keypoint of each individual across time.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data, assumed to be of shape (..., 2), where the last
-        dimension contains the x and y coordinates.
+        The input data containing `time` as a dimension.
 
     Returns
     -------
@@ -41,14 +42,13 @@ def compute_velocity(data: xr.DataArray) -> xr.DataArray:
 
 
 def compute_acceleration(data: xr.DataArray) -> xr.DataArray:
-    """Compute the acceleration between consecutive x, y
-    locations of each keypoint of each individual.
+    """Compute the acceleration between consecutive locations
+    of each keypoint of each individual.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data, assumed to be of shape (..., 2), where the last
-        dimension contains the x and y coordinates.
+        The input data containing `time` as a dimension.
 
     Returns
     -------
@@ -60,48 +60,58 @@ def compute_acceleration(data: xr.DataArray) -> xr.DataArray:
 
 
 def compute_approximate_derivative(
-    data: xr.DataArray, order: int = 1
+    data: xr.DataArray, order: int
 ) -> xr.DataArray:
     """Compute velocity or acceleration using numerical differentiation,
     assuming equidistant time spacing.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data, assumed to be of shape (..., 2), where the last
-        dimension contains data in the x and y dimensions.
+        The input data containing `time` as a dimension.
     order : int
         The order of the derivative. 1 for velocity, 2 for
-        acceleration. Default is 1.
+        acceleration. Value must be a positive integer.
 
     Returns
     -------
     xarray.DataArray
         An xarray DataArray containing the derived variable.
     """
+    if not isinstance(order, int):
+        raise log_error(
+            TypeError, f"Order must be an integer, but got {type(order)}."
+        )
     if order <= 0:
-        raise ValueError("order must be a positive integer.")
-    else:
-        result = data
-        dt = data["time"].values[1] - data["time"].values[0]
-        for _ in range(order):
-            result = xr.apply_ufunc(
-                np.gradient,
-                result,
-                dt,
-                kwargs={"axis": 0},
-            )
-        result = result.reindex_like(data)
+        raise log_error(ValueError, "Order must be a positive integer.")
+    _validate_time_dimension(data)
+    result = data
+    dt = data["time"].values[1] - data["time"].values[0]
+    for _ in range(order):
+        result = xr.apply_ufunc(
+            np.gradient,
+            result,
+            dt,
+            kwargs={"axis": 0},
+        )
+    result = result.reindex_like(data)
     return result
 
 
-# Locomotion Features
-# speed
-# speed_centroid
-# acceleration
-# acceleration_centroid
-# speed_fwd
-# radial_vel
-# tangential_vel
-# speed_centroid_w(s)
-# speed_(p)_w(s)
+def _validate_time_dimension(data: xr.DataArray) -> None:
+    """Validate the input data contains a 'time' dimension.
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+        The input data to validate.
+
+    Raises
+    ------
+    ValueError
+        If the input data does not contain a 'time' dimension.
+    """
+    if "time" not in data.dims:
+        raise log_error(
+            ValueError, "Input data must contain 'time' as a dimension."
+        )

tests/conftest.py

@@ -222,7 +222,7 @@ def _valid_tracks_array(array_type):
             n_keypoints = 1
         elif array_type == "single_track_array":
             n_individuals = 1
-        x_points = np.repeat(base * 3, n_individuals * n_keypoints)
+        x_points = np.repeat(base * base, n_individuals * n_keypoints)
         y_points = np.repeat(base * 4, n_individuals * n_keypoints)
         tracks_array = np.ravel(np.column_stack((x_points, y_points)))
         return tracks_array.reshape(n_frames, n_individuals, n_keypoints, 2)

tests/test_unit/test_kinematics.py

@@ -10,13 +10,41 @@ class TestKinematics:
 
     @pytest.fixture
     def expected_dataarray(self, valid_pose_dataset):
-        """Return an xarray.DataArray with default values and
-        the expected dimensions and coordinates."""
-        return xr.DataArray(
-            np.full((10, 2, 2, 2), [3.0, 4.0]),
-            dims=valid_pose_dataset.dims,
-            coords=valid_pose_dataset.coords,
-        )
+        """Return a function to generate the expected dataarray
+        for different kinematic properties."""
+
+        def _expected_dataarray(property):
+            """Return an xarray.DataArray with default values and
+            the expected dimensions and coordinates."""
+            # Expected x,y values for velocity
+            x_vals = np.array(
+                [1.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 17.0]
+            )
+            y_vals = np.full((10, 2, 2, 1), 4.0)
+            if property == "acceleration":
+                x_vals = np.array(
+                    [1.0, 1.5, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 1.5, 1.0]
+                )
+                y_vals = np.full((10, 2, 2, 1), 0)
+            elif property == "displacement":
+                x_vals = np.array(
+                    [0.0, 1.0, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0, 17.0]
+                )
+                y_vals[0] = 0
+
+            x_vals = x_vals.reshape(-1, 1, 1, 1)
+            # Repeat the x_vals to match the shape of the pose_tracks
+            x_vals = np.tile(x_vals, (1, 2, 2, 1))
+            return xr.DataArray(
+                np.concatenate(
+                    [x_vals, y_vals],
+                    axis=-1,
+                ),
+                dims=valid_pose_dataset.dims,
+                coords=valid_pose_dataset.coords,
+            )
+
+        return _expected_dataarray
 
     @pytest.fixture
     def expected_dataset(self, valid_pose_dataset):
@@ -61,25 +89,36 @@ def test_displacement(self, valid_pose_dataset, expected_dataarray):
         result = kinematics.compute_displacement(
             valid_pose_dataset.pose_tracks
         )
-        # Set the first displacement to zero
-        expected_dataarray[0, :, :, :] = 0
-        xr.testing.assert_allclose(result, expected_dataarray)
+        xr.testing.assert_allclose(result, expected_dataarray("displacement"))
 
     def test_velocity(self, valid_pose_dataset, expected_dataarray):
         """Test velocity computation."""
         result = kinematics.compute_velocity(valid_pose_dataset.pose_tracks)
-        xr.testing.assert_allclose(result, expected_dataarray)
+        xr.testing.assert_allclose(result, expected_dataarray("velocity"))
 
     def test_acceleration(self, valid_pose_dataset, expected_dataarray):
         """Test acceleration computation."""
         result = kinematics.compute_acceleration(
             valid_pose_dataset.pose_tracks
         )
-        expected_dataarray[:] = 0
-        xr.testing.assert_allclose(result, expected_dataarray)
+        xr.testing.assert_allclose(result, expected_dataarray("acceleration"))
 
-    def test_approximate_derivative_with_nonpositive_order(self):
+    @pytest.mark.parametrize("order", [0, -1, 1.0, "1"])
+    def test_approximate_derivative_with_invalid_order(self, order):
         """Test that an error is raised when the order is non-positive."""
         data = np.arange(10)
+        expected_exception = (
+            ValueError if isinstance(order, int) else TypeError
+        )
+        with pytest.raises(expected_exception):
+            kinematics.compute_approximate_derivative(data, order=order)
+
+    def test_compute_with_missing_time_dimension(
+        self, missing_dim_dataset, kinematic_property
+    ):
+        """Test that computing a property of a pose dataset with
+        missing 'time' dimension raises the appropriate error."""
         with pytest.raises(ValueError):
-            kinematics.compute_approximate_derivative(data, order=0)
+            eval(f"kinematics.compute_{kinematic_property}")(
+                missing_dim_dataset
+            )