parametrize test_vector, refactor sample_data fixture

Author: stellaprins

movement/plot.py

@@ -174,54 +174,3 @@ def vector(
     )
 
     return fig
-
-
-# # # FOR TESTING
-# from movement import sample_data
-# from movement.io import load_poses
-
-# ds_path = sample_data.fetch_dataset_paths(
-#     "SLEAP_single-mouse_EPM.analysis.h5"
-# )["poses"]
-# ds = load_poses.from_sleap_file(ds_path, fps=None)
-# # force time_unit = frames
-# frame_path = sample_data.fetch_dataset_paths(
-#     "SLEAP_single-mouse_EPM.analysis.h5"
-# )["frame"]
-
-# # head_trajectory = vector(
-# #     ds,
-# #     reference_points=("left_ear", "right_ear"),
-# #     vector_point="snout",
-# #     time_points=None,
-# #     x_lim=None,
-# #     y_lim=None,
-# #     individual=0,
-# # )
-
-# # plt.ion()
-# # head_trajectory.show()
-# # # user input to close window
-# # input("Press Enter to continue...")
-
-
-# # area of interest
-# xmin, ymin = 600.0, 665.0  # pixels
-# x_delta, y_delta = 125.0, 100.0  # pixels
-
-# # time window
-# time_window = tuple(range(1650, 1671))  # tuple of frame numbers
-
-# fig_head_vector = vector(
-#     ds,
-#     reference_points=["left_ear", "right_ear"],
-#     vector_point="snout",
-#     individual="individual_0",
-#     x_lim=(xmin, xmin + x_delta),
-#     y_lim=(ymin, ymin + y_delta),
-#     time_points=time_window,
-#     title="Zoomed in head vector (individual_0)",
-# )
-
-# fig_head_vector.show()
-# input("Press Enter to continue...")

tests/test_unit/test_plot.py

@@ -1,35 +1,23 @@
+"""Unit tests for the plot module."""
+
 import numpy as np
 import pytest
 import xarray as xr
 
 from movement.plot import vector
 
 
-@pytest.fixture
-def sample_data():
-    """Sample data for plot testing.
-
-    Data has three keypoints (left, centre, right) for one
-    individual that moves in a straight line along the y-axis with a
-    constant x-coordinate.
-
-    """
+def create_sample_data(keypoints, positions):
+    """Create sample data for testing."""
     time_steps = 4
     individuals = ["individual_0"]
-    keypoints = ["left", "centre", "right"]
     space = ["x", "y"]
-    positions = {
-        "left": {"x": -1, "y": np.arange(time_steps)},
-        "centre": {"x": 0, "y": np.arange(time_steps)},
-        "right": {"x": 1, "y": np.arange(time_steps)},
-    }
 
     time = np.arange(time_steps)
     position_data = np.zeros(
         (time_steps, len(space), len(keypoints), len(individuals))
     )
 
-    # Create x and y coordinates arrays
     x_coords = np.array([positions[key]["x"] for key in keypoints])
     y_coords = np.array([positions[key]["y"] for key in keypoints])
 
@@ -63,194 +51,86 @@ def sample_data():
 
 
 @pytest.fixture
