From 4a53405958c57c9e8d578850cee462071d1d78bf Mon Sep 17 00:00:00 2001 From: DanRyanIrish Date: Thu, 22 Mar 2018 10:37:58 -0400 Subject: [PATCH] Add NDCubeSequence ImageAnimator classes to plotting mixing. --- ndcube/mixins/sequence_plotting.py | 182 ++++++++++++++++++++++++++++- 1 file changed, 178 insertions(+), 4 deletions(-) diff --git a/ndcube/mixins/sequence_plotting.py b/ndcube/mixins/sequence_plotting.py index 09c354d4a..5e5fd4a64 100644 --- a/ndcube/mixins/sequence_plotting.py +++ b/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):