Plotting wrappers: Head Trajectory (#394)

* add plot module with trajectory function

* update example

* add plot kwargs and defaults

* add tests for trajectory plotting

* parametrize test_trajectory

Co-authored-by: Will Graham <[email protected]>

* add docstring to trajectory, remove test code

Co-authored-by: Will Graham <[email protected]>

* improve code coverage, make colorbar alpha resistant

Co-authored-by: Will Graham <[email protected]>

* Add Niko's plot trajectory suggestion, fix tests

* update example

* adjust example, make lines causing '\d' SyntaxWarning raw

* change trajectory inputs, replace keypoint and individual with selection, remove title

* change trajectory inputs, replace keypoint and individual with selection, remove title

* use selection dict for individuals and keypoints

* update trajectory examples, allow user to set marker colour

* fix examples, fix colorbar

* improve code coverage

* remove image_path from plot input, adjust example

* fix test after removing image_path input

* add more tests

* test trajectory without individuals and/or keypoints dimension

* fix logic test_trajectory_dropped_dim

* add tests, update examples, change selection to individuals and keypoints

Co-authored-by: Niko Sirmpilatze <[email protected]>

* change trajectory input in examples as well

* fix input load and explore poses example

* change folder structure

* add init file

Co-authored-by: Will Graham <[email protected]>

* fix test

* deal with drop deprecation (-> drop_vars)

* process review suggestions, adjust  drop dimension test

* reorder plot_trajectory inputs

* fix typo example

---------

Co-authored-by: Will Graham <[email protected]>
Co-authored-by: Niko Sirmpilatze <[email protected]>

Author: 3 people

examples/compute_kinematics.py

@@ -16,6 +16,7 @@
 from matplotlib import pyplot as plt
 
 from movement import sample_data
+from movement.plots import plot_trajectory
 from movement.utils.vector import compute_norm
 
 # %%
@@ -48,27 +49,36 @@
 # Visualise the data
 # ---------------------------
 # First, let's visualise the trajectories of the mice in the XY plane,
-# colouring them by individual.
+# colouring them by individual. We use the ``plot_trajectory`` function from
+# ``movement.plots`` which is a wrapper around
+# ``matplotlib.pyplot.scatter`` that simplifies plotting the trajectories of
+# individuals in the dataset. The fig and ax objects returned can be used to
+# further customise the plot.
 
+# Create a single figure and axes
 fig, ax = plt.subplots(1, 1)
+# Invert y-axis so (0,0) is in the top-left,
+# matching typical image coordinate systems
+ax.invert_yaxis()
+# Plot trajectories for each mouse on the same axes
 for mouse_name, col in zip(
-    position.individuals.values, ["r", "g", "b"], strict=False
+    position.individuals.values,
+    ["r", "g", "b"],  # colours
+    strict=False,
 ):
-    ax.plot(
-        position.sel(individuals=mouse_name, space="x"),
-        position.sel(individuals=mouse_name, space="y"),
-        linestyle="-",
-        marker=".",
-        markersize=2,
-        linewidth=0.5,
+    plot_trajectory(
+        position,
+        individual=mouse_name,
+        ax=ax,  # Use the same axes for all plots
         c=col,
+        marker="o",
+        s=10,
+        alpha=0.2,
         label=mouse_name,
     )
-    ax.invert_yaxis()
-    ax.set_xlabel("x (pixels)")
-    ax.set_ylabel("y (pixels)")
-    ax.axis("equal")
-    ax.legend()
+    ax.legend().set_alpha(1)
+ax.title.set_text("Trajectories of three mice")
+fig.show()
 
 # %%
 # We can see that the trajectories of the three mice are close to a circular
@@ -77,23 +87,19 @@
 # follows the convention for SLEAP and most image processing tools.
 
 # %%
-# We can also color the data points based on their timestamps:
+# By default the ``plot_trajectory`` function in ``movement.plots`` colours
+# data points based on their timestamps:
 fig, axes = plt.subplots(3, 1, sharey=True)
 for mouse_name, ax in zip(position.individuals.values, axes, strict=False):
-    sc = ax.scatter(
-        position.sel(individuals=mouse_name, space="x"),
-        position.sel(individuals=mouse_name, space="y"),
+    ax.invert_yaxis()
+    fig, ax = plot_trajectory(
+        position,
+        individual=mouse_name,
+        ax=ax,
         s=2,
-        c=position.time,
-        cmap="viridis",
     )
-    ax.invert_yaxis()
-    ax.set_title(mouse_name)
-    ax.set_xlabel("x (pixels)")
-    ax.set_ylabel("y (pixels)")
-    ax.axis("equal")
-    fig.colorbar(sc, ax=ax, label="time (s)")
 fig.tight_layout()
+fig.show()
 
 # %%
 # These plots show that for this snippet of the data,

examples/compute_polar_coordinates.py

@@ -16,6 +16,7 @@
 
 from movement import sample_data
 from movement.io import load_poses
+from movement.plots import plot_trajectory
 from movement.utils.vector import cart2pol, pol2cart
 
 # %%
@@ -69,65 +70,51 @@
 # We can plot the data to check that our computation of the head vector is
 # correct.
 #
-# We can start by plotting the trajectory of the midpoint between the ears. We
-# will refer to this as the head trajectory.
-
-fig, ax = plt.subplots(1, 1)
-mouse_name = ds.individuals.values[0]
-
-sc = ax.scatter(
-    midpoint_ears.sel(individuals=mouse_name, space="x"),
-    midpoint_ears.sel(individuals=mouse_name, space="y"),
-    s=15,
-    c=midpoint_ears.time,
-    cmap="viridis",
-    marker="o",
-)
-
-ax.axis("equal")
-ax.set_xlabel("x (pixels)")
-ax.set_ylabel("y (pixels)")
+# We can start by plotting the head trajectory with the ``plot_trajectory``
+# from ``movement.plots`` which creates a plot of the centroid of
+# the selected keypoints, for the head trajectory, we will use the midpoint
+# between the ears. By default, the trajectory of the first listed individual
+# is shown.
+
+fig, ax = plot_trajectory(position, keypoints=["left_ear", "right_ear"])
+# Invert y-axis so (0,0) is in the top-left,
+# matching typical image coordinate systems
 ax.invert_yaxis()
-ax.set_title(f"Head trajectory ({mouse_name})")
-fig.colorbar(sc, ax=ax, label=f"time ({ds.attrs['time_unit']})")
 fig.show()
 
+
 # %%
+# Overlay trajectory on Elevated Plus Maze
+# ----------------------------------------
 # We can see that the majority of the head trajectory data is within a
 # cruciform shape. This is because the dataset is of a mouse moving on an
 # `Elevated Plus Maze <https://en.wikipedia.org/wiki/Elevated_plus_maze>`_.
 # We can actually verify this is the case by overlaying the head
 # trajectory on the sample frame of the dataset.
 
-# read sample frame
+# Read sample frame
 frame_path = sample_data.fetch_dataset_paths(
     "SLEAP_single-mouse_EPM.analysis.h5"
 )["frame"]
-im = plt.imread(frame_path)
 
-
-# plot sample frame
+# Create figure and axis
 fig, ax = plt.subplots(1, 1)
-ax.imshow(im)
-
-# plot head trajectory with semi-transparent markers
-sc = ax.scatter(
-    midpoint_ears.sel(individuals=mouse_name, space="x"),
-    midpoint_ears.sel(individuals=mouse_name, space="y"),
-    s=15,
-    c=midpoint_ears.time,
+# Plot the frame using imshow
+ax.imshow(plt.imread(frame_path))
+# No need to invert the y-axis now, since the image is plotted
+# using a pixel coordinate system with origin on the top left of the image
+fig, ax = plot_trajectory(
+    ds.position,
+    individual="individual_0",
+    keypoints=["left_ear", "right_ear"],
+    ax=ax,
+    s=10,
     cmap="viridis",
     marker="o",
-    alpha=0.05,  # transparency
+    alpha=0.05,
 )
-
-ax.axis("equal")
-ax.set_xlabel("x (pixels)")
-ax.set_ylabel("y (pixels)")
-# No need to invert the y-axis now, since the image is plotted
-# using a pixel coordinate system with origin on the top left of the image
-ax.set_title(f"Head trajectory ({mouse_name})")
-
+# Adjust title
+ax.set_title("Head trajectory (individual_0)")
 fig.show()
 
 # %%

examples/load_and_explore_poses.py

@@ -7,10 +7,10 @@
 # %%
 # Imports
 # -------
-from matplotlib import pyplot as plt
 
 from movement import sample_data
 from movement.io import load_poses
+from movement.plots import plot_trajectory
 
 # %%
 # Define the file path
@@ -70,19 +70,10 @@
 # Trajectory plots
 # ----------------
 # We are not limited to ``xarray``'s built-in plots.
-# For example, we can use ``matplotlib`` to plot trajectories
-# (using scatter plots):
+# The ``movement.plots`` module provides some additional
+# visualisations, like ``plot_trajectory()``.
 
-mouse_name = "AEON3B_TP1"
 
-plt.scatter(
-    da.sel(individuals=mouse_name, space="x"),
-    da.sel(individuals=mouse_name, space="y"),
-    s=2,
-    c=da.time,
-    cmap="viridis",
-)
-plt.title(f"Trajectory of {mouse_name}")
-plt.xlabel("x")
-plt.ylabel("y")
-plt.colorbar(label="time (sec)")
+mouse_name = "AEON3B_TP1"
+fig, ax = plot_trajectory(position, individual=mouse_name)
+fig.show()

movement/plots/__init__.py

@@ -0,0 +1,3 @@
+from movement.plots.trajectory import plot_trajectory
+
+__all__ = ["plot_trajectory"]

movement/plots/trajectory.py

@@ -0,0 +1,114 @@
+"""Wrappers to plot movement data."""
+
+import xarray as xr
+from matplotlib import pyplot as plt
+
+DEFAULT_PLOTTING_ARGS = {
+    "s": 15,
+    "marker": "o",
+    "alpha": 1.0,
+}
+
+
+def plot_trajectory(
+    da: xr.DataArray,
+    individual: str | None = None,
+    keypoints: str | list[str] | None = None,
+    ax: plt.Axes | None = None,
+    **kwargs,
+) -> tuple[plt.Figure, plt.Axes]:
+    """Plot trajectory.
+
+    This function plots the trajectory of a specified keypoint or the centroid
+    of multiple keypoints for a given individual. By default, the first
+    is colored by time (using the default colormap). Pass a different colormap
+    through ``cmap`` if desired.
+
+    Parameters
+    ----------
+    da : xr.DataArray
+        A data array containing position information, with `time` and `space`
+        as required dimensions. Optionally, it may have `individuals` and/or
+        `keypoints` dimensions.
+    individual : str, optional
+        The name of the individual to be plotted. By default, the first
+        individual is plotted.
+    keypoints : str, list[str], optional
+        The name of the keypoint to be plotted, or a list of keypoint names
+        (their centroid will be plotted). By default, the centroid of all
+        keypoints is plotted.
+    ax : matplotlib.axes.Axes or None, optional
+        Axes object on which to draw the trajectory. If None, a new
+        figure and axes are created.
+    **kwargs : dict
+        Additional keyword arguments passed to
+        ``matplotlib.axes.Axes.scatter()``.
+
+    Returns
+    -------
+    (figure, axes) : tuple of (matplotlib.pyplot.Figure, matplotlib.axes.Axes)
+        The figure and axes containing the trajectory plot.
+
+    """
+    if isinstance(individual, list):
+        raise ValueError("Only one individual can be selected.")
+
+    selection = {}
+
+    if "individuals" in da.dims:
+        if individual is None:
+            selection["individuals"] = da.individuals.values[0]
+        else:
+            selection["individuals"] = individual
+
+    title_suffix = f" of {individual}" if "individuals" in da.dims else ""
+
+    if "keypoints" in da.dims:
+        if keypoints is None:
+            selection["keypoints"] = da.keypoints.values
+        else:
+            selection["keypoints"] = keypoints
+
+    plot_point = da.sel(**selection)
+
+    # If there are multiple selected keypoints, calculate the centroid
+    plot_point = (
+        plot_point.mean(dim="keypoints", skipna=True)
+        if "keypoints" in plot_point.dims and plot_point.sizes["keypoints"] > 1
+        else plot_point
+    )
+
+    plot_point = plot_point.squeeze()  # Only space and time should remain
+
+    fig, ax = plt.subplots(figsize=(6, 6)) if ax is None else (ax.figure, ax)
+
+    # Merge default plotting args with user-provided kwargs
+    for key, value in DEFAULT_PLOTTING_ARGS.items():
+        kwargs.setdefault(key, value)
+
+    colorbar = False
+    if "c" not in kwargs:
+        kwargs["c"] = plot_point.time
+        colorbar = True
+
+    # Plot the scatter, colouring by time or user-provided colour
+    sc = ax.scatter(
+        plot_point.sel(space="x"),
+        plot_point.sel(space="y"),
+        **kwargs,
+    )
+
+    space_unit = da.attrs.get("space_unit", "pixels")
+    ax.set_xlabel(f"x ({space_unit})")
+    ax.set_ylabel(f"y ({space_unit})")
+    ax.axis("equal")
+    ax.set_title(f"Trajectory{title_suffix}")
+
+    # Add 'colorbar' for time dimension if no colour was provided by user
+    time_unit = da.attrs.get("time_unit")
+    time_label = f"time ({time_unit})" if time_unit else "time steps (frames)"
+    fig.colorbar(sc, ax=ax, label=time_label).solids.set(
+        alpha=1.0
+    ) if colorbar else None
+
+    return fig, ax

tests/test_unit/test_load_bboxes.py

@@ -295,7 +295,7 @@ def test_from_via_tracks_file(
             r"_(0\d*)_$",
             AttributeError,
             "/crab_1/00000.jpg (row 0): "
-            "The provided frame regexp (_(0\d*)_$) did not return any "
+            r"The provided frame regexp (_(0\d*)_$) did not return any "
             "matches and a frame number could not be extracted from "
             "the filename.",
         ),
@@ -304,7 +304,7 @@ def test_from_via_tracks_file(
             ValueError,
             "/crab_1/00000.jpg (row 0): "
             "The frame number extracted from the filename "
-            "using the provided regexp ((0\d*\.\w+)$) "
+            r"using the provided regexp ((0\d*\.\w+)$) "
             "could not be cast as an integer.",
         ),
     ],