Check that new / existing axes are respected

Author: willGraham01

movement/plot.py

@@ -16,8 +16,9 @@ def occupancy_histogram(
     keypoint: int | str = 0,
     individual: int | str = 0,
     title: str | None = None,
+    ax: plt.Axes | None = None,
     **kwargs: Any,
-) -> tuple[plt.Figure, dict[HistInfoKeys, np.ndarray]]:
+) -> tuple[plt.Figure, plt.Axes, dict[HistInfoKeys, np.ndarray]]:
     """Create a 2D histogram of the occupancy data given.
 
     Time-points whose corresponding spatial coordinates have NaN values
@@ -35,13 +36,20 @@ def occupancy_histogram(
     title : str, optional
         Title to give to the plot. Default will be generated from the
         ``keypoint`` and ``individual``
+    ax : matplotlib.axes.Axes, optional
+        Axes object on which to draw the histogram. If not provided, a new
+        figure and axes are created and returned.
     kwargs : Any
         Keyword arguments passed to ``matplotlib.pyplot.hist2d``
 
     Returns
     -------
     matplotlib.pyplot.Figure
-        Plot handle containing the rendered 2D histogram.
+        Plot handle containing the rendered 2D histogram. If ``ax`` is
+        supplied, this will be the figure that ``ax`` belongs to.
+    matplotlib.axes.Axes
+        Axes on which the histogram was drawn. If ``ax`` was supplied,
+        the input will be directly modified and returned in this value.
     dict[str, numpy.ndarray]
         Information about the created histogram (see Notes).
 
@@ -99,7 +107,10 @@ def occupancy_histogram(
         if key not in kwargs:
             kwargs[key] = value
     # Now it should just be a case of creating the histogram
-    fig, ax = plt.subplots()
+    if ax is not None:
+        fig = ax.get_figure()
+    else:
+        fig, ax = plt.subplots()
     counts, xedges, yedges, hist_image = ax.hist2d(
         data.sel(space=x_coord), data.sel(space=y_coord), **kwargs
     )
@@ -114,4 +125,4 @@ def occupancy_histogram(
     ax.set_xlabel(x_coord)
     ax.set_ylabel(y_coord)
 
-    return fig, {"counts": counts, "xedges": xedges, "yedges": yedges}
+    return fig, ax, {"counts": counts, "xedges": xedges, "yedges": yedges}

tests/test_unit/test_plot.py

@@ -1,6 +1,8 @@
+import matplotlib.pyplot as plt
 import numpy as np
 import pytest
 import xarray as xr
+from matplotlib.collections import QuadMesh
 from numpy.random import RandomState
 
 from movement.plot import occupancy_histogram
@@ -197,7 +199,7 @@ def test_occupancy_histogram(
         # when examining the xy data later in the test.
         kwargs_to_select_xy_data.pop(d, None)
 
-    _, histogram_info = occupancy_histogram(
+    _, _, histogram_info = occupancy_histogram(
         data, individual=individual, keypoint=keypoint, bins=n_bins
     )
     plotted_values = histogram_info["counts"]
@@ -265,3 +267,35 @@ def test_occupancy_histogram(
     else:
         # No non-nan values were given
         assert plotted_values.sum() == 0
+
+
+def test_respects_axes(histogram_data: xr.DataArray) -> None:
+    """Check that existing axes objects are respected if passed."""
+    # Plotting on existing axes
+    existing_fig, existing_ax = plt.subplots(1, 2)
+
+    existing_ax[0].plot(
+        np.linspace(0.0, 10.0, num=100), np.linspace(0.0, 10.0, num=100)
+    )
+
+    _, _, hist_info_existing = occupancy_histogram(
+        histogram_data, ax=existing_ax[1]
+    )
+    hist_plots_added = [
+        qm for qm in existing_ax[1].get_children() if isinstance(qm, QuadMesh)
+    ]
+    assert len(hist_plots_added) == 1
+
+    # Plot on new axis and create a new figure
+    new_fig, new_ax, hist_info_new = occupancy_histogram(histogram_data)
+    hist_plots_created = [
+        qm for qm in new_ax.get_children() if isinstance(qm, QuadMesh)
+    ]
+    assert len(hist_plots_created) == 1
+
+    # Check that the same plot was made for each
+    assert set(hist_info_new.keys()) == set(hist_info_existing.keys())
+    for key, new_ax_value in hist_info_new.items():
+        existing_ax_value = hist_info_existing[key]
+
+        assert np.allclose(new_ax_value, existing_ax_value)