From dbd07dfe1b5ff9f4be018ad5ab286e08a52279a1 Mon Sep 17 00:00:00 2001 From: Michael Hardman <29800382+mrhardman@users.noreply.github.com> Date: Mon, 10 Feb 2025 15:24:09 +0000 Subject: [PATCH 1/5] WIP: animate function based on plot2d_polygon. --- xbout/plotting/animate.py | 202 +++++++++++++++++++++++++++++++++++++- 1 file changed, 201 insertions(+), 1 deletion(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index e223983f..95912793 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -2,7 +2,7 @@ import matplotlib import matplotlib.pyplot as plt import warnings - +from mpl_toolkits.axes_grid1 import make_axes_locatable import animatplot as amp from .utils import ( @@ -679,3 +679,203 @@ def animate_line( return anim return line_block + +def animate_polygon( + da, + ax=None, + cax=None, + cmap="viridis", + norm=None, + logscale=False, + antialias=False, + vmin=None, + vmax=None, + extend="neither", + add_colorbar=True, + colorbar_label=None, + separatrix=True, + separatrix_kwargs={"color": "white", "linestyle": "-", "linewidth": 2}, + targets=False, + add_limiter_hatching=False, + grid_only=False, + linewidth=0, + linecolor="black", +): + """ + Nice looking 2D plots which have no visual artifacts around the X-point. + + Parameters + ---------- + da : xarray.DataArray + A 2D (x,y) DataArray of data to plot + ax : Axes, optional + Axes to plot on. If not provided, will make its own. + cax : Axes, optional + Axes to plot colorbar on. If not provided, will plot on the same axes as the plot. + cmap : str or matplotlib.colors.Colormap, default "viridis" + Colormap to use for the plot + norm : matplotlib.colors.Normalize, optional + Normalization to use for the color scale + logscale : bool, default False + If True, use a symlog color scale + antialias : bool, default False + Enables antialiasing. Note: this also shows mesh cell edges - it's unclear how to disable this. + vmin : float, optional + Minimum value for the color scale + vmax : float, optional + Maximum value for the color scale + extend : str, optional, default "neither" + Extend the colorbar. Options are "neither", "both", "min", "max" + add_colorbar : bool, default True + Enable colorbar in figure? + colorbar_label : str, optional + Label for the colorbar + separatrix : bool, default True + Add lines showing separatrices + separatrix_kwargs : dict + Keyword arguments to pass custom style to the separatrices plot + targets : bool, default True + Draw solid lines at the target surfaces + add_limiter_hatching : bool, default True + Draw hatched areas at the targets + grid_only : bool, default False + Only plot the grid, not the data. This sets all the polygons to have a white face. + linewidth : float, default 0 + Width of the gridlines on cell edges + linecolor : str, default "black" + Color of the gridlines on cell edges + """ + + if ax is None: + fig, ax = plt.subplots(figsize=(3, 6), dpi=120) + else: + fig = ax.get_figure() + + if cax is None: + cax = ax + + if vmin is None: + vmin = np.nanmin(da.values) + + if vmax is None: + vmax = np.nanmax(da.max().values) + + if colorbar_label is None: + if "short_name" in da.attrs: + colorbar_label = da.attrs["short_name"] + elif da.name is not None: + colorbar_label = da.name + else: + colorbar_label = "" + + if "units" in da.attrs: + colorbar_label += f" [{da.attrs['units']}]" + + if "Rxy_lower_right_corners" in da.coords: + r_nodes = [ + "R", + "Rxy_lower_left_corners", + "Rxy_lower_right_corners", + "Rxy_upper_left_corners", + "Rxy_upper_right_corners", + ] + z_nodes = [ + "Z", + "Zxy_lower_left_corners", + "Zxy_lower_right_corners", + "Zxy_upper_left_corners", + "Zxy_upper_right_corners", + ] + cell_r = np.concatenate( + [np.expand_dims(da[x], axis=2) for x in r_nodes], axis=2 + ) + cell_z = np.concatenate( + [np.expand_dims(da[x], axis=2) for x in z_nodes], axis=2 + ) + else: + raise Exception("Cell corners not present in mesh, cannot do polygon plot") + + Nx = len(cell_r) + Ny = len(cell_r[0]) + patches = [] + + # https://matplotlib.org/2.0.2/examples/api/patch_collection.html + + idx = [np.array([1, 2, 4, 3, 1])] + patches = [] + for i in range(Nx): + for j in range(Ny): + p = matplotlib.patches.Polygon( + np.concatenate((cell_r[i][j][tuple(idx)], cell_z[i][j][tuple(idx)])) + .reshape(2, 5) + .T, + fill=False, + closed=True, + facecolor=None, + ) + patches.append(p) + + norm = _create_norm(logscale, norm, vmin, vmax) + + if grid_only is True: + cmap = matplotlib.colors.ListedColormap(["white"]) + polys = matplotlib.collections.PatchCollection( + patches, + alpha=1, + norm=norm, + cmap=cmap, + antialiaseds=antialias, + edgecolors=linecolor, + linewidths=linewidth, + joinstyle="bevel", + ) + + colors = da.data[0,:,:].flatten() + polys.set_array(colors) + ax.add_collection(polys) + + #def update(frame): + # # for each frame, update the data stored on each artist. + # colors = da.data[frame,:,:].flatten() + # polys.set_array(colors) + # return polys + def update(frame): + colors = da.data[frame,:,:].flatten() + polys.set_array(colors) + print(da.data[0,0,0].compute()) + print(frame) + #update(0) + + if add_colorbar: + # This produces a "foolproof" colorbar which + # is always the height of the plot + # From https://joseph-long.com/writing/colorbars/ + divider = make_axes_locatable(ax) + cax = divider.append_axes("right", size="5%", pad=0.05) + fig.colorbar(polys, cax=cax, label=colorbar_label, extend=extend) + cax.grid(which="both", visible=False) + + + + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel("R [m]") + ax.set_ylabel("Z [m]") + ax.set_ylim(cell_z.min(), cell_z.max()) + ax.set_xlim(cell_r.min(), cell_r.max()) + ax.set_title(da.name) + + if separatrix: + plot_separatrices(da, ax, x="R", y="Z", **separatrix_kwargs) + + if targets: + plot_targets(da, ax, x="R", y="Z", hatching=add_limiter_hatching) + + #print(np.shape(da.data)) + #colors = da.data[0,:,:].flatten() + #polys.set_array(colors) + #ax.add_collection(polys) + + + ani = matplotlib.animation.FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=3000) + return ani + From 367c225ebc8294c3b35a76c03772b905f808f9ef Mon Sep 17 00:00:00 2001 From: Michael Hardman <29800382+mrhardman@users.noreply.github.com> Date: Mon, 10 Feb 2025 16:09:33 +0000 Subject: [PATCH 2/5] WIP: tidy up animate function based on plot2d_polygon, add comments. --- xbout/plotting/animate.py | 24 ++++++------------------ 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 95912793..412c52ea 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -833,19 +833,12 @@ def animate_polygon( colors = da.data[0,:,:].flatten() polys.set_array(colors) ax.add_collection(polys) - - #def update(frame): - # # for each frame, update the data stored on each artist. - # colors = da.data[frame,:,:].flatten() - # polys.set_array(colors) - # return polys + # function to update the data plotted + # assuming data in shape (t,x,y) def update(frame): colors = da.data[frame,:,:].flatten() polys.set_array(colors) - print(da.data[0,0,0].compute()) - print(frame) - #update(0) - + if add_colorbar: # This produces a "foolproof" colorbar which # is always the height of the plot @@ -869,13 +862,8 @@ def update(frame): if targets: plot_targets(da, ax, x="R", y="Z", hatching=add_limiter_hatching) - - #print(np.shape(da.data)) - #colors = da.data[0,:,:].flatten() - #polys.set_array(colors) - #ax.add_collection(polys) - - - ani = matplotlib.animation.FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=3000) + + # make the animation by using FuncAnimation and update() to generate frames + ani = matplotlib.animation.FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=30) return ani From bc4d9ebdd91c475238b11af96e8ba25ed8ba00c1 Mon Sep 17 00:00:00 2001 From: Michael Hardman <29800382+mrhardman@users.noreply.github.com> Date: Fri, 14 Feb 2025 10:44:28 +0000 Subject: [PATCH 3/5] Working, but perhaps slow, version of the animate_polygon for animating lists of data with Mike Kryjak's polygon plot. --- xbout/boutdataset.py | 255 ++++++++++++++++++++++++++++++++++++++ xbout/plotting/animate.py | 15 ++- 2 files changed, 266 insertions(+), 4 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 94987b4b..1a5a5a60 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -11,6 +11,7 @@ import animatplot as amp from matplotlib import pyplot as plt from matplotlib.animation import PillowWriter +from matplotlib.animation import FuncAnimation from mpl_toolkits.axes_grid1 import make_axes_locatable @@ -22,6 +23,7 @@ animate_poloidal, animate_pcolormesh, animate_line, + animate_polygon, _add_controls, _normalise_time_coord, _parse_coord_option, @@ -1345,6 +1347,259 @@ def is_list(variable): return anim + def animate_polygon_list( + self, + variables, + animate_over=None, + save_as=None, + show=False, + fps=10, + nrows=None, + ncols=None, + poloidal_plot=False, + axis_coords=None, + subplots_adjust=None, + vmin=None, + vmax=None, + logscale=None, + titles=None, + aspect=None, + extend=None, + #controls="both", + tight_layout=True, + **kwargs, + ): + """ + Parameters + ---------- + variables : list of str or BoutDataArray + The variables to plot. For any string passed, the corresponding + variable in this DataSet is used - then the calling DataSet must + have only 3 dimensions. It is possible to pass BoutDataArrays to + allow more flexible plots, e.g. with different variables being + plotted against different axes. + animate_over : str, optional + Dimension over which to animate, defaults to the time dimension + save_as : str, optional + If passed, a gif is created with this filename + show : bool, optional + Call pyplot.show() to display the animation + fps : float, optional + Indicates the number of frames per second to play + nrows : int, optional + Specify the number of rows of plots + ncols : int, optional + Specify the number of columns of plots + poloidal_plot : bool or sequence of bool, optional + If set to True, make all 2D animations in the poloidal plane instead of using + grid coordinates, per variable if sequence is given + axis_coords : None, str, dict or list of None, str or dict + Coordinates to use for axis labelling. + + - None: Use the dimension coordinate for each axis, if it exists. + - "index": Use the integer index values. + - dict: keys are dimension names, values set axis_coords for each axis + separately. Values can be: None, "index", the name of a 1d variable or + coordinate (which must have the dimension given by 'key'), or a 1d + numpy array, dask array or DataArray whose length matches the length of + the dimension given by 'key'. + + Only affects time coordinate for plots with poloidal_plot=True. + If a list is passed, it must have the same length as 'variables' and gives + the axis_coords setting for each plot individually. + The setting to use for the 'animate_over' coordinate can be passed in one or + more dict values, but must be the same in all dicts if given more than once. + subplots_adjust : dict, optional + Arguments passed to fig.subplots_adjust()() + vmin : float or sequence of floats + Minimum value for color scale, per variable if a sequence is given + vmax : float or sequence of floats + Maximum value for color scale, per variable if a sequence is given + logscale : bool or float, sequence of bool or float, optional + If True, default to a logarithmic color scale instead of a linear one. + If a non-bool type is passed it is treated as a float used to set the linear + threshold of a symmetric logarithmic scale as + linthresh=min(abs(vmin),abs(vmax))*logscale, defaults to 1e-5 if True is + passed. + Per variable if sequence is given. + titles : sequence of str or None, optional + Custom titles for each plot. Pass None in the sequence to use the default for + a certain variable + aspect : str or None, or sequence of str or None, optional + Argument to set_aspect() for each plot. Defaults to "equal" for poloidal + plots and "auto" for others. + extend : str or None, optional + Passed to fig.colorbar() + controls : string or None, default "both" + By default, add both the timeline and play/pause toggle to the animation. If + "timeline" is passed add only the timeline, if "toggle" is passed add only + the play/pause toggle. If None or an empty string is passed, add neither. + tight_layout : bool or dict, optional + If set to False, don't call tight_layout() on the figure. + If a dict is passed, the dict entries are passed as arguments to + tight_layout() + **kwargs : dict, optional + Additional keyword arguments are passed on to each animation function, per + variable if a sequence is given. + + Returns + ------- + animation + An animatplot.Animation object. + """ + + if animate_over is None: + animate_over = self.metadata.get("bout_tdim", "t") + + nvars = len(variables) + + if nrows is None and ncols is None: + ncols = int(np.ceil(np.sqrt(nvars))) + nrows = int(np.ceil(nvars / ncols)) + elif nrows is None: + nrows = int(np.ceil(nvars / ncols)) + elif ncols is None: + ncols = int(np.ceil(nvars / nrows)) + else: + if nrows * ncols < nvars: + raise ValueError("Not enough rows*columns to fit all variables") + + fig, axes = plt.subplots(nrows, ncols, squeeze=False) + axes = axes.flatten() + + ncells = nrows * ncols + + if nvars < ncells: + for index in range(ncells - nvars): + fig.delaxes(axes[ncells - index - 1]) + + if subplots_adjust is not None: + fig.subplots_adjust(**subplots_adjust) + + def _expand_list_arg(arg, arg_name): + if isinstance(arg, collections.abc.Sequence) and not isinstance(arg, str): + if len(arg) != len(variables): + raise ValueError( + "if %s is a sequence, it must have the same " + 'number of elements as "variables"' % arg_name + ) + else: + arg = [arg] * len(variables) + return arg + + poloidal_plot = _expand_list_arg(poloidal_plot, "poloidal_plot") + vmin = _expand_list_arg(vmin, "vmin") + vmax = _expand_list_arg(vmax, "vmax") + logscale = _expand_list_arg(logscale, "logscale") + titles = _expand_list_arg(titles, "titles") + aspect = _expand_list_arg(aspect, "aspect") + extend = _expand_list_arg(extend, "extend") + axis_coords = _expand_list_arg(axis_coords, "axis_coords") + for k in kwargs: + kwargs[k] = _expand_list_arg(kwargs[k], k) + + animate_data = [] + + def is_list(variable): + return ( + isinstance(variable, list) + or isinstance(variable, tuple) + or isinstance(variable, set) + ) + + for i, subplot_args in enumerate( + zip( + variables, + axes, + poloidal_plot, + vmin, + vmax, + logscale, + titles, + ) + ): + ( + v, + ax, + this_poloidal_plot, + this_vmin, + this_vmax, + this_logscale, + this_title, + ) = subplot_args + + this_kwargs = {k: v[i] for k, v in kwargs.items()} + + divider = make_axes_locatable(ax) + cax = divider.append_axes("right", size="5%", pad=0.1) + + if isinstance(v, str): + v = self.data[v] + + data = v.bout.data + ndims = len(data.dims) + ax.set_title(data.name) + + if ndims == 3: + if this_poloidal_plot: + polys, da, update_func = animate_polygon( + data, + ax=ax, + cax=cax, + vmin=this_vmin, + vmax=this_vmax, + logscale=this_logscale, + animate=False, + **this_kwargs, + ) + animate_data.append([polys,da,update_func]) + else: + raise ValueError( + "Unsupported option " + + ". this_poloidal_plot " + + str(this_poloidal_plot) + ) + else: + raise ValueError( + "Unsupported number of dimensions " + + str(ndims) + + ". Dims are " + + str(v.dims) + ) + + if this_title is not None: + # Replace default title with user-specified one + ax.set_title(this_title) + + def update(frame): + for list in animate_data: + (polys, da, update_func) = list + # call update function for each axes + update_func(frame,polys,da) + + # make the animation for all the subplots simultaneously + # use the last data array da to choose the number of frames + # assumes time dimension same length for all variables + anim = FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=30) + if tight_layout: + if subplots_adjust is not None: + warnings.warn( + "tight_layout argument to animate_list() is True, but " + "subplots_adjust argument is not None. subplots_adjust " + "is being ignored." + ) + if not isinstance(tight_layout, dict): + tight_layout = {} + fig.tight_layout(**tight_layout) + + if save_as is not None: + anim.save(save_as + ".gif", writer=PillowWriter(fps=fps)) + + if show: + plt.show() + + return anim + def with_cherab_grid(self): """ Returns a new DataSet with a 'cherab_grid' attribute. diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 412c52ea..9fb8770a 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -700,6 +700,7 @@ def animate_polygon( grid_only=False, linewidth=0, linecolor="black", + animate=True, ): """ Nice looking 2D plots which have no visual artifacts around the X-point. @@ -862,8 +863,14 @@ def update(frame): if targets: plot_targets(da, ax, x="R", y="Z", hatching=add_limiter_hatching) - - # make the animation by using FuncAnimation and update() to generate frames - ani = matplotlib.animation.FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=30) - return ani + if animate: + # make the animation by using FuncAnimation and update() to generate frames + ani = matplotlib.animation.FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=30) + return ani + else: + # return function and data for making the animation + def update_out(frame,polys,da): + colors = da.data[frame,:,:].flatten() + polys.set_array(colors) + return polys, da, update_out From 108989e7ea9595de06a0a4111c6e13ce038c2bd4 Mon Sep 17 00:00:00 2001 From: Michael Hardman <29800382+mrhardman@users.noreply.github.com> Date: Fri, 14 Feb 2025 12:01:29 +0000 Subject: [PATCH 4/5] Make the update function going to animate_polygon_list from animate_polygon fully encapsulated, use the t dimension to set the number of frames. --- xbout/boutdataset.py | 13 ++++++------- xbout/plotting/animate.py | 7 ++----- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 1a5a5a60..ec1bedc6 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -1542,7 +1542,7 @@ def is_list(variable): if ndims == 3: if this_poloidal_plot: - polys, da, update_func = animate_polygon( + update_func = animate_polygon( data, ax=ax, cax=cax, @@ -1552,7 +1552,7 @@ def is_list(variable): animate=False, **this_kwargs, ) - animate_data.append([polys,da,update_func]) + animate_data.append(update_func) else: raise ValueError( "Unsupported option " @@ -1572,15 +1572,14 @@ def is_list(variable): ax.set_title(this_title) def update(frame): - for list in animate_data: - (polys, da, update_func) = list + for update_func in animate_data: # call update function for each axes - update_func(frame,polys,da) + update_func(frame) # make the animation for all the subplots simultaneously - # use the last data array da to choose the number of frames + # use time data array "t" to choose the number of frames # assumes time dimension same length for all variables - anim = FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=30) + anim = FuncAnimation(fig=fig, func=update, frames=self.data["t"].data.size, interval=30) if tight_layout: if subplots_adjust is not None: warnings.warn( diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 9fb8770a..8a52681d 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -868,9 +868,6 @@ def update(frame): ani = matplotlib.animation.FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=30) return ani else: - # return function and data for making the animation - def update_out(frame,polys,da): - colors = da.data[frame,:,:].flatten() - polys.set_array(colors) - return polys, da, update_out + # return function for making the animation + return update From dd068de7e0e8b09ed243a9a24144ca1fd2832f60 Mon Sep 17 00:00:00 2001 From: mrhardman Date: Thu, 20 Feb 2025 16:30:01 +0000 Subject: [PATCH 5/5] Apply black formatting --- xbout/boutdataset.py | 20 +++++++++++--------- xbout/plotting/animate.py | 21 +++++++++++---------- 2 files changed, 22 insertions(+), 19 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index ec1bedc6..501a962a 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -1365,7 +1365,7 @@ def animate_polygon_list( titles=None, aspect=None, extend=None, - #controls="both", + # controls="both", tight_layout=True, **kwargs, ): @@ -1555,10 +1555,10 @@ def is_list(variable): animate_data.append(update_func) else: raise ValueError( - "Unsupported option " - + ". this_poloidal_plot " - + str(this_poloidal_plot) - ) + "Unsupported option " + + ". this_poloidal_plot " + + str(this_poloidal_plot) + ) else: raise ValueError( "Unsupported number of dimensions " @@ -1570,16 +1570,18 @@ def is_list(variable): if this_title is not None: # Replace default title with user-specified one ax.set_title(this_title) - + def update(frame): for update_func in animate_data: # call update function for each axes update_func(frame) - + # make the animation for all the subplots simultaneously # use time data array "t" to choose the number of frames # assumes time dimension same length for all variables - anim = FuncAnimation(fig=fig, func=update, frames=self.data["t"].data.size, interval=30) + anim = FuncAnimation( + fig=fig, func=update, frames=self.data["t"].data.size, interval=30 + ) if tight_layout: if subplots_adjust is not None: warnings.warn( @@ -1598,7 +1600,7 @@ def update(frame): plt.show() return anim - + def with_cherab_grid(self): """ Returns a new DataSet with a 'cherab_grid' attribute. diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 8a52681d..ebbd8829 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -680,6 +680,7 @@ def animate_line( return line_block + def animate_polygon( da, ax=None, @@ -700,7 +701,7 @@ def animate_polygon( grid_only=False, linewidth=0, linecolor="black", - animate=True, + animate=True, ): """ Nice looking 2D plots which have no visual artifacts around the X-point. @@ -830,16 +831,17 @@ def animate_polygon( linewidths=linewidth, joinstyle="bevel", ) - - colors = da.data[0,:,:].flatten() + + colors = da.data[0, :, :].flatten() polys.set_array(colors) ax.add_collection(polys) + # function to update the data plotted # assuming data in shape (t,x,y) def update(frame): - colors = da.data[frame,:,:].flatten() + colors = da.data[frame, :, :].flatten() polys.set_array(colors) - + if add_colorbar: # This produces a "foolproof" colorbar which # is always the height of the plot @@ -849,8 +851,6 @@ def update(frame): fig.colorbar(polys, cax=cax, label=colorbar_label, extend=extend) cax.grid(which="both", visible=False) - - ax.set_aspect("equal", adjustable="box") ax.set_xlabel("R [m]") ax.set_ylabel("Z [m]") @@ -864,10 +864,11 @@ def update(frame): if targets: plot_targets(da, ax, x="R", y="Z", hatching=add_limiter_hatching) if animate: - # make the animation by using FuncAnimation and update() to generate frames - ani = matplotlib.animation.FuncAnimation(fig=fig, func=update, frames=np.shape(da.data)[0], interval=30) + # make the animation by using FuncAnimation and update() to generate frames + ani = matplotlib.animation.FuncAnimation( + fig=fig, func=update, frames=np.shape(da.data)[0], interval=30 + ) return ani else: # return function for making the animation return update -