Start of 3D plots: Basic plot_surface

mpl_interactions/helpers.py

@@ -39,6 +39,7 @@
     "gogogo_display",
     "create_mpl_controls_fig",
     "eval_xy",
+    "eval_xyz",
     "choose_fmt_str",
 ]
 
@@ -255,6 +256,23 @@ def eval_xy(x_, y_, params, cache=None):
         y = y_
     return np.asanyarray(x), np.asanyarray(y)
 
+def eval_xyz(x, y, z, params, cache=None):
+    # maybe should allow for `y` to not need `x`
+    x_, y_ = eval_xy(x, y, params, cache)
+
+    if isinstance(z, Callable):
+        if cache is not None:
+            if z in cache:
+                z_ = cache[z]
+            else:
+                z_ = z(x, y, **params)
+        else:
+            z_ = z(x, y, **params)
+    else:
+        z_ = z
+    return x_, y_, np.asanyarray(z_)
+
+
 
 def kwarg_to_ipywidget(key, val, update, slider_format_string, play_button=None):
     """

mpl_interactions/ipyplot.py

@@ -7,3 +7,4 @@
 from .pyplot import interactive_title as title
 from .pyplot import interactive_xlabel as xlabel
 from .pyplot import interactive_ylabel as ylabel
+from .pyplot import interactive_plot_surface as plot_surface

mpl_interactions/mpl_kwargs.py

@@ -1,6 +1,8 @@
 from matplotlib.artist import ArtistInspector
 from matplotlib.collections import Collection
 from matplotlib.image import AxesImage
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
 
 # this is a list of options to Line2D partially taken from
 # https://github.com/matplotlib/matplotlib/blob/f9d29189507cfe4121a231f6ab63539d216c37bd/lib/matplotlib/lines.py#L271
@@ -63,6 +65,7 @@
 
 imshow_kwargs_list = ArtistInspector(AxesImage).get_setters()
 collection_kwargs_list = ArtistInspector(Collection).get_setters()
+Poly3D_collection_kwargs_list = ArtistInspector(Poly3DCollection).get_setters()
 
 Text_kwargs_list = [
     "agg_filter",

mpl_interactions/pyplot.py

@@ -17,6 +17,7 @@
     callable_else_value_no_cast,
     eval_xy,
     create_slider_format_dict,
+    eval_xyz,
     gogogo_display,
     gogogo_figure,
     kwarg_to_ipywidget,
@@ -29,6 +30,7 @@
     imshow_kwargs_list,
     Text_kwargs_list,
     collection_kwargs_list,
+    Poly3D_collection_kwargs_list,
     kwarg_popper,
 )
 
@@ -42,6 +44,7 @@
     "interactive_title",
     "interactive_xlabel",
     "interactive_ylabel",
+    "interactive_plot_surface",
 ]
 
 
@@ -1229,3 +1232,81 @@ def update(params, indices, cache):
         **text_kwargs,
     )
     return controls
+
+
+def interactive_plot_surface(
+    X,
+    Y,
+    Z,
+    *args,
+    norm=None,
+    vmin=None,
+    vmax=None,
+    vmin_vmax=None,
+    lightsource=None,
+    ax=None,
+    controls=None,
+    slider_formats=None,
+    display_controls=True,
+    play_buttons=False,
+    force_ipywidgets=False,
+    **kwargs,
+):
+    """
+    beep boop docstrings are for future ian
+    """
+    ipympl = notebook_backend()
+    fig, ax = gogogo_figure(ipympl, ax)
+    use_ipywidgets = ipympl or force_ipywidgets
+    slider_formats = create_slider_format_dict(slider_formats)
+    kwargs, poly3d_kwargs = kwarg_popper(kwargs, Poly3D_collection_kwargs_list)
+
+    funcs, extra_ctrls, param_excluder = prep_scalars(kwargs, vmin=vmin, vmax=vmax)
+    vmin = funcs["vmin"]
+    vmax = funcs["vmax"]
+
+    if vmin_vmax is not None:
+        if isinstance(vmin_vmax, tuple) and not isinstance(vmin_vmax[0], str):
+            vmin_vmax = ("r", *vmin_vmax)
+        kwargs["vmin_vmax"] = vmin_vmax
+
+    controls, params = gogogo_controls(
+        kwargs, controls, display_controls, slider_formats, play_buttons, extra_ctrls
+    )
+
+    if vmin_vmax is not None:
+        params.pop("vmin_vmax")
+        params["vmin"] = controls.params["vmin"]
+        params["vmax"] = controls.params["vmax"]
+
+        def vmin(**kwargs):
+            return kwargs["vmin"]
+
+        def vmax(**kwargs):
+            return kwargs["vmax"]
+
+    if "color" not in poly3d_kwargs:
+        # Call ourselves and keep around in order
+        # to avoid modifying the color-cycle
+        # https://stackoverflow.com/questions/13831549/get-matplotlib-color-cycle-state
+        poly3d_kwargs["color"] = ax._get_lines.get_next_color()
+
+    def update(params, indices, cache):
+        nonlocal artist
+        artist.remove()
+        artist = make_new_surface()
+
+    controls._register_function(update, fig, params)
+
+    def make_new_surface():
+        X_, Y_, Z_ = eval_xyz(X, Y, Z, params)
+        return ax.plot_surface(
+            X=X_,
+            Y=Y_,
+            Z=Z_,
+            vmin=callable_else_value(vmin, param_excluder(params, "vmin")),
+            vmax=callable_else_value(vmax, param_excluder(params, "vmax")),
+            **poly3d_kwargs
+        )
+
+    artist = make_new_surface()