Simplify and expand kinematic tests for bboxes

movement/analysis/kinematics.py

@@ -6,24 +6,40 @@
 
 
 def compute_displacement(data: xr.DataArray) -> xr.DataArray:
-    """Compute displacement between consecutive positions.
+    """Compute displacement array in cartesian coordinates.
 
-    Displacement is the difference between consecutive positions
-    of each keypoint for each individual across ``time``.
-    At each time point ``t``, it is defined as a vector in
-    cartesian ``(x,y)`` coordinates, pointing from the previous
-    ``(t-1)`` to the current ``(t)`` position.
+    The displacement array is defined as the difference between the position
+    array at time point ``t`` and the position array at time point ``t-1``.
+
+    As a result, for a given individual and keypoint, the displacement vector
+    at time point ``t``, is the vector pointing from the previous
+    ``(t-1)`` to the current ``(t)`` position, in cartesian coordinates.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data containing position information, with
-        ``time`` as a dimension.
+        The input data array containing position vectors in cartesian
+        coordinates, with ``time`` as a dimension.
 
     Returns
     -------
     xarray.DataArray
-        An xarray DataArray containing the computed displacement.
+        An xarray DataArray containing displacement vectors in cartesian
+        coordinates.
+
+    Notes
+    -----
+    For the ``position`` array of a ``poses`` dataset, the ``displacement``
+    array will hold the displacement vectors for every keypoint and every
+    individual.
+
+    For the ``position`` array of a ``bboxes`` dataset, the ``displacement``
+    array will hold the displacement vectors for the centroid of every
+    individual bounding box.
+
+    For the ``shape`` array of a ``bboxes`` dataset, the
+    ``displacement`` array will hold vectors with the change in width and
+    height per bounding box, between consecutive time points.
 
     """
     _validate_time_dimension(data)
@@ -33,52 +49,73 @@ def compute_displacement(data: xr.DataArray) -> xr.DataArray:
 
 
 def compute_velocity(data: xr.DataArray) -> xr.DataArray:
-    """Compute the velocity in cartesian ``(x,y)`` coordinates.
+    """Compute velocity array in cartesian coordinates.
 
-    Velocity is the first derivative of position for each keypoint
-    and individual across ``time``, computed with the second order
-    accurate central differences.
+    The velocity array is the first time-derivative of the position
+    array. It is computed by applying the second-order accurate central
+    differences method on the position array.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data containing position information, with
-        ``time`` as a dimension.
+        The input data array containing position vectors in cartesian
+        coordinates, with ``time`` as a dimension.
 
     Returns
     -------
     xarray.DataArray
-        An xarray DataArray containing the computed velocity.
+        An xarray DataArray containing velocity vectors in cartesian
+        coordinates.
+
+    Notes
+    -----
+    For the ``position`` array of a ``poses`` dataset, the ``velocity`` array
+    will hold the velocity vectors for every keypoint and every individual.
+
+    For the ``position`` array of a ``bboxes`` dataset, the ``velocity`` array
+    will hold the velocity vectors for the centroid of every individual
+    bounding box.
 
     See Also
     --------
-    :py:meth:`xarray.DataArray.differentiate` : The underlying method used.
+    :meth:`xarray.DataArray.differentiate` : The underlying method used.
 
     """
     return _compute_approximate_time_derivative(data, order=1)
 
 
 def compute_acceleration(data: xr.DataArray) -> xr.DataArray:
-    """Compute acceleration in cartesian ``(x,y)`` coordinates.
+    """Compute acceleration array in cartesian coordinates.
 
-    Acceleration is the second derivative of position for each keypoint
-    and individual across ``time``, computed with the second order
-    accurate central differences.
+    The acceleration array is the second time-derivative of the
+    position array. It is computed by applying the second-order accurate
+    central differences method on the velocity array.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data containing position information, with
-        ``time`` as a dimension.
+        The input data array containing position vectors in cartesian
+        coordinates, with``time`` as a dimension.
 
     Returns
     -------
     xarray.DataArray
-        An xarray DataArray containing the computed acceleration.
+        An xarray DataArray containing acceleration vectors in cartesian
+        coordinates.
+
+    Notes
+    -----
+    For the ``position`` array of a ``poses`` dataset, the ``acceleration``
+    array will hold the acceleration vectors for every keypoint and every
+    individual.
+
+    For the ``position`` array of a ``bboxes`` dataset, the ``acceleration``
+    array will hold the acceleration vectors for the centroid of every
+    individual bounding box.
 
     See Also
     --------
