Merge pull request #348 from European-XFEL/hyperslicer2

Implement hyperslicer2() for plotting image arrays

Author: JamesWrigley

docs/changelog.md

@@ -28,6 +28,8 @@ Added:
   (AGIPD1M, AGIPD500K, DSSC1M, JUNGFRAU, LPD1M) (!177).
 - [imshow2][extra.utils.imshow2] now supports plotting 2D
   [DataArray][xarray.DataArray]s properly (!333).
+- Added [hyperslicer2()][extra.utils.hyperslicer2] to make plotting image arrays
+  easier (!348).
 
 Changed:
 - [Timepix3.spatial_bins()] is now a static method.

docs/images/hyperslicer2.gif

@@ -20,6 +20,8 @@ else.
 
 ::: extra.utils.imshow2
 
+::: extra.utils.hyperslicer2
+
 ## Fitting functions
 
 ::: extra.utils.fit_gaussian

docs/utilities.md

@@ -88,6 +88,7 @@ plugins:
           - https://docs.xarray.dev/en/stable/objects.inv
           - https://pint.readthedocs.io/en/stable/objects.inv
           - https://docs.scipy.org/doc/scipy/objects.inv
+          - https://mpl-interactions.readthedocs.io/en/stable/objects.inv
           options:
             docstring_style: google
             separate_signature: yes

mkdocs.yml

@@ -31,6 +31,7 @@ dependencies = [
     "oauth2-xfel-client >=6.1",
     "pandas",
     "xarray",
+    "mpl_interactions",
     "pasha",
     "pint",
     "requests",

pyproject.toml

@@ -90,3 +90,73 @@ def imshow2(image, *args, lognorm=False, ax=None, **kwargs):
         if ax is None:
             ax = plt
         return ax.imshow(image, *args, **kwargs)
+
+def hyperslicer2(arr, *args, ax=None, lognorm=False, colorbar=True, **kwargs):
+    """Interactively visualize arrays of images.
+
+    This is a lightweight wrapper around
+    [hyperslicer()][mpl_interactions.generic.hyperslicer] with some useful defaults:
+
+    - Try to set `vmin`/`vmax` to reasonable values. Note that setting
+      `vmin`/`vmax` is incompatible with the `norm` argument, so they will only
+      be set if `norm` is not passed.
+    - Set `interpolation="none"`.
+    - Enable the play buttons.
+    - Draw a colorbar.
+
+    Example usage:
+    ```python
+    plt.figure()
+    # Note the trailing semi-colon to swallow the return value. hyperslicer2()
+    # returns a `controls` object by default that displays the play buttons, so
+    # returning it from a notebook cell will end up displaying the play buttons
+    # twice.
+    hyperslicer2(images);
+    ```
+    ![](../images/hyperslicer2.gif)
+
+    All arguments other than the ones listed below are passed to
+    [hyperslicer()][mpl_interactions.generic.hyperslicer], and explicitly
+    passing any of `vmin`/`vmax`/`interpolation`/`play_buttons` will override
+    the defaults.
+
+    Args:
+        arr (array_like): The array of images to display. Should have at least
+            three dimensions.
+        ax (matplotlib.axes.Axes): The axis to plot the image in.
+        lognorm (bool): Whether to display the images in a log color scale.
+        colorbar (bool): Whether to display a colorbar.
+    """
+    import matplotlib.pyplot as plt
+    from mpl_interactions import hyperslicer
+
+    # Enable the controls by default
+    if "play_buttons" not in kwargs:
+        kwargs["play_buttons"] = True
+
+    # Disable interpolation by default
+    if "interpolation" not in kwargs:
+        kwargs["interpolation"] = "none"
+
+    # Enable log color scale if requested and `norm` is not already set
+    if lognorm and "norm" not in kwargs:
+        from matplotlib.colors import LogNorm
+        kwargs["norm"] = LogNorm()
+
+    # Set the vmin/vmax if we're not using `norm`
+    if "norm" not in kwargs and np.issubdtype(arr.dtype, np.number):
+        if "vmin" not in kwargs:
+            kwargs["vmin"] = np.nanquantile(arr, 0.01)
+        if "vmax" not in kwargs:
+            kwargs["vmax"] = np.nanquantile(arr, 0.99)
+
+    if ax is None:
+        ax = plt.gca()
+    fig = ax.get_figure()
+
+    controls = hyperslicer(arr, *args, ax=ax, **kwargs)
+
+    if colorbar:
+        fig.colorbar(ax.get_images()[-1], ax=ax)
+
+    return controls

src/extra/utils/misc.py

@@ -1,7 +1,7 @@
 import numpy as np
 import xarray as xr
 
-from extra.utils import imshow2, fit_gaussian, gaussian
+from extra.utils import imshow2, hyperslicer2, fit_gaussian, gaussian
 
 
 def test_imshow2():
@@ -13,6 +13,11 @@ def test_imshow2():
     imshow2(image)
 
 
+def test_hyperslicer2():
+    # Smoke test
+    images = np.random.rand(10, 100, 100)
+    hyperslicer2(images)
+
 def test_fit_gaussian():
     # Test with auto-generated xdata and nans/infs
     params = [0, 1, 20, 5]