From eb152a7a80b7dc6a6848156ba2f0a97838520f0d Mon Sep 17 00:00:00 2001 From: niksirbi Date: Tue, 23 Jul 2024 13:43:00 +0100 Subject: [PATCH 01/12] added utility for vector magnitude --- movement/utils/vector.py | 41 ++++++++++++++++++++++++++++++++++ tests/test_unit/test_vector.py | 26 +++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index c35990ebe..bdf235034 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -82,6 +82,47 @@ def pol2cart(data: xr.DataArray) -> xr.DataArray: ).transpose(*dims) +def magnitude(data: xr.DataArray) -> xr.DataArray: + """Compute the magnitude in space. + + The magnitude is computed as the Euclidean norm of a vector + with spatial components ``x`` and ``y`` in Cartesian coordinates. + If the input data contains polar coordinates, the magnitude + is the same as the radial distance ``rho``. + + Parameters + ---------- + data : xarray.DataArray + The input data containing either ``space`` or ``space_pol`` + as a dimension. + + Returns + ------- + xarray.DataArray + An xarray DataArray representing the magnitude of the vector + in space. The output has no spatial dimension. + + + """ + if "space" in data.dims: + _validate_dimension_coordinates(data, {"space": ["x", "y"]}) + return xr.apply_ufunc( + np.linalg.norm, + data, + input_core_dims=[["space"]], + kwargs={"axis": -1}, + ) + elif "space_pol" in data.dims: + _validate_dimension_coordinates(data, {"space_pol": ["rho", "phi"]}) + return data.sel(space_pol="rho", drop=True) + else: + raise log_error( + ValueError, + "Input data must contain either 'space' or 'space_pol' " + "as dimensions.", + ) + + def _validate_dimension_coordinates( data: xr.DataArray, required_dim_coords: dict ) -> None: diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index 88dd85b57..dba66ff4c 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -121,3 +121,29 @@ def test_pol2cart(self, ds, expected_exception, request): with expected_exception: result = vector.pol2cart(ds.pol) xr.testing.assert_allclose(result, ds.cart) + + @pytest.mark.parametrize( + "ds, expected_exception", + [ + ("cart_pol_dataset", does_not_raise()), + ("cart_pol_dataset_with_nan", does_not_raise()), + ("cart_pol_dataset_missing_cart_dim", pytest.raises(ValueError)), + ( + "cart_pol_dataset_missing_cart_coords", + pytest.raises(ValueError), + ), + ], + ) + def test_magnitude(self, ds, expected_exception, request): + """Test vector magnitude with known values.""" + ds = request.getfixturevalue(ds) + with expected_exception: + result = vector.magnitude(ds.cart) + expected = np.sqrt( + ds.cart.sel(space="x") ** 2 + ds.cart.sel(space="y") ** 2 + ) + xr.testing.assert_allclose(result, expected) + # result should be the same from Cartesian and polar coordinates + xr.testing.assert_allclose(result, vector.magnitude(ds.pol)) + # The result should only contain the time dimension. + assert result.dims == ("time",) From 47d7b534e71ff112f06453e7a5ebbd38bc373fc8 Mon Sep 17 00:00:00 2001 From: niksirbi Date: Tue, 23 Jul 2024 15:05:57 +0100 Subject: [PATCH 02/12] added utility for normalising vectors --- movement/utils/vector.py | 25 +++++++++++++++++++++++++ tests/test_unit/test_vector.py | 25 +++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index bdf235034..f6f3da82f 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -123,6 +123,31 @@ def magnitude(data: xr.DataArray) -> xr.DataArray: ) +def normalize(data: xr.DataArray) -> xr.DataArray: + """Normalize data by the magnitude in space. + + Parameters + ---------- + data : xarray.DataArray + The input data containing ``space`` as a dimension, + with ``x`` and ``y`` in the dimension coordinate. + + Returns + ------- + xarray.DataArray + An xarray DataArray representing the normalized data, + having the same dimensions as the input data. + + Notes + ----- + Where the input values are 0 for both ``x`` and ``y``, the normalized + values will be NaN, because of zero-division. + + """ + _validate_dimension_coordinates(data, {"space": ["x", "y"]}) + return data / magnitude(data) + + def _validate_dimension_coordinates( data: xr.DataArray, required_dim_coords: dict ) -> None: diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index dba66ff4c..8e185455c 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -147,3 +147,28 @@ def test_magnitude(self, ds, expected_exception, request): xr.testing.assert_allclose(result, vector.magnitude(ds.pol)) # The result should only contain the time dimension. assert result.dims == ("time",) + + @pytest.mark.parametrize( + "ds, expected_exception", + [ + ("cart_pol_dataset", does_not_raise()), + ("cart_pol_dataset_with_nan", does_not_raise()), + ("cart_pol_dataset_missing_cart_dim", pytest.raises(ValueError)), + ], + ) + def test_normalize(self, ds, expected_exception, request): + """Test data normalization (division by magnitude).""" + ds = request.getfixturevalue(ds) + with expected_exception: + normalized = vector.normalize(ds.cart) + # the normalized data should have the same dimensions as the input + assert normalized.dims == ds.cart.dims + # the first time point is NaN because the input vector is [0, 0] + # (zero-division during normalization). + assert normalized.sel(time=0).isnull().all() + # the magnitude of the normalized vector should be 1 for all + # time points except for the expected NaNs. + normalized_mag = vector.magnitude(normalized).values + expected_mag = np.ones_like(normalized_mag) + expected_mag[normalized.isnull().any("space")] = np.nan + np.testing.assert_allclose(normalized_mag, expected_mag) From 49e5c671be07732979314b6ef1ad48fede1f0f2a Mon Sep 17 00:00:00 2001 From: niksirbi Date: Tue, 23 Jul 2024 15:39:24 +0100 Subject: [PATCH 03/12] use magnitude utility in kinematics example --- examples/compute_kinematics.py | 34 +++++++++++++--------------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/examples/compute_kinematics.py b/examples/compute_kinematics.py index 1ab2731df..2b5d8aa56 100644 --- a/examples/compute_kinematics.py +++ b/examples/compute_kinematics.py @@ -10,14 +10,13 @@ # Imports # ------- -import numpy as np - # For interactive plots: install ipympl with `pip install ipympl` and uncomment # the following line in your notebook # %matplotlib widget from matplotlib import pyplot as plt from movement import sample_data +from movement.utils.vector import magnitude # %% # Load sample dataset @@ -255,13 +254,12 @@ # mouse along its trajectory. # length of each displacement vector -displacement_vectors_lengths = np.linalg.norm( - displacement.sel(individuals=mouse_name, space=["x", "y"]).squeeze(), - axis=1, +displacement_vectors_lengths = magnitude( + displacement.sel(individuals=mouse_name) ) -# sum of all displacement vectors -total_displacement = np.sum(displacement_vectors_lengths, axis=0) # in pixels +# sum of all displacement vectors lengths (in pixels) +total_displacement = displacement_vectors_lengths.sum(dim="time").values[0] print( f"The mouse {mouse_name}'s trajectory is {total_displacement:.2f} " @@ -299,14 +297,12 @@ # uses second order central differences. # %% -# We can also visualise the speed, as the norm of the velocity vector: +# We can also visualise the speed, as the magnitude (norm) +# of the velocity vector: fig, axes = plt.subplots(3, 1, sharex=True, sharey=True) for mouse_name, ax in zip(velocity.individuals.values, axes, strict=False): - # compute the norm of the velocity vector for one mouse - speed_one_mouse = np.linalg.norm( - velocity.sel(individuals=mouse_name, space=["x", "y"]).squeeze(), - axis=1, - ) + # compute the magnitude of the velocity vector for one mouse + speed_one_mouse = magnitude(velocity.sel(individuals=mouse_name)) # plot speed against time ax.plot(speed_one_mouse) ax.set_title(mouse_name) @@ -379,16 +375,12 @@ fig.tight_layout() # %% -# The norm of the acceleration vector is the magnitude of the -# acceleration. -# We can also represent this for each individual. +# The can also represent the magnitude (norm) of the acceleration vector +# for each individual: fig, axes = plt.subplots(3, 1, sharex=True, sharey=True) for mouse_name, ax in zip(accel.individuals.values, axes, strict=False): - # compute norm of the acceleration vector for one mouse - accel_one_mouse = np.linalg.norm( - accel.sel(individuals=mouse_name, space=["x", "y"]).squeeze(), - axis=1, - ) + # compute magnitude of the acceleration vector for one mouse + accel_one_mouse = magnitude(accel.sel(individuals=mouse_name)) # plot acceleration against time ax.plot(accel_one_mouse) From 400a60777739cac320f286ee9b1fe94e067774de Mon Sep 17 00:00:00 2001 From: niksirbi Date: Mon, 12 Aug 2024 14:56:17 +0100 Subject: [PATCH 04/12] use `magnitude` in `cart2pol` --- movement/utils/vector.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index f6f3da82f..2893039b3 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -25,12 +25,7 @@ def cart2pol(data: xr.DataArray) -> xr.DataArray: """ _validate_dimension_coordinates(data, {"space": ["x", "y"]}) - rho = xr.apply_ufunc( - np.linalg.norm, - data, - input_core_dims=[["space"]], - kwargs={"axis": -1}, - ) + rho = magnitude(data) phi = xr.apply_ufunc( np.arctan2, data.sel(space="y"), From cf50d2668ea6a29ba4ec4d4320b2b5054bde635a Mon Sep 17 00:00:00 2001 From: niksirbi Date: Mon, 12 Aug 2024 14:58:14 +0100 Subject: [PATCH 05/12] reorder functions in vector.py --- movement/utils/vector.py | 132 +++++++++++++++++++-------------------- 1 file changed, 66 insertions(+), 66 deletions(-) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index 2893039b3..804ac976e 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -6,6 +6,72 @@ from movement.utils.logging import log_error +def magnitude(data: xr.DataArray) -> xr.DataArray: + """Compute the magnitude in space. + + The magnitude is computed as the Euclidean norm of a vector + with spatial components ``x`` and ``y`` in Cartesian coordinates. + If the input data contains polar coordinates, the magnitude + is the same as the radial distance ``rho``. + + Parameters + ---------- + data : xarray.DataArray + The input data containing either ``space`` or ``space_pol`` + as a dimension. + + Returns + ------- + xarray.DataArray + An xarray DataArray representing the magnitude of the vector + in space. The output has no spatial dimension. + + + """ + if "space" in data.dims: + _validate_dimension_coordinates(data, {"space": ["x", "y"]}) + return xr.apply_ufunc( + np.linalg.norm, + data, + input_core_dims=[["space"]], + kwargs={"axis": -1}, + ) + elif "space_pol" in data.dims: + _validate_dimension_coordinates(data, {"space_pol": ["rho", "phi"]}) + return data.sel(space_pol="rho", drop=True) + else: + raise log_error( + ValueError, + "Input data must contain either 'space' or 'space_pol' " + "as dimensions.", + ) + + +def normalize(data: xr.DataArray) -> xr.DataArray: + """Normalize data by the magnitude in space. + + Parameters + ---------- + data : xarray.DataArray + The input data containing ``space`` as a dimension, + with ``x`` and ``y`` in the dimension coordinate. + + Returns + ------- + xarray.DataArray + An xarray DataArray representing the normalized data, + having the same dimensions as the input data. + + Notes + ----- + Where the input values are 0 for both ``x`` and ``y``, the normalized + values will be NaN, because of zero-division. + + """ + _validate_dimension_coordinates(data, {"space": ["x", "y"]}) + return data / magnitude(data) + + def cart2pol(data: xr.DataArray) -> xr.DataArray: """Transform Cartesian coordinates to polar. @@ -77,72 +143,6 @@ def pol2cart(data: xr.DataArray) -> xr.DataArray: ).transpose(*dims) -def magnitude(data: xr.DataArray) -> xr.DataArray: - """Compute the magnitude in space. - - The magnitude is computed as the Euclidean norm of a vector - with spatial components ``x`` and ``y`` in Cartesian coordinates. - If the input data contains polar coordinates, the magnitude - is the same as the radial distance ``rho``. - - Parameters - ---------- - data : xarray.DataArray - The input data containing either ``space`` or ``space_pol`` - as a dimension. - - Returns - ------- - xarray.DataArray - An xarray DataArray representing the magnitude of the vector - in space. The output has no spatial dimension. - - - """ - if "space" in data.dims: - _validate_dimension_coordinates(data, {"space": ["x", "y"]}) - return xr.apply_ufunc( - np.linalg.norm, - data, - input_core_dims=[["space"]], - kwargs={"axis": -1}, - ) - elif "space_pol" in data.dims: - _validate_dimension_coordinates(data, {"space_pol": ["rho", "phi"]}) - return data.sel(space_pol="rho", drop=True) - else: - raise log_error( - ValueError, - "Input data must contain either 'space' or 'space_pol' " - "as dimensions.", - ) - - -def normalize(data: xr.DataArray) -> xr.DataArray: - """Normalize data by the magnitude in space. - - Parameters - ---------- - data : xarray.DataArray - The input data containing ``space`` as a dimension, - with ``x`` and ``y`` in the dimension coordinate. - - Returns - ------- - xarray.DataArray - An xarray DataArray representing the normalized data, - having the same dimensions as the input data. - - Notes - ----- - Where the input values are 0 for both ``x`` and ``y``, the normalized - values will be NaN, because of zero-division. - - """ - _validate_dimension_coordinates(data, {"space": ["x", "y"]}) - return data / magnitude(data) - - def _validate_dimension_coordinates( data: xr.DataArray, required_dim_coords: dict ) -> None: From d7f81d010fe32c31724a0ee7f41dd1b04f6cb666 Mon Sep 17 00:00:00 2001 From: niksirbi Date: Tue, 13 Aug 2024 16:25:40 +0100 Subject: [PATCH 06/12] rename `magnitude` to `compute_norm` --- examples/compute_kinematics.py | 8 ++++---- movement/utils/vector.py | 32 +++++++++++++++++++++----------- tests/test_unit/test_vector.py | 10 +++++----- 3 files changed, 30 insertions(+), 20 deletions(-) diff --git a/examples/compute_kinematics.py b/examples/compute_kinematics.py index 2b5d8aa56..68797510d 100644 --- a/examples/compute_kinematics.py +++ b/examples/compute_kinematics.py @@ -16,7 +16,7 @@ from matplotlib import pyplot as plt from movement import sample_data -from movement.utils.vector import magnitude +from movement.utils.vector import compute_norm # %% # Load sample dataset @@ -254,7 +254,7 @@ # mouse along its trajectory. # length of each displacement vector -displacement_vectors_lengths = magnitude( +displacement_vectors_lengths = compute_norm( displacement.sel(individuals=mouse_name) ) @@ -302,7 +302,7 @@ fig, axes = plt.subplots(3, 1, sharex=True, sharey=True) for mouse_name, ax in zip(velocity.individuals.values, axes, strict=False): # compute the magnitude of the velocity vector for one mouse - speed_one_mouse = magnitude(velocity.sel(individuals=mouse_name)) + speed_one_mouse = compute_norm(velocity.sel(individuals=mouse_name)) # plot speed against time ax.plot(speed_one_mouse) ax.set_title(mouse_name) @@ -380,7 +380,7 @@ fig, axes = plt.subplots(3, 1, sharex=True, sharey=True) for mouse_name, ax in zip(accel.individuals.values, axes, strict=False): # compute magnitude of the acceleration vector for one mouse - accel_one_mouse = magnitude(accel.sel(individuals=mouse_name)) + accel_one_mouse = compute_norm(accel.sel(individuals=mouse_name)) # plot acceleration against time ax.plot(accel_one_mouse) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index 804ac976e..7d73d6120 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -6,25 +6,35 @@ from movement.utils.logging import log_error -def magnitude(data: xr.DataArray) -> xr.DataArray: - """Compute the magnitude in space. +def compute_norm(data: xr.DataArray) -> xr.DataArray: + """Compute the norm of the vectors along the spatial dimension. - The magnitude is computed as the Euclidean norm of a vector - with spatial components ``x`` and ``y`` in Cartesian coordinates. - If the input data contains polar coordinates, the magnitude - is the same as the radial distance ``rho``. + The norm of a vector is its magnitude, also called Euclidean norm, 2-norm + or Euclidean length. Note that if the input data is expressed in polar + coordinates, the magnitude of a vector is the same as its radial coordinate + ``rho``. Parameters ---------- data : xarray.DataArray - The input data containing either ``space`` or ``space_pol`` + The input data array containing either ``space`` or ``space_pol`` as a dimension. Returns ------- xarray.DataArray - An xarray DataArray representing the magnitude of the vector - in space. The output has no spatial dimension. + A data array holding the norm of the input vectors. + Note that this output array has no spatial dimension but preserves + all other dimensions of the input data array (see Notes). + + Notes + ----- + If the input data array is a ``position`` array, this function will compute + the magnitude of the position vectors, for every individual and keypoint, + at every timestep. If the input data array is a ``shape`` array of a + bounding boxes dataset, it will compute the magnitude of the shape + vectors (i.e., the diagonal of the bounding box), + for every individual and at every timestep. """ @@ -69,7 +79,7 @@ def normalize(data: xr.DataArray) -> xr.DataArray: """ _validate_dimension_coordinates(data, {"space": ["x", "y"]}) - return data / magnitude(data) + return data / compute_norm(data) def cart2pol(data: xr.DataArray) -> xr.DataArray: @@ -91,7 +101,7 @@ def cart2pol(data: xr.DataArray) -> xr.DataArray: """ _validate_dimension_coordinates(data, {"space": ["x", "y"]}) - rho = magnitude(data) + rho = compute_norm(data) phi = xr.apply_ufunc( np.arctan2, data.sel(space="y"), diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index 8e185455c..357ce5e01 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -134,17 +134,17 @@ def test_pol2cart(self, ds, expected_exception, request): ), ], ) - def test_magnitude(self, ds, expected_exception, request): - """Test vector magnitude with known values.""" + def test_compute_norm(self, ds, expected_exception, request): + """Test vector norm computation with known values.""" ds = request.getfixturevalue(ds) with expected_exception: - result = vector.magnitude(ds.cart) + result = vector.compute_norm(ds.cart) expected = np.sqrt( ds.cart.sel(space="x") ** 2 + ds.cart.sel(space="y") ** 2 ) xr.testing.assert_allclose(result, expected) # result should be the same from Cartesian and polar coordinates - xr.testing.assert_allclose(result, vector.magnitude(ds.pol)) + xr.testing.assert_allclose(result, vector.compute_norm(ds.pol)) # The result should only contain the time dimension. assert result.dims == ("time",) @@ -168,7 +168,7 @@ def test_normalize(self, ds, expected_exception, request): assert normalized.sel(time=0).isnull().all() # the magnitude of the normalized vector should be 1 for all # time points except for the expected NaNs. - normalized_mag = vector.magnitude(normalized).values + normalized_mag = vector.compute_norm(normalized).values expected_mag = np.ones_like(normalized_mag) expected_mag[normalized.isnull().any("space")] = np.nan np.testing.assert_allclose(normalized_mag, expected_mag) From 526cb0c6ac22018adadc455e267fb6031a4d7f5d Mon Sep 17 00:00:00 2001 From: niksirbi Date: Tue, 13 Aug 2024 18:26:25 +0100 Subject: [PATCH 07/12] define normalisation in polar coordinates too --- movement/utils/vector.py | 36 +++++++++++++++++++++++++--------- tests/test_unit/test_vector.py | 12 ++++++++++++ 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index 7d73d6120..3aaebf696 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -58,28 +58,46 @@ def compute_norm(data: xr.DataArray) -> xr.DataArray: def normalize(data: xr.DataArray) -> xr.DataArray: - """Normalize data by the magnitude in space. + """Convert the vectors along the spatial dimension into unit vectors. + + A unit vector is a vector pointing in the same direction as the original + vector but with norm = 1. Parameters ---------- data : xarray.DataArray - The input data containing ``space`` as a dimension, - with ``x`` and ``y`` in the dimension coordinate. + The input data array containing either ``space`` or ``space_pol`` + as a dimension. Returns ------- xarray.DataArray - An xarray DataArray representing the normalized data, - having the same dimensions as the input data. + A data array holding the unit vectors of the input data array + (all input dimensions are preserved). Notes ----- - Where the input values are 0 for both ``x`` and ``y``, the normalized - values will be NaN, because of zero-division. + Note that the unit vector for the null vector is undefined, since the null + vector has 0 norm and no direction associated with it. """ - _validate_dimension_coordinates(data, {"space": ["x", "y"]}) - return data / compute_norm(data) + if "space" in data.dims: + _validate_dimension_coordinates(data, {"space": ["x", "y"]}) + return data / compute_norm(data) + elif "space_pol" in data.dims: + _validate_dimension_coordinates(data, {"space_pol": ["rho", "phi"]}) + # Set both rho and phi values to NaN at null vectors (where rho = 0) + new_data = xr.where(data.sel(space_pol="rho") == 0, np.nan, data) + # Set the rho values to 1 for non-null vectors (phi is preserved) + new_data.loc[{"space_pol": "rho"}] = xr.where( + new_data.sel(space_pol="rho").isnull(), np.nan, 1 + ) + return new_data + raise log_error( + ValueError, + "Input data must contain either 'space' or 'space_pol' " + "as dimensions.", + ) def cart2pol(data: xr.DataArray) -> xr.DataArray: diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index 357ce5e01..e25b07fe9 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -172,3 +172,15 @@ def test_normalize(self, ds, expected_exception, request): expected_mag = np.ones_like(normalized_mag) expected_mag[normalized.isnull().any("space")] = np.nan np.testing.assert_allclose(normalized_mag, expected_mag) + # Normalising the polar data should yield the same result + # as converting the normalised Cartesian data to polar. + xr.testing.assert_allclose( + vector.normalize(ds.pol), + vector.cart2pol(vector.normalize(ds.cart)), + ) + # Normalising the Cartesian data should yield the same result + # as converting the normalised polar data to Cartesian. + xr.testing.assert_allclose( + vector.normalize(ds.cart), + vector.pol2cart(vector.normalize(ds.pol)), + ) From 1b22b6033f5acffdb1b745460b651ecf374fb74d Mon Sep 17 00:00:00 2001 From: Niko Sirmpilatze Date: Tue, 13 Aug 2024 18:32:53 +0100 Subject: [PATCH 08/12] Update comment phrasing Co-authored-by: sfmig <33267254+sfmig@users.noreply.github.com> --- examples/compute_kinematics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/compute_kinematics.py b/examples/compute_kinematics.py index 68797510d..b3fefd4a4 100644 --- a/examples/compute_kinematics.py +++ b/examples/compute_kinematics.py @@ -258,7 +258,7 @@ displacement.sel(individuals=mouse_name) ) -# sum of all displacement vectors lengths (in pixels) +# sum the lengths of all displacement vectors (in pixels) total_displacement = displacement_vectors_lengths.sum(dim="time").values[0] print( From 968f5fd5676777ba10eb11b3a789e9ab94176ff4 Mon Sep 17 00:00:00 2001 From: niksirbi Date: Tue, 13 Aug 2024 19:00:40 +0100 Subject: [PATCH 09/12] generalise assertion about null vector normalisation --- tests/test_unit/test_vector.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index e25b07fe9..265ace168 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -138,13 +138,16 @@ def test_compute_norm(self, ds, expected_exception, request): """Test vector norm computation with known values.""" ds = request.getfixturevalue(ds) with expected_exception: + # validate the norm computation result = vector.compute_norm(ds.cart) expected = np.sqrt( ds.cart.sel(space="x") ** 2 + ds.cart.sel(space="y") ** 2 ) xr.testing.assert_allclose(result, expected) + # result should be the same from Cartesian and polar coordinates xr.testing.assert_allclose(result, vector.compute_norm(ds.pol)) + # The result should only contain the time dimension. assert result.dims == ("time",) @@ -161,24 +164,29 @@ def test_normalize(self, ds, expected_exception, request): ds = request.getfixturevalue(ds) with expected_exception: normalized = vector.normalize(ds.cart) + # the normalized data should have the same dimensions as the input assert normalized.dims == ds.cart.dims - # the first time point is NaN because the input vector is [0, 0] - # (zero-division during normalization). - assert normalized.sel(time=0).isnull().all() + + # normalising null vectors (where x,y = 0,0) should yield NaNs + is_null_vector = (ds.cart == 0).all("space") + assert normalized.where(is_null_vector).isnull().all() + # the magnitude of the normalized vector should be 1 for all # time points except for the expected NaNs. normalized_mag = vector.compute_norm(normalized).values expected_mag = np.ones_like(normalized_mag) expected_mag[normalized.isnull().any("space")] = np.nan np.testing.assert_allclose(normalized_mag, expected_mag) - # Normalising the polar data should yield the same result + + # normalising the polar data should yield the same result # as converting the normalised Cartesian data to polar. xr.testing.assert_allclose( vector.normalize(ds.pol), vector.cart2pol(vector.normalize(ds.cart)), ) - # Normalising the Cartesian data should yield the same result + + # normalising the Cartesian data should yield the same result # as converting the normalised polar data to Cartesian. xr.testing.assert_allclose( vector.normalize(ds.cart), From 3f2318f70130d784b853153e6b41140e94d78d3c Mon Sep 17 00:00:00 2001 From: niksirbi Date: Tue, 13 Aug 2024 19:23:34 +0100 Subject: [PATCH 10/12] renamed `normalize` to `convert_to_unit` --- movement/utils/vector.py | 2 +- tests/test_unit/test_vector.py | 35 ++++++++++++++++------------------ 2 files changed, 17 insertions(+), 20 deletions(-) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index 3aaebf696..40d190a33 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -57,7 +57,7 @@ def compute_norm(data: xr.DataArray) -> xr.DataArray: ) -def normalize(data: xr.DataArray) -> xr.DataArray: +def convert_to_unit(data: xr.DataArray) -> xr.DataArray: """Convert the vectors along the spatial dimension into unit vectors. A unit vector is a vector pointing in the same direction as the original diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index 265ace168..381143fa4 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -159,36 +159,33 @@ def test_compute_norm(self, ds, expected_exception, request): ("cart_pol_dataset_missing_cart_dim", pytest.raises(ValueError)), ], ) - def test_normalize(self, ds, expected_exception, request): - """Test data normalization (division by magnitude).""" + def test_convert_to_unit(self, ds, expected_exception, request): + """Test conversion to unit vectors (normalisation).""" ds = request.getfixturevalue(ds) with expected_exception: - normalized = vector.normalize(ds.cart) + unit_cart = vector.convert_to_unit(ds.cart) # normalise Cartesian + unit_pol = vector.convert_to_unit(ds.pol) # normalise polar - # the normalized data should have the same dimensions as the input - assert normalized.dims == ds.cart.dims + # the normalised data should have the same dimensions as the input + assert unit_cart.dims == ds.cart.dims + assert unit_pol.dims == ds.pol.dims # normalising null vectors (where x,y = 0,0) should yield NaNs is_null_vector = (ds.cart == 0).all("space") - assert normalized.where(is_null_vector).isnull().all() + assert unit_cart.where(is_null_vector).isnull().all() + assert unit_pol.where(is_null_vector).isnull().all() - # the magnitude of the normalized vector should be 1 for all + # the norm of the unit vectors should be 1 for all # time points except for the expected NaNs. - normalized_mag = vector.compute_norm(normalized).values - expected_mag = np.ones_like(normalized_mag) - expected_mag[normalized.isnull().any("space")] = np.nan - np.testing.assert_allclose(normalized_mag, expected_mag) + unit_cart_norm = vector.compute_norm(unit_cart).values + expected_norm = np.ones_like(unit_cart_norm) + expected_norm[unit_cart.isnull().any("space")] = np.nan + np.testing.assert_allclose(unit_cart_norm, expected_norm) # normalising the polar data should yield the same result # as converting the normalised Cartesian data to polar. - xr.testing.assert_allclose( - vector.normalize(ds.pol), - vector.cart2pol(vector.normalize(ds.cart)), - ) + xr.testing.assert_allclose(unit_pol, vector.cart2pol(unit_cart)) # normalising the Cartesian data should yield the same result # as converting the normalised polar data to Cartesian. - xr.testing.assert_allclose( - vector.normalize(ds.cart), - vector.pol2cart(vector.normalize(ds.pol)), - ) + xr.testing.assert_allclose(unit_cart, vector.pol2cart(unit_pol)) From 16a5087aa46e49d1f9ac22ee634a23eb7576897a Mon Sep 17 00:00:00 2001 From: niksirbi Date: Wed, 14 Aug 2024 16:32:19 +0100 Subject: [PATCH 11/12] extend test for conversion to unit vectors --- tests/test_unit/test_vector.py | 43 +++++++++++++++++----------------- 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py index 381143fa4..8787a4682 100644 --- a/tests/test_unit/test_vector.py +++ b/tests/test_unit/test_vector.py @@ -163,29 +163,28 @@ def test_convert_to_unit(self, ds, expected_exception, request): """Test conversion to unit vectors (normalisation).""" ds = request.getfixturevalue(ds) with expected_exception: - unit_cart = vector.convert_to_unit(ds.cart) # normalise Cartesian - unit_pol = vector.convert_to_unit(ds.pol) # normalise polar + # normalise both the Cartesian and the polar data to unit vectors + unit_cart = vector.convert_to_unit(ds.cart) + unit_pol = vector.convert_to_unit(ds.pol) + # they should yield the same result, just in different coordinates + xr.testing.assert_allclose(unit_cart, vector.pol2cart(unit_pol)) + xr.testing.assert_allclose(unit_pol, vector.cart2pol(unit_cart)) + + # since we established that polar vs Cartesian unit vectors are + # equivalent, it's enough to do other assertions on either one # the normalised data should have the same dimensions as the input assert unit_cart.dims == ds.cart.dims - assert unit_pol.dims == ds.pol.dims - - # normalising null vectors (where x,y = 0,0) should yield NaNs - is_null_vector = (ds.cart == 0).all("space") - assert unit_cart.where(is_null_vector).isnull().all() - assert unit_pol.where(is_null_vector).isnull().all() - - # the norm of the unit vectors should be 1 for all - # time points except for the expected NaNs. - unit_cart_norm = vector.compute_norm(unit_cart).values - expected_norm = np.ones_like(unit_cart_norm) - expected_norm[unit_cart.isnull().any("space")] = np.nan - np.testing.assert_allclose(unit_cart_norm, expected_norm) - - # normalising the polar data should yield the same result - # as converting the normalised Cartesian data to polar. - xr.testing.assert_allclose(unit_pol, vector.cart2pol(unit_cart)) - # normalising the Cartesian data should yield the same result - # as converting the normalised polar data to Cartesian. - xr.testing.assert_allclose(unit_cart, vector.pol2cart(unit_pol)) + # unit vector should be NaN if the input vector was null or NaN + is_null_vec = (ds.cart == 0).all("space") # null vec: x=0, y=0 + is_nan_vec = ds.cart.isnull().any("space") # any NaN in x or y + expected_nan_idxs = is_null_vec | is_nan_vec + assert unit_cart.where(expected_nan_idxs).isnull().all() + + # For non-NaN unit vectors in polar coordinates, the rho values + # should be 1 and the phi values should be the same as the input + expected_unit_pol = ds.pol.copy() + expected_unit_pol.loc[{"space_pol": "rho"}] = 1 + expected_unit_pol = expected_unit_pol.where(~expected_nan_idxs) + xr.testing.assert_allclose(unit_pol, expected_unit_pol) From 5d8d42cf1746964dfc99d90d27056f3a4a7de670 Mon Sep 17 00:00:00 2001 From: niksirbi Date: Wed, 14 Aug 2024 16:37:07 +0100 Subject: [PATCH 12/12] refactor erro raising into separate func --- movement/utils/vector.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/movement/utils/vector.py b/movement/utils/vector.py index 40d190a33..0d5d88c83 100644 --- a/movement/utils/vector.py +++ b/movement/utils/vector.py @@ -50,11 +50,7 @@ def compute_norm(data: xr.DataArray) -> xr.DataArray: _validate_dimension_coordinates(data, {"space_pol": ["rho", "phi"]}) return data.sel(space_pol="rho", drop=True) else: - raise log_error( - ValueError, - "Input data must contain either 'space' or 'space_pol' " - "as dimensions.", - ) + _raise_error_for_missing_spatial_dim() def convert_to_unit(data: xr.DataArray) -> xr.DataArray: @@ -93,11 +89,8 @@ def convert_to_unit(data: xr.DataArray) -> xr.DataArray: new_data.sel(space_pol="rho").isnull(), np.nan, 1 ) return new_data - raise log_error( - ValueError, - "Input data must contain either 'space' or 'space_pol' " - "as dimensions.", - ) + else: + _raise_error_for_missing_spatial_dim() def cart2pol(data: xr.DataArray) -> xr.DataArray: @@ -211,3 +204,11 @@ def _validate_dimension_coordinates( ) if error_message: raise log_error(ValueError, error_message) + + +def _raise_error_for_missing_spatial_dim() -> None: + raise log_error( + ValueError, + "Input data array must contain either 'space' or 'space_pol' " + "as dimensions.", + )