-    :py:meth:`xarray.DataArray.differentiate` : The underlying method used.
+    :meth:`xarray.DataArray.differentiate` : The underlying method used.
 
     """
     return _compute_approximate_time_derivative(data, order=2)
@@ -87,24 +124,26 @@ def compute_acceleration(data: xr.DataArray) -> xr.DataArray:
 def _compute_approximate_time_derivative(
     data: xr.DataArray, order: int
 ) -> xr.DataArray:
-    """Compute the derivative using numerical differentiation.
+    """Compute the time-derivative of an array using numerical differentiation.
 
-    This function uses :py:meth:`xarray.DataArray.differentiate`,
-    which differentiates the array with the second order
-    accurate central differences.
+    This function uses :meth:`xarray.DataArray.differentiate`,
+    which differentiates the array with the second-order
+    accurate central differences method.
 
     Parameters
     ----------
     data : xarray.DataArray
-        The input data containing ``time`` as a dimension.
+        The input data array containing ``time`` as a dimension.
     order : int
-        The order of the derivative. 1 for velocity, 2 for
-        acceleration. Value must be a positive integer.
+        The order of the time-derivative. For an input containing position
+        data, use 1 to compute velocity, and 2 to compute acceleration. Value
+        must be a positive integer.
 
     Returns
     -------
     xarray.DataArray
-        An xarray DataArray containing the derived variable.
+        An xarray DataArray containing the time-derivative of the
+        input data.
 
     """
     if not isinstance(order, int):
@@ -113,7 +152,9 @@ def _compute_approximate_time_derivative(
         )
     if order <= 0:
         raise log_error(ValueError, "Order must be a positive integer.")
+
     _validate_time_dimension(data)
+
     result = data
     for _ in range(order):
         result = result.differentiate("time")

tests/conftest.py

@@ -238,12 +238,18 @@ def valid_bboxes_arrays_all_zeros():
 
 # --------------------- Bboxes dataset fixtures ----------------------------
 @pytest.fixture
-def valid_bboxes_array():
-    """Return a dictionary of valid non-zero arrays for a
-    ValidBboxesDataset.
-
-    Contains realistic data for 10 frames, 2 individuals, in 2D
-    with 5 low confidence bounding boxes.
+def valid_bboxes_arrays():
+    """Return a dictionary of valid arrays for a
+    ValidBboxesDataset representing a uniform linear motion.
+
+    It represents 2 individuals for 10 frames, in 2D space.
+    - Individual 0 moves along the x=y line from the origin.
+    - Individual 1 moves along the x=-y line line from the origin.
+
+    All confidence values are set to 0.9 except the following which are set
+    to 0.1:
+    - Individual 0 at frames 2, 3, 4
+    - Individual 1 at frames 2, 3
     """
     # define the shape of the arrays
     n_frames, n_individuals, n_space = (10, 2, 2)
@@ -276,22 +282,21 @@ def valid_bboxes_array():
         "position": position,
         "shape": shape,
         "confidence": confidence,
-        "individual_names": ["id_" + str(id) for id in range(n_individuals)],
     }
 
 
 @pytest.fixture
 def valid_bboxes_dataset(
-    valid_bboxes_array,
+    valid_bboxes_arrays,
 ):
-    """Return a valid bboxes dataset with low confidence values and
-    time in frames.
+    """Return a valid bboxes dataset for two individuals moving in uniform
+    linear motion, with 5 frames with low confidence values and time in frames.
     """
     dim_names = MovementDataset.dim_names["bboxes"]
 
-    position_array = valid_bboxes_array["position"]
-    shape_array = valid_bboxes_array["shape"]
-    confidence_array = valid_bboxes_array["confidence"]
+    position_array = valid_bboxes_arrays["position"]
+    shape_array = valid_bboxes_arrays["shape"]
+    confidence_array = valid_bboxes_arrays["confidence"]
 
     n_frames, n_individuals, _ = position_array.shape
 
@@ -409,12 +414,110 @@ def valid_poses_dataset(valid_position_array, request):
 @pytest.fixture
 def valid_poses_dataset_with_nan(valid_poses_dataset):
     """Return a valid pose tracks dataset with NaN values."""
+    # Sets position for all keypoints in individual ind1 to NaN
+    # at timepoints 3, 7, 8
     valid_poses_dataset.position.loc[
         {"individuals": "ind1", "time": [3, 7, 8]}
     ] = np.nan
     return valid_poses_dataset
 
 
