Additional return values to help extract histogram information

Author: willGraham01

movement/plot.py

@@ -1,10 +1,13 @@
 """Wrappers to plot movement data."""
 
-from typing import Any
+from typing import Any, Literal, TypeAlias
 
 import matplotlib.pyplot as plt
+import numpy as np
 import xarray as xr
 
+HistInfoKeys: TypeAlias = Literal["counts", "xedges", "yedges"]
+
 DEFAULT_HIST_ARGS = {"alpha": 1.0, "bins": 30, "cmap": "viridis"}
 
 
@@ -14,9 +17,13 @@ def occupancy_histogram(
     individual: int | str = 0,
     title: str | None = None,
     **kwargs: Any,
-) -> plt.Figure:
+) -> tuple[plt.Figure, dict[HistInfoKeys, np.ndarray]]:
     """Create a 2D histogram of the occupancy data given.
 
+    Time-points whose corresponding spatial coordinates have NaN values
+    are ignored. Histogram information is returned as the second output
+    value (see Notes).
+
     Parameters
     ----------
     da : xarray.DataArray
@@ -35,6 +42,35 @@ def occupancy_histogram(
     -------
     matplotlib.pyplot.Figure
         Plot handle containing the rendered 2D histogram.
+    dict[str, numpy.ndarray]
+        Information about the created histogram (see Notes).
+
+    Notes
+    -----
+    In instances where the counts or information about the histogram bins is
+    desired, the ``return_hist_info`` argument should be provided as ``True``.
+    This will force the function to return a second output value, which is a
+    dictionary containing the bin edges and bin counts that were used to create
+    the histogram.
+
+    For data with ``N`` time-points, the dictionary output has key-value pairs;
+    - ``xedges``, an ``(N+1,)`` ``numpy`` array specifying the bin edges in the
+    first spatial dimension.
+    - ``yedges``, same as ``xedges`` but for the second spatial dimension.
+    - ``counts``, an ``(N, N)`` ``numpy`` array with the count for each bin.
+
+    ``counts[x, y]`` is the number of datapoints in the
+    ``(xedges[x], xedges[x+1]), (yedges[y], yedges[y+1])`` bin. These values
+    are those returned from ``matplotlib.pyplot.Axes.hist2d``.
+
+    Note that the ``counts`` values do not necessarily match the mappable
+    values that one gets from extracting the data from the
+    ``matplotlib.collections.QuadMesh`` object (that represents the rendered
+    histogram) via its ``get_array()`` attribute.
+
+    See Also
+    --------
+    matplotlib.pyplot.Axes.hist2d : The underlying plotting function.
 
     """
     data = da.position if isinstance(da, xr.Dataset) else da
