Repository: neuroinformatics-unit/movement

Issue:
Add norm function to the vector utils
**Is your feature request related to a problem? Please describe.**
We often need to get the magnitude of a vector.

**Describe the solution you'd like**
A `norm()` utility function that would function similar to existing `cart2pol` and `pol2cart` utilites.

This can facilitate #147 


Discussion:
sonarcloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarCloud](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243)


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/neuroinformatics-unit/movement/pull/243?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.74%. Comparing base [(`242f532`)](https://app.codecov.io/gh/neuroinformatics-unit/movement/commit/242f5328c11f21ae605169a8dc780b6d855ff478?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit) to head [(`3c078a3`)](https://app.codecov.io/gh/neuroinformatics-unit/movement/commit/3c078a395b70349ae922692a998523326d31015d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main     #243   +/-   ##
=======================================
  Coverage   99.74%   99.74%           
=======================================
  Files          13       13           
  Lines         771      782   +11     
=======================================
+ Hits          769      780   +11     
  Misses          2        2           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/neuroinformatics-unit/movement/pull/243?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=neuroinformatics-unit).

sonarcloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarCloud](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243)


niksirbi: @sfmig I'll give you a chance to comment on this before merging, because we've discussed this together. No rush
sonarcloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarCloud](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243)


sonarcloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarCloud](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243)


sonarcloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarCloud](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243)


niksirbi: Thanks for the thorough review @sfmig. I believe I've addressed all your points by basically adopting all your suggestions (see edited part of the PR description). Re-requesting your review.
sonarcloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=neuroinformatics-unit_movement&pullRequest=243&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=neuroinformatics-unit_movement&pullRequest=243&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarCloud](https://sonarcloud.io/dashboard?id=neuroinformatics-unit_movement&pullRequest=243)



Patch:
diff --git a/examples/compute_kinematics.py b/examples/compute_kinematics.py
index 1ab2731..b3fefd4 100644
--- a/examples/compute_kinematics.py
+++ b/examples/compute_kinematics.py
@@ -10,14 +10,13 @@ visualise the results.
 # 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 compute_norm
 
 # %%
 # Load sample dataset
@@ -255,13 +254,12 @@ fig.colorbar(sc, ax=ax, label="time (s)")
 # 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 = compute_norm(
+    displacement.sel(individuals=mouse_name)
 )
 
-# sum of all displacement vectors
-total_displacement = np.sum(displacement_vectors_lengths, axis=0)  # in pixels
+# sum the lengths of all displacement vectors (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 @@ plt.gcf().show()
 # 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 = compute_norm(velocity.sel(individuals=mouse_name))
     # plot speed against time
     ax.plot(speed_one_mouse)
     ax.set_title(mouse_name)
@@ -379,16 +375,12 @@ for mouse_name, ax in zip(accel.individuals.values, axes, strict=False):
 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 = 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 c35990e..0d5d88c 100644
--- a/movement/utils/vector.py
+++ b/movement/utils/vector.py
@@ -6,6 +6,93 @@ import xarray as xr
 from movement.utils.logging import log_error
 
 
+def compute_norm(data: xr.DataArray) -> xr.DataArray:
+    """Compute the norm of the vectors along the spatial dimension.
+
+    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 array containing either ``space`` or ``space_pol``
+        as a dimension.
+
+    Returns
+    -------
+    xarray.DataArray
+         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.
+
+
+    """
+    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_error_for_missing_spatial_dim()
+
+
+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
+    vector but with norm = 1.
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+        The input data array containing either ``space`` or ``space_pol``
+        as a dimension.
+
+    Returns
+    -------
+    xarray.DataArray
+        A data array holding the unit vectors of the input data array
+        (all input dimensions are preserved).
+
+    Notes
+    -----
+    Note that the unit vector for the null vector is undefined, since the null
+    vector has 0 norm and no direction associated with it.
+
+    """
+    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
+    else:
+        _raise_error_for_missing_spatial_dim()
+
+
 def cart2pol(data: xr.DataArray) -> xr.DataArray:
     """Transform Cartesian coordinates to polar.
 
@@ -25,12 +112,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 = compute_norm(data)
     phi = xr.apply_ufunc(
         np.arctan2,
         data.sel(space="y"),
@@ -122,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.",
+    )


Test patch:
diff --git a/tests/test_unit/test_vector.py b/tests/test_unit/test_vector.py
index 88dd85b..8787a46 100644
--- a/tests/test_unit/test_vector.py
+++ b/tests/test_unit/test_vector.py
@@ -121,3 +121,70 @@ class TestVector:
         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_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",)
+
+    @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_convert_to_unit(self, ds, expected_exception, request):
+        """Test conversion to unit vectors (normalisation)."""
+        ds = request.getfixturevalue(ds)
+        with expected_exception:
+            # 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
+
+            # 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)