+@pytest.fixture
+def valid_poses_array_uniform_linear_motion():
+    """Return a dictionary of valid arrays for a
+    ValidPosesDataset representing a uniform linear motion.
+
+    It represents 2 individuals with 3 keypoints, for 10 frames, in 2D space.
+    - Individual 0 moves along the x=y line from the origin.
+    - Individual 1 moves along the x=-y line line from the origin.
+
+    All confidence values for all keypoints are set to 0.9 except
+    for the keypoints at the following frames which are set to 0.1:
+    - Individual 0 at frames 2, 3, 4
+    - Individual 1 at frames 2, 3
+    """
+    # define the shape of the arrays
+    n_frames, n_individuals, n_keypoints, n_space = (10, 2, 3, 2)
+
+    # define centroid (index=0) trajectory in position array
+    # for each individual, the centroid moves along
+    # the x=+/-y line, starting from the origin.
+    # - individual 0 moves along x = y line
+    # - individual 1 moves along x = -y line
+    # They move one unit along x and y axes in each frame
+    frames = np.arange(n_frames)
+    position = np.empty((n_frames, n_individuals, n_keypoints, n_space))
+    position[:, :, 0, 0] = frames[:, None]  # reshape to (n_frames, 1)
+    position[:, 0, 0, 1] = frames
+    position[:, 1, 0, 1] = -frames
+
+    # define trajectory of left and right keypoints
+    # for individual 0, at each timepoint:
+    # - the left keypoint (index=1) is at x_centroid, y_centroid + 1
+    # - the right keypoint (index=2) is at x_centroid + 1, y_centroid
+    # for individual 1, at each timepoint:
+    # - the left keypoint (index=1) is at x_centroid - 1, y_centroid
+    # - the right keypoint (index=2) is at x_centroid, y_centroid + 1
+    offsets = [
+        [(0, 1), (1, 0)],  # individual 0: left, right keypoints (x,y) offsets
+        [(-1, 0), (0, 1)],  # individual 1: left, right keypoints (x,y) offsets
+    ]
+    for i in range(n_individuals):
+        for kpt in range(1, n_keypoints):
+            position[:, i, kpt, 0] = (
+                position[:, i, 0, 0] + offsets[i][kpt - 1][0]
+            )
+            position[:, i, kpt, 1] = (
+                position[:, i, 0, 1] + offsets[i][kpt - 1][1]
+            )
+
+    # build an array of confidence values, all 0.9
+    confidence = np.full((n_frames, n_individuals, n_keypoints), 0.9)
+    # set 5 low-confidence values
+    # - set 3 confidence values for individual id_0's centroid to 0.1
+    # - set 2 confidence values for individual id_1's centroid to 0.1
+    idx_start = 2
+    confidence[idx_start : idx_start + 3, 0, 0] = 0.1
+    confidence[idx_start : idx_start + 2, 1, 0] = 0.1
+
+    return {"position": position, "confidence": confidence}
+
+
+@pytest.fixture
+def valid_poses_dataset_uniform_linear_motion(
+    valid_poses_array_uniform_linear_motion,
+):
+    """Return a valid poses dataset for two individuals moving in uniform
+    linear motion, with 5 frames with low confidence values and time in frames.
+    """
+    dim_names = MovementDataset.dim_names["poses"]
+
+    position_array = valid_poses_array_uniform_linear_motion["position"]
+    confidence_array = valid_poses_array_uniform_linear_motion["confidence"]
+
+    n_frames, n_individuals, _, _ = position_array.shape
+
+    return xr.Dataset(
+        data_vars={
+            "position": xr.DataArray(position_array, dims=dim_names),
+            "confidence": xr.DataArray(confidence_array, dims=dim_names[:-1]),
+        },
+        coords={
+            dim_names[0]: np.arange(n_frames),
+            dim_names[1]: [f"id_{i}" for i in range(1, n_individuals + 1)],
+            dim_names[2]: ["centroid", "left", "right"],
+            dim_names[3]: ["x", "y"],
+        },
+        attrs={
+            "fps": None,
+            "time_unit": "frames",
+            "source_software": "test",
+            "source_file": "test_poses.h5",
+            "ds_type": "poses",
+        },
+    )
+
+
 # -------------------- Invalid datasets fixtures ------------------------------
 @pytest.fixture
 def not_a_dataset():