@@ -64,9 +100,9 @@ def occupancy_histogram(
             kwargs[key] = value
     # Now it should just be a case of creating the histogram
     fig, ax = plt.subplots()
-    _, _, _, hist_image = ax.hist2d(
+    counts, xedges, yedges, hist_image = ax.hist2d(
         data.sel(space=x_coord), data.sel(space=y_coord), **kwargs
-    )  # counts, xedges, yedges, image
+    )
     colourbar = fig.colorbar(hist_image, ax=ax)
     colourbar.solids.set(alpha=1.0)
 
@@ -78,4 +114,4 @@ def occupancy_histogram(
     ax.set_xlabel(x_coord)
     ax.set_ylabel(y_coord)
 
-    return fig
+    return fig, {"counts": counts, "xedges": xedges, "yedges": yedges}

tests/test_unit/test_plot.py

@@ -1,22 +1,11 @@
-from itertools import product
-
-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
 
 
-def get_histogram_binning_data(fig: plt.Figure) -> list[QuadMesh]:
-    """Fetch 2D array data from a histogram plot."""
-    return [
-        qm for qm in fig.axes[0].get_children() if isinstance(qm, QuadMesh)
-    ]
-
-
 @pytest.fixture
 def seed() -> int:
     return 0
@@ -72,19 +61,30 @@ def histogram_data_with_nans(
 ) -> xr.DataArray:
     """DataArray whose data is the ``normal_dist_2d`` points.
 
-    Each datapoint has a chance of being turned into a NaN value.
-
     Axes 2 and 3 are the individuals and keypoints axes, respectively.
     These dimensions are given coordinates {i,k}{0,1,2,3,4,5,...} for
     the purposes of indexing.
+
+    For individual i0, keypoint k0, the following (time, space) values are
+    converted into NaNs:
+    - (100, "x")
+    - (200, "y")
+    - (150, "x")
+    - (150, "y")
+
     """
+    individual_0 = "i0"
+    keypoint_0 = "k0"
     data_with_nans = histogram_data.copy(deep=True)
-    data_shape = data_with_nans.shape
-    nan_chance = 1.0 / 25.0
-    index_ranges = [range(dim_length) for dim_length in data_shape]
-    for multiindex in product(*index_ranges):
-        if rng.uniform() < nan_chance:
-            data_with_nans[*multiindex] = float("nan")
+    for time_index, space_coord in [
+        (100, "x"),
+        (200, "y"),
+        (150, "x"),
+        (150, "y"),
+    ]:
+        data_with_nans.loc[
+            time_index, space_coord, individual_0, keypoint_0
+        ] = float("nan")
     return data_with_nans
 
 
@@ -97,7 +97,29 @@ def histogram_data_with_nans(
 
 @pytest.mark.parametrize(
     ["data", "individual", "keypoint", "n_bins"],
-    [pytest.param("histogram_data", "i0", "k0", 30, id="30 bins each axis")],
+    [
+        pytest.param(
+            "histogram_data",
+            "i0",
+            "k0",
+            30,
+            id="30 bins each axis",
+        ),
+        pytest.param(
+            "histogram_data",
+            "i1",
+            "k0",
+            (20, 30),
+            id="(20, 30) bins",
+        ),
+        pytest.param(
+            "histogram_data_with_nans",
+            "i0",
+            "k0",
+            30,
+            id="NaNs should be removed",
+        ),
+    ],
 )
 def test_occupancy_histogram(
     data: xr.DataArray,
@@ -110,62 +132,61 @@ def test_occupancy_histogram(
     if isinstance(data, str):
         data = request.getfixturevalue(data)
 
-    plotted_hist = occupancy_histogram(
+    _, histogram_info = occupancy_histogram(
         data, individual=individual, keypoint=keypoint, bins=n_bins
     )
-
-    # Confirm that a histogram was made
-    plotted_data = get_histogram_binning_data(plotted_hist)
-    assert len(plotted_data) == 1
-    plotted_data = plotted_data[0]
-    plotting_coords = plotted_data.get_coordinates()
-    plotted_values = plotted_data.get_array()
+    plotted_values = histogram_info["counts"]
 
     # Confirm the binned array has the correct size
     if not isinstance(n_bins, tuple):
         n_bins = (n_bins, n_bins)
-    assert plotted_data.get_array().shape == n_bins
+    assert plotted_values.shape == n_bins
 
     # Confirm that each bin has the correct number of assignments
     data_time_xy = data.sel(individuals=individual, keypoints=keypoint)
-    x_values = data_time_xy.sel(space="x").values
-    y_values = data_time_xy.sel(space="y").values
+    data_time_xy = data_time_xy.dropna(dim="time", how="any")
+    plotted_x_values = data_time_xy.sel(space="x").values
+    plotted_y_values = data_time_xy.sel(space="y").values
+    assert plotted_x_values.shape == plotted_y_values.shape
+    # This many non-NaN values were plotted
+    n_non_nan_values = plotted_x_values.shape[0]
+
     reconstructed_bins_limits_x = np.linspace(
-        x_values.min(),
-        x_values.max(),
+        plotted_x_values.min(),
+        plotted_x_values.max(),
         num=n_bins[0] + 1,
         endpoint=True,
     )
-    assert all(
-        np.allclose(reconstructed_bins_limits_x, plotting_coords[i, :, 0])
-        for i in range(n_bins[0])
-    )
+    assert np.allclose(reconstructed_bins_limits_x, histogram_info["xedges"])
     reconstructed_bins_limits_y = np.linspace(
-        y_values.min(),
-        y_values.max(),
+        plotted_y_values.min(),
+        plotted_y_values.max(),
         num=n_bins[1] + 1,
         endpoint=True,
     )
-    assert all(
-        np.allclose(reconstructed_bins_limits_y, plotting_coords[:, j, 1])
-        for j in range(n_bins[1])
-    )
+    assert np.allclose(reconstructed_bins_limits_y, histogram_info["yedges"])
 
     reconstructed_bin_counts = np.zeros(shape=n_bins, dtype=float)
     for i, xi in enumerate(reconstructed_bins_limits_x[:-1]):
         xi_p1 = reconstructed_bins_limits_x[i + 1]
 
-        x_pts_in_range = (x_values >= xi) & (x_values <= xi_p1)
+        x_pts_in_range = (plotted_x_values >= xi) & (plotted_x_values <= xi_p1)
         for j, yj in enumerate(reconstructed_bins_limits_y[:-1]):
             yj_p1 = reconstructed_bins_limits_y[j + 1]
 
-            y_pts_in_range = (y_values >= yj) & (y_values <= yj_p1)
+            y_pts_in_range = (plotted_y_values >= yj) & (
+                plotted_y_values <= yj_p1
+            )
 
             pts_in_this_bin = (x_pts_in_range & y_pts_in_range).sum()
             reconstructed_bin_counts[i, j] = pts_in_this_bin
 
             if pts_in_this_bin != plotted_values[i, j]:
                 pass
 
+    # We agree with a manual count
     assert reconstructed_bin_counts.sum() == plotted_values.sum()
+    # All non-NaN values were plotted
+    assert n_non_nan_values == plotted_values.sum()
+    # The counts were actually correct
     assert np.all(reconstructed_bin_counts == plotted_values)