Add NDCubeSequence ImageAnimator classes to plotting mixing.

ndcube/mixins/sequence_plotting.py

@@ -1,11 +1,10 @@
 import numpy as np
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-
 import astropy.units as u
+from sunpy.visualization.imageanimator import ImageAnimatorWCS
 
 from ndcube import utils
-from ndcube.visualization import animation as ani
 
 __all__ = ['NDCubePlotMixin']
 
@@ -426,7 +425,7 @@ def _animate_ND_sequence(self, cubesequence, *args, **kwargs):
         Visualizes an NDCubeSequence of >2D NDCubes as 2D an animation with N-2 sliders.
 
         """
-        return ani.ImageAnimatorNDCubeSequence(cubesequence, *args, **kwargs)
+        return ImageAnimatorNDCubeSequence(cubesequence, *args, **kwargs)
 
     def _animate_ND_sequence_as_Nminus1Danimation(self, cubesequence, *args, **kwargs):
         """
@@ -435,7 +434,182 @@ def _animate_ND_sequence_as_Nminus1Danimation(self, cubesequence, *args, **kwarg
         Called if plot_as_cube=True.
 
         """
-        return ani.ImageAnimatorCommonAxisNDCubeSequence(cubesequence, *args, **kwargs)
+        return ImageAnimatorCommonAxisNDCubeSequence(cubesequence, *args, **kwargs)
+
+
+class ImageAnimatorNDCubeSequence(ImageAnimatorWCS):
+    """
+    Animates N-dimensional data with the associated astropy WCS object.
+
+    The following keyboard shortcuts are defined in the viewer:
+
+    left': previous step on active slider
+    right': next step on active slider
+    top': change the active slider up one
+    bottom': change the active slider down one
+    'p': play/pause active slider
+
+    This viewer can have user defined buttons added by specifying the labels
+    and functions called when those buttons are clicked as keyword arguments.
+
+    Parameters
+    ----------
+    seq: `ndcube.datacube.CubeSequence`
+        The list of cubes.
+
+    image_axes: `list`
+        The two axes that make the image
+
+    fig: `matplotlib.figure.Figure`
+        Figure to use
+
+    axis_ranges: list of physical coordinates for array or None
+        If None array indices will be used for all axes.
+        If a list it should contain one element for each axis of the numpy array.
+        For the image axes a [min, max] pair should be specified which will be
+        passed to :func:`matplotlib.pyplot.imshow` as extent.
+        For the slider axes a [min, max] pair can be specified or an array the
+        same length as the axis which will provide all values for that slider.
+        If None is specified for an axis then the array indices will be used
+        for that axis.
+
+    interval: `int`
+        Animation interval in ms
+
+    colorbar: `bool`
+        Plot colorbar
+
+    button_labels: `list`
+        List of strings to label buttons
+
+    button_func: `list`
+        List of functions to map to the buttons
+
+    unit_x_axis: `astropy.units.Unit`
+        The unit of x axis.
+
+    unit_y_axis: `astropy.units.Unit`
+        The unit of y axis.
+
+    Extra keywords are passed to imshow.
+
+    """
+    def __init__(self, seq, wcs=None, **kwargs):
+        if wcs is None:
+            wcs = seq[0].wcs
+        self.sequence = seq.data
+        self.cumul_cube_lengths = np.cumsum(np.ones(len(self.sequence)))
+        data_concat = np.stack([cube.data for cube in seq.data])
+        # Add dimensions of length 1 of concatenated data array
+        # shape for an missing axes.
+        if seq[0].wcs.naxis != len(seq.dimensions) - 1:
+            new_shape = list(data_concat.shape)
+            for i in np.arange(seq[0].wcs.naxis)[seq[0].missing_axis[::-1]]:
+                new_shape.insert(i+1, 1)
+            data_concat = data_concat.reshape(new_shape)
+        # Add dummy axis to WCS object to represent sequence axis.
+        new_wcs = utils.wcs.append_sequence_axis_to_wcs(wcs)
+
+        super(ImageAnimatorNDCubeSequence, self).__init__(data_concat, wcs=new_wcs, **kwargs)
+
+
+class ImageAnimatorCommonAxisNDCubeSequence(ImageAnimatorWCS):
+    """
+    Animates N-dimensional data with the associated astropy WCS object.
+
+    The following keyboard shortcuts are defined in the viewer:
+
+    left': previous step on active slider
+    right': next step on active slider
+    top': change the active slider up one
+    bottom': change the active slider down one
+    'p': play/pause active slider
+
+    This viewer can have user defined buttons added by specifying the labels
+    and functions called when those buttons are clicked as keyword arguments.
+
+    Parameters
+    ----------
+    seq: `ndcube.datacube.CubeSequence`
+        The list of cubes.
+
+    image_axes: `list`
+        The two axes that make the image
+
+    fig: `matplotlib.figure.Figure`
+        Figure to use
+
+    axis_ranges: list of physical coordinates for array or None
+        If None array indices will be used for all axes.
+        If a list it should contain one element for each axis of the numpy array.
+        For the image axes a [min, max] pair should be specified which will be
+        passed to :func:`matplotlib.pyplot.imshow` as extent.
+        For the slider axes a [min, max] pair can be specified or an array the
+        same length as the axis which will provide all values for that slider.
+        If None is specified for an axis then the array indices will be used
+        for that axis.
+
+    interval: `int`
+        Animation interval in ms
+
+    colorbar: `bool`
+        Plot colorbar
+
+    button_labels: `list`
+        List of strings to label buttons
+
+    button_func: `list`
+        List of functions to map to the buttons
+
+    unit_x_axis: `astropy.units.Unit`
+        The unit of x axis.
+
+    unit_y_axis: `astropy.units.Unit`
+        The unit of y axis.
+
+    Extra keywords are passed to imshow.
+
+    """
+    def __init__(self, seq, wcs=None, **kwargs):
+        if seq._common_axis is None:
+            raise ValueError("Common axis must be set to use this class. "
+                             "Use ImageAnimatorNDCubeSequence.")
+        if wcs is None:
+            wcs = seq[0].wcs
+        self.sequence = seq.data
+        self.cumul_cube_lengths = np.cumsum(np.array(
+            [c.dimensions[0].value for c in self.sequence], dtype=int))
+        data_concat = np.concatenate([cube.data for cube in seq.data], axis=seq._common_axis)
+        # Add dimensions of length 1 of concatenated data array
+        # shape for an missing axes.
+        if seq[0].wcs.naxis != len(seq.dimensions) - 1:
+            new_shape = list(data_concat.shape)
+            for i in np.arange(seq[0].wcs.naxis)[seq[0].missing_axis[::-1]]:
+                new_shape.insert(i, 1)
+            data_concat = data_concat.reshape(new_shape)
+
+        super(ImageAnimatorCommonAxisNDCubeSequence, self).__init__(
+            data_concat, wcs=wcs, **kwargs)
+
+    def update_plot(self, val, im, slider):
+        val = int(val)
+        ax_ind = self.slider_axes[slider.slider_ind]
+        ind = np.argmin(np.abs(self.axis_ranges[ax_ind] - val))
+        self.frame_slice[ax_ind] = ind
+        list_slices_wcsaxes = list(self.slices_wcsaxes)
+        sequence_slice = utils.sequence._convert_cube_like_index_to_sequence_slice(
+            val, self.cumul_cube_lengths)
+        sequence_index = sequence_slice.sequence_index
+        cube_index = sequence_slice.common_axis_item
+        list_slices_wcsaxes[self.wcs.naxis-ax_ind-1] = cube_index
+        self.slices_wcsaxes = list_slices_wcsaxes
+        if val != slider.cval:
+            self.axes.reset_wcs(
+                wcs=self.sequence[sequence_index].wcs, slices=self.slices_wcsaxes)
+            self._set_unit_in_axis(self.axes)
+            im.set_array(self.data[self.frame_slice])
+            slider.cval = val
+
 
 
 def _determine_sequence_units(cubesequence_data, unit=None):