-def sample_data_quiver1():
+def sample_data():
     """Sample data for plot testing.
 
-    Data has three keypoints (left, centre, right) for one
-    individual that moves in a straight line along the y-axis with a
-    constant x-coordinate.
+    Data has six keypoints for one individual that moves in a straight line
+    along the y-axis. All keypoints have a constant x-coordinate and move in
+    steps of 1 along the y-axis.
 
-    """
-    time_steps = 4
-    individuals = ["individual_0"]
-    keypoints = ["left", "centre", "right"]
-    space = ["x", "y"]
-    positions = {
-        "left": {"x": -1, "y": np.arange(time_steps)},
-        "centre": {"x": 0, "y": np.arange(time_steps) + 1},
-        "right": {"x": 1, "y": np.arange(time_steps)},
-    }
-
-    time = np.arange(time_steps)
-    position_data = np.zeros(
-        (time_steps, len(space), len(keypoints), len(individuals))
-    )
-
-    # Create x and y coordinates arrays
-    x_coords = np.array([positions[key]["x"] for key in keypoints])
-    y_coords = np.array([positions[key]["y"] for key in keypoints])
-
-    for i, _ in enumerate(keypoints):
-        position_data[:, 0, i, 0] = x_coords[i]  # x-coordinates
-        position_data[:, 1, i, 0] = y_coords[i]  # y-coordinates
-
-    confidence_data = np.full(
-        (time_steps, len(keypoints), len(individuals)), 0.90
-    )
-
-    ds = xr.Dataset(
-        {
-            "position": (
-                ["time", "space", "keypoints", "individuals"],
-                position_data,
-            ),
-            "confidence": (
-                ["time", "keypoints", "individuals"],
-                confidence_data,
-            ),
-        },
-        coords={
-            "time": time,
-            "space": space,
-            "keypoints": keypoints,
-            "individuals": individuals,
-        },
-    )
-    return ds
-
-
-@pytest.fixture
-def sample_data_quiver2():
-    """Sample data for plot testing.
-
-    Data has three keypoints (left, centre, right) for one
-    individual that moves in a straight line along the y-axis with a
-    constant x-coordinate.
+    Keypoint starting positions:
+    - left1: (-1, 0)
+    - right1: (1, 0)
+    - left2: (-2, 0)
+    - right2: (2, 0)
+    - centre0: (0, 0)
+    - centre1: (0, 1)
 
     """
-    time_steps = 4
-    individuals = ["individual_0"]
-    keypoints = ["left1", "right1", "left2", "right2"]
-    space = ["x", "y"]
+    keypoints = ["left1", "right1", "left2", "right2", "centre0", "centre1"]
     positions = {
-        "left1": {"x": -1, "y": np.arange(time_steps)},
-        "right1": {"x": 1, "y": np.arange(time_steps)},
-        "left2": {"x": -1, "y": np.arange(time_steps)},
-        "right2": {"x": 1, "y": np.arange(time_steps)},
+        "left1": {"x": -1, "y": np.arange(4)},
+        "right1": {"x": 1, "y": np.arange(4)},
+        "left2": {"x": -2, "y": np.arange(4)},
+        "right2": {"x": 2, "y": np.arange(4)},
+        "centre0": {"x": 0, "y": np.arange(4)},
+        "centre1": {"x": 0, "y": np.arange(4) + 1},
     }
