Implement 1D animations

.travis.yml

@@ -1,7 +1,7 @@
 language: python
 python:
-  - "3.5"
   - "3.6"
+  - "3.7"
 install:
   - pip install --upgrade setuptools pip pytest pytest-cov coverage codecov
   - pip install -r requirements.txt

requirements.txt

@@ -1,6 +1,7 @@
 xarray >= 0.12.2
 dask[array] >= 1.0.0
 natsort >= 5.5.0
-matplotlib >= 2.2
+matplotlib >= 3.0.3
 animatplot >= 0.3
 netcdf4 >= 1.4.0
+Pillow >= 6.1.0

setup.py

@@ -19,14 +19,15 @@
     author_email="thomas.nicholas@york.ac.uk",
     description='Collect data from BOUT++ runs in python using xarray',
     license="Apache",
-    python_requires='>=3.5',
+    python_requires='>=3.6',
     install_requires=[
         'xarray>=v0.12.2',
         'dask[array]>=1.0.0',
         'natsort>=5.5.0',
-        'matplotlib>=2.2',
+        'matplotlib>=3.0.3',
         'animatplot>=0.3',
         'netcdf4>=1.4.0',
+        'Pillow>=6.1.0'
     ],
     extras_require={
         'tests': ['pytest >= 3.3.0'],
@@ -42,7 +43,6 @@
         "License :: OSI Approved :: Apache License",
         "Natural Language :: English",
         "Operating System :: POSIX :: Linux",
-        "Programming Language :: Python :: 3.5",
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Topic :: Scientific/Engineering :: Visualization"

xbout/boutdataarray.py

@@ -97,7 +97,6 @@ def animate1D(self, animate_over='t', x='x', y='y', animate=True,
             print("{} data passed has {} dimensions - will use "
                   "animatplot.blocks.Line()".format(variable, str(n_dims)))
             line_block = animate_line(data=data, animate_over=animate_over,
-                                      x=x, y=y, sep_pos=sep_pos,
-                                      animate=animate, fps=fps,
+                                      sep_pos=sep_pos, animate=animate, fps=fps,
                                       save_as=save_as, ax=ax, **kwargs)
             return line_block

xbout/plotting/animate.py

@@ -4,10 +4,10 @@
 import animatplot as amp
 
 from .utils import plot_separatrix
-
+from matplotlib.animation import PillowWriter
 
 def animate_imshow(data, animate_over='t', x='x', y='y', animate=True,
-                   vmin='min', vmax='max', fps=10, save_as=None,
+                   vmin=None, vmax=None, fps=10, save_as=None,
                    sep_pos=None, ax=None, **kwargs):
     """
     Plots a color plot which is animated with time over the specified
@@ -60,9 +60,9 @@ def animate_imshow(data, animate_over='t', x='x', y='y', animate=True,
     image_data = data.values
 
     # If not specified, determine max and min values across entire data series
-    if vmax is 'max':
+    if vmax is None:
         vmax = np.max(image_data)
-    if vmin is 'min':
+    if vmin is None:
         vmin = np.min(image_data)
 
     if not ax:
@@ -93,16 +93,89 @@ def animate_imshow(data, animate_over='t', x='x', y='y', animate=True,
 
         if not save_as:
             save_as = "{}_over_{}".format(variable, animate_over)
-        # TODO save using PillowWriter instead once matplotlib 3.1 comes out
-        # see https://github.com/t-makaro/animatplot/issues/24
-        anim.save(save_as + '.gif', writer='imagemagick')
+        anim.save(save_as + '.gif', writer=PillowWriter(fps=fps))
 
     return imshow_block
 
 
-def animate_line(data, animate_over='t', x='x', y='y', animate=True,
-                 fps=10, save_as=None, sep_pos=None, ax=None, **kwargs):
+def animate_line(data, animate_over='t', animate=True,
+                 vmin=None, vmax=None, fps=10, save_as=None, sep_pos=None, ax=None,
+                 **kwargs):
+    """
+    Plots a line plot which is animated with time.
+
+    Currently only supports 1D+1 data, which it plots with xarray's
+    wrapping of matplotlib's plot.
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+    animate_over : str, optional
+        Dimension over which to animate
+    vmin : float, optional
+        Minimum value to use for colorbar. Default is to use minimum value of
+        data across whole timeseries.
+    vmax : float, optional
+        Maximum value to use for colorbar. Default is to use maximum value of
+        data across whole timeseries.
+    sep_pos : int, optional
+        Radial position at which to plot the separatrix
+    save_as: str, optional
+        Filename to give to the resulting gif
+    fps : int, optional
+        Frames per second of resulting gif
+    kwargs : dict, optional
+        Additional keyword arguments are passed on to the plotting function
+        (e.g. imshow for 2D plots).
+    """
+
     variable = data.name
 
+    # Check plot is the right orientation
     t_read, x_read = data.dims
-    raise NotImplementedError
+    if (t_read is animate_over):
+        pass
+    else:
+        data = data.transpose(x_read, animate_over)
+
+    # Load values eagerly otherwise for some reason the plotting takes
+    # 100's of times longer - for some reason animatplot does not deal
+    # well with dask arrays!
+    image_data = data.values
+
+    # If not specified, determine max and min values across entire data series
+    if vmax is None:
+        vmax = np.max(image_data)
+    if vmin is None:
+        vmin = np.min(image_data)
+
+    if not ax:
+        fig, ax = plt.subplots()
+
+    # set range of plot
+    ax.set_ylim([vmin, vmax])
+
+    line_block = amp.blocks.Line(image_data, ax=ax, **kwargs)
+
+    timeline = amp.Timeline(np.arange(data.sizes[animate_over]), fps=fps)
+
+    if animate:
+        anim = amp.Animation([line_block], timeline)
+
+    # Add title and axis labels
+    ax.set_title("{} variation over {}".format(variable, animate_over))
+    ax.set_xlabel(x_read)
+    ax.set_ylabel(variable)
+
+    # Plot separatrix
+    if sep_pos:
+        ax.plot_vline(sep_pos, '--')
+
+    if animate:
+        anim.controls(timeline_slider_args={'text': animate_over})
+
+        if not save_as:
+            save_as = "{}_over_{}".format(variable, animate_over)
+        anim.save(save_as + '.gif', writer=PillowWriter(fps=fps))
+
+    return line_block

xbout/tests/test_animate.py

@@ -0,0 +1,44 @@
+import pytest
+
+from xbout import open_boutdataset
+from xbout.boutdataarray import BoutDataArrayAccessor
+from .test_load import create_bout_ds_list
+
+from animatplot.blocks import Imshow, Line
+
+
+@pytest.fixture
+def create_test_file(tmpdir_factory):
+
+    # Create temp dir for output of animate1D/2D
+    save_dir = tmpdir_factory.mktemp("test_data")
+
+    # Generate some test data
+    ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=3, nype=3,
+                                             syn_data_type="linear")
+    for ds, file_name in zip(ds_list, file_list):
+        ds.to_netcdf(str(save_dir.join(str(file_name))))
+
+    ds = open_boutdataset(save_dir.join("BOUT.dmp.*.nc")) # Open test data
+
+    return save_dir, ds
+
+
+class TestAnimate:
+    """
+    Set of tests to check whether animate1D() and animate2D() are running properly
+    and PillowWriter is saving each animation correctly
+    """
+    def test_animate2D(self, create_test_file):
+
+        save_dir, ds = create_test_file
+        animation = ds['n'].isel(y=1).bout.animate2D(y='z', save_as="%s/test" % save_dir)
+
+        assert isinstance(animation, Imshow)
+
+    def test_animate1D(self, create_test_file):
+
+        save_dir, ds = create_test_file
+        animation = ds['n'].isel(y=2, z=0).bout.animate1D(save_as="%s/test" % save_dir)
+
+        assert isinstance(animation, Line)