compute_forward_vector returns empty DataArray with wrong dims and coords

[niksirbi]

Describe the bug

compute_forward_vector (and also it's alias compute_head_direction_vector) return an empty array with incorrect dimensions and coordinates.

To Reproduce

from movement import sample_data
from movement.kinematics import compute_forward_vector

# Load sample data
ds = sample_data.fetch_dataset("DLC_single-mouse_EPM.predictions.csv")

# Compute the forward vector
forward_vector = compute_forward_vector(
    ds.position,
    left_keypoint="left_ear",
    right_keypoint="right_ear",
    camera_view="top_down",
)

print(forward_vector)

Returns:

<xarray.DataArray 'position' (time: 18485, space: 3, individuals: 0)> Size: 0B

Coordinates:
  * time         (time) float64 148kB 0.0 0.03333 0.06667 ... 616.1 616.1 616.1
  * space        (space) object 24B 'x' 'y' nan
  * individuals  (individuals) <U12 0B

Expected behaviour
There should have been only 2 spatial dimensions 'x' and 'y' (instead of 'x' 'y' nan) and the individual's coordinate (ID) should have been preserved (instead of being lost).

Likely culprit
This line which intends to get rid of the extra spatial coordinate that was introduced during the cross product computation. This was correct to begin with, but we forgot to change the order of dimensions during #236 (both me and @lochhh missed this line). As a consequence, this line now gets rid of the "individuals" coordinate (the animal ID and all data that goes with it) instead of the 3rd spatial coordinate 🤦🏼 .

Proposed fix
The fix is easy, we just have to replace )[:, :, :-1] with [:, -1, :] in the above quoted line.
The more involved part is figuring out why the relevant unit tests missed this bug and adjusting them to catch such errors in future.

[willGraham01]

The suggested fix doesn't appear to generate the expected result; we still get a NaN spatial dimension in the output (and a 3-spatial dimension output):

>>> print(input)
<xarray.DataArray (time: 4, individuals: 1, keypoints: 3, space: 2)> Size: 192B
1.0 0.0 -1.0 0.0 0.0 -1.0 nan nan 0.0 ... 0.0 0.0 1.0 0.0 -1.0 0.0 1.0 -1.0 0.0
Coordinates:
  * time         (time) int64 32B 0 1 2 3
  * individuals  (individuals) <U12 48B 'individual_0'
  * keypoints    (keypoints) <U9 108B 'left_ear' 'right_ear' 'nose'
  * space        (space) <U1 8B 'x' 'y'
>>> forward_vector = compute_forward_vector(input, 'left_ear', 'right_ear')
>>> print(forward_vector)
<xarray.DataArray (time: 4, space: 3)> Size: 96B
-0.0 1.0 0.0 nan nan nan -0.0 -1.0 -0.0 1.0 0.0 0.0
Coordinates:
  * time         (time) int64 32B 0 1 2 3
    individuals  <U12 48B 'individual_0'
  * space        (space) object 24B 'x' 'y' nan

though we are preserving the individuals axis now. In fact, the individuals axis appears to be preserved even if we remove the [:, -1, :] portion of the code on the suggested line. I think the issue might be that we're just not dropping what we expect to drop? Replacing these lines with a more explicit construction and drop seems to do the trick;

upward_vector = xr.DataArray(
    np.tile(upward_vector.reshape(1, -1), [len(data.time), 1]),
    dims=["time", "space"],
    coords={
        "space": ["x", "y", "z"],
    },
)
# Compute forward direction as the cross product
# (right-to-left) cross (forward) = up
forward_vector = xr.cross(
    right_to_left_vector, upward_vector, dim="space"
).drop_sel(space="z")  # keep only the first 2 spatial dimensions of the result

Then for the same input as above,

>>> forward_vector = compute_forward_vector(input, 'left_ear', 'right_ear')
>>> print(forward_vector)
<xarray.DataArray (time: 4, individuals: 1, space: 2)> Size: 64B
-0.0 1.0 nan nan -0.0 -1.0 1.0 0.0
Coordinates:
  * space        (space) <U1 8B 'x' 'y'
  * time         (time) int64 32B 0 1 2 3
  * individuals  (individuals) <U12 48B 'individual_0'

[niksirbi]

upward_vector = xr.DataArray(
    np.tile(upward_vector.reshape(1, -1), [len(data.time), 1]),
    dims=["time", "space"],
    coords={
        "space": ["x", "y", "z"],
    },
)
# Compute forward direction as the cross product
# (right-to-left) cross (forward) = up
forward_vector = xr.cross(
    right_to_left_vector, upward_vector, dim="space"
).drop_sel(space="z")  # keep only the first 2 spatial dimensions of the result

I like this solution. Using drop_sel has the additional advantage of being robust to future reorderings of dimensions.
I'm happy for this solution to be implemented and tested.

[niksirbi]

And sorry, my suggested fix had a typo, it should have been [:, :-1, :], but regardless your alternative fix is better for the reason mentioned above.