From 7015cf587f152ac6975069ad0bad81115188b105 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 22 Dec 2023 18:38:26 +0000 Subject: [PATCH 01/30] Draft compute velocity --- movement/analysis/kinematics.py | 97 ++++++++++++++++++++++++++++++ tests/conftest.py | 16 +++-- tests/test_unit/test_kinematics.py | 27 +++++++++ 3 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 movement/analysis/kinematics.py create mode 100644 tests/test_unit/test_kinematics.py diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py new file mode 100644 index 000000000..2c6681ad6 --- /dev/null +++ b/movement/analysis/kinematics.py @@ -0,0 +1,97 @@ +import numpy as np +import xarray as xr + + +def compute_velocity( + data: xr.DataArray, method: str = "euclidean" +) -> np.ndarray: + """Compute the instantaneous velocity of a single keypoint from + a single individual. + + Parameters + ---------- + data : xarray.DataArray + The input data, assumed to be of shape (..., 2), where the last + dimension contains the x and y coordinates. + method : str + The method to use for computing velocity. Can be "euclidean" or + "numerical". + + Returns + ------- + numpy.ndarray + A numpy array containing the computed velocity. + """ + if method == "euclidean": + return compute_euclidean_velocity(data) + return approximate_derivative(data, order=1) + + +def compute_euclidean_velocity(data: xr.DataArray) -> np.ndarray: + """Compute velocity based on the Euclidean norm (magnitude) of the + differences between consecutive points, i.e. the straight-line + distance travelled. + + Parameters + ---------- + data : xarray.DataArray + The input data, assumed to be of shape (..., 2), where the last + dimension contains the x and y coordinates. + + Returns + ------- + numpy.ndarray + A numpy array containing the computed velocity. + """ + time_diff = data["time"].diff(dim="time") + space_diff = np.linalg.norm(np.diff(data.values, axis=0), axis=1) + velocity = space_diff / time_diff + # Pad with zero to match the original shape of the data + velocity = np.concatenate([np.zeros((1,) + velocity.shape[1:]), velocity]) + return velocity + + +def approximate_derivative(data: xr.DataArray, order: int = 0) -> np.ndarray: + """Compute displacement, 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 the x and y coordinates. + order : int + The order of the derivative. 0 for displacement, 1 for velocity, 2 for + acceleration. + + Returns + ------- + numpy.ndarray + A numpy array containing the computed kinematic variable. + """ + if order == 0: # Compute displacement + result = np.diff(data, axis=0) + # Pad with zeros to match the original shape of the data + result = np.concatenate([np.zeros((1,) + result.shape[1:]), result]) + else: + result = data + dt = data["time"].diff(dim="time").values[0] + for _ in range(order): + result = np.gradient(result, dt, axis=0) + magnitude = np.linalg.norm(result, axis=-1) + # Pad with zero to match the output of compute_euclidean_velocity + magnitude = np.pad(magnitude[:-1], (1, 0), "constant") + # direction = np.arctan2(result[..., 1], result[..., 0]) + return magnitude + + +# Locomotion Features +# speed +# speed_centroid +# acceleration +# acceleration_centroid +# speed_fwd +# radial_vel +# tangential_vel +# speed_centroid_w(s) +# speed_(p)_w(s) diff --git a/tests/conftest.py b/tests/conftest.py index 361305dd6..5108d6a34 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -212,12 +212,20 @@ def valid_tracks_array(): def _valid_tracks_array(array_type): """Return a valid tracks array.""" + # Unless specified, default is a multi_track_array with + # 10 frames, 2 individuals, and 2 keypoints. + n_frames = 10 + n_individuals = 2 + n_keypoints = 2 + base = np.arange(n_frames)[:, np.newaxis, np.newaxis, np.newaxis] if array_type == "single_keypoint_array": - return np.zeros((10, 2, 1, 2)) + n_keypoints = 1 elif array_type == "single_track_array": - return np.zeros((10, 1, 2, 2)) - else: # "multi_track_array": - return np.zeros((10, 2, 2, 2)) + n_individuals = 1 + x_points = np.repeat(base * 3, 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) return _valid_tracks_array diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py new file mode 100644 index 000000000..03756de58 --- /dev/null +++ b/tests/test_unit/test_kinematics.py @@ -0,0 +1,27 @@ +import numpy as np +import pytest + +from movement.analysis import kinematics + + +class TestKinematics: + """Test suite for the kinematics module.""" + + @pytest.mark.parametrize("method", ["euclidean", "numerical"]) + def test_compute_velocity(self, valid_pose_dataset, method): + """Test the `compute_velocity` function.""" + # Select a single keypoint from a single individual + data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + # Compute velocity + velocity = kinematics.compute_velocity(data, method=method) + expected = np.pad([5.0] * 9, (1, 0), "constant") + assert np.allclose(velocity, expected) + + def test_compute_acceleration(self, valid_pose_dataset): + """Test the `approximate_derivative` function for + calculating acceleration.""" + # Select a single keypoint from a single individual + data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + # Compute acceleration + acceleration = kinematics.approximate_derivative(data, order=2) + assert np.allclose(acceleration, np.zeros(10)) From 6121be2ee97d4b740c09cfb1727797b960e432d0 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 22 Dec 2023 22:00:23 +0000 Subject: [PATCH 02/30] Add test for displacement --- movement/analysis/kinematics.py | 4 ++-- tests/test_unit/test_kinematics.py | 10 ++++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 2c6681ad6..d59cd2a95 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -78,9 +78,9 @@ def approximate_derivative(data: xr.DataArray, order: int = 0) -> np.ndarray: dt = data["time"].diff(dim="time").values[0] for _ in range(order): result = np.gradient(result, dt, axis=0) + # Pad with zeros to match the output of compute_euclidean_velocity + result = np.pad(result[1:], ((1, 0), (0, 0)), "constant") magnitude = np.linalg.norm(result, axis=-1) - # Pad with zero to match the output of compute_euclidean_velocity - magnitude = np.pad(magnitude[:-1], (1, 0), "constant") # direction = np.arctan2(result[..., 1], result[..., 0]) return magnitude diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 03756de58..67566a25b 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -7,6 +7,16 @@ class TestKinematics: """Test suite for the kinematics module.""" + def test_compute_displacement(self, valid_pose_dataset): + """Test the `approximate_derivative` function for + calculating displacement.""" + # Select a single keypoint from a single individual + data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + # Compute displacement + displacement = kinematics.approximate_derivative(data) + expected = np.pad([5.0] * 9, (1, 0), "constant") + assert np.allclose(displacement, expected) + @pytest.mark.parametrize("method", ["euclidean", "numerical"]) def test_compute_velocity(self, valid_pose_dataset, method): """Test the `compute_velocity` function.""" From 44ff2b0d13b84f2190f44c66b4ed1fbf1bc1feda Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 8 Jan 2024 10:44:58 +0000 Subject: [PATCH 03/30] Fix confidence values in `valid_pose_dataset` fixture --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 5108d6a34..3f762a248 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -245,7 +245,7 @@ def valid_pose_dataset(valid_tracks_array, request): data_vars={ "pose_tracks": xr.DataArray(tracks_array, dims=dim_names), "confidence": xr.DataArray( - tracks_array[..., 0], + np.ones(tracks_array.shape[:-1]), dims=dim_names[:-1], ), }, From a4d1633e952efcf2dd9b5638a7d606cabe388c63 Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 8 Jan 2024 17:59:26 +0000 Subject: [PATCH 04/30] Refactor kinematics test and functions --- movement/analysis/kinematics.py | 107 +++++++++++++++++++++-------- tests/test_unit/test_kinematics.py | 55 ++++++++++----- 2 files changed, 115 insertions(+), 47 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index d59cd2a95..9602b5eb1 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -2,10 +2,50 @@ import xarray as xr -def compute_velocity( - data: xr.DataArray, method: str = "euclidean" -) -> np.ndarray: - """Compute the instantaneous velocity of a single keypoint from +def displacement(data: xr.DataArray) -> np.ndarray: + """Compute the displacement between consecutive locations + of a single keypoint from a single individual. + + Parameters + ---------- + data : xarray.DataArray + The input data, assumed to be of shape (..., 2), where the last + dimension contains the x and y coordinates. + + Returns + ------- + numpy.ndarray + A numpy array containing the computed magnitude and + direction of the displacement. + """ + displacement_vector = np.diff(data, axis=0, prepend=data[0:1]) + magnitude = np.linalg.norm(displacement_vector, axis=1) + direction = np.arctan2( + displacement_vector[..., 1], displacement_vector[..., 0] + ) + return np.stack((magnitude, direction), axis=1) + + +def distance(data: xr.DataArray) -> np.ndarray: + """Compute the distances between consecutive locations of + a single keypoint from a single individual. + + Parameters + ---------- + data : xarray.DataArray + The input data, assumed to be of shape (..., 2), where the last + dimension contains the x and y coordinates. + + Returns + ------- + numpy.ndarray + A numpy array containing the computed distance. + """ + return displacement(data)[:, 0] + + +def velocity(data: xr.DataArray) -> np.ndarray: + """Compute the velocity of a single keypoint from a single individual. Parameters @@ -13,24 +53,19 @@ def compute_velocity( data : xarray.DataArray The input data, assumed to be of shape (..., 2), where the last dimension contains the x and y coordinates. - method : str - The method to use for computing velocity. Can be "euclidean" or - "numerical". Returns ------- numpy.ndarray A numpy array containing the computed velocity. """ - if method == "euclidean": - return compute_euclidean_velocity(data) return approximate_derivative(data, order=1) -def compute_euclidean_velocity(data: xr.DataArray) -> np.ndarray: +def speed(data: xr.DataArray) -> np.ndarray: """Compute velocity based on the Euclidean norm (magnitude) of the differences between consecutive points, i.e. the straight-line - distance travelled. + distance travelled, assuming equidistant time spacing. Parameters ---------- @@ -43,17 +78,30 @@ def compute_euclidean_velocity(data: xr.DataArray) -> np.ndarray: numpy.ndarray A numpy array containing the computed velocity. """ - time_diff = data["time"].diff(dim="time") - space_diff = np.linalg.norm(np.diff(data.values, axis=0), axis=1) - velocity = space_diff / time_diff - # Pad with zero to match the original shape of the data - velocity = np.concatenate([np.zeros((1,) + velocity.shape[1:]), velocity]) - return velocity + return velocity(data)[:, 0] + + +def acceleration(data: xr.DataArray) -> np.ndarray: + """Compute the acceleration of a single keypoint from + a single individual. + + Parameters + ---------- + data : xarray.DataArray + The input data, assumed to be of shape (..., 2), where the last + dimension contains the x and y coordinates. + + Returns + ------- + numpy.ndarray + A numpy array containing the computed acceleration. + """ + return approximate_derivative(data, order=2) -def approximate_derivative(data: xr.DataArray, order: int = 0) -> np.ndarray: - """Compute displacement, velocity, or acceleration using numerical - differentiation, assuming equidistant time spacing. +def approximate_derivative(data: xr.DataArray, order: int = 1) -> np.ndarray: + """Compute velocity or acceleration using numerical differentiation, + assuming equidistant time spacing. Parameters ---------- @@ -61,28 +109,27 @@ def approximate_derivative(data: xr.DataArray, order: int = 0) -> np.ndarray: The input data, assumed to be of shape (..., 2), where the last dimension contains the x and y coordinates. order : int - The order of the derivative. 0 for displacement, 1 for velocity, 2 for - acceleration. + The order of the derivative. 1 for velocity, 2 for + acceleration. Default is 1. Returns ------- numpy.ndarray - A numpy array containing the computed kinematic variable. + A numpy array containing the computed magnitudes and directions of + the kinematic variable. """ - if order == 0: # Compute displacement - result = np.diff(data, axis=0) - # Pad with zeros to match the original shape of the data - result = np.concatenate([np.zeros((1,) + result.shape[1:]), result]) + if order <= 0: + raise ValueError("order must be a positive integer.") else: result = data dt = data["time"].diff(dim="time").values[0] for _ in range(order): result = np.gradient(result, dt, axis=0) - # Pad with zeros to match the output of compute_euclidean_velocity + # Prepend with zeros to match match output to the input shape result = np.pad(result[1:], ((1, 0), (0, 0)), "constant") magnitude = np.linalg.norm(result, axis=-1) - # direction = np.arctan2(result[..., 1], result[..., 0]) - return magnitude + direction = np.arctan2(result[..., 1], result[..., 0]) + return np.stack((magnitude, direction), axis=1) # Locomotion Features diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 67566a25b..c9ce36647 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -7,31 +7,52 @@ class TestKinematics: """Test suite for the kinematics module.""" - def test_compute_displacement(self, valid_pose_dataset): - """Test the `approximate_derivative` function for - calculating displacement.""" + def test_distance(self, valid_pose_dataset): + """Test distance calculation.""" # Select a single keypoint from a single individual data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) - # Compute displacement - displacement = kinematics.approximate_derivative(data) + result = kinematics.distance(data) expected = np.pad([5.0] * 9, (1, 0), "constant") - assert np.allclose(displacement, expected) + assert np.allclose(result, expected) - @pytest.mark.parametrize("method", ["euclidean", "numerical"]) - def test_compute_velocity(self, valid_pose_dataset, method): - """Test the `compute_velocity` function.""" + def test_displacement(self, valid_pose_dataset): + """Test displacement calculation.""" + # Select a single keypoint from a single individual + data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + result = kinematics.displacement(data) + expected_magnitude = np.pad([5.0] * 9, (1, 0), "constant") + expected_direction = np.concatenate(([0], np.full(9, 0.92729522))) + expected = np.stack((expected_magnitude, expected_direction), axis=1) + assert np.allclose(result, expected) + + def test_velocity(self, valid_pose_dataset): + """Test velocity calculation.""" # Select a single keypoint from a single individual data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) # Compute velocity - velocity = kinematics.compute_velocity(data, method=method) + result = kinematics.velocity(data) + expected_magnitude = np.pad([5.0] * 9, (1, 0), "constant") + expected_direction = np.concatenate(([0], np.full(9, 0.92729522))) + expected = np.stack((expected_magnitude, expected_direction), axis=1) + assert np.allclose(result, expected) + + def test_speed(self, valid_pose_dataset): + """Test velocity calculation.""" + # Select a single keypoint from a single individual + data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + result = kinematics.speed(data) expected = np.pad([5.0] * 9, (1, 0), "constant") - assert np.allclose(velocity, expected) + assert np.allclose(result, expected) - def test_compute_acceleration(self, valid_pose_dataset): - """Test the `approximate_derivative` function for - calculating acceleration.""" + def test_acceleration(self, valid_pose_dataset): + """Test acceleration calculation.""" # Select a single keypoint from a single individual data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) - # Compute acceleration - acceleration = kinematics.approximate_derivative(data, order=2) - assert np.allclose(acceleration, np.zeros(10)) + result = kinematics.acceleration(data) + assert np.allclose(result, np.zeros((10, 2))) + + def test_approximate_derivative_with_nonpositive_order(self): + """Test that an error is raised when the order is non-positive.""" + data = np.arange(10) + with pytest.raises(ValueError): + kinematics.approximate_derivative(data, order=0) From e21fbb3b5175799c143c7203653c415387f0e166 Mon Sep 17 00:00:00 2001 From: lochhh Date: Thu, 25 Jan 2024 18:27:14 +0000 Subject: [PATCH 05/30] Vectorise kinematic functions --- movement/analysis/kinematics.py | 99 +++++++++++++++++++----------- tests/test_unit/test_kinematics.py | 90 +++++++++++++++++++-------- 2 files changed, 128 insertions(+), 61 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 9602b5eb1..f5b176158 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -2,49 +2,58 @@ import xarray as xr -def displacement(data: xr.DataArray) -> np.ndarray: +def displacement(data: xr.DataArray) -> xr.Dataset: """Compute the displacement between consecutive locations - of a single keypoint from a single individual. + 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, assumed to be of shape (..., 2), where + the last dimension contains the x and y coordinates. Returns ------- - numpy.ndarray - A numpy array containing the computed magnitude and + xarray.Dataset + An xarray Dataset containing the computed magnitude and direction of the displacement. """ - displacement_vector = np.diff(data, axis=0, prepend=data[0:1]) - magnitude = np.linalg.norm(displacement_vector, axis=1) - direction = np.arctan2( - displacement_vector[..., 1], displacement_vector[..., 0] + displacement_da = data.diff(dim="time") + magnitude = xr.apply_ufunc( + np.linalg.norm, + displacement_da, + input_core_dims=[["space"]], + kwargs={"axis": -1}, ) - return np.stack((magnitude, direction), axis=1) + magnitude = magnitude.reindex_like(data.sel(space="x")) + direction = xr.apply_ufunc( + np.arctan2, + displacement_da[..., 1], + displacement_da[..., 0], + ) + direction = direction.reindex_like(data.sel(space="x")) + return xr.Dataset({"magnitude": magnitude, "direction": direction}) -def distance(data: xr.DataArray) -> np.ndarray: - """Compute the distances between consecutive locations of - a single keypoint from a single individual. +def distance(data: xr.DataArray) -> xr.DataArray: + """Compute the Euclidean distances 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, assumed to be of shape (..., 2), where + the last dimension contains the x and y coordinates. Returns ------- - numpy.ndarray - A numpy array containing the computed distance. + xarray.DataArray + An xarray DataArray containing the magnitude of displacement. """ - return displacement(data)[:, 0] + return displacement(data).magnitude -def velocity(data: xr.DataArray) -> np.ndarray: +def velocity(data: xr.DataArray) -> xr.Dataset: """Compute the velocity of a single keypoint from a single individual. @@ -56,14 +65,15 @@ def velocity(data: xr.DataArray) -> np.ndarray: Returns ------- - numpy.ndarray - A numpy array containing the computed velocity. + xarray.Dataset + An xarray Dataset containing the computed magnitude and + direction of the velocity. """ return approximate_derivative(data, order=1) -def speed(data: xr.DataArray) -> np.ndarray: - """Compute velocity based on the Euclidean norm (magnitude) of the +def speed(data: xr.DataArray) -> xr.DataArray: + """Compute speed based on the Euclidean norm (magnitude) of the differences between consecutive points, i.e. the straight-line distance travelled, assuming equidistant time spacing. @@ -75,10 +85,10 @@ def speed(data: xr.DataArray) -> np.ndarray: Returns ------- - numpy.ndarray - A numpy array containing the computed velocity. + xarray.DataArray + An xarray DataArray containing the magnitude of velocity. """ - return velocity(data)[:, 0] + return velocity(data).magnitude def acceleration(data: xr.DataArray) -> np.ndarray: @@ -99,7 +109,7 @@ def acceleration(data: xr.DataArray) -> np.ndarray: return approximate_derivative(data, order=2) -def approximate_derivative(data: xr.DataArray, order: int = 1) -> np.ndarray: +def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.Dataset: """Compute velocity or acceleration using numerical differentiation, assuming equidistant time spacing. @@ -114,9 +124,9 @@ def approximate_derivative(data: xr.DataArray, order: int = 1) -> np.ndarray: Returns ------- - numpy.ndarray - A numpy array containing the computed magnitudes and directions of - the kinematic variable. + xarray.Dataset + An xarray Dataset containing the computed magnitudes and + directions of the derived variable. """ if order <= 0: raise ValueError("order must be a positive integer.") @@ -124,12 +134,27 @@ def approximate_derivative(data: xr.DataArray, order: int = 1) -> np.ndarray: result = data dt = data["time"].diff(dim="time").values[0] for _ in range(order): - result = np.gradient(result, dt, axis=0) - # Prepend with zeros to match match output to the input shape - result = np.pad(result[1:], ((1, 0), (0, 0)), "constant") - magnitude = np.linalg.norm(result, axis=-1) - direction = np.arctan2(result[..., 1], result[..., 0]) - return np.stack((magnitude, direction), axis=1) + result = xr.apply_ufunc( + np.gradient, + result, + dt, + kwargs={"axis": 0}, + ) + result = result.reindex_like(data) + magnitude = xr.apply_ufunc( + np.linalg.norm, + result, + input_core_dims=[["space"]], + kwargs={"axis": -1}, + ) + magnitude = magnitude.reindex_like(data.sel(space="x")) + direction = xr.apply_ufunc( + np.arctan2, + result[..., 1], + result[..., 0], + ) + direction = direction.reindex_like(data.sel(space="x")) + return xr.Dataset({"magnitude": magnitude, "direction": direction}) # Locomotion Features diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index c9ce36647..cb8624ced 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -1,5 +1,6 @@ import numpy as np import pytest +import xarray as xr from movement.analysis import kinematics @@ -9,47 +10,88 @@ class TestKinematics: def test_distance(self, valid_pose_dataset): """Test distance calculation.""" - # Select a single keypoint from a single individual - data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + data = valid_pose_dataset.pose_tracks result = kinematics.distance(data) - expected = np.pad([5.0] * 9, (1, 0), "constant") - assert np.allclose(result, expected) + expected = np.full((10, 2, 2), 5.0) + expected[0, :, :] = np.nan + np.testing.assert_allclose(result.values, expected) def test_displacement(self, valid_pose_dataset): """Test displacement calculation.""" - # Select a single keypoint from a single individual - data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + data = valid_pose_dataset.pose_tracks result = kinematics.displacement(data) - expected_magnitude = np.pad([5.0] * 9, (1, 0), "constant") - expected_direction = np.concatenate(([0], np.full(9, 0.92729522))) - expected = np.stack((expected_magnitude, expected_direction), axis=1) - assert np.allclose(result, expected) + expected_magnitude = np.full((10, 2, 2), 5.0) + expected_magnitude[0, :, :] = np.nan + expected_direction = np.full((10, 2, 2), 0.92729522) + expected_direction[0, :, :] = np.nan + expected = xr.Dataset( + data_vars={ + "magnitude": xr.DataArray( + expected_magnitude, dims=data.dims[:-1] + ), + "direction": xr.DataArray( + expected_direction, dims=data.dims[:-1] + ), + }, + coords={ + "time": data.time, + "keypoints": data.keypoints, + "individuals": data.individuals, + }, + ) + xr.testing.assert_allclose(result, expected) def test_velocity(self, valid_pose_dataset): """Test velocity calculation.""" - # Select a single keypoint from a single individual - data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + data = valid_pose_dataset.pose_tracks # Compute velocity result = kinematics.velocity(data) - expected_magnitude = np.pad([5.0] * 9, (1, 0), "constant") - expected_direction = np.concatenate(([0], np.full(9, 0.92729522))) - expected = np.stack((expected_magnitude, expected_direction), axis=1) - assert np.allclose(result, expected) + expected_magnitude = np.full((10, 2, 2), 5.0) + expected_direction = np.full((10, 2, 2), 0.92729522) + expected = xr.Dataset( + data_vars={ + "magnitude": xr.DataArray( + expected_magnitude, dims=data.dims[:-1] + ), + "direction": xr.DataArray( + expected_direction, dims=data.dims[:-1] + ), + }, + coords={ + "time": data.time, + "keypoints": data.keypoints, + "individuals": data.individuals, + }, + ) + xr.testing.assert_allclose(result, expected) def test_speed(self, valid_pose_dataset): - """Test velocity calculation.""" - # Select a single keypoint from a single individual - data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + """Test speed calculation.""" + data = valid_pose_dataset.pose_tracks result = kinematics.speed(data) - expected = np.pad([5.0] * 9, (1, 0), "constant") - assert np.allclose(result, expected) + expected = np.full((10, 2, 2), 5.0) + np.testing.assert_allclose(result.values, expected) def test_acceleration(self, valid_pose_dataset): """Test acceleration calculation.""" - # Select a single keypoint from a single individual - data = valid_pose_dataset.pose_tracks.isel(keypoints=0, individuals=0) + data = valid_pose_dataset.pose_tracks result = kinematics.acceleration(data) - assert np.allclose(result, np.zeros((10, 2))) + expected = xr.Dataset( + data_vars={ + "magnitude": xr.DataArray( + np.zeros((10, 2, 2)), dims=data.dims[:-1] + ), + "direction": xr.DataArray( + np.zeros((10, 2, 2)), dims=data.dims[:-1] + ), + }, + coords={ + "time": data.time, + "keypoints": data.keypoints, + "individuals": data.individuals, + }, + ) + xr.testing.assert_allclose(result, expected) def test_approximate_derivative_with_nonpositive_order(self): """Test that an error is raised when the order is non-positive.""" From 630d7a08e5ee326176c60ba54a43d1daef3fc40a Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 26 Jan 2024 14:23:18 +0000 Subject: [PATCH 06/30] Refactor repeated calls to compute magnitude + direction --- movement/analysis/kinematics.py | 55 ++++++++++++++++++--------------- 1 file changed, 30 insertions(+), 25 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index f5b176158..5cef7d5f5 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -16,23 +16,11 @@ def displacement(data: xr.DataArray) -> xr.Dataset: ------- xarray.Dataset An xarray Dataset containing the computed magnitude and - direction of the displacement. + direction of displacement. """ - displacement_da = data.diff(dim="time") - magnitude = xr.apply_ufunc( - np.linalg.norm, - displacement_da, - input_core_dims=[["space"]], - kwargs={"axis": -1}, - ) - magnitude = magnitude.reindex_like(data.sel(space="x")) - direction = xr.apply_ufunc( - np.arctan2, - displacement_da[..., 1], - displacement_da[..., 0], - ) - direction = direction.reindex_like(data.sel(space="x")) - return xr.Dataset({"magnitude": magnitude, "direction": direction}) + displacement_xy = data.diff(dim="time") + displacement_xy = displacement_xy.reindex_like(data) + return compute_vector_magnitude_direction(displacement_xy) def distance(data: xr.DataArray) -> xr.DataArray: @@ -67,7 +55,7 @@ def velocity(data: xr.DataArray) -> xr.Dataset: ------- xarray.Dataset An xarray Dataset containing the computed magnitude and - direction of the velocity. + direction of velocity. """ return approximate_derivative(data, order=1) @@ -91,7 +79,7 @@ def speed(data: xr.DataArray) -> xr.DataArray: return velocity(data).magnitude -def acceleration(data: xr.DataArray) -> np.ndarray: +def acceleration(data: xr.DataArray) -> xr.Dataset: """Compute the acceleration of a single keypoint from a single individual. @@ -103,8 +91,9 @@ def acceleration(data: xr.DataArray) -> np.ndarray: Returns ------- - numpy.ndarray - A numpy array containing the computed acceleration. + xarray.Dataset + An xarray Dataset containing the magnitude and direction + of acceleration. """ return approximate_derivative(data, order=2) @@ -141,19 +130,35 @@ def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.Dataset: kwargs={"axis": 0}, ) result = result.reindex_like(data) + return compute_vector_magnitude_direction(result) + + +def compute_vector_magnitude_direction(input: xr.DataArray) -> xr.Dataset: + """Compute the magnitude and direction of a vector. + + Parameters + ---------- + input : xarray.DataArray + The input data, assumed to be of shape (..., 2), where the last + dimension contains the x and y coordinates. + + Returns + ------- + xarray.Dataset + An xarray Dataset containing the computed magnitude and + direction. + """ magnitude = xr.apply_ufunc( np.linalg.norm, - result, + input, input_core_dims=[["space"]], kwargs={"axis": -1}, ) - magnitude = magnitude.reindex_like(data.sel(space="x")) direction = xr.apply_ufunc( np.arctan2, - result[..., 1], - result[..., 0], + input[..., 1], + input[..., 0], ) - direction = direction.reindex_like(data.sel(space="x")) return xr.Dataset({"magnitude": magnitude, "direction": direction}) From 684930ffab4a8904e552a093d12cf032517c8e63 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 26 Jan 2024 14:25:18 +0000 Subject: [PATCH 07/30] Displacement to return 0 instead of NaN --- movement/analysis/kinematics.py | 2 +- tests/test_unit/test_kinematics.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 5cef7d5f5..288c1ac16 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -19,7 +19,7 @@ def displacement(data: xr.DataArray) -> xr.Dataset: direction of displacement. """ displacement_xy = data.diff(dim="time") - displacement_xy = displacement_xy.reindex_like(data) + displacement_xy = displacement_xy.reindex(data.coords, fill_value=0) return compute_vector_magnitude_direction(displacement_xy) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index cb8624ced..f2182ceba 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -13,7 +13,7 @@ def test_distance(self, valid_pose_dataset): data = valid_pose_dataset.pose_tracks result = kinematics.distance(data) expected = np.full((10, 2, 2), 5.0) - expected[0, :, :] = np.nan + expected[0, :, :] = 0 np.testing.assert_allclose(result.values, expected) def test_displacement(self, valid_pose_dataset): @@ -21,9 +21,9 @@ def test_displacement(self, valid_pose_dataset): data = valid_pose_dataset.pose_tracks result = kinematics.displacement(data) expected_magnitude = np.full((10, 2, 2), 5.0) - expected_magnitude[0, :, :] = np.nan + expected_magnitude[0, :, :] = 0 expected_direction = np.full((10, 2, 2), 0.92729522) - expected_direction[0, :, :] = np.nan + expected_direction[0, :, :] = 0 expected = xr.Dataset( data_vars={ "magnitude": xr.DataArray( From 0fa3acc03cccc46d580b73f1adb37584c3d16d55 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 26 Jan 2024 18:27:06 +0000 Subject: [PATCH 08/30] Return x y components in kinematic functions --- movement/analysis/kinematics.py | 88 ++++++++++++++++++------------ tests/test_unit/test_kinematics.py | 51 ++++++++++++----- 2 files changed, 88 insertions(+), 51 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 288c1ac16..1da73f1c4 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -2,9 +2,9 @@ import xarray as xr -def displacement(data: xr.DataArray) -> xr.Dataset: - """Compute the displacement between consecutive locations - of each keypoint of each individual. +def displacement(data: xr.DataArray) -> xr.DataArray: + """Compute the displacement between consecutive x, y + locations of each keypoint of each individual. Parameters ---------- @@ -14,18 +14,17 @@ def displacement(data: xr.DataArray) -> xr.Dataset: Returns ------- - xarray.Dataset - An xarray Dataset containing the computed magnitude and - direction of displacement. + 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 compute_vector_magnitude_direction(displacement_xy) + return displacement_xy -def distance(data: xr.DataArray) -> xr.DataArray: - """Compute the Euclidean distances between consecutive - locations of each keypoint of each individual. +def displacement_vector(data: xr.DataArray) -> xr.Dataset: + """Compute the Euclidean magnitude (distance) and direction + (angle relative to the positive x-axis) of displacement. Parameters ---------- @@ -35,15 +34,16 @@ def distance(data: xr.DataArray) -> xr.DataArray: Returns ------- - xarray.DataArray - An xarray DataArray containing the magnitude of displacement. + xarray.Dataset + An xarray Dataset containing the magnitude and direction + of the computed displacement. """ - return displacement(data).magnitude + return compute_vector_magnitude_direction(displacement(data)) -def velocity(data: xr.DataArray) -> xr.Dataset: - """Compute the velocity of a single keypoint from - a single individual. +def velocity(data: xr.DataArray) -> xr.DataArray: + """Compute the velocity between consecutive x, y locations + of each keypoint of each individual. Parameters ---------- @@ -53,35 +53,34 @@ def velocity(data: xr.DataArray) -> xr.Dataset: Returns ------- - xarray.Dataset - An xarray Dataset containing the computed magnitude and - direction of velocity. + xarray.DataArray + An xarray Dataset containing the computed velocity. """ return approximate_derivative(data, order=1) -def speed(data: xr.DataArray) -> xr.DataArray: - """Compute speed based on the Euclidean norm (magnitude) of the - differences between consecutive points, i.e. the straight-line - distance travelled, assuming equidistant time spacing. +def velocity_vector(data: xr.DataArray) -> xr.Dataset: + """Compute the Euclidean magnitude (speed) and direction + (angle relative to the positive x-axis) of velocity. 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, assumed to be of shape (..., 2), where + the last dimension contains the x and y coordinates. Returns ------- - xarray.DataArray - An xarray DataArray containing the magnitude of velocity. + xarray.Dataset + An xarray Dataset containing the magnitude and direction + of the computed velocity. """ - return velocity(data).magnitude + return compute_vector_magnitude_direction(velocity(data)) -def acceleration(data: xr.DataArray) -> xr.Dataset: - """Compute the acceleration of a single keypoint from - a single individual. +def acceleration(data: xr.DataArray) -> xr.DataArray: + """Compute the acceleration between consecutive x, y + locations of each keypoint of each individual. Parameters ---------- @@ -98,7 +97,25 @@ def acceleration(data: xr.DataArray) -> xr.Dataset: return approximate_derivative(data, order=2) -def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.Dataset: +def acceleration_vector(data: xr.DataArray) -> xr.Dataset: + """Compute the Euclidean magnitude and direction of acceleration. + + Parameters + ---------- + data : xarray.DataArray + The input data, assumed to be of shape (..., 2), where + the last dimension contains the x and y coordinates. + + Returns + ------- + xarray.Dataset + An xarray Dataset containing the magnitude and direction + of the computed acceleration. + """ + return compute_vector_magnitude_direction(acceleration(data)) + + +def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.DataArray: """Compute velocity or acceleration using numerical differentiation, assuming equidistant time spacing. @@ -113,9 +130,8 @@ def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.Dataset: Returns ------- - xarray.Dataset - An xarray Dataset containing the computed magnitudes and - directions of the derived variable. + xarray.DataArray + An xarray DataArray containing the derived variable. """ if order <= 0: raise ValueError("order must be a positive integer.") @@ -130,7 +146,7 @@ def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.Dataset: kwargs={"axis": 0}, ) result = result.reindex_like(data) - return compute_vector_magnitude_direction(result) + return result def compute_vector_magnitude_direction(input: xr.DataArray) -> xr.Dataset: diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index f2182ceba..0a1c690da 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -8,18 +8,23 @@ class TestKinematics: """Test suite for the kinematics module.""" - def test_distance(self, valid_pose_dataset): - """Test distance calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.distance(data) - expected = np.full((10, 2, 2), 5.0) - expected[0, :, :] = 0 - np.testing.assert_allclose(result.values, expected) - def test_displacement(self, valid_pose_dataset): """Test displacement calculation.""" data = valid_pose_dataset.pose_tracks result = kinematics.displacement(data) + expected_values = np.full((10, 2, 2, 2), [3.0, 4.0]) + expected_values[0, :, :, :] = 0 + expected = xr.DataArray( + expected_values, + dims=data.dims, + coords=data.coords, + ) + xr.testing.assert_allclose(result, expected) + + def test_displacement_vector(self, valid_pose_dataset): + """Test displacement vector calculation.""" + data = valid_pose_dataset.pose_tracks + result = kinematics.displacement_vector(data) expected_magnitude = np.full((10, 2, 2), 5.0) expected_magnitude[0, :, :] = 0 expected_direction = np.full((10, 2, 2), 0.92729522) @@ -46,6 +51,18 @@ def test_velocity(self, valid_pose_dataset): data = valid_pose_dataset.pose_tracks # Compute velocity result = kinematics.velocity(data) + expected_values = np.full((10, 2, 2, 2), [3.0, 4.0]) + expected = xr.DataArray( + expected_values, + dims=data.dims, + coords=data.coords, + ) + xr.testing.assert_allclose(result, expected) + + def test_velocity_vector(self, valid_pose_dataset): + """Test velocity vector calculation.""" + data = valid_pose_dataset.pose_tracks + result = kinematics.velocity_vector(data) expected_magnitude = np.full((10, 2, 2), 5.0) expected_direction = np.full((10, 2, 2), 0.92729522) expected = xr.Dataset( @@ -65,17 +82,21 @@ def test_velocity(self, valid_pose_dataset): ) xr.testing.assert_allclose(result, expected) - def test_speed(self, valid_pose_dataset): - """Test speed calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.speed(data) - expected = np.full((10, 2, 2), 5.0) - np.testing.assert_allclose(result.values, expected) - def test_acceleration(self, valid_pose_dataset): """Test acceleration calculation.""" data = valid_pose_dataset.pose_tracks result = kinematics.acceleration(data) + expected = xr.DataArray( + np.zeros((10, 2, 2, 2)), + dims=data.dims, + coords=data.coords, + ) + xr.testing.assert_allclose(result, expected) + + def test_acceleration_vector(self, valid_pose_dataset): + """Test acceleration vector calculation.""" + data = valid_pose_dataset.pose_tracks + result = kinematics.acceleration_vector(data) expected = xr.Dataset( data_vars={ "magnitude": xr.DataArray( From 39489c030dd8d6be2ad1b71988588c3b8be93476 Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 29 Jan 2024 12:30:38 +0000 Subject: [PATCH 09/30] Refactor kinematics tests --- tests/test_unit/test_kinematics.py | 128 +++++++++++------------------ 1 file changed, 50 insertions(+), 78 deletions(-) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 0a1c690da..34d6fe755 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -8,111 +8,83 @@ class TestKinematics: """Test suite for the kinematics module.""" - def test_displacement(self, valid_pose_dataset): - """Test displacement calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.displacement(data) - expected_values = np.full((10, 2, 2, 2), [3.0, 4.0]) - expected_values[0, :, :, :] = 0 - expected = xr.DataArray( - expected_values, - dims=data.dims, - coords=data.coords, + @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, ) - xr.testing.assert_allclose(result, expected) - def test_displacement_vector(self, valid_pose_dataset): - """Test displacement vector calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.displacement_vector(data) - expected_magnitude = np.full((10, 2, 2), 5.0) - expected_magnitude[0, :, :] = 0 - expected_direction = np.full((10, 2, 2), 0.92729522) - expected_direction[0, :, :] = 0 - expected = xr.Dataset( + @pytest.fixture + def expected_dataset(self, valid_pose_dataset): + """Return an xarray.Dataset with default `magnitude` and + `direction` data variables, and the expected dimensions + and coordinates.""" + dims = valid_pose_dataset.pose_tracks.dims[:-1] + return xr.Dataset( data_vars={ "magnitude": xr.DataArray( - expected_magnitude, dims=data.dims[:-1] + np.full((10, 2, 2), 5.0), + dims=dims, ), "direction": xr.DataArray( - expected_direction, dims=data.dims[:-1] + np.full((10, 2, 2), 0.92729522), + dims=dims, ), }, coords={ - "time": data.time, - "keypoints": data.keypoints, - "individuals": data.individuals, + "time": valid_pose_dataset.time, + "individuals": valid_pose_dataset.individuals, + "keypoints": valid_pose_dataset.keypoints, }, ) - xr.testing.assert_allclose(result, expected) - def test_velocity(self, valid_pose_dataset): + def test_displacement(self, valid_pose_dataset, expected_dataarray): + """Test displacement calculation.""" + data = valid_pose_dataset.pose_tracks + result = kinematics.displacement(data) + # Set the first displacement to zero + expected_dataarray[0, :, :, :] = 0 + xr.testing.assert_allclose(result, expected_dataarray) + + def test_displacement_vector(self, valid_pose_dataset, expected_dataset): + """Test displacement vector calculation.""" + data = valid_pose_dataset.pose_tracks + result = kinematics.displacement_vector(data) + # Set the first displacement to zero + expected_dataset.magnitude[0, :, :] = 0 + expected_dataset.direction[0, :, :] = 0 + xr.testing.assert_allclose(result, expected_dataset) + + def test_velocity(self, valid_pose_dataset, expected_dataarray): """Test velocity calculation.""" data = valid_pose_dataset.pose_tracks - # Compute velocity result = kinematics.velocity(data) - expected_values = np.full((10, 2, 2, 2), [3.0, 4.0]) - expected = xr.DataArray( - expected_values, - dims=data.dims, - coords=data.coords, - ) - xr.testing.assert_allclose(result, expected) + xr.testing.assert_allclose(result, expected_dataarray) - def test_velocity_vector(self, valid_pose_dataset): + def test_velocity_vector(self, valid_pose_dataset, expected_dataset): """Test velocity vector calculation.""" data = valid_pose_dataset.pose_tracks result = kinematics.velocity_vector(data) - expected_magnitude = np.full((10, 2, 2), 5.0) - expected_direction = np.full((10, 2, 2), 0.92729522) - expected = xr.Dataset( - data_vars={ - "magnitude": xr.DataArray( - expected_magnitude, dims=data.dims[:-1] - ), - "direction": xr.DataArray( - expected_direction, dims=data.dims[:-1] - ), - }, - coords={ - "time": data.time, - "keypoints": data.keypoints, - "individuals": data.individuals, - }, - ) - xr.testing.assert_allclose(result, expected) + xr.testing.assert_allclose(result, expected_dataset) - def test_acceleration(self, valid_pose_dataset): + def test_acceleration(self, valid_pose_dataset, expected_dataarray): """Test acceleration calculation.""" data = valid_pose_dataset.pose_tracks result = kinematics.acceleration(data) - expected = xr.DataArray( - np.zeros((10, 2, 2, 2)), - dims=data.dims, - coords=data.coords, - ) - xr.testing.assert_allclose(result, expected) + expected_dataarray[:] = 0 + xr.testing.assert_allclose(result, expected_dataarray) - def test_acceleration_vector(self, valid_pose_dataset): + def test_acceleration_vector(self, valid_pose_dataset, expected_dataset): """Test acceleration vector calculation.""" data = valid_pose_dataset.pose_tracks result = kinematics.acceleration_vector(data) - expected = xr.Dataset( - data_vars={ - "magnitude": xr.DataArray( - np.zeros((10, 2, 2)), dims=data.dims[:-1] - ), - "direction": xr.DataArray( - np.zeros((10, 2, 2)), dims=data.dims[:-1] - ), - }, - coords={ - "time": data.time, - "keypoints": data.keypoints, - "individuals": data.individuals, - }, - ) - xr.testing.assert_allclose(result, expected) + expected_dataset.magnitude[:] = 0 + expected_dataset.direction[:] = 0 + xr.testing.assert_allclose(result, expected_dataset) def test_approximate_derivative_with_nonpositive_order(self): """Test that an error is raised when the order is non-positive.""" From 6e55ab6a7d21f87ff6716e50945b308ff114373f Mon Sep 17 00:00:00 2001 From: lochhh Date: Wed, 31 Jan 2024 10:45:28 +0000 Subject: [PATCH 10/30] Remove unnecessary instantiations --- tests/test_unit/test_kinematics.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 34d6fe755..669ef6dbc 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -44,16 +44,14 @@ def expected_dataset(self, valid_pose_dataset): def test_displacement(self, valid_pose_dataset, expected_dataarray): """Test displacement calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.displacement(data) + result = kinematics.displacement(valid_pose_dataset.pose_tracks) # Set the first displacement to zero expected_dataarray[0, :, :, :] = 0 xr.testing.assert_allclose(result, expected_dataarray) def test_displacement_vector(self, valid_pose_dataset, expected_dataset): """Test displacement vector calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.displacement_vector(data) + result = kinematics.displacement_vector(valid_pose_dataset.pose_tracks) # Set the first displacement to zero expected_dataset.magnitude[0, :, :] = 0 expected_dataset.direction[0, :, :] = 0 @@ -61,27 +59,23 @@ def test_displacement_vector(self, valid_pose_dataset, expected_dataset): def test_velocity(self, valid_pose_dataset, expected_dataarray): """Test velocity calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.velocity(data) + result = kinematics.velocity(valid_pose_dataset.pose_tracks) xr.testing.assert_allclose(result, expected_dataarray) def test_velocity_vector(self, valid_pose_dataset, expected_dataset): """Test velocity vector calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.velocity_vector(data) + result = kinematics.velocity_vector(valid_pose_dataset.pose_tracks) xr.testing.assert_allclose(result, expected_dataset) def test_acceleration(self, valid_pose_dataset, expected_dataarray): """Test acceleration calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.acceleration(data) + result = kinematics.acceleration(valid_pose_dataset.pose_tracks) expected_dataarray[:] = 0 xr.testing.assert_allclose(result, expected_dataarray) def test_acceleration_vector(self, valid_pose_dataset, expected_dataset): """Test acceleration vector calculation.""" - data = valid_pose_dataset.pose_tracks - result = kinematics.acceleration_vector(data) + result = kinematics.acceleration_vector(valid_pose_dataset.pose_tracks) expected_dataset.magnitude[:] = 0 expected_dataset.direction[:] = 0 xr.testing.assert_allclose(result, expected_dataset) From ed10d94d207283dcc3248457b4a804f6aad5f824 Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 5 Feb 2024 15:03:18 +0000 Subject: [PATCH 11/30] Improve time diff calculation --- movement/analysis/kinematics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 1da73f1c4..973971240 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -137,7 +137,7 @@ def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.DataArray: raise ValueError("order must be a positive integer.") else: result = data - dt = data["time"].diff(dim="time").values[0] + dt = data["time"].values[1] - data["time"].values[0] for _ in range(order): result = xr.apply_ufunc( np.gradient, From 26351c753cfcf5999f2c86b38177dab068497c74 Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 5 Feb 2024 15:07:07 +0000 Subject: [PATCH 12/30] Prefix kinematics methods with `compute_` --- movement/analysis/kinematics.py | 26 ++++++++++++++------------ tests/test_unit/test_kinematics.py | 24 +++++++++++++++++------- 2 files changed, 31 insertions(+), 19 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 973971240..512c091e2 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -2,7 +2,7 @@ import xarray as xr -def displacement(data: xr.DataArray) -> xr.DataArray: +def compute_displacement(data: xr.DataArray) -> xr.DataArray: """Compute the displacement between consecutive x, y locations of each keypoint of each individual. @@ -22,7 +22,7 @@ def displacement(data: xr.DataArray) -> xr.DataArray: return displacement_xy -def displacement_vector(data: xr.DataArray) -> xr.Dataset: +def compute_displacement_vector(data: xr.DataArray) -> xr.Dataset: """Compute the Euclidean magnitude (distance) and direction (angle relative to the positive x-axis) of displacement. @@ -38,10 +38,10 @@ def displacement_vector(data: xr.DataArray) -> xr.Dataset: An xarray Dataset containing the magnitude and direction of the computed displacement. """ - return compute_vector_magnitude_direction(displacement(data)) + return compute_vector_magnitude_direction(compute_displacement(data)) -def velocity(data: xr.DataArray) -> xr.DataArray: +def compute_velocity(data: xr.DataArray) -> xr.DataArray: """Compute the velocity between consecutive x, y locations of each keypoint of each individual. @@ -56,10 +56,10 @@ def velocity(data: xr.DataArray) -> xr.DataArray: xarray.DataArray An xarray Dataset containing the computed velocity. """ - return approximate_derivative(data, order=1) + return compute_approximate_derivative(data, order=1) -def velocity_vector(data: xr.DataArray) -> xr.Dataset: +def compute_velocity_vector(data: xr.DataArray) -> xr.Dataset: """Compute the Euclidean magnitude (speed) and direction (angle relative to the positive x-axis) of velocity. @@ -75,10 +75,10 @@ def velocity_vector(data: xr.DataArray) -> xr.Dataset: An xarray Dataset containing the magnitude and direction of the computed velocity. """ - return compute_vector_magnitude_direction(velocity(data)) + return compute_vector_magnitude_direction(compute_velocity(data)) -def acceleration(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. @@ -94,10 +94,10 @@ def acceleration(data: xr.DataArray) -> xr.DataArray: An xarray Dataset containing the magnitude and direction of acceleration. """ - return approximate_derivative(data, order=2) + return compute_approximate_derivative(data, order=2) -def acceleration_vector(data: xr.DataArray) -> xr.Dataset: +def compute_acceleration_vector(data: xr.DataArray) -> xr.Dataset: """Compute the Euclidean magnitude and direction of acceleration. Parameters @@ -112,10 +112,12 @@ def acceleration_vector(data: xr.DataArray) -> xr.Dataset: An xarray Dataset containing the magnitude and direction of the computed acceleration. """ - return compute_vector_magnitude_direction(acceleration(data)) + return compute_vector_magnitude_direction(compute_acceleration(data)) -def approximate_derivative(data: xr.DataArray, order: int = 1) -> xr.DataArray: +def compute_approximate_derivative( + data: xr.DataArray, order: int = 1 +) -> xr.DataArray: """Compute velocity or acceleration using numerical differentiation, assuming equidistant time spacing. diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 669ef6dbc..ab7f43d54 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -44,14 +44,18 @@ def expected_dataset(self, valid_pose_dataset): def test_displacement(self, valid_pose_dataset, expected_dataarray): """Test displacement calculation.""" - result = kinematics.displacement(valid_pose_dataset.pose_tracks) + 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) def test_displacement_vector(self, valid_pose_dataset, expected_dataset): """Test displacement vector calculation.""" - result = kinematics.displacement_vector(valid_pose_dataset.pose_tracks) + result = kinematics.compute_displacement_vector( + valid_pose_dataset.pose_tracks + ) # Set the first displacement to zero expected_dataset.magnitude[0, :, :] = 0 expected_dataset.direction[0, :, :] = 0 @@ -59,23 +63,29 @@ def test_displacement_vector(self, valid_pose_dataset, expected_dataset): def test_velocity(self, valid_pose_dataset, expected_dataarray): """Test velocity calculation.""" - result = kinematics.velocity(valid_pose_dataset.pose_tracks) + result = kinematics.compute_velocity(valid_pose_dataset.pose_tracks) xr.testing.assert_allclose(result, expected_dataarray) def test_velocity_vector(self, valid_pose_dataset, expected_dataset): """Test velocity vector calculation.""" - result = kinematics.velocity_vector(valid_pose_dataset.pose_tracks) + result = kinematics.compute_velocity_vector( + valid_pose_dataset.pose_tracks + ) xr.testing.assert_allclose(result, expected_dataset) def test_acceleration(self, valid_pose_dataset, expected_dataarray): """Test acceleration calculation.""" - result = kinematics.acceleration(valid_pose_dataset.pose_tracks) + result = kinematics.compute_acceleration( + valid_pose_dataset.pose_tracks + ) expected_dataarray[:] = 0 xr.testing.assert_allclose(result, expected_dataarray) def test_acceleration_vector(self, valid_pose_dataset, expected_dataset): """Test acceleration vector calculation.""" - result = kinematics.acceleration_vector(valid_pose_dataset.pose_tracks) + result = kinematics.compute_acceleration_vector( + valid_pose_dataset.pose_tracks + ) expected_dataset.magnitude[:] = 0 expected_dataset.direction[:] = 0 xr.testing.assert_allclose(result, expected_dataset) @@ -84,4 +94,4 @@ def test_approximate_derivative_with_nonpositive_order(self): """Test that an error is raised when the order is non-positive.""" data = np.arange(10) with pytest.raises(ValueError): - kinematics.approximate_derivative(data, order=0) + kinematics.compute_approximate_derivative(data, order=0) From 425d4e895d25346f31509c2800b60a363686544d Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 5 Feb 2024 15:28:33 +0000 Subject: [PATCH 13/30] Add kinematic properties to `PosesAccessor` --- movement/move_accessor.py | 32 ++++++++++++++++++++++++++ tests/test_unit/test_poses_accessor.py | 18 +++++++++++++++ 2 files changed, 50 insertions(+) create mode 100644 tests/test_unit/test_poses_accessor.py diff --git a/movement/move_accessor.py b/movement/move_accessor.py index 3186fe3bd..11b65e72e 100644 --- a/movement/move_accessor.py +++ b/movement/move_accessor.py @@ -3,6 +3,7 @@ import xarray as xr +from movement.analysis import kinematics from movement.io.validators import ValidPoseTracks logger = logging.getLogger(__name__) @@ -70,6 +71,37 @@ class MoveAccessor: def __init__(self, ds: xr.Dataset): self._obj = ds + @property + def displacement(self) -> xr.DataArray: + """Return the displacement between consecutive x, y + locations of each keypoint of each individual. + """ + pose_tracks = self._obj[self.var_names[0]] + self._obj["displacement"] = kinematics.compute_displacement( + pose_tracks + ) + return self._obj["displacement"] + + @property + def velocity(self) -> xr.DataArray: + """Return the velocity between consecutive x, y locations + of each keypoint of each individual. + """ + pose_tracks = self._obj[self.var_names[0]] + self._obj["velocity"] = kinematics.compute_velocity(pose_tracks) + return self._obj["velocity"] + + @property + def acceleration(self) -> xr.DataArray: + """Return the acceleration between consecutive x, y locations + of each keypoint of each individual. + """ + pose_tracks = self._obj[self.var_names[0]] + self._obj["acceleration"] = kinematics.compute_acceleration( + pose_tracks + ) + return self._obj["acceleration"] + def validate(self) -> None: """Validate the PoseTracks dataset.""" fps = self._obj.attrs.get("fps", None) diff --git a/tests/test_unit/test_poses_accessor.py b/tests/test_unit/test_poses_accessor.py new file mode 100644 index 000000000..46debb980 --- /dev/null +++ b/tests/test_unit/test_poses_accessor.py @@ -0,0 +1,18 @@ +import pytest +import xarray as xr + + +class TestMoveAccessor: + """Test suite for the move_accessor module.""" + + @pytest.mark.parametrize( + "property_name", ["displacement", "velocity", "acceleration"] + ) + def test_property(self, valid_pose_dataset, property_name): + """Test that the property returns an instance of xr.DataArray + with the correct name, and that the input xr.Dataset contains + the newly-added data variable with the same name.""" + result = getattr(valid_pose_dataset.move, property_name) + assert isinstance(result, xr.DataArray) + assert result.name == property_name + assert property_name in valid_pose_dataset.data_vars From 79141567da15c8d45921d9345e3e3c4c450c0b6d Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 5 Feb 2024 15:31:28 +0000 Subject: [PATCH 14/30] Update `test_property` docstring --- tests/test_unit/test_poses_accessor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_unit/test_poses_accessor.py b/tests/test_unit/test_poses_accessor.py index 46debb980..99820a65f 100644 --- a/tests/test_unit/test_poses_accessor.py +++ b/tests/test_unit/test_poses_accessor.py @@ -9,9 +9,9 @@ class TestMoveAccessor: "property_name", ["displacement", "velocity", "acceleration"] ) def test_property(self, valid_pose_dataset, property_name): - """Test that the property returns an instance of xr.DataArray - with the correct name, and that the input xr.Dataset contains - the newly-added data variable with the same name.""" + """Test that accessing a property returns an instance of + xr.DataArray with the correct name, and that the input xr.Dataset + now contains the property as a data variable.""" result = getattr(valid_pose_dataset.move, property_name) assert isinstance(result, xr.DataArray) assert result.name == property_name From 7664e6e794ebc007688c9cc55ed090abfebc3e47 Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 5 Feb 2024 18:29:44 +0000 Subject: [PATCH 15/30] Refactor `_vector` methods and kinematics tests --- movement/analysis/kinematics.py | 98 +++++++---------------- tests/conftest.py | 7 ++ tests/test_unit/test_kinematics.py | 104 +++++++++++++------------ tests/test_unit/test_poses_accessor.py | 12 +-- 4 files changed, 94 insertions(+), 127 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 512c091e2..6b9bb5e3c 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -22,25 +22,6 @@ def compute_displacement(data: xr.DataArray) -> xr.DataArray: return displacement_xy -def compute_displacement_vector(data: xr.DataArray) -> xr.Dataset: - """Compute the Euclidean magnitude (distance) and direction - (angle relative to the positive x-axis) of displacement. - - Parameters - ---------- - data : xarray.DataArray - The input data, assumed to be of shape (..., 2), where - the last dimension contains the x and y coordinates. - - Returns - ------- - xarray.Dataset - An xarray Dataset containing the magnitude and direction - of the computed displacement. - """ - return compute_vector_magnitude_direction(compute_displacement(data)) - - def compute_velocity(data: xr.DataArray) -> xr.DataArray: """Compute the velocity between consecutive x, y locations of each keypoint of each individual. @@ -59,25 +40,6 @@ def compute_velocity(data: xr.DataArray) -> xr.DataArray: return compute_approximate_derivative(data, order=1) -def compute_velocity_vector(data: xr.DataArray) -> xr.Dataset: - """Compute the Euclidean magnitude (speed) and direction - (angle relative to the positive x-axis) of velocity. - - Parameters - ---------- - data : xarray.DataArray - The input data, assumed to be of shape (..., 2), where - the last dimension contains the x and y coordinates. - - Returns - ------- - xarray.Dataset - An xarray Dataset containing the magnitude and direction - of the computed velocity. - """ - return compute_vector_magnitude_direction(compute_velocity(data)) - - def compute_acceleration(data: xr.DataArray) -> xr.DataArray: """Compute the acceleration between consecutive x, y locations of each keypoint of each individual. @@ -97,24 +59,6 @@ def compute_acceleration(data: xr.DataArray) -> xr.DataArray: return compute_approximate_derivative(data, order=2) -def compute_acceleration_vector(data: xr.DataArray) -> xr.Dataset: - """Compute the Euclidean magnitude and direction of acceleration. - - Parameters - ---------- - data : xarray.DataArray - The input data, assumed to be of shape (..., 2), where - the last dimension contains the x and y coordinates. - - Returns - ------- - xarray.Dataset - An xarray Dataset containing the magnitude and direction - of the computed acceleration. - """ - return compute_vector_magnitude_direction(compute_acceleration(data)) - - def compute_approximate_derivative( data: xr.DataArray, order: int = 1 ) -> xr.DataArray: @@ -125,7 +69,7 @@ def compute_approximate_derivative( ---------- data : xarray.DataArray The input data, assumed to be of shape (..., 2), where the last - dimension contains the x and y coordinates. + dimension contains data in the x and y dimensions. order : int The order of the derivative. 1 for velocity, 2 for acceleration. Default is 1. @@ -151,33 +95,47 @@ def compute_approximate_derivative( return result -def compute_vector_magnitude_direction(input: xr.DataArray) -> xr.Dataset: - """Compute the magnitude and direction of a vector. +def compute_norm(data: xr.DataArray) -> xr.DataArray: + """Compute the Euclidean norm (magnitude) of a vector. Parameters ---------- - input : xarray.DataArray + data : xarray.DataArray The input data, assumed to be of shape (..., 2), where the last - dimension contains the x and y coordinates. + dimension contains data in the x and y dimensions. Returns ------- - xarray.Dataset - An xarray Dataset containing the computed magnitude and - direction. + xarray.DataArray + An xarray DataArray containing the computed Euclidean norm. """ - magnitude = xr.apply_ufunc( + return xr.apply_ufunc( np.linalg.norm, - input, + data, input_core_dims=[["space"]], kwargs={"axis": -1}, ) - direction = xr.apply_ufunc( + + +def compute_theta(data: xr.DataArray) -> xr.DataArray: + """Compute the theta (direction) of a vector. + + 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. + + Returns + ------- + xarray.DataArray + An xarray DataArray containing the computed theta. + """ + return xr.apply_ufunc( np.arctan2, - input[..., 1], - input[..., 0], + data[..., 1], + data[..., 0], ) - return xr.Dataset({"magnitude": magnitude, "direction": direction}) # Locomotion Features diff --git a/tests/conftest.py b/tests/conftest.py index 3f762a248..47958fe3e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -297,4 +297,11 @@ def missing_dim_dataset(valid_pose_dataset): ] ) def invalid_pose_dataset(request): + """Return an invalid pose tracks dataset.""" return request.getfixturevalue(request.param) + + +@pytest.fixture(params=["displacement", "velocity", "acceleration"]) +def kinematic_property(request): + """Return a kinematic property.""" + return request.param diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index ab7f43d54..44c38a002 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -20,30 +20,42 @@ def expected_dataarray(self, valid_pose_dataset): @pytest.fixture def expected_dataset(self, valid_pose_dataset): - """Return an xarray.Dataset with default `magnitude` and - `direction` data variables, and the expected dimensions - and coordinates.""" - dims = valid_pose_dataset.pose_tracks.dims[:-1] - return xr.Dataset( - data_vars={ - "magnitude": xr.DataArray( - np.full((10, 2, 2), 5.0), - dims=dims, - ), - "direction": xr.DataArray( - np.full((10, 2, 2), 0.92729522), - dims=dims, - ), - }, - coords={ - "time": valid_pose_dataset.time, - "individuals": valid_pose_dataset.individuals, - "keypoints": valid_pose_dataset.keypoints, - }, - ) + """Return a function to generate an expected dataset.""" + + def _expected_dataset(name): + dims = valid_pose_dataset.pose_tracks.dims[:-1] + ds = xr.Dataset( + data_vars={ + "norm": xr.DataArray( + np.full((10, 2, 2), 5.0), + dims=dims, + ), + "theta": xr.DataArray( + np.full((10, 2, 2), 0.92729522), + dims=dims, + ), + }, + coords={ + "time": valid_pose_dataset.time, + "individuals": valid_pose_dataset.individuals, + "keypoints": valid_pose_dataset.keypoints, + }, + ) + if name == "displacement": + # Set the first values to zero + ds.norm[0, :, :] = 0 + ds.theta[0, :, :] = 0 + elif name == "acceleration": + # Set all values to zero + ds.norm[:] = 0 + ds.theta[:] = 0 + ds.attrs["name"] = name + return ds + + return _expected_dataset def test_displacement(self, valid_pose_dataset, expected_dataarray): - """Test displacement calculation.""" + """Test displacement computation.""" result = kinematics.compute_displacement( valid_pose_dataset.pose_tracks ) @@ -51,44 +63,38 @@ def test_displacement(self, valid_pose_dataset, expected_dataarray): expected_dataarray[0, :, :, :] = 0 xr.testing.assert_allclose(result, expected_dataarray) - def test_displacement_vector(self, valid_pose_dataset, expected_dataset): - """Test displacement vector calculation.""" - result = kinematics.compute_displacement_vector( - valid_pose_dataset.pose_tracks - ) - # Set the first displacement to zero - expected_dataset.magnitude[0, :, :] = 0 - expected_dataset.direction[0, :, :] = 0 - xr.testing.assert_allclose(result, expected_dataset) - def test_velocity(self, valid_pose_dataset, expected_dataarray): - """Test velocity calculation.""" + """Test velocity computation.""" result = kinematics.compute_velocity(valid_pose_dataset.pose_tracks) xr.testing.assert_allclose(result, expected_dataarray) - def test_velocity_vector(self, valid_pose_dataset, expected_dataset): - """Test velocity vector calculation.""" - result = kinematics.compute_velocity_vector( - valid_pose_dataset.pose_tracks - ) - xr.testing.assert_allclose(result, expected_dataset) - def test_acceleration(self, valid_pose_dataset, expected_dataarray): - """Test acceleration calculation.""" + """Test acceleration computation.""" result = kinematics.compute_acceleration( valid_pose_dataset.pose_tracks ) expected_dataarray[:] = 0 xr.testing.assert_allclose(result, expected_dataarray) - def test_acceleration_vector(self, valid_pose_dataset, expected_dataset): - """Test acceleration vector calculation.""" - result = kinematics.compute_acceleration_vector( - valid_pose_dataset.pose_tracks - ) - expected_dataset.magnitude[:] = 0 - expected_dataset.direction[:] = 0 - xr.testing.assert_allclose(result, expected_dataset) + def test_compute_norm( + self, valid_pose_dataset, kinematic_property, expected_dataset + ): + """Test Euclidean norm (magnitude) computation for different + kinematic properties.""" + data = getattr(valid_pose_dataset.poses, kinematic_property) + result = kinematics.compute_norm(data) + expected = expected_dataset(kinematic_property).norm + xr.testing.assert_allclose(result, expected) + + def test_compute_theta( + self, valid_pose_dataset, kinematic_property, expected_dataset + ): + """Test theta (direction) computation for different + kinematic properties.""" + data = getattr(valid_pose_dataset.poses, kinematic_property) + result = kinematics.compute_theta(data) + expected = expected_dataset(kinematic_property).theta + xr.testing.assert_allclose(result, expected) def test_approximate_derivative_with_nonpositive_order(self): """Test that an error is raised when the order is non-positive.""" diff --git a/tests/test_unit/test_poses_accessor.py b/tests/test_unit/test_poses_accessor.py index 99820a65f..4af50a2d0 100644 --- a/tests/test_unit/test_poses_accessor.py +++ b/tests/test_unit/test_poses_accessor.py @@ -1,18 +1,14 @@ -import pytest import xarray as xr class TestMoveAccessor: """Test suite for the move_accessor module.""" - @pytest.mark.parametrize( - "property_name", ["displacement", "velocity", "acceleration"] - ) - def test_property(self, valid_pose_dataset, property_name): + def test_property(self, valid_pose_dataset, kinematic_property): """Test that accessing a property returns an instance of xr.DataArray with the correct name, and that the input xr.Dataset now contains the property as a data variable.""" - result = getattr(valid_pose_dataset.move, property_name) + result = getattr(valid_pose_dataset.move, kinematic_property) assert isinstance(result, xr.DataArray) - assert result.name == property_name - assert property_name in valid_pose_dataset.data_vars + assert result.name == kinematic_property + assert kinematic_property in valid_pose_dataset.data_vars From 8d96421670baaf531e8ab7acd92ca386fca7709a Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 5 Feb 2024 18:42:12 +0000 Subject: [PATCH 16/30] Update `expected_dataset` docstring --- tests/test_unit/test_kinematics.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 44c38a002..d73ba8b0e 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -20,9 +20,12 @@ def expected_dataarray(self, valid_pose_dataset): @pytest.fixture def expected_dataset(self, valid_pose_dataset): - """Return a function to generate an expected dataset.""" + """Return a function to generate the expected dataset + for different kinematic properties.""" def _expected_dataset(name): + """Return an xarray.Dataset with expected norm and + theta values.""" dims = valid_pose_dataset.pose_tracks.dims[:-1] ds = xr.Dataset( data_vars={ @@ -49,7 +52,6 @@ def _expected_dataset(name): # Set all values to zero ds.norm[:] = 0 ds.theta[:] = 0 - ds.attrs["name"] = name return ds return _expected_dataset From f0fd4694ce4bb3c8b27ff7a47931e4eacc11dc08 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 9 Feb 2024 17:01:46 +0000 Subject: [PATCH 17/30] Rename `poses` to `move` in `PosesAccessor` --- tests/test_unit/test_kinematics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index d73ba8b0e..abde66191 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -83,7 +83,7 @@ def test_compute_norm( ): """Test Euclidean norm (magnitude) computation for different kinematic properties.""" - data = getattr(valid_pose_dataset.poses, kinematic_property) + data = getattr(valid_pose_dataset.move, kinematic_property) result = kinematics.compute_norm(data) expected = expected_dataset(kinematic_property).norm xr.testing.assert_allclose(result, expected) @@ -93,7 +93,7 @@ def test_compute_theta( ): """Test theta (direction) computation for different kinematic properties.""" - data = getattr(valid_pose_dataset.poses, kinematic_property) + data = getattr(valid_pose_dataset.move, kinematic_property) result = kinematics.compute_theta(data) expected = expected_dataset(kinematic_property).theta xr.testing.assert_allclose(result, expected) From 56451c1f5d1ba7b7618693a79e55b702647aec4c Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 9 Feb 2024 17:10:52 +0000 Subject: [PATCH 18/30] Refactor properties in `PosesAccessor` --- movement/move_accessor.py | 42 ++++++++++++++++++++++++++++----------- 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/movement/move_accessor.py b/movement/move_accessor.py index 11b65e72e..cecdb8ce4 100644 --- a/movement/move_accessor.py +++ b/movement/move_accessor.py @@ -1,5 +1,5 @@ import logging -from typing import ClassVar +from typing import Callable, ClassVar import xarray as xr @@ -71,36 +71,54 @@ class MoveAccessor: def __init__(self, ds: xr.Dataset): self._obj = ds + def _compute_property( + self, + property: str, + compute_function: Callable[[xr.DataArray], xr.DataArray], + ) -> xr.DataArray: + """Compute a kinematic property and store it in the dataset. + + Parameters + ---------- + property : str + The name of the property to compute. + compute_function : Callable[[xarray.DataArray], xarray.DataArray] + The function to compute the property. + + Returns + ------- + xarray.DataArray + The computed property. + """ + if property not in self._obj: + pose_tracks = self._obj[self.var_names[0]] + self._obj[property] = compute_function(pose_tracks) + return self._obj[property] + @property def displacement(self) -> xr.DataArray: """Return the displacement between consecutive x, y locations of each keypoint of each individual. """ - pose_tracks = self._obj[self.var_names[0]] - self._obj["displacement"] = kinematics.compute_displacement( - pose_tracks + return self._compute_property( + "displacement", kinematics.compute_displacement ) - return self._obj["displacement"] @property def velocity(self) -> xr.DataArray: """Return the velocity between consecutive x, y locations of each keypoint of each individual. """ - pose_tracks = self._obj[self.var_names[0]] - self._obj["velocity"] = kinematics.compute_velocity(pose_tracks) - return self._obj["velocity"] + return self._compute_property("velocity", kinematics.compute_velocity) @property def acceleration(self) -> xr.DataArray: """Return the acceleration between consecutive x, y locations of each keypoint of each individual. """ - pose_tracks = self._obj[self.var_names[0]] - self._obj["acceleration"] = kinematics.compute_acceleration( - pose_tracks + return self._compute_property( + "acceleration", kinematics.compute_acceleration ) - return self._obj["acceleration"] def validate(self) -> None: """Validate the PoseTracks dataset.""" From 344f8c5c65c9fcacfca03e1d5c999bb78e6f7519 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 9 Feb 2024 17:13:11 +0000 Subject: [PATCH 19/30] Remove vector util functions and tests --- movement/analysis/kinematics.py | 43 ------------------------------ tests/test_unit/test_kinematics.py | 20 -------------- 2 files changed, 63 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 6b9bb5e3c..6c7cdf271 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -95,49 +95,6 @@ def compute_approximate_derivative( return result -def compute_norm(data: xr.DataArray) -> xr.DataArray: - """Compute the Euclidean norm (magnitude) of a vector. - - 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. - - Returns - ------- - xarray.DataArray - An xarray DataArray containing the computed Euclidean norm. - """ - return xr.apply_ufunc( - np.linalg.norm, - data, - input_core_dims=[["space"]], - kwargs={"axis": -1}, - ) - - -def compute_theta(data: xr.DataArray) -> xr.DataArray: - """Compute the theta (direction) of a vector. - - 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. - - Returns - ------- - xarray.DataArray - An xarray DataArray containing the computed theta. - """ - return xr.apply_ufunc( - np.arctan2, - data[..., 1], - data[..., 0], - ) - - # Locomotion Features # speed # speed_centroid diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index abde66191..190813e2c 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -78,26 +78,6 @@ def test_acceleration(self, valid_pose_dataset, expected_dataarray): expected_dataarray[:] = 0 xr.testing.assert_allclose(result, expected_dataarray) - def test_compute_norm( - self, valid_pose_dataset, kinematic_property, expected_dataset - ): - """Test Euclidean norm (magnitude) computation for different - kinematic properties.""" - data = getattr(valid_pose_dataset.move, kinematic_property) - result = kinematics.compute_norm(data) - expected = expected_dataset(kinematic_property).norm - xr.testing.assert_allclose(result, expected) - - def test_compute_theta( - self, valid_pose_dataset, kinematic_property, expected_dataset - ): - """Test theta (direction) computation for different - kinematic properties.""" - data = getattr(valid_pose_dataset.move, kinematic_property) - result = kinematics.compute_theta(data) - expected = expected_dataset(kinematic_property).theta - xr.testing.assert_allclose(result, expected) - def test_approximate_derivative_with_nonpositive_order(self): """Test that an error is raised when the order is non-positive.""" data = np.arange(10) From d8102fcc617f93ebe43995e187ac68b0fa54ac0a Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 19 Feb 2024 14:59:49 +0000 Subject: [PATCH 20/30] Update `not_a_dataset` fixture description --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 47958fe3e..04dd31df0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -266,7 +266,7 @@ def valid_pose_dataset(valid_tracks_array, request): @pytest.fixture def not_a_dataset(): - """Return an invalid pose tracks dataset.""" + """Return data that is not a pose tracks dataset.""" return [1, 2, 3] From 4838837fb004b6b1a9c4c7492dbbccbd951ed869 Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 19 Feb 2024 15:03:37 +0000 Subject: [PATCH 21/30] Validate dataset upon accessor property access --- movement/move_accessor.py | 1 + tests/test_unit/test_poses_accessor.py | 25 +++++++++++++++++++++---- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/movement/move_accessor.py b/movement/move_accessor.py index cecdb8ce4..6bf44e99b 100644 --- a/movement/move_accessor.py +++ b/movement/move_accessor.py @@ -90,6 +90,7 @@ def _compute_property( xarray.DataArray The computed property. """ + self.validate() if property not in self._obj: pose_tracks = self._obj[self.var_names[0]] self._obj[property] = compute_function(pose_tracks) diff --git a/tests/test_unit/test_poses_accessor.py b/tests/test_unit/test_poses_accessor.py index 4af50a2d0..5cdfb2992 100644 --- a/tests/test_unit/test_poses_accessor.py +++ b/tests/test_unit/test_poses_accessor.py @@ -1,14 +1,31 @@ +import pytest import xarray as xr class TestMoveAccessor: """Test suite for the move_accessor module.""" - def test_property(self, valid_pose_dataset, kinematic_property): - """Test that accessing a property returns an instance of - xr.DataArray with the correct name, and that the input xr.Dataset - now contains the property as a data variable.""" + def test_property_with_valid_dataset( + self, valid_pose_dataset, kinematic_property + ): + """Test that accessing a property of a valid pose dataset + returns an instance of xr.DataArray with the correct name, + and that the input xr.Dataset now contains the property as + a data variable.""" result = getattr(valid_pose_dataset.move, kinematic_property) assert isinstance(result, xr.DataArray) assert result.name == kinematic_property assert kinematic_property in valid_pose_dataset.data_vars + + def test_property_with_invalid_dataset( + self, invalid_pose_dataset, kinematic_property + ): + """Test that accessing a property of an invalid pose dataset + raises a ValueError.""" + expected_exception = ( + ValueError + if isinstance(invalid_pose_dataset, xr.Dataset) + else AttributeError + ) + with pytest.raises(expected_exception): + getattr(invalid_pose_dataset.move, kinematic_property) From a945999ee233be5ec6911bf04731fb8bbe89781e Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 19 Feb 2024 15:05:10 +0000 Subject: [PATCH 22/30] Update `poses_accessor` test description --- tests/test_unit/test_poses_accessor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_unit/test_poses_accessor.py b/tests/test_unit/test_poses_accessor.py index 5cdfb2992..94e410ca9 100644 --- a/tests/test_unit/test_poses_accessor.py +++ b/tests/test_unit/test_poses_accessor.py @@ -21,7 +21,7 @@ def test_property_with_invalid_dataset( self, invalid_pose_dataset, kinematic_property ): """Test that accessing a property of an invalid pose dataset - raises a ValueError.""" + raises the appropriate error.""" expected_exception = ( ValueError if isinstance(invalid_pose_dataset, xr.Dataset) From 15ca8a5e3d5a98b83f35568bba865ec9f2cbfc2c Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 19 Feb 2024 17:33:04 +0000 Subject: [PATCH 23/30] Validate input data in kinematic functions --- movement/analysis/kinematics.py | 92 +++++++++++++++++------------- tests/conftest.py | 2 +- tests/test_unit/test_kinematics.py | 69 +++++++++++++++++----- 3 files changed, 106 insertions(+), 57 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 6c7cdf271..8914872b3 100644 --- a/movement/analysis/kinematics.py +++ b/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,7 +60,7 @@ 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. @@ -68,40 +68,50 @@ def compute_approximate_derivative( 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." + ) diff --git a/tests/conftest.py b/tests/conftest.py index 04dd31df0..93a40a0c5 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 190813e2c..7be6bf0aa 100644 --- a/tests/test_unit/test_kinematics.py +++ b/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 + ) From d3a6e0f58645618392c8e630ba962e621a48c694 Mon Sep 17 00:00:00 2001 From: lochhh Date: Mon, 19 Feb 2024 17:34:06 +0000 Subject: [PATCH 24/30] Remove unused fixture --- tests/test_unit/test_kinematics.py | 38 ------------------------------ 1 file changed, 38 deletions(-) diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index 7be6bf0aa..c434c3e4c 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -46,44 +46,6 @@ def _expected_dataarray(property): return _expected_dataarray - @pytest.fixture - def expected_dataset(self, valid_pose_dataset): - """Return a function to generate the expected dataset - for different kinematic properties.""" - - def _expected_dataset(name): - """Return an xarray.Dataset with expected norm and - theta values.""" - dims = valid_pose_dataset.pose_tracks.dims[:-1] - ds = xr.Dataset( - data_vars={ - "norm": xr.DataArray( - np.full((10, 2, 2), 5.0), - dims=dims, - ), - "theta": xr.DataArray( - np.full((10, 2, 2), 0.92729522), - dims=dims, - ), - }, - coords={ - "time": valid_pose_dataset.time, - "individuals": valid_pose_dataset.individuals, - "keypoints": valid_pose_dataset.keypoints, - }, - ) - if name == "displacement": - # Set the first values to zero - ds.norm[0, :, :] = 0 - ds.theta[0, :, :] = 0 - elif name == "acceleration": - # Set all values to zero - ds.norm[:] = 0 - ds.theta[:] = 0 - return ds - - return _expected_dataset - def test_displacement(self, valid_pose_dataset, expected_dataarray): """Test displacement computation.""" result = kinematics.compute_displacement( From 9c92ae8c9acc561b80f0e7953da3c50b2e2a0b83 Mon Sep 17 00:00:00 2001 From: lochhh Date: Tue, 20 Feb 2024 16:46:32 +0000 Subject: [PATCH 25/30] Parametrise kinematics tests --- tests/conftest.py | 13 +++++- tests/test_unit/test_kinematics.py | 70 ++++++++++++++++++++---------- 2 files changed, 59 insertions(+), 24 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 93a40a0c5..c952d44f4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -217,7 +217,9 @@ def _valid_tracks_array(array_type): n_frames = 10 n_individuals = 2 n_keypoints = 2 - base = np.arange(n_frames)[:, np.newaxis, np.newaxis, np.newaxis] + base = np.arange(n_frames, dtype=float)[ + :, np.newaxis, np.newaxis, np.newaxis + ] if array_type == "single_keypoint_array": n_keypoints = 1 elif array_type == "single_track_array": @@ -264,6 +266,15 @@ def valid_pose_dataset(valid_tracks_array, request): ) +@pytest.fixture +def valid_pose_dataset_with_nan(valid_pose_dataset): + """Return a valid pose tracks dataset with NaN values.""" + valid_pose_dataset.pose_tracks.loc[ + {"individuals": "ind1", "time": [3, 7, 8]} + ] = np.nan + return valid_pose_dataset + + @pytest.fixture def not_a_dataset(): """Return data that is not a pose tracks dataset.""" diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index c434c3e4c..d1b251474 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -1,3 +1,5 @@ +from contextlib import nullcontext as does_not_raise + import numpy as np import pytest import xarray as xr @@ -46,24 +48,56 @@ def _expected_dataarray(property): return _expected_dataarray - def test_displacement(self, valid_pose_dataset, expected_dataarray): + kinematic_test_params = [ + ("valid_pose_dataset", does_not_raise()), + ("valid_pose_dataset_with_nan", does_not_raise()), + ("missing_dim_dataset", pytest.raises(ValueError)), + ] + + @pytest.mark.parametrize("ds, expected_exception", kinematic_test_params) + def test_displacement( + self, ds, expected_exception, expected_dataarray, request + ): """Test displacement computation.""" - result = kinematics.compute_displacement( - valid_pose_dataset.pose_tracks - ) - xr.testing.assert_allclose(result, expected_dataarray("displacement")) + ds = request.getfixturevalue(ds) + with expected_exception: + result = kinematics.compute_displacement(ds.pose_tracks) + expected = expected_dataarray("displacement") + if ds.pose_tracks.isnull().any(): + expected.loc[ + {"individuals": "ind1", "time": [3, 4, 7, 8, 9]} + ] = np.nan + xr.testing.assert_allclose(result, expected) - def test_velocity(self, valid_pose_dataset, expected_dataarray): + @pytest.mark.parametrize("ds, expected_exception", kinematic_test_params) + def test_velocity( + self, ds, expected_exception, expected_dataarray, request + ): """Test velocity computation.""" - result = kinematics.compute_velocity(valid_pose_dataset.pose_tracks) - xr.testing.assert_allclose(result, expected_dataarray("velocity")) + ds = request.getfixturevalue(ds) + with expected_exception: + result = kinematics.compute_velocity(ds.pose_tracks) + expected = expected_dataarray("velocity") + if ds.pose_tracks.isnull().any(): + expected.loc[ + {"individuals": "ind1", "time": [2, 4, 6, 7, 8, 9]} + ] = np.nan + xr.testing.assert_allclose(result, expected) - def test_acceleration(self, valid_pose_dataset, expected_dataarray): + @pytest.mark.parametrize("ds, expected_exception", kinematic_test_params) + def test_acceleration( + self, ds, expected_exception, expected_dataarray, request + ): """Test acceleration computation.""" - result = kinematics.compute_acceleration( - valid_pose_dataset.pose_tracks - ) - xr.testing.assert_allclose(result, expected_dataarray("acceleration")) + ds = request.getfixturevalue(ds) + with expected_exception: + result = kinematics.compute_acceleration(ds.pose_tracks) + expected = expected_dataarray("acceleration") + if ds.pose_tracks.isnull().any(): + expected.loc[ + {"individuals": "ind1", "time": [1, 3, 5, 6, 7, 8, 9]} + ] = np.nan + xr.testing.assert_allclose(result, expected) @pytest.mark.parametrize("order", [0, -1, 1.0, "1"]) def test_approximate_derivative_with_invalid_order(self, order): @@ -74,13 +108,3 @@ def test_approximate_derivative_with_invalid_order(self, order): ) 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): - eval(f"kinematics.compute_{kinematic_property}")( - missing_dim_dataset - ) From 9f1c464a1d8504c95fb6969cb8d9251ec4bd11c1 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 23 Feb 2024 10:21:56 +0000 Subject: [PATCH 26/30] Set `compute_derivative` as internal function --- movement/analysis/kinematics.py | 6 +++--- tests/test_unit/test_kinematics.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 8914872b3..0487102d2 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -38,7 +38,7 @@ def compute_velocity(data: xr.DataArray) -> xr.DataArray: xarray.DataArray An xarray Dataset containing the computed velocity. """ - return compute_approximate_derivative(data, order=1) + return _compute_approximate_derivative(data, order=1) def compute_acceleration(data: xr.DataArray) -> xr.DataArray: @@ -56,10 +56,10 @@ def compute_acceleration(data: xr.DataArray) -> xr.DataArray: An xarray Dataset containing the magnitude and direction of acceleration. """ - return compute_approximate_derivative(data, order=2) + return _compute_approximate_derivative(data, order=2) -def compute_approximate_derivative( +def _compute_approximate_derivative( data: xr.DataArray, order: int ) -> xr.DataArray: """Compute velocity or acceleration using numerical differentiation, diff --git a/tests/test_unit/test_kinematics.py b/tests/test_unit/test_kinematics.py index d1b251474..ca9a78d55 100644 --- a/tests/test_unit/test_kinematics.py +++ b/tests/test_unit/test_kinematics.py @@ -107,4 +107,4 @@ def test_approximate_derivative_with_invalid_order(self, order): ValueError if isinstance(order, int) else TypeError ) with pytest.raises(expected_exception): - kinematics.compute_approximate_derivative(data, order=order) + kinematics._compute_approximate_derivative(data, order=order) From b3421006aade1c88724ed149458df3b59305b29c Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 23 Feb 2024 10:29:29 +0000 Subject: [PATCH 27/30] Update `kinematics.py` docstrings --- movement/analysis/kinematics.py | 41 ++++++++++++++++++++------------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py index 0487102d2..a89ffdf47 100644 --- a/movement/analysis/kinematics.py +++ b/movement/analysis/kinematics.py @@ -5,13 +5,13 @@ def compute_displacement(data: xr.DataArray) -> xr.DataArray: - """Compute the displacement between consecutive locations - of each keypoint of each individual across time. + """Compute the displacement between consecutive positions + of each keypoint for each individual across time. Parameters ---------- data : xarray.DataArray - The input data containing `time` as a dimension. + The input data containing ``time`` as a dimension. Returns ------- @@ -25,36 +25,45 @@ def compute_displacement(data: xr.DataArray) -> xr.DataArray: def compute_velocity(data: xr.DataArray) -> xr.DataArray: - """Compute the velocity between consecutive locations - of each keypoint of each individual across time. + """Compute the velocity between consecutive positions + of each keypoint for each individual across time. Parameters ---------- data : xarray.DataArray - The input data containing `time` as a dimension. + The input data containing ``time`` as a dimension. Returns ------- xarray.DataArray - An xarray Dataset containing the computed velocity. + An xarray DataArray containing the computed velocity. + + Notes + ----- + This function computes velocity using numerical differentiation + and assumes equidistant time spacing. """ return _compute_approximate_derivative(data, order=1) def compute_acceleration(data: xr.DataArray) -> xr.DataArray: - """Compute the acceleration between consecutive locations - of each keypoint of each individual. + """Compute the acceleration between consecutive positions + of each keypoint for each individual across time. Parameters ---------- data : xarray.DataArray - The input data containing `time` as a dimension. + The input data containing ``time`` as a dimension. Returns ------- - xarray.Dataset - An xarray Dataset containing the magnitude and direction - of acceleration. + xarray.DataArray + An xarray DataArray containing the computed acceleration. + + Notes + ----- + This function computes acceleration using numerical differentiation + and assumes equidistant time spacing. """ return _compute_approximate_derivative(data, order=2) @@ -68,7 +77,7 @@ def _compute_approximate_derivative( Parameters ---------- data : xarray.DataArray - The input data containing `time` as a dimension. + The input data containing ``time`` as a dimension. order : int The order of the derivative. 1 for velocity, 2 for acceleration. Value must be a positive integer. @@ -99,7 +108,7 @@ def _compute_approximate_derivative( def _validate_time_dimension(data: xr.DataArray) -> None: - """Validate the input data contains a 'time' dimension. + """Validate the input data contains a ``time`` dimension. Parameters ---------- @@ -109,7 +118,7 @@ def _validate_time_dimension(data: xr.DataArray) -> None: Raises ------ ValueError - If the input data does not contain a 'time' dimension. + If the input data does not contain a ``time`` dimension. """ if "time" not in data.dims: raise log_error( From ac0a579bd8328ff5b01990f818d31382125204e3 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 23 Feb 2024 10:36:00 +0000 Subject: [PATCH 28/30] Add new modules to API docs --- docs/source/api_index.rst | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/docs/source/api_index.rst b/docs/source/api_index.rst index 78e27563b..146c7fa1f 100644 --- a/docs/source/api_index.rst +++ b/docs/source/api_index.rst @@ -44,6 +44,24 @@ Sample Data fetch_sample_data_path fetch_sample_data +Analysis +----------- +.. currentmodule:: movement.analysis.kinematics +.. autosummary:: + :toctree: api + + compute_displacement + compute_velocity + compute_acceleration + +Move Accessor +------------- +.. currentmodule:: movement.move_accessor +.. autosummary:: + :toctree: api + + MoveAccessor + Logging ------- .. currentmodule:: movement.logging From e895de0fe4f59616e7fafe07e9ec37e8d52c9985 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 23 Feb 2024 14:51:54 +0000 Subject: [PATCH 29/30] Update `move_accessor` docstrings --- movement/move_accessor.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/movement/move_accessor.py b/movement/move_accessor.py index 6bf44e99b..36a3c6eee 100644 --- a/movement/move_accessor.py +++ b/movement/move_accessor.py @@ -14,9 +14,10 @@ @xr.register_dataset_accessor("move") class MoveAccessor: - """An accessor that extends an xarray Dataset object. + """An accessor that extends an xarray Dataset by implementing + `movement`-specific properties and methods. - The xarray Dataset has the following dimensions: + The xarray Dataset contains the following expected dimensions: - ``time``: the number of frames in the video - ``individuals``: the number of individuals in the video - ``keypoints``: the number of keypoints in the skeleton @@ -27,11 +28,18 @@ class MoveAccessor: ['x','y',('z')] for ``space``. The coordinates of the ``time`` dimension are in seconds if ``fps`` is provided, otherwise they are in frame numbers. - The dataset contains two data variables (xarray DataArray objects): + The dataset contains two expected data variables (xarray DataArrays): - ``pose_tracks``: with shape (``time``, ``individuals``, ``keypoints``, ``space``) - ``confidence``: with shape (``time``, ``individuals``, ``keypoints``) + When accessing a ``.move`` property (e.g. ``displacement``, ``velocity``, + ``acceleration``) for the first time, the property is computed and stored + as a data variable with the same name in the dataset. The ``.move`` + accessor can be omitted in subsequent accesses, i.e. + ``ds.move.displacement`` and ``ds.displacement`` will return the same data + variable. + The dataset may also contain following attributes as metadata: - ``fps``: the number of frames per second in the video - ``time_unit``: the unit of the ``time`` coordinates, frames or @@ -46,7 +54,7 @@ class MoveAccessor: Using an accessor is the recommended way to extend xarray objects. See [1]_ for more details. - Methods/properties that are specific to this class can be used via + Methods/properties that are specific to this class can be accessed via the ``.move`` accessor, e.g. ``ds.move.validate()``. References @@ -98,8 +106,8 @@ def _compute_property( @property def displacement(self) -> xr.DataArray: - """Return the displacement between consecutive x, y - locations of each keypoint of each individual. + """Return the displacement between consecutive positions + of each keypoint for each individual across time. """ return self._compute_property( "displacement", kinematics.compute_displacement @@ -107,15 +115,15 @@ def displacement(self) -> xr.DataArray: @property def velocity(self) -> xr.DataArray: - """Return the velocity between consecutive x, y locations - of each keypoint of each individual. + """Return the velocity between consecutive positions + of each keypoint for each individual across time. """ return self._compute_property("velocity", kinematics.compute_velocity) @property def acceleration(self) -> xr.DataArray: - """Return the acceleration between consecutive x, y locations - of each keypoint of each individual. + """Return the acceleration between consecutive positions + of each keypoint for each individual across time. """ return self._compute_property( "acceleration", kinematics.compute_acceleration From e5cb36ce851505a47c3ec62b3729e8907d73e282 Mon Sep 17 00:00:00 2001 From: lochhh Date: Fri, 23 Feb 2024 15:01:03 +0000 Subject: [PATCH 30/30] Rename `test_move_accessor` filename --- tests/test_unit/{test_poses_accessor.py => test_move_accessor.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/test_unit/{test_poses_accessor.py => test_move_accessor.py} (100%) diff --git a/tests/test_unit/test_poses_accessor.py b/tests/test_unit/test_move_accessor.py similarity index 100% rename from tests/test_unit/test_poses_accessor.py rename to tests/test_unit/test_move_accessor.py