diff --git a/animatplot/__init__.py b/animatplot/__init__.py index 06ffa3a..fa99c80 100644 --- a/animatplot/__init__.py +++ b/animatplot/__init__.py @@ -2,3 +2,4 @@ from .timeline import Timeline from .animations.animation import Animation from . import blocks, util +from . import xanim diff --git a/animatplot/examples/__init__.py b/animatplot/examples/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/animatplot/examples/xanim/__init__.py b/animatplot/examples/xanim/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/animatplot/examples/xanim/imshow.py b/animatplot/examples/xanim/imshow.py new file mode 100644 index 0000000..6fb9b45 --- /dev/null +++ b/animatplot/examples/xanim/imshow.py @@ -0,0 +1,21 @@ +import numpy as np +import xarray as xr +import matplotlib.pyplot as plt +from animatplot.xanim import animated_plot + + +# Create data set +x = np.linspace(-2, 2, 41) +y = np.linspace(-2, 2, 41) +t = np.linspace(0, 2*np.pi, 30) +X, Y, T = np.meshgrid(x, y, t) +data = np.sin(X*X+Y*Y-T) +da = xr.DataArray(data, coords=[('horizontal position', x), + ('vertical position', y), + ('time', t)]) + +# Create animated 2D plot +anim, block, timeline = animated_plot(da) +anim.controls() +anim.save_gif('../xarray_imshow') +plt.show() diff --git a/animatplot/examples/xarray_imshow.gif b/animatplot/examples/xarray_imshow.gif new file mode 100644 index 0000000..352e633 Binary files /dev/null and b/animatplot/examples/xarray_imshow.gif differ diff --git a/animatplot/xanim.py b/animatplot/xanim.py new file mode 100644 index 0000000..e3938a5 --- /dev/null +++ b/animatplot/xanim.py @@ -0,0 +1,61 @@ +import animatplot as amp +import xarray as xr + + +def animated_plot(da, anim_over='__guess__', x='x', y='y', plot_type='imshow', fps=10): + """ + Function to plot an animated plot from an xarray DataArray. + + The intention is that this would form the basis of a xr.DataArray.plot.animate() method. + """ + + if not isinstance(da, xr.DataArray): + raise TypeError('First argument must be an xarray DataArray object') + + timeline, evolving_coord = _timeline_from_coords(da.coords, anim_over, fps) + t_axis = da.dims.index(evolving_coord) + + if len(da.dims) is not 3: + raise NotImplementedError('Currently only plots 2D dataarrays') + + if plot_type is 'imshow': + # TODO ideally the blocks method should call the xarray plotting method da.plot.imshow() + print(da.values.shape) + print(t_axis) + block = amp.blocks.Imshow(da.values, t_axis=t_axis) + else: + # TODO deal with other kinds of plot (lines, quiver...) here + raise NotImplementedError('Currently only plots using imshow') + + # TODO if we can't call the xarray plotting method directly then add titles, axes labels etc here + anim = amp.Animation([block]) + + return anim, block, timeline + + +def _timeline_from_coords(coords, anim_over, fps): + """Create the animatplot Timeline object from the information in the coordinates attribute of the DataArray.""" + + # Determine coordinate over which plot should be animated + if anim_over is '__guess__': + # Attempt to automatically determine coordinate over which to animate plot + guesses = {'t', 'T', 'time', 'Time'} + matches = list(guesses & set(coords)) + if len(matches) is 1: + evolving_coord = matches[0] + else: + raise ValueError('Could not automatically determine coordinate to animate over - ' + 'multiple possibilities found: ' + str(matches)) + elif anim_over in coords: + evolving_coord = anim_over + else: + raise ValueError('Could not determine coordinate to animate over') + + print(evolving_coord) + t = coords[evolving_coord].values + + # TODO Attempt to determine units from metadata in dataarray attributes, according to CF conventions + + timeline = amp.Timeline(t, units=None, fps=fps) + + return timeline, evolving_coord