diff --git a/movement/analysis/navigation.py b/movement/analysis/navigation.py index 89bb10bea..022d713f0 100644 --- a/movement/analysis/navigation.py +++ b/movement/analysis/navigation.py @@ -50,6 +50,12 @@ def compute_head_direction_vector( raise log_error( ValueError, "The left and right keypoints may not be identical." ) + if len(data.space) != 2: + raise log_error( + ValueError, + "Input data must have 2 (and only 2) spatial dimensions, but " + f"currently has {len(data.space)}.", + ) # Select the right and left keypoints head_left = data.sel(keypoints=left_keypoint, drop=True) diff --git a/tests/test_unit/test_navigation.py b/tests/test_unit/test_navigation.py index 2d84a7076..a8a90e028 100644 --- a/tests/test_unit/test_navigation.py +++ b/tests/test_unit/test_navigation.py @@ -36,7 +36,7 @@ def test_compute_head_direction_vector(mock_dataset): are computed from a basic mock dataset. """ # Test that validators work - with pytest.raises(TypeError): + with pytest.raises(TypeError): # Catch incorrect datatype np_array = [ [[[1, 0], [-1, 0], [0, -1]]], [[[0, 1], [0, -1], [1, 0]]], @@ -47,17 +47,38 @@ def test_compute_head_direction_vector(mock_dataset): np_array, "left_ear", "right_ear" ) - with pytest.raises(AttributeError): + with pytest.raises(AttributeError): # Catch incorrect dimensions mock_dataset_keypoint = mock_dataset.sel(keypoints="nose", drop=True) navigation.compute_head_direction_vector( mock_dataset_keypoint, "left_ear", "right_ear" ) - with pytest.raises(ValueError): + with pytest.raises(ValueError): # Catch identical left and right keypoints navigation.compute_head_direction_vector( mock_dataset, "left_ear", "left_ear" ) + with pytest.raises(ValueError): # Catch incorrect spatial dimensions + time = np.array([0]) + individuals = np.array(["individual_0"]) + keypoints = np.array(["left", "right"]) + space = np.array(["x", "y", "z"]) + + ds = xr.DataArray( + [ + [[[1, 0, 0], [-1, 0, 0]]], + ], + dims=["time", "individuals", "keypoints", "space"], + coords={ + "time": time, + "individuals": individuals, + "keypoints": keypoints, + "space": space, + }, + ) + + navigation.compute_head_direction_vector(ds, "left", "right") + # Test that output contains correct datatype, dimensions, and values head_vector = navigation.compute_head_direction_vector( mock_dataset, "left_ear", "right_ear"