From 16f5e0c246c659308e2778c0fad4b9c602081514 Mon Sep 17 00:00:00 2001 From: Ian Hunt-Isaak Date: Mon, 19 Jul 2021 17:20:59 -0400 Subject: [PATCH] basic plot_surface --- mpl_interactions/helpers.py | 18 ++++++++ mpl_interactions/ipyplot.py | 1 + mpl_interactions/mpl_kwargs.py | 3 ++ mpl_interactions/pyplot.py | 81 ++++++++++++++++++++++++++++++++++ 4 files changed, 103 insertions(+) diff --git a/mpl_interactions/helpers.py b/mpl_interactions/helpers.py index 80b15cee..93076cb3 100644 --- a/mpl_interactions/helpers.py +++ b/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): """ diff --git a/mpl_interactions/ipyplot.py b/mpl_interactions/ipyplot.py index 1f556be8..acd96bfe 100644 --- a/mpl_interactions/ipyplot.py +++ b/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 \ No newline at end of file diff --git a/mpl_interactions/mpl_kwargs.py b/mpl_interactions/mpl_kwargs.py index 5bf30685..cc780c0e 100644 --- a/mpl_interactions/mpl_kwargs.py +++ b/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", diff --git a/mpl_interactions/pyplot.py b/mpl_interactions/pyplot.py index df067bcf..8d009efb 100644 --- a/mpl_interactions/pyplot.py +++ b/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()