-
-    time = np.arange(time_steps)
-    position_data = np.zeros(
-        (time_steps, len(space), len(keypoints), len(individuals))
-    )
-
-    # Create x and y coordinates arrays
-    x_coords = np.array([positions[key]["x"] for key in keypoints])
-    y_coords = np.array([positions[key]["y"] for key in keypoints])
-
-    for i, _ in enumerate(keypoints):
-        position_data[:, 0, i, 0] = x_coords[i]  # x-coordinates
-        position_data[:, 1, i, 0] = y_coords[i]  # y-coordinates
-
-    confidence_data = np.full(
-        (time_steps, len(keypoints), len(individuals)), 0.90
-    )
-
-    ds = xr.Dataset(
-        {
-            "position": (
-                ["time", "space", "keypoints", "individuals"],
-                position_data,
-            ),
-            "confidence": (
-                ["time", "keypoints", "individuals"],
-                confidence_data,
-            ),
-        },
-        coords={
-            "time": time,
-            "space": space,
-            "keypoints": keypoints,
-            "individuals": individuals,
-        },
-    )
-    return ds
-
-
-def test_vector_no_quiver(sample_data):
-    """Test midpoint between left and right keypoints."""
+    return create_sample_data(keypoints, positions)
+
+
+@pytest.mark.parametrize(
+    ["vector_point", "expected_u", "expected_v"],
+    [
+        pytest.param(
+            "centre0",
+            [0.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0],
+            id="u = 0, v = 0",
+        ),
+        pytest.param(
+            "centre1",
+            [0.0, 0.0, 0.0, 0.0],
+            [1.0, 1.0, 1.0, 1.0],
+            id="u = 0, v = 1",
+        ),
+        pytest.param(
+            "right2",
+            [2.0, 2.0, 2.0, 2.0],
+            [0.0, 0.0, 0.0, 0.0],
+            id="u = 2, v = 0",
+        ),
+        pytest.param(
+            "left2",
+            [-2.0, -2.0, -2.0, -2.0],
+            [0.0, 0.0, 0.0, 0.0],
+            id="u = -2, v = 0",
+        ),
+    ],
+)
+def test_vector(sample_data, vector_point, expected_u, expected_v):
+    """Test vector plot.
+
+    Test the vector plot for different vector points. The U and V values
+    represent the horizontal (x) and vertical (y) displacement of the vector
+
+    The reference points are "left1" and "right1".
+    """
     vector_fig = vector(
         sample_data,
-        reference_points=["left", "right"],
-        vector_point="centre",
-    )
-
-    quiver = vector_fig.axes[0].collections[-1]
-
-    # Extract the X, Y, U, V data
-    x = quiver.X
-    y = quiver.Y
-    u = quiver.U
-    v = quiver.V
-
-    expected_x = np.array([0.0, 0.0, 0.0, 0.0])
-    expected_y = np.array([0.0, 1.0, 2.0, 3.0])
-    expected_u = np.array([0.0, 0.0, 0.0, 0.0])
-    expected_v = np.array([0.0, 0.0, 0.0, 0.0])
-
-    assert np.allclose(x, expected_x)
-    assert np.allclose(y, expected_y)
-    assert np.allclose(u, expected_u)
-    assert np.allclose(v, expected_v)
-
-
-def test_vector_quiver2(sample_data_quiver2):
-    """Test midpoint between left and right keypoints."""
-    vector_fig = vector(
-        sample_data_quiver2,
         reference_points=["left1", "right1"],
-        vector_point="right2",
-    )
-
-    quiver = vector_fig.axes[0].collections[-1]
-
-    # Extract the X, Y, U, V data
-    x = quiver.X
-    y = quiver.Y
-    u = quiver.U
-    v = quiver.V
-
-    expected_x = np.array([0.0, 0.0, 0.0, 0.0])
-    expected_y = np.array([0.0, 1.0, 2.0, 3.0])
-    expected_u = np.array([1.0, 1.0, 1.0, 1.0])
-    expected_v = np.array([0.0, 0.0, 0.0, 0.0])
-
-    assert np.allclose(x, expected_x)
-    assert np.allclose(y, expected_y)
-    assert np.allclose(u, expected_u)
-    assert np.allclose(v, expected_v)
-
-
-def test_vector_quiver1(sample_data_quiver1):
-    """Test midpoint between left and right keypoints."""
-    vector_fig = vector(
-        sample_data_quiver1,
-        reference_points=["left", "right"],
-        vector_point="centre",
+        vector_point=vector_point,
     )
 
     quiver = vector_fig.axes[0].collections[-1]
 
-    # Extract the X, Y, U, V data
     x = quiver.X
     y = quiver.Y
     u = quiver.U
     v = quiver.V
 
     expected_x = np.array([0.0, 0.0, 0.0, 0.0])
     expected_y = np.array([0.0, 1.0, 2.0, 3.0])
-    expected_u = np.array([0.0, 0.0, 0.0, 0.0])
-    expected_v = np.array([1.0, 1.0, 1.0, 1.0])
 
     assert np.allclose(x, expected_x)
     assert np.allclose(y, expected_y)