roughly working

Author: ddkohler

CHANGELOG.md

@@ -5,6 +5,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 
 ## [Unreleased]
 
+### Added
+- new artists submodule `animate` as a convenient wrapper for matplotlib's `FuncAnimate`
+
+### Changed
+- `interact2D`: replaced SimpleNamespace object with a dataclass for more explicit typing
+
 ## [3.6.2]
 
 ### Fixed

WrightTools/artists/__init__.py

@@ -2,9 +2,14 @@
 
 # flake8: noqa
 
+import pathlib
+for x in pathlib.Path(__file__).parent.walk():
+    print(x)
 
+
+from ._animate import *
 from ._base import *
 from ._colors import *
 from ._helpers import *
-from ._quick import *
 from ._interact import *
+from ._quick import *

WrightTools/artists/_animate.py

@@ -0,0 +1,90 @@
+"""library for making animations"""
+"""TODO: push these functions into WrightTools"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from functools import partial
+from matplotlib.animation import FuncAnimation
+
+# typing
+from ._interact import interact2D_fig
+
+__all__ = ["animate2D", "animate_interact2D"]
+
+
+def animate2D(
+    d,
+    cmap, norm,
+    snake=False,
+    back_and_forth=False,
+    **ani_kwargs,
+):
+    """
+    animate pcolormesh of a nd dataset (ndim >=2), 
+    mesh plots last two axes of the dataset (use `Data.transform` if needed)
+    uses first channel in dataset (use `bring_to_front` if needed)
+
+    Note: snake, back_and_forth are not yet implemented are no-ops
+    """
+
+    # initialize canvas
+    # need fig, ax, art, and norm to get this working
+    fig, ax = plt.subplots(subplot_kw=dict(projection="wright"), dpi=140, layout="constrained")
+    art = ax.pcolormesh(d[tuple([0 for i in d.shape[:-2]])], cmap=cmap, norm=norm)
+    fig.colorbar(art, ax=ax)
+
+    # funcs for updating
+    def title(ind):
+        parts = [
+            f"{var.natural_name} = {var[:].squeeze()[ind]:.2f} {var.units}" 
+            for var in map(lambda a: a.variables[0], d.axes[:-2])
+        ]
+        return "\n".join(parts)
+
+    ax.set_title(title(0))
+
+    def update2D(frame, data, fig, ax, mesh, norm):
+        print(frame)
+        # for ind, axis in zip(frame, data.axes[:-2]):
+        mesh.set_array(data.channels[0][frame])
+        ax.set_title(title(frame))
+        mesh.set_norm(norm)
+        fig.canvas.draw_idle()
+        return mesh
+
+    frames = list(np.ndindex(d.shape[:-2]))
+
+    return FuncAnimation(
+        fig=fig,
+        func=partial(update2D, data=d, mesh=art, fig=fig, ax=ax, norm=norm),
+        frames=frames,
+        **ani_kwargs,
+    )
+
+
+def animate_interact2D(
+    interact2D:interact2D_fig,
+    snake=False, 
+    back_and_forth=False,
+    **kwargs
+):
+    """
+    Take an interact2D figure and create an animation by moving the sliders.
+
+    Note: snake, back_and_forth are not yet implemented are no-ops
+    """
+
+    def update(frame):
+        print(frame)
+        for ind, slider in zip(frame, interact2D.sliders.values()):
+            slider.set_val(ind)
+
+    frames = list(np.ndindex(tuple([s.valmax+1 for s in interact2D.sliders.values()])))
+    print(f"beginning animation: {len(frames)} frames to write")
+    return FuncAnimation(
+        fig=interact2D.fig,
+        func=update,
+        frames=frames,
+    )
+

WrightTools/artists/_interact.py

@@ -4,16 +4,29 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 from matplotlib.widgets import Slider, RadioButtons
+from typing import Any
 from types import SimpleNamespace
 
+from dataclasses import dataclass
+
 from ._helpers import create_figure, plot_colorbar, add_sideplot
 from ._base import _order_for_imshow
 from ._colors import colormaps
 from ..exceptions import DimensionalityError
 from .. import kit as wt_kit
 from .. import data as wt_data
 
-__all__ = ["interact2D"]
+__all__ = ["interact2D", "interact2D_fig"]
+
+
+@dataclass
+class interact2D_fig:
+    fig:plt.figure
+    image:Any
+    sliders:dict[str, Slider]
+    crosshairs:list
+    radio:RadioButtons
+    cax: plt.axes
 
 
 class Focus:
@@ -206,7 +219,7 @@ def interact2D(
 
     Returns
     -------
-    out : types.SimpleNamespace
+    out : wt.artists.interact2D_fig
         container for important interactive elements of the plot.
 
         Properties
@@ -583,11 +596,13 @@ def update_key_press(info):
     for slider in sliders.values():
         slider.on_changed(update_slider)
 
-    out = SimpleNamespace(
+    return interact2D_fig(
+        fig=fig,
         image=obj2D,
         sliders=sliders,
         crosshairs=[crosshair_hline, crosshair_vline],
         radio=radio,
-        colorbar=colorbar,
+        cax=cax,
     )
-    return out
+
+