Ensure derived data variables have correct name (#523)

* Set name attribute for derived variables

* Add names for forward_vector and signed angles

* Add tests to check if name is correctly set

* Add remaining tests

* Rename head_direction to head_direction_vector

Author: Ishaan0132

movement/kinematics.py

@@ -60,6 +60,7 @@ def compute_displacement(data: xr.DataArray) -> xr.DataArray:
     validate_dims_coords(data, {"time": [], "space": []})
     result = data.diff(dim="time")
     result = result.reindex(data.coords, fill_value=0)
+    result.name = "displacement"
     return result
 
 
@@ -99,7 +100,9 @@ def compute_velocity(data: xr.DataArray) -> xr.DataArray:
     # validate only presence of Cartesian space dimension
     # (presence of time dimension will be checked in compute_time_derivative)
     validate_dims_coords(data, {"space": []})
-    return compute_time_derivative(data, order=1)
+    result = compute_time_derivative(data, order=1)
+    result.name = "velocity"
+    return result
 
 
 def compute_acceleration(data: xr.DataArray) -> xr.DataArray:
@@ -139,7 +142,9 @@ def compute_acceleration(data: xr.DataArray) -> xr.DataArray:
     # validate only presence of Cartesian space dimension
     # (presence of time dimension will be checked in compute_time_derivative)
     validate_dims_coords(data, {"space": []})
-    return compute_time_derivative(data, order=2)
+    result = compute_time_derivative(data, order=2)
+    result.name = "acceleration"
+    return result
 
 
 def compute_time_derivative(data: xr.DataArray, order: int) -> xr.DataArray:
@@ -202,7 +207,9 @@ def compute_speed(data: xr.DataArray) -> xr.DataArray:
         except ``space`` is removed.
 
     """
-    return compute_norm(compute_velocity(data))
+    result = compute_norm(compute_velocity(data))
+    result.name = "speed"
+    return result
 
 
 def compute_forward_vector(
@@ -312,7 +319,9 @@ def compute_forward_vector(
         space="z"
     )  # keep only the first 2 spatal dimensions of the result
     # Return unit vector
-    return convert_to_unit(forward_vector)
+    result = convert_to_unit(forward_vector)
+    result.name = "forward_vector"
+    return result
 
 
 def compute_head_direction_vector(
@@ -353,9 +362,11 @@ def compute_head_direction_vector(
         ``keypoints`` dimension.
 
     """
-    return compute_forward_vector(
+    result = compute_forward_vector(
         data, left_keypoint, right_keypoint, camera_view=camera_view
     )
+    result.name = "head_direction_vector"
+    return result
 
 
 def compute_forward_vector_angle(
@@ -439,6 +450,7 @@ def compute_forward_vector_angle(
     if in_degrees:
         heading_array = np.rad2deg(heading_array)
 
+    heading_array.name = "forward_vector_angle"
     return heading_array
 
 
@@ -534,6 +546,7 @@ def _cdist(
             elem2: getattr(a, labels_dim).values,
         }
     )
+    result.name = "distance"
     # Drop any squeezed coordinates
     return result.squeeze(drop=True)
 
@@ -873,13 +886,13 @@ def compute_path_length(
     _warn_about_nan_proportion(data, nan_warn_threshold)
 
     if nan_policy == "ffill":
-        return compute_norm(
+        result = compute_norm(
             compute_displacement(data.ffill(dim="time")).isel(
                 time=slice(1, None)
             )  # skip first displacement (always 0)
         ).sum(dim="time", min_count=1)  # return NaN if no valid segment
     elif nan_policy == "scale":
-        return _compute_scaled_path_length(data)
+        result = _compute_scaled_path_length(data)
     else:
         raise logger.error(
             ValueError(
@@ -888,6 +901,9 @@ def compute_path_length(
             )
         )
 
+    result.name = "path_length"
+    return result
+
 
 def _warn_about_nan_proportion(
     data: xr.DataArray, nan_warn_threshold: float

movement/utils/vector.py

@@ -274,6 +274,7 @@ def compute_signed_angle_2d(
     # arctan2 returns values in [-pi, pi].
     # We need to map -pi angles to pi, to stay in the (-pi, pi] range
     angles.values[angles <= -np.pi] = np.pi
+    angles.name = "signed_angle"
     return angles
 
 

tests/test_integration/test_kinematics_vector_transform.py

@@ -61,6 +61,7 @@ def test_cart2pol_transform_on_kinematics(
     kinematic_array_cart = getattr(kin, f"compute_{kinematic_variable}")(
         ds.position
     )
+    assert kinematic_array_cart.name == kinematic_variable
     kinematic_array_pol = vector.cart2pol(kinematic_array_cart)
 
     # Build expected data array

tests/test_unit/test_kinematics.py

@@ -44,6 +44,7 @@ def test_kinematics(self, valid_dataset, kinematic_variable, request):
         kinematic_array = getattr(kinematics, f"compute_{kinematic_variable}")(
             position
         )
+        assert kinematic_array.name == kinematic_variable
         # Figure out which dimensions to expect in kinematic_array
         # and in the final xarray.DataArray
         expected_dims = ["time", "individuals"]
@@ -112,6 +113,7 @@ def test_kinematics_with_dataset_with_nans(
         kinematic_array = getattr(kinematics, f"compute_{kinematic_variable}")(
             position
         )
+        assert kinematic_array.name == kinematic_variable
         # compute n nans in kinematic array per individual
         n_nans_kinematics_per_indiv = [
             helpers.count_nans(kinematic_array.isel(individuals=i))
@@ -203,6 +205,7 @@ def test_path_length_across_time_ranges(
         path_length = kinematics.compute_path_length(
             position, start=start, stop=stop
         )
+        assert path_length.name == "path_length"
 
         # Expected number of segments (displacements) in selected time range
         num_segments = 9  # full time range: 10 frames - 1
@@ -266,6 +269,7 @@ def test_path_length_with_nan(
         path_length = kinematics.compute_path_length(
             position, nan_policy=nan_policy
         )
+        assert path_length.name == "path_length"
         # Get path_length for individual "id_0" as a numpy array
         path_length_id_0 = path_length.sel(individuals="id_0").values
         # Check them against the expected values
@@ -312,9 +316,10 @@ def test_path_length_nan_warn_threshold(
     """
     position = valid_poses_dataset_with_nan.position
     with expected_exception:
-        kinematics.compute_path_length(
+        result = kinematics.compute_path_length(
             position, nan_warn_threshold=nan_warn_threshold
         )
+        assert result.name == "path_length"
 
 
 @pytest.fixture
@@ -408,6 +413,10 @@ def test_compute_forward_vector(valid_data_array_for_forward_vector):
         "right_ear",
         camera_view="bottom_up",
     )
+    assert forward_vector.name == "forward_vector"
+    assert forward_vector_flipped.name == "forward_vector"
+    assert head_vector.name == "head_direction_vector"
+
     known_vectors = np.array([[[0, -1]], [[1, 0]], [[0, 1]], [[-1, 0]]])
 
     for output_array in [forward_vector, forward_vector_flipped, head_vector]:
@@ -480,6 +489,7 @@ def test_nan_behavior_forward_vector(
     forward_vector = kinematics.compute_forward_vector(
         valid_data_array_for_forward_vector_with_nan, "left_ear", "right_ear"
     )
+    assert forward_vector.name == "forward_vector"
     # Check coord preservation
     for preserved_coord in ["time", "space", "individuals"]:
         assert np.all(
@@ -552,6 +562,7 @@ def test_cdist_with_known_values(dim, expected_data, valid_poses_dataset):
     a = input_dataarray.sel({dim: pairs[0]})
     b = input_dataarray.sel({dim: pairs[1]})
     result = kinematics._cdist(a, b, dim)
+    assert result.name == "distance"
     xr.testing.assert_equal(
         result,
         expected,
@@ -624,7 +635,9 @@ def test_cdist_with_single_dim_inputs(valid_dataset, selection_fn, request):
         valid_dataset = request.getfixturevalue(valid_dataset)
         position = valid_dataset.position
         a, b = selection_fn(position)
-        assert isinstance(kinematics._cdist(a, b, "individuals"), xr.DataArray)
+        result = kinematics._cdist(a, b, "individuals")
+        assert result.name == "distance"
+        assert isinstance(result, xr.DataArray)
 
 
 @pytest.mark.parametrize(
@@ -662,12 +675,16 @@ def test_compute_pairwise_distances_with_valid_pairs(
         valid_poses_dataset.position, dim, pairs
     )
     if isinstance(result, dict):
+        for _, value in result.items():
+            assert isinstance(value, xr.DataArray)
+            assert value.name == "distance"
         expected_data_vars = [
             f"dist_{pair[0]}_{pair[1]}" for pair in expected_data_vars
         ]
         assert set(result.keys()) == set(expected_data_vars)
     else:  # expect single DataArray
         assert isinstance(result, xr.DataArray)
+        assert result.name == "distance"
 
 
 @pytest.mark.parametrize(
@@ -803,6 +820,8 @@ def test_antisymmetry_properties(
             right_keypoint=right_keypoint,
             reference_vector=reference_vector,
         )
+        assert without_orientations_swapped.name == "forward_vector_angle"
+        assert with_orientations_swapped.name == "forward_vector_angle"
 
         expected_orientations = without_orientations_swapped.copy(deep=True)
         if swap_left_right:
@@ -841,6 +860,8 @@ def test_in_degrees_toggle(
             reference_vector=reference_vector,
             in_degrees=True,
         )
+        assert in_radians.name == "forward_vector_angle"
+        assert in_degrees.name == "forward_vector_angle"
 
         xr.testing.assert_allclose(in_degrees, np.rad2deg(in_radians))
 
@@ -901,6 +922,9 @@ def test_transformation_invariance(
             reference_vector=reference_vector,
         )
 
+        assert untranslated_output.name == "forward_vector_angle"
+        assert translated_output.name == "forward_vector_angle"
+
         xr.testing.assert_allclose(untranslated_output, translated_output)
 
     def test_casts_from_tuple(

tests/test_unit/test_vector.py

@@ -322,6 +322,9 @@ def test_compute_signed_angle_2d(
                 left_vector, right_vector, v_as_left_operand=True
             )
 
+            assert computed_angles.name == "signed_angle"
+            assert computed_angles_reversed.name == "signed_angle"
+
             xr.testing.assert_allclose(computed_angles, expected_angles)
             xr.testing.assert_allclose(
                 computed_angles_reversed, expected_angles_reversed
@@ -390,4 +393,5 @@ def test_multidimensional_input(
         )
 
         computed_angles = vector.compute_signed_angle_2d(v_left, v_right)
+        assert computed_angles.name == "signed_angle"
         xr.testing.assert_allclose(computed_angles, expected_angles)