Method to make multiple animations in one figure

xbout/__init__.py

@@ -6,7 +6,7 @@
 from .boutdataset import BoutDatasetAccessor
 from .boutdataarray import BoutDataArrayAccessor
 
-from .plotting.animate import animate_imshow
+from .plotting.animate import animate_pcolormesh, animate_poloidal
 from .plotting.utils import plot_separatrix
 
 from ._version import __version__

xbout/boutdataarray.py

@@ -4,8 +4,9 @@
 import xarray as xr
 from xarray import register_dataarray_accessor
 
-from .plotting.animate import animate_imshow, animate_line
+from .plotting.animate import animate_poloidal, animate_pcolormesh, animate_line
 from .plotting import plotfuncs
+from .plotting.utils import _create_norm
 
 
 @register_dataarray_accessor('bout')
@@ -38,14 +39,14 @@ def __str__(self):
             text += "Options:\n{}".format(styled(self.options))
         return text
 
-    def animate2D(self, animate_over='t', x=None, y=None, animate=True,
-                  fps=10, save_as=None, sep_pos=None, ax=None, **kwargs):
+    def animate2D(self, animate_over='t', x=None, y=None, animate=True, fps=10,
+                  save_as=None, ax=None, poloidal_plot=False, logscale=None, **kwargs):
         """
         Plots a color plot which is animated with time over the specified
         coordinate.
 
-        Currently only supports 2D+1 data, which it plots with xarray's
-        wrapping of matplotlib's imshow.
+        Currently only supports 2D+1 data, which it plots with animatplot's
+        wrapping of matplotlib's pcolormesh.
 
         Parameters
         ----------
@@ -63,10 +64,17 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True,
             Frames per second of resulting gif
         save_as : str, optional
             Filename to give to the resulting gif
-        sep_pos : int, optional
-            Position along the 'x' dimension to plot the separatrix
         ax : matplotlib.pyplot.axes object, optional
             Axis on which to plot the gif
+        poloidal_plot : bool, optional
+            Use animate_poloidal to make a plot in R-Z coordinates (input field must be
+            (t,x,y))
+        logscale : 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.
         kwargs : dict, optional
             Additional keyword arguments are passed on to the plotting function
             (e.g. imshow for 2D plots).
@@ -75,14 +83,31 @@ def animate2D(self, animate_over='t', x=None, y=None, animate=True,
         data = self.data
         variable = data.name
         n_dims = len(data.dims)
+
         if n_dims == 3:
-            print("{} data passed has {} dimensions - will use "
-                  "animatplot.blocks.Imshow()".format(variable, str(n_dims)))
-            imshow_block = animate_imshow(data=data, animate_over=animate_over,
-                                          x=x, y=y, sep_pos=sep_pos,
-                                          animate=animate, fps=fps,
-                                          save_as=save_as, ax=ax, **kwargs)
-            return imshow_block
+            vmin = kwargs['vmin'] if 'vmin' in kwargs else data.min().values
+            vmax = kwargs['vmax'] if 'vmax' in kwargs else data.max().values
+            kwargs['norm'] = _create_norm(logscale, kwargs.get('norm', None), vmin, vmax)
+
+            if poloidal_plot:
+                print("{} data passed has {} dimensions - making poloidal plot with "
+                      "animate_poloidal()".format(variable, str(n_dims)))
+                if x is not None:
+                    kwargs['x'] = x
+                if y is not None:
+                    kwargs['y'] = y
+                poloidal_blocks = animate_poloidal(data, animate_over=animate_over,
+                                                   animate=animate, fps=fps,
+                                                   save_as=save_as, ax=ax, **kwargs)
+                return poloidal_blocks
+            else:
+                print("{} data passed has {} dimensions - will use "
+                      "animatplot.blocks.Pcolormesh()".format(variable, str(n_dims)))
+                pcolormesh_block = animate_pcolormesh(data=data,
+                                                      animate_over=animate_over, x=x, y=y,
+                                                      animate=animate, fps=fps,
+                                                      save_as=save_as, ax=ax, **kwargs)
+                return pcolormesh_block
         else:
             raise ValueError(
                 "Data passed has an unsupported number of dimensions "

xbout/boutdataset.py

@@ -2,8 +2,15 @@
 from functools import partial
 
 from xarray import register_dataset_accessor, save_mfdataset, merge
+
+import animatplot as amp
+from matplotlib import pyplot as plt
+import numpy as np
 from dask.diagnostics import ProgressBar
 
+from .plotting.animate import animate_poloidal, animate_pcolormesh, animate_line
+from .plotting.utils import _create_norm
+
 
 @register_dataset_accessor('bout')
 class BoutDatasetAccessor:
@@ -158,6 +165,132 @@ def to_restart(self, savepath='.', nxpe=None, nype=None,
             save_mfdataset(restart_datasets, paths, compute=True)
         return
 
+    def animate_list(self, variables, animate_over='t', save_as=None, show=False, fps=10,
+                     nrows=None, ncols=None, poloidal_plot=False, subplots_adjust=None,
+                     vmin=None, vmax=None, logscale=None, **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
+        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, optional
+            If set to True, make all 2D animations in the poloidal plane instead of using
+            grid coordinates
+        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, 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.
+        **kwargs : dict, optional
+            Additional keyword arguments are passed on to each animation function
+        """
+
+        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)
+
+        if subplots_adjust is not None:
+            fig.subplots_adjust(**subplots_adjust)
+
+        try:
+            if len(vmin) != len(variables):
+                raise ValueError('if vmin is a sequence, it must have the same number '
+                                 'of elements as "variables"')
+        except TypeError:
+            vmin = [vmin] * len(variables)
+
+        try:
+            if len(vmax) != len(variables):
+                raise ValueError('if vmin is a sequence, it must have the same number '
+                                 'of elements as "variables"')
+        except TypeError:
+            vmax = [vmax] * len(variables)
+
+        blocks = []
+        for v, ax, this_vmin, this_vmax in zip(variables, axes.flatten(), vmin, vmax):
+
+            if isinstance(v, str):
+                v = self.data[v]
+
+            data = v.bout.data
+            ndims = len(data.dims)
+            ax.set_title(data.name)
+
+            if ndims == 2:
+                blocks.append(animate_line(data=data, ax=ax, animate_over=animate_over,
+                                           animate=False, **kwargs))
+            elif ndims == 3:
+                if this_vmin is None:
+                    this_vmin = data.min().values
+                if this_vmax is None:
+                    this_vmax = data.max().values
+
+                norm = _create_norm(logscale, kwargs.get('norm', None), this_vmin,
+                                    this_vmax)
+
+                if poloidal_plot:
+                    var_blocks = animate_poloidal(data, ax=ax,
+                                                  animate_over=animate_over,
+                                                  animate=False, vmin=this_vmin,
+                                                  vmax=this_vmax, norm=norm, **kwargs)
+                    for block in var_blocks:
+                        blocks.append(block)
+                else:
+                    blocks.append(animate_pcolormesh(data=data, ax=ax,
+                                                     animate_over=animate_over,
+                                                     animate=False, vmin=this_vmin,
+                                                     vmax=this_vmax, norm=norm,
+                                                     **kwargs))
+            else:
+                raise ValueError("Unsupported number of dimensions "
+                                 + str(ndims) + ". Dims are " + str(v.dims))
+
+        timeline = amp.Timeline(np.arange(v.sizes[animate_over]), fps=fps)
+        anim = amp.Animation(blocks, timeline)
+        anim.controls(timeline_slider_args={'text': animate_over})
+
+        if save_as is not None:
+            anim.save(save_as + '.gif', writer='imagemagick')
+
+        if show:
+            plt.show()
+
+        return anim
+
 
 def _find_time_dependent_vars(data):
     evolving_vars = set(var for var in data.data_vars if 't' in data[var].dims)