diff --git a/.travis.yml b/.travis.yml index 284cfe99c..b0e33ae37 100644 --- a/.travis.yml +++ b/.travis.yml @@ -54,6 +54,9 @@ matrix: include: + # Add a matplotlib backend for generating plots in tests. + - env: MPLBACKEND='Agg' + # Do a basic test. - os: linux stage: Initial tests @@ -99,8 +102,13 @@ install: - git clone git://github.com/astropy/ci-helpers.git - source ci-helpers/travis/setup_conda.sh +before_script: + - "export DISPLAY=:99.0" + - "sh -e /etc/init.d/xvfb start" + - sleep 3 # give xvfb some time to start + script: - - $MAIN_CMD $SETUP_CMD + - $MAIN_CMD $SETUP_CMD after_success: - - if [[ $SETUP_CMD = *"cov"* ]]; then pip install codecov; codecov; fi + - if [[ $SETUP_CMD = *"cov"* ]]; then pip install codecov; codecov; fi diff --git a/ndcube/mixins/sequence_plotting.py b/ndcube/mixins/sequence_plotting.py new file mode 100644 index 000000000..f304dd403 --- /dev/null +++ b/ndcube/mixins/sequence_plotting.py @@ -0,0 +1,1001 @@ +import numpy as np +import matplotlib as mpl +import matplotlib.pyplot as plt +import astropy.units as u +from sunpy.visualization.imageanimator import ImageAnimatorWCS, LineAnimator + +from ndcube import utils + +__all__ = ['NDCubePlotMixin'] + +NON_COMPATIBLE_UNIT_MESSAGE = \ + "All sequence sub-cubes' unit attribute are not compatible with data_unit set by user." +AXES_UNIT_ERRONESLY_SET_MESSAGE = \ + "axes_unit element must be None unless corresponding axes_coordinate is None or a Quantity." + + +class NDCubeSequencePlotMixin: + def plot(self, axes=None, plot_axis_indices=None, + axes_coordinates=None, axes_units=None, data_unit=None, **kwargs): + """ + Visualizes data in the NDCubeSequence with the sequence axis as a separate dimension. + + Based on the dimensionality of the sequence and value of plot_axis_indices kwarg, + a Line/Image Animation/Plot is produced. + + Parameters + ---------- + axes: `astropy.visualization.wcsaxes.core.WCSAxes` or ??? or None. + The axes to plot onto. If None the current axes will be used. + + plot_axis_indices: `int` or iterable of one or two `int`. + If two axis indices are given, the sequence is visualized as an image or + 2D animation, assuming the sequence has at least 2 dimensions. + The dimension indicated by the 0th index is displayed on the + x-axis while the dimension indicated by the 1st index is displayed on the y-axis. + If only one axis index is given (either as an int or a list of one int), + then a 1D line animation is produced with the indicated dimension on the x-axis + and other dimensions represented by animations sliders. + Default=[-1, -2]. If sequence only has one dimension, + plot_axis_indices is ignored and a static 1D line plot is produced. + + axes_coordinates: `None` or `list` of `None` `astropy.units.Quantity` `numpy.ndarray` `str` + Denotes physical coordinates for plot and slider axes. + If None coordinates derived from the WCS objects will be used for all axes. + If a list, its length should equal either the number sequence dimensions or + the length of plot_axis_indices. + If the length equals the number of sequence dimensions, each element describes + the coordinates of the corresponding sequence dimension. + If the length equals the length of plot_axis_indices, + the 0th entry describes the coordinates of the x-axis + while (if length is 2) the 1st entry describes the coordinates of the y-axis. + Slider axes are implicitly set to None. + If the number of sequence dimensions equals the length of plot_axis_indices, + the latter convention takes precedence. + The value of each entry should be either + `None` (implies derive the coordinates from the WCS objects), + an `astropy.units.Quantity` or a `numpy.ndarray` of coordinates for each pixel, + or a `str` denoting a valid extra coordinate. + + axes_units: `None or `list` of `None`, `astropy.units.Unit` and/or `str` + If None units derived from the WCS objects will be used for all axes. + If a list, its length should equal either the number sequence dimensions or + the length of plot_axis_indices. + If the length equals the number of sequence dimensions, each element gives the + unit in which the coordinates along the corresponding sequence dimension should + displayed whether they be a plot axes or a slider axes. + If the length equals the length of plot_axis_indices, + the 0th entry describes the unit in which the x-axis coordinates should be displayed + while (if length is 2) the 1st entry describes the unit in which the y-axis should + be displayed. Slider axes are implicitly set to None. + If the number of sequence dimensions equals the length of plot_axis_indices, + the latter convention takes precedence. + The value of each entry should be either + `None` (implies derive the unit from the WCS object of the 0th sub-cube), + `astropy.units.Unit` or a valid unit `str`. + + data_unit: `astropy.unit.Unit` or valid unit `str` or None + Unit in which data be displayed. If the length of plot_axis_indices is 2, + a 2D image/animation is produced and data_unit determines the unit represented by + the color table. If the length of plot_axis_indices is 1, + a 1D plot/animation is produced and data_unit determines the unit in which the + y-axis is displayed. + + Returns + ------- + ax: `matplotlib.axes.Axes`, `ndcube.mixins.sequence_plotting.ImageAnimatorNDCubeSequence` or `ndcube.mixins.sequence_plotting.ImageAnimatorCubeLikeNDCubeSequence` + Axes or animation object depending on dimensionality of NDCubeSequence + + """ + # Check kwargs are in consistent formats and set default values if not done so by user. + naxis = len(self.dimensions) + plot_axis_indices, axes_coordinates, axes_units = _prep_axes_kwargs( + naxis, plot_axis_indices, axes_coordinates, axes_units) + if naxis == 1: + # Make 1D line plot. + ax = self._plot_1D_sequence(axes_coordinates, + axes_units, data_unit, **kwargs) + else: + if len(plot_axis_indices) == 1: + # Since sequence has more than 1 dimension and number of plot axes is 1, + # produce a 1D line animation. + raise NotImplementedError() + elif len(plot_axis_indices) == 2: + if naxis == 2: + # Since sequence has 2 dimensions and number of plot axes is 2, + # produce a 2D image. + ax = self._plot_2D_sequence(plot_axis_indices, axes_coordinates, + axes_units, data_unit, **kwargs) + else: + # Since sequence has more than 2 dimensions and number of plot axes is 2, + # produce a 2D animation. + ax = ImageAnimatorNDCubeSequence( + self, plot_axis_indices=plot_axis_indices, + axes_coordinates=axes_coordinates, axes_units=axes_units, **kwargs) + + return ax + + def plot_as_cube(self, axes=None, plot_axis_indices=None, + axes_coordinates=None, axes_units=None, data_unit=None, **kwargs): + """ + Visualizes data in the NDCubeSequence with the sequence axis folded into the common axis. + + Based on the cube-like dimensionality of the sequence and value of plot_axis_indices + kwarg, a Line/Image Plot/Animation is produced. + + Parameters + ---------- + axes: `astropy.visualization.wcsaxes.core.WCSAxes` or ??? or None. + The axes to plot onto. If None the current axes will be used. + + plot_axis_indices: `int` or iterable of one or two `int`. + If two axis indices are given, the sequence is visualized as an image or + 2D animation, assuming the sequence has at least 2 cube-like dimensions. + The cube-like dimension indicated by the 0th index is displayed on the + x-axis while the cube-like dimension indicated by the 1st index is + displayed on the y-axis. If only one axis index is given (either as an int + or a list of one int), then a 1D line animation is produced with the indicated + cube-like dimension on the x-axis and other cube-like dimensions represented + by animations sliders. + Default=[-1, -2]. If sequence only has one cube-like dimension, + plot_axis_indices is ignored and a static 1D line plot is produced. + + axes_coordinates: `None` or `list` of `None` `astropy.units.Quantity` `numpy.ndarray` `str` + Denotes physical coordinates for plot and slider axes. + If None coordinates derived from the WCS objects will be used for all axes. + If a list, its length should equal either the number cube-like dimensions or + the length of plot_axis_indices. + If the length equals the number of cube-like dimensions, each element describes + the coordinates of the corresponding cube-like dimension. + If the length equals the length of plot_axis_indices, + the 0th entry describes the coordinates of the x-axis + while (if length is 2) the 1st entry describes the coordinates of the y-axis. + Slider axes are implicitly set to None. + If the number of cube-like dimensions equals the length of plot_axis_indices, + the latter convention takes precedence. + The value of each entry should be either + `None` (implies derive the coordinates from the WCS objects), + an `astropy.units.Quantity` or a `numpy.ndarray` of coordinates for each pixel, + or a `str` denoting a valid extra coordinate. + + axes_units: `None or `list` of `None`, `astropy.units.Unit` and/or `str` + If None units derived from the WCS objects will be used for all axes. + If a list, its length should equal either the number cube-like dimensions or + the length of plot_axis_indices. + If the length equals the number of cube-like dimensions, each element gives the + unit in which the coordinates along the corresponding cube-like dimension should + displayed whether they be a plot axes or a slider axes. + If the length equals the length of plot_axis_indices, + the 0th entry describes the unit in which the x-axis coordinates should be displayed + while (if length is 2) the 1st entry describes the unit in which the y-axis should + be displayed. Slider axes are implicitly set to None. + If the number of cube-like dimensions equals the length of plot_axis_indices, + the latter convention takes precedence. + The value of each entry should be either + `None` (implies derive the unit from the WCS object of the 0th sub-cube), + `astropy.units.Unit` or a valid unit `str`. + + data_unit: `astropy.unit.Unit` or valid unit `str` or None + Unit in which data be displayed. If the length of plot_axis_indices is 2, + a 2D image/animation is produced and data_unit determines the unit represented by + the color table. If the length of plot_axis_indices is 1, + a 1D plot/animation is produced and data_unit determines the unit in which the + y-axis is displayed. + + Returns + ------- + ax: ax: `matplotlib.axes.Axes`, `ndcube.mixins.sequence_plotting.ImageAnimatorNDCubeSequence` or `ndcube.mixins.sequence_plotting.ImageAnimatorCubeLikeNDCubeSequence` + Axes or animation object depending on dimensionality of NDCubeSequence + + """ + # Verify common axis is set. + if self._common_axis is None: + raise TypeError("Common axis must be set.") + # Check kwargs are in consistent formats and set default values if not done so by user. + naxis = len(self.cube_like_dimensions) + plot_axis_indices, axes_coordinates, axes_units = _prep_axes_kwargs( + naxis, plot_axis_indices, axes_coordinates, axes_units) + # Produce plot/image/animation based on cube-like dimensions of sequence. + if naxis == 1: + # Since sequence has 1 cube-like dimension, produce a 1D line plot. + ax = self._plot_2D_sequence_as_1Dline(axes_coordinates, axes_units, data_unit, + **kwargs) + else: + if len(plot_axis_indices) == 1: + # Since sequence has more than 1 cube-like dimension and + # number of plot axes is 1, produce a 1D line animation. + raise NotImplementedError() + elif len(plot_axis_indices) == 2: + if naxis == 2: + # Since sequence has 2 cube-like dimensions and + # number of plot axes is 2, produce a 2D image. + ax = self._plot_3D_sequence_as_2Dimage(axes, plot_axis_indices, + axes_coordinates, axes_units, + data_unit, **kwargs) + else: + # Since sequence has more than 2 cube-like dimensions and + # number of plot axes is 2, produce a 2D animation. + ax = ImageAnimatorCubeLikeNDCubeSequence( + self, plot_axis_indices=plot_axis_indices, + axes_coordinates=axes_coordinates, axes_units=axes_units, **kwargs) + return ax + + def _plot_1D_sequence(self, axes_coordinates=None, + axes_units=None, data_unit=None, **kwargs): + """ + Visualizes an NDCubeSequence of scalar NDCubes as a line plot. + + A scalar NDCube is one whose NDCube.data is a scalar rather than an array. + + Parameters + ---------- + axes_coordinates: `numpy.ndarray` `astropy.unit.Quantity` `str` `None` or length 1 `list` + Denotes the physical coordinates of the x-axis. + If list, must be of length 1 containing one object of one of the other allowed types. + If None, coordinates are derived from the WCS objects. + If an `astropy.units.Quantity` or a `numpy.ndarray` gives the coordinates for + each pixel along the x-axis. + If a `str`, denotes the extra coordinate to be used. The extra coordinate must + correspond to the sequence axis. + + axes_units: `astropy.unit.Unit` or valid unit `str` or length 1 `list` of those types. + Unit in which X-axis should be displayed. Must be compatible with the unit of + the coordinate denoted by x_axis_range. Not used if x_axis_range is a + `numpy.ndarray` or the designated extra coordinate is a `numpy.ndarray` + + data_unit: `astropy.units.unit` or valid unit `str` + The units into which the y-axis should be displayed. The unit attribute of all + the sub-cubes must be compatible to set this kwarg. + + """ + # Derive x-axis coordinates and unit from inputs. + x_axis_coordinates, unit_x_axis = _derive_1D_coordinates_and_units(axes_coordinates, + axes_units) + # Check that the unit attribute is a set in all cubes and derive unit_y_axis if not set. + unit_y_axis = data_unit + sequence_units, unit_y_axis = _determine_sequence_units(self.data, unit_y_axis) + # If not all cubes have their unit set, create a data array from cube's data. + if sequence_units is None: + ydata = np.array([cube.data for cube in self.data]) + else: + # If all cubes have unit set, create a data quantity from cubes' data. + ydata = u.Quantity([cube.data * sequence_units[i] + for i, cube in enumerate(self.data)], unit=unit_y_axis).value + # Determine uncertainties. + sequence_uncertainty_nones = [] + for i, cube in enumerate(self.data): + if cube.uncertainty is None: + sequence_uncertainty_nones.append(i) + if sequence_uncertainty_nones == list(range(len(self.data))): + # If all cube uncertainties are None, make yerror also None. + yerror = None + else: + # Else determine uncertainties, giving 0 uncertainty for + # cubes with uncertainty of None. + if sequence_units is None: + yerror = np.array([cube.uncertainty.array for cube in self.data]) + yerror[sequence_uncertainty_nones] = 0. + else: + # If all cubes have compatible units, ensure uncertainties are in the same unit. + yerror = [] + for i, cube in enumerate(self.data): + if i in sequence_uncertainty_nones: + yerror.append(0. * sequence_units[i]) + else: + yerror.append(cube.uncertainty.array * sequence_units[i]) + yerror = u.Quantity(yerror, unit=unit_y_axis).value + # Define x-axis data. + if x_axis_coordinates is None: + # Since scalar NDCubes have no array/pixel indices, WCS translations don't work. + # Therefore x-axis values will be unitless sequence indices unless supplied by user + # or an extra coordinate is designated. + xdata = np.arange(int(self.dimensions[0].value)) + xname = self.world_axis_physical_types[0] + elif isinstance(x_axis_coordinates, str): + xdata = self.sequence_axis_extra_coords[x_axis_coordinates] + xname = x_axis_coordinates + else: + xdata = x_axis_coordinates + xname = self.world_axis_physical_types[0] + if isinstance(xdata, u.Quantity): + if unit_x_axis is None: + unit_x_axis = xdata.unit + else: + xdata = xdata.to(unit_x_axis) + else: + unit_x_axis = None + default_xlabel = "{0} [{1}]".format(xname, unit_x_axis) + fig, ax = _make_1D_sequence_plot(xdata, ydata, yerror, unit_y_axis, default_xlabel, kwargs) + return ax + + def _plot_2D_sequence_as_1Dline(self, axes_coordinates=None, + axes_units=None, data_unit=None, **kwargs): + """ + Visualizes an NDCubeSequence of 1D NDCubes with a common axis as a line plot. + + Called if plot_as_cube=True. + + Parameters + ---------- + Same as _plot_1D_sequence + + """ + # Derive x-axis coordinates and unit from inputs. + x_axis_coordinates, unit_x_axis = _derive_1D_coordinates_and_units(axes_coordinates, + axes_units) + # Check that the unit attribute is set of all cubes and derive unit_y_axis if not set. + unit_y_axis = data_unit + sequence_units, unit_y_axis = _determine_sequence_units(self.data, unit_y_axis) + # If all cubes have unit set, create a y data quantity from cube's data. + if sequence_units is None: + ydata = np.concatenate([cube.data for cube in self.data]) + else: + # If all cubes have unit set, create a data quantity from cubes' data. + ydata = np.concatenate([(cube.data * sequence_units[i]).to(unit_y_axis).value + for i, cube in enumerate(self.data)]) + # Determine uncertainties. + # Check which cubes don't have uncertainties. + sequence_uncertainty_nones = [] + for i, cube in enumerate(self.data): + if cube.uncertainty is None: + sequence_uncertainty_nones.append(i) + if sequence_uncertainty_nones == list(range(len(self.data))): + # If no sub-cubes have uncertainty, set overall yerror to None. + yerror = None + else: + # Else determine uncertainties, giving 0 uncertainty for + # cubes with uncertainty of None. + yerror = [] + if sequence_units is None: + for i, cube in enumerate(self.data): + if i in sequence_uncertainty_nones: + yerror.append(np.zeros(cube.data.shape)) + else: + yerror.append(cube.uncertainty.array) + else: + for i, cube in enumerate(self.data): + if i in sequence_uncertainty_nones: + yerror.append((np.zeros(cube.data.shape) * sequence_units[i]).to( + unit_y_axis).value) + else: + yerror.append((cube.uncertainty.array * sequence_units[i]).to( + unit_y_axis).value) + yerror = np.concatenate(yerror) + # Define x-axis data. + if x_axis_coordinates is None: + print('a', unit_x_axis) + if unit_x_axis is None: + print('b', unit_x_axis) + unit_x_axis = np.asarray(self[0].wcs.wcs.cunit)[ + np.invert(self[0].missing_axis)][0] + print('c', unit_x_axis) + xdata = u.Quantity(np.concatenate([cube.axis_world_coords().to(unit_x_axis).value + for cube in self.data]), unit=unit_x_axis) + xname = self.cube_like_world_axis_physical_types[0] + print('d', unit_x_axis) + elif isinstance(x_axis_coordinates, str): + xdata = self.common_axis_extra_coords[x_axis_coordinates] + xname = x_axis_coordinates + else: + xdata = x_axis_coordinates + xname = "" + if isinstance(xdata, u.Quantity): + if unit_x_axis is None: + unit_x_axis = xdata.unit + else: + xdata = xdata.to(unit_x_axis) + else: + unit_x_axis = None + default_xlabel = "{0} [{1}]".format(xname, unit_x_axis) + # For consistency, make xdata an array if a Quantity. Wait until now + # because if xdata is a Quantity, its unit is needed until now. + if isinstance(xdata, u.Quantity): + xdata = xdata.value + # Plot data + fig, ax = _make_1D_sequence_plot(xdata, ydata, yerror, unit_y_axis, default_xlabel, kwargs) + return ax + + def _plot_2D_sequence(self, plot_axis_indices=None, axes_coordinates=None, + axes_units=None, data_unit=None, **kwargs): + """ + Visualizes an NDCubeSequence of 1D NDCubes as a 2D image. + + **kwargs are fed into matplotlib.image.NonUniformImage. + + Parameters + ---------- + Same as self.plot() + + """ + # Set default values of kwargs if not set. + if axes_coordinates is None: + axes_coordinates = [None, None] + if axes_units is None: + axes_units = [None, None] + # Convert plot_axis_indices to array for function operations. + plot_axis_indices = np.asarray(plot_axis_indices) + # Check that the unit attribute is set of all cubes and derive unit_y_axis if not set. + sequence_units, data_unit = _determine_sequence_units(self.data, data_unit) + # If all cubes have unit set, create a data quantity from cube's data. + if sequence_units is not None: + data = np.stack([(cube.data * sequence_units[i]).to(data_unit).value + for i, cube in enumerate(self.data)]) + else: + data = np.stack([cube.data for i, cube in enumerate(self.data)]) + if plot_axis_indices[0] < plot_axis_indices[1]: + # Transpose data if user-defined images_axes require it. + data = data.transpose() + # Determine index of above axes variables corresponding to sequence and cube axes. + # Since the axes variables have been re-oriented before this function was called + # so the 0th element corresponds to the sequence axis, and the 1st to the cube axis, + # determining this is trivial. + sequence_axis_index = 0 + cube_axis_index = 1 + # Derive the coordinates, unit, and default label of the cube axis. + cube_axis_unit = axes_units[cube_axis_index] + if axes_coordinates[cube_axis_index] is None: + if cube_axis_unit is None: + cube_axis_unit = np.array(self[0].wcs.wcs.cunit)[ + np.invert(self[0].missing_axis)][0] + cube_axis_coords = self[0].axis_world_coords().to(cube_axis_unit).value + cube_axis_name = self.world_axis_physical_types[1] + else: + if isinstance(axes_coordinates[cube_axis_index], str): + cube_axis_coords = \ + self[0].extra_coords[axes_coordinates[cube_axis_index]]["value"] + cube_axis_name = axes_coordinates[cube_axis_index] + else: + cube_axis_coords = axes_coordinates[cube_axis_index] + cube_axis_name = "" + if isinstance(cube_axis_coords, u.Quantity): + if cube_axis_unit is None: + cube_axis_unit = cube_axis_coords.unit + cube_axis_coords = cube_axis_coords.value + else: + cube_axis_coords = cube_axis_coords.to(cube_axis_unit).value + else: + if cube_axis_unit is not None: + raise ValueError(AXES_UNIT_ERRONESLY_SET_MESSAGE) + default_cube_axis_label = "{0} [{1}]".format(cube_axis_name, cube_axis_unit) + axes_coordinates[cube_axis_index] = cube_axis_coords + axes_units[cube_axis_index] = cube_axis_unit + # Derive the coordinates, unit, and default label of the sequence axis. + sequence_axis_unit = axes_units[sequence_axis_index] + if axes_coordinates[sequence_axis_index] is None: + sequence_axis_coords = np.arange(len(self.data)) + sequence_axis_name = self.world_axis_physical_types[0] + elif isinstance(axes_coordinates[sequence_axis_index], str): + sequence_axis_coords = \ + self.sequence_axis_extra_coords[axes_coordinates[sequence_axis_index]] + sequence_axis_name = axes_coordinates[sequence_axis_index] + else: + sequence_axis_coords = axes_coordinates[sequence_axis_index] + sequence_axis_name = self.world_axis_physical_types[0] + if isinstance(sequence_axis_coords, u.Quantity): + if sequence_axis_unit is None: + sequence_axis_unit = sequence_axis_coords.unit + sequence_axis_coords = sequence_axis_coords.value + else: + sequence_axis_coords = sequence_axis_coords.to(sequence_axis_unit).value + else: + if sequence_axis_unit is not None: + raise ValueError(AXES_UNIT_ERRONESLY_SET_MESSAGE) + default_sequence_axis_label = "{0} [{1}]".format(sequence_axis_name, sequence_axis_unit) + axes_coordinates[sequence_axis_index] = sequence_axis_coords + axes_units[sequence_axis_index] = sequence_axis_unit + axes_labels = [None, None] + axes_labels[cube_axis_index] = default_cube_axis_label + axes_labels[sequence_axis_index] = default_sequence_axis_label + # Plot image. + # Create figure and axes objects. + fig, ax = plt.subplots(1, 1) + # Since we can't assume the x-axis will be uniform, create NonUniformImage + # axes and add it to the axes object. + im_ax = mpl.image.NonUniformImage(ax, + extent=(axes_coordinates[plot_axis_indices[0]][0], + axes_coordinates[plot_axis_indices[0]][-1], + axes_coordinates[plot_axis_indices[1]][0], + axes_coordinates[plot_axis_indices[1]][-1]), + **kwargs) + im_ax.set_data(axes_coordinates[plot_axis_indices[0]], + axes_coordinates[plot_axis_indices[1]], data) + ax.add_image(im_ax) + # Set the limits, labels, etc. of the axes. + ax.set_xlim((axes_coordinates[plot_axis_indices[0]][0], + axes_coordinates[plot_axis_indices[0]][-1])) + ax.set_ylim((axes_coordinates[plot_axis_indices[1]][0], + axes_coordinates[plot_axis_indices[1]][-1])) + ax.set_xlabel(axes_labels[plot_axis_indices[0]]) + ax.set_ylabel(axes_labels[plot_axis_indices[1]]) + + return ax + + def _plot_3D_sequence_as_2Dimage(self, axes=None, plot_axis_indices=None, + axes_coordinates=None, axes_units=None, data_unit=None, + **kwargs): + """ + Visualizes an NDCubeSequence of 2D NDCubes with a common axis as a 2D image. + + Called if plot_as_cube=True. + + """ + # Set default values of kwargs if not set. + if axes_coordinates is None: + axes_coordinates = [None, None] + if axes_units is None: + axes_units = [None, None] + # Convert plot_axis_indices to array for function operations. + plot_axis_indices = np.asarray(plot_axis_indices) + # Check that the unit attribute is set of all cubes and derive unit_y_axis if not set. + sequence_units, data_unit = _determine_sequence_units(self.data, data_unit) + # If all cubes have unit set, create a data quantity from cube's data. + if sequence_units is not None: + data = np.concatenate([(cube.data * sequence_units[i]).to(data_unit).value + for i, cube in enumerate(self.data)], + axis=self._common_axis) + else: + data = np.concatenate([cube.data for cube in self.data], + axis=self._common_axis) + if plot_axis_indices[0] < plot_axis_indices[1]: + data = data.transpose() + # Determine index of common axis and other cube axis. + common_axis_index = self._common_axis + cube_axis_index = [0, 1] + cube_axis_index.pop(common_axis_index) + cube_axis_index = cube_axis_index[0] + # Derive the coordinates, unit, and default label of the cube axis. + cube_axis_unit = axes_units[cube_axis_index] + if axes_coordinates[cube_axis_index] is None: + if cube_axis_unit is None: + cube_axis_unit = np.array(self[0].wcs.wcs.cunit)[ + np.invert(self[0].missing_axis)][0] + cube_axis_coords = \ + self[0].axis_world_coords()[cube_axis_index].to(cube_axis_unit).value + cube_axis_name = self.cube_like_world_axis_physical_types[1] + else: + if isinstance(axes_coordinates[cube_axis_index], str): + cube_axis_coords = \ + self[0].extra_coords[axes_coordinates[cube_axis_index]]["value"] + cube_axis_name = axes_coordinates[cube_axis_index] + else: + cube_axis_coords = axes_coordinates[cube_axis_index] + cube_axis_name = "" + if isinstance(cube_axis_coords, u.Quantity): + if cube_axis_unit is None: + cube_axis_unit = cube_axis_coords.unit + cube_axis_coords = cube_axis_coords.value + else: + cube_axis_coords = cube_axis_coords.to(cube_axis_unit).value + else: + if cube_axis_unit is not None: + raise ValueError(AXES_UNIT_ERRONESLY_SET_MESSAGE) + default_cube_axis_label = "{0} [{1}]".format(cube_axis_name, cube_axis_unit) + axes_coordinates[cube_axis_index] = cube_axis_coords + axes_units[cube_axis_index] = cube_axis_unit + # Derive the coordinates, unit, and default label of the common axis. + common_axis_unit = axes_units[common_axis_index] + if axes_coordinates[common_axis_index] is None: + # Concatenate values along common axis for each cube. + if common_axis_unit is None: + wcs_common_axis_index = utils.cube.data_axis_to_wcs_axis( + common_axis_index, self[0].missing_axis) + common_axis_unit = np.array(self[0].wcs.wcs.cunit)[wcs_common_axis_index] + common_axis_coords = u.Quantity(np.concatenate( + [cube.axis_world_coords()[common_axis_index].to(common_axis_unit).value + for cube in self.data]), unit=common_axis_unit) + common_axis_name = self.cube_like_world_axis_physical_types[common_axis_index] + elif isinstance(axes_coordinates[common_axis_index], str): + common_axis_coords = \ + self.common_axis_extra_coords[axes_coordinates[common_axis_index]] + common_axis_name = axes_coordinates[common_axis_index] + else: + common_axis_coords = axes_coordinates[common_axis_index] + common_axis_name = "" + if isinstance(common_axis_coords, u.Quantity): + if common_axis_unit is None: + common_axis_unit = common_axis_coords.unit + common_axis_coords = common_axis_coords.value + else: + common_axis_coords = common_axis_coords.to(common_axis_unit).value + else: + if common_axis_unit is not None: + raise ValueError(AXES_UNIT_ERRONESLY_SET_MESSAGE) + default_common_axis_label = "{0} [{1}]".format(common_axis_name, common_axis_unit) + axes_coordinates[common_axis_index] = common_axis_coords + axes_units[common_axis_index] = common_axis_unit + axes_labels = [None, None] + axes_labels[cube_axis_index] = default_cube_axis_label + axes_labels[common_axis_index] = default_common_axis_label + # Plot image. + # Create figure and axes objects. + fig, ax = plt.subplots(1, 1) + # Since we can't assume the x-axis will be uniform, create NonUniformImage + # axes and add it to the axes object. + im_ax = mpl.image.NonUniformImage( + ax, extent=(axes_coordinates[plot_axis_indices[0]][0], + axes_coordinates[plot_axis_indices[0]][-1], + axes_coordinates[plot_axis_indices[1]][0], + axes_coordinates[plot_axis_indices[1]][-1]), + **kwargs) + im_ax.set_data(axes_coordinates[plot_axis_indices[0]], + axes_coordinates[plot_axis_indices[1]], data) + ax.add_image(im_ax) + # Set the limits, labels, etc. of the axes. + ax.set_xlim((axes_coordinates[plot_axis_indices[0]][0], + axes_coordinates[plot_axis_indices[0]][-1])) + ax.set_ylim((axes_coordinates[plot_axis_indices[1]][0], + axes_coordinates[plot_axis_indices[1]][-1])) + ax.set_xlabel(axes_labels[plot_axis_indices[0]]) + ax.set_ylabel(axes_labels[plot_axis_indices[1]]) + + return ax + + +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.NDCubeSequence` + 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, axes=None, plot_axis_indices=None, + axes_coordinates=None, axes_units=None, data_unit=None, **kwargs): + self.sequence = seq.data # Required by parent class. + # Set default values of kwargs if not set. + if wcs is None: + wcs = seq[0].wcs + if axes_coordinates is None: + axes_coordinates = [None] * len(seq.dimensions) + if axes_units is None: + axes_units = [None] * len(seq.dimensions) + # Determine units of each cube in sequence. + sequence_units, data_unit = _determine_sequence_units(seq.data, data_unit) + # If all cubes have unit set, create a data quantity from cube's data. + if sequence_units is None: + data_stack = np.stack([cube.data for i, cube in enumerate(seq.data)]) + else: + data_stack = np.stack([(cube.data * sequence_units[i]).to(data_unit).value + for i, cube in enumerate(seq.data)]) + self.cumul_cube_lengths = np.cumsum(np.ones(len(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_stack.shape) + for i in np.arange(seq[0].wcs.naxis)[seq[0].missing_axis[::-1]]: + new_shape.insert(i+1, 1) + # Also insert dummy coordinates and units. + axes_coordinates.insert(i+1, None) + axes_units.insert(i+1, None) + data_stack = data_stack.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_stack, wcs=new_wcs, image_axes=plot_axis_indices, axis_ranges=axes_coordinates, + unit_x_axis=axes_units[plot_axis_indices[0]], + unit_y_axis=axes_units[plot_axis_indices[1]], **kwargs) + + +class ImageAnimatorCubeLikeNDCubeSequence(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, axes=None, plot_axis_indices=None, + axes_coordinates=None, axes_units=None, data_unit=None, **kwargs): + if seq._common_axis is None: + raise TypeError("Common axis must be set to use this class. " + "Use ImageAnimatorNDCubeSequence.") + self.sequence = seq.data # Required by parent class. + # Set default values of kwargs if not set. + if wcs is None: + wcs = seq[0].wcs + if axes_coordinates is None: + axes_coordinates = [None] * len(seq.cube_like_dimensions) + if axes_units is None: + axes_units = [None] * len(seq.cube_like_dimensions) + # Determine units of each cube in sequence. + sequence_units, data_unit = _determine_sequence_units(seq.data, data_unit) + # If all cubes have unit set, create a data quantity from cube's data. + if sequence_units is None: + data_concat = np.concatenate([cube.data for cube in seq.data], axis=seq._common_axis) + else: + data_concat = np.concatenate( + [(cube.data * sequence_units[i]).to(data_unit).value + for i, cube in enumerate(seq.data)], axis=seq._common_axis) + self.cumul_cube_lengths = np.cumsum(np.array( + [c.dimensions[0].value for c in seq.data], dtype=int)) + # 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) + # Also insert dummy coordinates and units. + axes_coordinates.insert(i, None) + axes_units.insert(i, None) + data_concat = data_concat.reshape(new_shape) + + super(ImageAnimatorCubeLikeNDCubeSequence, self).__init__( + data_concat, wcs=wcs, image_axes=plot_axis_indices, axis_ranges=axes_coordinates, + unit_x_axis=axes_units[plot_axis_indices[0]], + unit_y_axis=axes_units[plot_axis_indices[1]], **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): + """ + Returns units of cubes in sequence and derives data unit if not set. + + If not all cubes have their unit attribute set, an error is raised. + + Parameters + ---------- + cubesequence_data: `list` of `ndcube.NDCube` + Taken from NDCubeSequence.data attribute. + + unit: `astropy.units.Unit` or `None` + If None, an appropriate unit is derived from first cube in sequence. + + Returns + ------- + sequence_units: `list` of `astropy.units.Unit` + Unit of each cube. + + unit: `astropy.units.Unit` + If input unit is not None, then the same as input. Otherwise it is + the unit of the first cube in the sequence. + + """ + # Check that the unit attribute is set of all cubes. If not, unit_y_axis + sequence_units = [] + for i, cube in enumerate(cubesequence_data): + if cube.unit is None: + break + else: + sequence_units.append(cube.unit) + if len(sequence_units) != len(cubesequence_data): + sequence_units = None + # If all cubes have unit set, create a data quantity from cube's data. + if sequence_units is None: + if unit is not None: + raise ValueError(NON_COMPATIBLE_UNIT_MESSAGE) + else: + if unit is None: + unit = sequence_units[0] + return sequence_units, unit + + +def _make_1D_sequence_plot(xdata, ydata, yerror, unit_y_axis, default_xlabel, kwargs): + # Define plot settings if not set in kwargs. + xlabel = kwargs.pop("xlabel", default_xlabel) + ylabel = kwargs.pop("ylabel", "Data [{0}]".format(unit_y_axis)) + title = kwargs.pop("title", "") + xlim = kwargs.pop("xlim", None) + ylim = kwargs.pop("ylim", None) + # Plot data + fig, ax = plt.subplots(1, 1) + ax.errorbar(xdata, ydata, yerror, **kwargs) + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.set_xlim(xlim) + ax.set_ylim(ylim) + return fig, ax + + +def _prep_axes_kwargs(naxis, plot_axis_indices, axes_coordinates, axes_units): + """ + Checks input values are correct based on number of sequence dimensions and sets defaults. + + Parameters + ---------- + plot_axis_indices: As for NDCubeSequencePlotMixin.plot or NDCubeSequencePlotMixin.plot_as_cube + + axes_coordinates: As for NDCubeSequencePlotMixin.plot or NDCubeSequencePlotMixin.plot_as_cube + + axes_units: As for NDCubeSequencePlotMixin.plot or NDCubeSequencePlotMixin.plot_as_cube + + Returns + ------- + plot_axis_indices: None or `list` of `int` of length 1 or 2. + + axes_coordinates: `None` or `list` of `None` `astropy.units.Quantity` `numpy.ndarray` `str` + Length of list equals number of sequence axes. + + axes_units: None or `list` of `None` `astropy.units.Unit` or `str` + Length of list equals number of sequence axes. + + """ + # If plot_axis_indices, axes_coordinates, axes_units are not None and not lists, + # convert to lists for consistent indexing behaviour. + if (not isinstance(plot_axis_indices, list)) and (plot_axis_indices is not None): + plot_axis_indices = [plot_axis_indices] + if (not isinstance(axes_coordinates, list)) and (axes_coordinates is not None): + axes_coordinates = [axes_coordinates] + if (not isinstance(axes_units, list)) and (axes_units is not None): + axes_units = [axes_units] + # Set default value of plot_axis_indices if not set by user. + if plot_axis_indices is None: + plot_axis_indices = [-1, -2] + else: + # If number of sequence dimensions is greater than 1, + # ensure length of plot_axis_indices is 1 or 2. + # No need to check case where number of sequence dimensions is 1 + # as plot_axis_indices is ignored in that case. + if naxis > 1 and len(plot_axis_indices) not in [1, 2]: + raise ValueError("plot_axis_indices can have at most length 2.") + if axes_coordinates is not None: + if naxis > 1: + # If convention of axes_coordinates and axes_units being length of + # plot_axis_index is being used, convert to convention where their + # length equals sequence dimensions. Only do this if number of dimensions if + # greater than 1 as the conventions are equivalent if there is only one dimension. + if len(axes_coordinates) == len(plot_axis_indices): + none_axes_coordinates = np.array([None] * naxis) + none_axes_coordinates[plot_axis_indices] = axes_coordinates + axes_coordinates = list(none_axes_coordinates) + # Now axes_coordinates have been converted to a consistent convention, + # ensure their length equals the number of sequence dimensions. + if len(axes_coordinates) != naxis: + raise ValueError("length of axes_coordinates must be {0}.".format(naxis)) + # Ensure all elements in axes_coordinates are of correct types. + ax_coord_types = (u.Quantity, np.ndarray, str) + for axis_coordinate in axes_coordinates: + if axis_coordinate is not None and not isinstance(axis_coordinate, ax_coord_types): + raise TypeError("axes_coordinates must be one of {0} or list of those.".format( + [None] + list(ax_coord_types))) + if axes_units is not None: + if naxis > 1: + if len(axes_units) == len(plot_axis_indices): + none_axes_units = np.array([None] * naxis) + none_axes_units[plot_axis_indices] = axes_units + axes_units = list(none_axes_units) + # Now axes_units have been converted to a consistent convention, + # ensure their length equals the number of sequence dimensions. + if len(axes_units) != naxis: + raise ValueError("length of axes_units must be {0}.".format(naxis)) + # Ensure all elements in axes_units are of correct types. + ax_unit_types = (u.UnitBase, str) + for axis_unit in axes_units: + if axis_unit is not None and not isinstance(axis_unit, ax_unit_types): + raise TypeError("axes_units must be one of {0} or list of {0}.".format( + ax_unit_types)) + + return plot_axis_indices, axes_coordinates, axes_units + + +def _derive_1D_coordinates_and_units(axes_coordinates, axes_units): + if axes_coordinates is None: + x_axis_coordinates = axes_coordinates + else: + if not isinstance(axes_coordinates, list): + axes_coordinates = [axes_coordinates] + x_axis_coordinates = axes_coordinates[0] + if axes_units is None: + unit_x_axis = axes_units + else: + if not isinstance(axes_units, list): + axes_units = [axes_units] + unit_x_axis = axes_units[0] + return x_axis_coordinates, unit_x_axis diff --git a/ndcube/ndcube_sequence.py b/ndcube/ndcube_sequence.py index 305e33628..50deadb66 100644 --- a/ndcube/ndcube_sequence.py +++ b/ndcube/ndcube_sequence.py @@ -5,12 +5,12 @@ from sunpy.map import MapCube from ndcube import utils -from ndcube.visualization import animation as ani +from ndcube.mixins.sequence_plotting import NDCubeSequencePlotMixin __all__ = ['NDCubeSequence'] -class NDCubeSequence: +class NDCubeSequenceBase: """ Class representing list of cubes. @@ -171,13 +171,6 @@ def sequence_axis_extra_coords(self): sequence_extra_coords = None return sequence_extra_coords - def plot(self, *args, **kwargs): - if self._common_axis is None: - i = ani.ImageAnimatorNDCubeSequence(self, *args, **kwargs) - else: - i = ani.ImageAnimatorCommonAxisNDCubeSequence(self, *args, **kwargs) - return i - def explode_along_axis(self, axis): """ Separates slices of NDCubes in sequence along a given cube axis into (N-1)DCubes. @@ -228,6 +221,10 @@ def _new_instance(cls, data_list, meta=None, common_axis=None): return cls(data_list, meta=meta, common_axis=common_axis) +class NDCubeSequence(NDCubeSequenceBase, NDCubeSequencePlotMixin): + pass + + """ Cube Sequence Helpers """ diff --git a/ndcube/tests/test_sequence_plotting.py b/ndcube/tests/test_sequence_plotting.py new file mode 100644 index 000000000..b85e5ca2b --- /dev/null +++ b/ndcube/tests/test_sequence_plotting.py @@ -0,0 +1,578 @@ +# -*- coding: utf-8 -*- +import pytest +import datetime + +import numpy as np +import astropy.units as u +import matplotlib + +from ndcube import NDCube, NDCubeSequence +from ndcube.utils.wcs import WCS +import ndcube.mixins.sequence_plotting + +# Set matplotlib display for testing +#matplotlib.use('Agg') + +# sample data for tests +# TODO: use a fixture reading from a test file. file TBD. +data = np.array([[[1, 2, 3, 4], [2, 4, 5, 3], [0, -1, 2, 3]], + [[2, 4, 5, 1], [10, 5, 2, 2], [10, 3, 3, 0]]]) + +data2 = np.array([[[11, 22, 33, 44], [22, 44, 55, 33], [0, -1, 22, 33]], + [[22, 44, 55, 11], [10, 55, 22, 22], [10, 33, 33, 0]]]) + +ht = {'CTYPE3': 'HPLT-TAN', 'CUNIT3': 'deg', 'CDELT3': 0.5, 'CRPIX3': 0, 'CRVAL3': 0, 'NAXIS3': 2, + 'CTYPE2': 'WAVE ', 'CUNIT2': 'Angstrom', 'CDELT2': 0.2, 'CRPIX2': 0, 'CRVAL2': 0, + 'NAXIS2': 3, + 'CTYPE1': 'TIME ', 'CUNIT1': 'min', 'CDELT1': 0.4, 'CRPIX1': 0, 'CRVAL1': 0, 'NAXIS1': 4} + +wt = WCS(header=ht, naxis=3) + +cube1 = NDCube( + data, wt, missing_axis=[False, False, False, True], + extra_coords=[ + ('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)), + ('hi', 1, u.Quantity(range(data.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(0, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 0))]) + +cube1_with_unit = NDCube( + data, wt, missing_axis=[False, False, False, True], + unit=u.km, + extra_coords=[ + ('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)), + ('hi', 1, u.Quantity(range(data.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(0, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 0))]) + +cube1_with_mask = NDCube( + data, wt, missing_axis=[False, False, False, True], + mask=np.zeros_like(data, dtype=bool), + extra_coords=[ + ('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)), + ('hi', 1, u.Quantity(range(data.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(0, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 0))]) + +cube1_with_uncertainty = NDCube( + data, wt, missing_axis=[False, False, False, True], + uncertainty=np.sqrt(data), + extra_coords=[ + ('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)), + ('hi', 1, u.Quantity(range(data.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(0, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 0))]) + +cube1_with_unit_and_uncertainty = NDCube( + data, wt, missing_axis=[False, False, False, True], + unit=u.km, uncertainty=np.sqrt(data), + extra_coords=[ + ('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)), + ('hi', 1, u.Quantity(range(data.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(0, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 0))]) + +cube3 = NDCube( + data2, wt, missing_axis=[False, False, False, True], + extra_coords=[ + ('pix', 0, u.Quantity(np.arange(1, data2.shape[0]+1), unit=u.pix) + + cube1.extra_coords['pix']['value'][-1]), + ('hi', 1, u.Quantity(range(data2.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(2, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 2))]) + +cube3_with_unit = NDCube( + data2, wt, missing_axis=[False, False, False, True], + unit=u.m, + extra_coords=[ + ('pix', 0, u.Quantity(np.arange(1, data2.shape[0]+1), unit=u.pix) + + cube1.extra_coords['pix']['value'][-1]), + ('hi', 1, u.Quantity(range(data2.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(2, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 2))]) + +cube3_with_mask = NDCube( + data2, wt, missing_axis=[False, False, False, True], + mask=np.zeros_like(data2, dtype=bool), + extra_coords=[ + ('pix', 0, u.Quantity(np.arange(1, data2.shape[0]+1), unit=u.pix) + + cube1.extra_coords['pix']['value'][-1]), + ('hi', 1, u.Quantity(range(data2.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(2, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 2))]) + +cube3_with_uncertainty = NDCube( + data2, wt, missing_axis=[False, False, False, True], + uncertainty=np.sqrt(data2), + extra_coords=[ + ('pix', 0, u.Quantity(np.arange(1, data2.shape[0]+1), unit=u.pix) + + cube1.extra_coords['pix']['value'][-1]), + ('hi', 1, u.Quantity(range(data2.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(2, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 2))]) + +cube3_with_unit_and_uncertainty = NDCube( + data2, wt, missing_axis=[False, False, False, True], + unit=u.m, uncertainty=np.sqrt(data2), + extra_coords=[ + ('pix', 0, u.Quantity(np.arange(1, data2.shape[0]+1), unit=u.pix) + + cube1.extra_coords['pix']['value'][-1]), + ('hi', 1, u.Quantity(range(data2.shape[1]), unit=u.s)), + ('distance', None, u.Quantity(2, unit=u.cm)), + ('time', None, datetime.datetime(2000, 1, 1, 0, 2))]) + +# Define some test NDCubeSequences. +common_axis = 0 +seq = NDCubeSequence(data_list=[cube1, cube3, cube1, cube3], common_axis=common_axis) + +seq_no_common_axis = NDCubeSequence(data_list=[cube1, cube3, cube1, cube3]) + +seq_with_units = NDCubeSequence( + data_list=[cube1_with_unit, cube3_with_unit, cube1_with_unit, cube3_with_unit], + common_axis=common_axis) + +seq_with_masks = NDCubeSequence( + data_list=[cube1_with_mask, cube3_with_mask, cube1_with_mask, cube3_with_mask], + common_axis=common_axis) + +seq_with_unit0 = NDCubeSequence(data_list=[cube1_with_unit, cube3, + cube1_with_unit, cube3], common_axis=common_axis) + +seq_with_mask0 = NDCubeSequence(data_list=[cube1_with_mask, cube3, + cube1_with_mask, cube3], common_axis=common_axis) + +seq_with_uncertainty = NDCubeSequence(data_list=[cube1_with_uncertainty, cube3_with_uncertainty, + cube1_with_uncertainty, cube3_with_uncertainty], + common_axis=common_axis) + +seq_with_some_uncertainty = NDCubeSequence( + data_list=[cube1_with_uncertainty, cube3, cube1, cube3_with_uncertainty], + common_axis=common_axis) + +seq_with_units_and_uncertainty = NDCubeSequence( + data_list=[cube1_with_unit_and_uncertainty, cube3_with_unit_and_uncertainty, + cube1_with_unit_and_uncertainty, cube3_with_unit_and_uncertainty], + common_axis=common_axis) + +seq_with_units_and_some_uncertainty = NDCubeSequence( + data_list=[cube1_with_unit_and_uncertainty, cube3_with_unit, + cube1_with_unit, cube3_with_unit_and_uncertainty], + common_axis=common_axis) + +# Derive some expected data arrays in plot objects. +seq_data_stack = np.stack([cube.data for cube in seq_with_masks.data]) +seq_mask_stack = np.stack([cube.mask for cube in seq_with_masks.data]) + +seq_stack = np.ma.masked_array(seq_data_stack, seq_mask_stack) +seq_stack_km = np.ma.masked_array( + np.stack([(cube.data * cube.unit).to(u.km).value for cube in seq_with_units.data]), + seq_mask_stack) + +seq_data_concat = np.concatenate([cube.data for cube in seq_with_masks.data], axis=common_axis) +seq_mask_concat = np.concatenate([cube.mask for cube in seq_with_masks.data], axis=common_axis) + +seq_concat = np.ma.masked_array(seq_data_concat, seq_mask_concat) +seq_concat_km = np.ma.masked_array( + np.concatenate([(cube.data * cube.unit).to(u.km).value + for cube in seq_with_units.data], axis=common_axis), + seq_mask_concat) + +# Derive expected axis_ranges +x_axis_coords = np.array([0.4, 0.8, 1.2, 1.6]).reshape((1, 1, 4)) +new_x_axis_coords_shape = u.Quantity(seq.dimensions, unit=u.pix).value.astype(int) +new_x_axis_coords_shape[-1] = 1 +none_axis_ranges_axis3 = [np.arange(len(seq.data)), np.array([0., 2.]), np.array([0., 1.5, 3.]), + np.tile(np.array(x_axis_coords), new_x_axis_coords_shape)] + +# Derive expected extents +seq_axis1_lim_deg = [0.49998731, 0.99989848] +seq_axis1_lim_arcsec = [(axis1_xlim*u.deg).to(u.arcsec).value for axis1_xlim in seq_axis1_lim_deg] +seq_axis2_lim_m = [seq[:, :, :, 0].data[0].axis_world_coords()[-1][0].value, + seq[:, :, :, 0].data[0].axis_world_coords()[-1][-1].value] + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_values", [ + (seq[:, 0, 0, 0], {}, + (np.arange(len(seq.data)), np.array([1, 11, 1, 11]), + "meta.obs.sequence [None]", "Data [None]", (0, len(seq[:, 0, 0, 0].data)-1), + (min([cube.data.min() for cube in seq[:, 0, 0, 0].data]), + max([cube.data.max() for cube in seq[:, 0, 0, 0].data])))), + + (seq_with_units[:, 0, 0, 0], {}, + (np.arange(len(seq_with_units.data)), np.array([1, 0.011, 1, 0.011]), + "meta.obs.sequence [None]", "Data [km]", (0, len(seq_with_units[:, 0, 0, 0].data)-1), + (min([(cube.data * cube.unit).to(seq_with_units[:, 0, 0, 0].data[0].unit).value + for cube in seq_with_units[:, 0, 0, 0].data]), + max([(cube.data * cube.unit).to(seq_with_units[:, 0, 0, 0].data[0].unit).value + for cube in seq_with_units[:, 0, 0, 0].data])))), + + (seq_with_uncertainty[:, 0, 0, 0], {}, + (np.arange(len(seq_with_uncertainty.data)), np.array([1, 11, 1, 11]), + "meta.obs.sequence [None]", "Data [None]", (0, len(seq_with_uncertainty[:, 0, 0, 0].data)-1), + (min([cube.data for cube in seq_with_uncertainty[:, 0, 0, 0].data]), + max([cube.data for cube in seq_with_uncertainty[:, 0, 0, 0].data])))), + + (seq_with_units_and_uncertainty[:, 0, 0, 0], {}, + (np.arange(len(seq_with_units_and_uncertainty.data)), np.array([1, 0.011, 1, 0.011]), + "meta.obs.sequence [None]", "Data [km]", + (0, len(seq_with_units_and_uncertainty[:, 0, 0, 0].data)-1), + (min([(cube.data*cube.unit).to(seq_with_units_and_uncertainty[:, 0, 0, 0].data[0].unit).value + for cube in seq_with_units_and_uncertainty[:, 0, 0, 0].data]), + max([(cube.data*cube.unit).to(seq_with_units_and_uncertainty[:, 0, 0, 0].data[0].unit).value + for cube in seq_with_units_and_uncertainty[:, 0, 0, 0].data])))), + + (seq_with_units_and_some_uncertainty[:, 0, 0, 0], {}, + (np.arange(len(seq_with_units_and_some_uncertainty.data)), np.array([1, 0.011, 1, 0.011]), + "meta.obs.sequence [None]", "Data [km]", + (0, len(seq_with_units_and_some_uncertainty[:, 0, 0, 0].data)-1), + (min([(cube.data*cube.unit).to( + seq_with_units_and_some_uncertainty[:, 0, 0, 0].data[0].unit).value + for cube in seq_with_units_and_some_uncertainty[:, 0, 0, 0].data]), + max([(cube.data*cube.unit).to( + seq_with_units_and_some_uncertainty[:, 0, 0, 0].data[0].unit).value + for cube in seq_with_units_and_some_uncertainty[:, 0, 0, 0].data])))), + + (seq[:, 0, 0, 0], {"axes_coordinates": "distance"}, + ((seq.sequence_axis_extra_coords["distance"]), np.array([1, 11, 1, 11]), + "distance [{0}]".format(seq.sequence_axis_extra_coords["distance"].unit), "Data [None]", + (min(seq.sequence_axis_extra_coords["distance"].value), + max(seq.sequence_axis_extra_coords["distance"].value)), + (min([cube.data.min() for cube in seq[:, 0, 0, 0].data]), + max([cube.data.max() for cube in seq[:, 0, 0, 0].data])))), + + (seq[:, 0, 0, 0], {"axes_coordinates": u.Quantity(np.arange(len(seq.data)), unit=u.cm), + "axes_units": u.km}, + (u.Quantity(np.arange(len(seq.data)), unit=u.cm).to(u.km), np.array([1, 11, 1, 11]), + "meta.obs.sequence [km]", "Data [None]", + (min((u.Quantity(np.arange(len(seq.data)), unit=u.cm).to(u.km).value)), + max((u.Quantity(np.arange(len(seq.data)), unit=u.cm).to(u.km).value))), + (min([cube.data.min() for cube in seq[:, 0, 0, 0].data]), + max([cube.data.max() for cube in seq[:, 0, 0, 0].data])))) + ]) +def test_sequence_plot_1D_plot(test_input, test_kwargs, expected_values): + # Unpack expected values + expected_x_data, expected_y_data, expected_xlabel, expected_ylabel, \ + expected_xlim, expected_ylim = expected_values + # Run plot method + output = test_input.plot(**test_kwargs) + # Check values are correct + assert isinstance(output, matplotlib.axes.Axes) + np.testing.assert_array_equal(output.lines[0].get_xdata(), expected_x_data) + np.testing.assert_array_equal(output.lines[0].get_ydata(), expected_y_data) + assert output.axes.get_xlabel() == expected_xlabel + assert output.axes.get_ylabel() == expected_ylabel + output_xlim = output.axes.get_xlim() + assert output_xlim[0] <= expected_xlim[0] + assert output_xlim[1] >= expected_xlim[1] + output_ylim = output.axes.get_ylim() + assert output_ylim[0] <= expected_ylim[0] + assert output_ylim[1] >= expected_ylim[1] + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_values", [ + (seq[:, :, 0, 0], {}, + (np.array([0.49998731, 0.99989848, 0.49998731, 0.99989848, + 0.49998731, 0.99989848, 0.49998731, 0.99989848]), + np.array([1, 2, 11, 22, 1, 2, 11, 22]), + "{0} [{1}]".format(seq[:, :, 0, 0].cube_like_world_axis_physical_types[common_axis], "deg"), + "Data [None]", tuple(seq_axis1_lim_deg), + (min([cube.data.min() for cube in seq[:, :, 0, 0].data]), + max([cube.data.max() for cube in seq[:, :, 0, 0].data])))), + + (seq_with_units[:, :, 0, 0], {}, + (np.array([0.49998731, 0.99989848, 0.49998731, 0.99989848, + 0.49998731, 0.99989848, 0.49998731, 0.99989848]), + np.array([1, 2, 0.011, 0.022, 1, 2, 0.011, 0.022]), + "{0} [{1}]".format(seq[:, :, 0, 0].cube_like_world_axis_physical_types[common_axis], "deg"), + "Data [km]", tuple(seq_axis1_lim_deg), + (min([min((cube.data * cube.unit).to(u.km).value) + for cube in seq_with_units[:, :, 0, 0].data]), + max([max((cube.data * cube.unit).to(u.km).value) + for cube in seq_with_units[:, :, 0, 0].data])))), + + (seq_with_uncertainty[:, :, 0, 0], {}, + (np.array([0.49998731, 0.99989848, 0.49998731, 0.99989848, + 0.49998731, 0.99989848, 0.49998731, 0.99989848]), + np.array([1, 2, 11, 22, 1, 2, 11, 22]), + "{0} [{1}]".format( + seq_with_uncertainty[:, :, 0, 0].cube_like_world_axis_physical_types[ + common_axis], "deg"), + "Data [None]", tuple(seq_axis1_lim_deg), + (min([cube.data.min() for cube in seq_with_uncertainty[:, :, 0, 0].data]), + max([cube.data.max() for cube in seq_with_uncertainty[:, :, 0, 0].data])))), + + (seq_with_some_uncertainty[:, :, 0, 0], {}, + (np.array([0.49998731, 0.99989848, 0.49998731, 0.99989848, + 0.49998731, 0.99989848, 0.49998731, 0.99989848]), + np.array([1, 2, 11, 22, 1, 2, 11, 22]), + "{0} [{1}]".format( + seq_with_some_uncertainty[:, :, 0, 0].cube_like_world_axis_physical_types[ + common_axis], "deg"), + "Data [None]", tuple(seq_axis1_lim_deg), + (min([cube.data.min() for cube in seq_with_some_uncertainty[:, :, 0, 0].data]), + max([cube.data.max() for cube in seq_with_some_uncertainty[:, :, 0, 0].data])))), + + (seq_with_units_and_uncertainty[:, :, 0, 0], {}, + (np.array([0.49998731, 0.99989848, 0.49998731, 0.99989848, + 0.49998731, 0.99989848, 0.49998731, 0.99989848]), + np.array([1, 2, 0.011, 0.022, 1, 2, 0.011, 0.022]), + "{0} [{1}]".format( + seq_with_units_and_uncertainty[:, :, 0, 0].cube_like_world_axis_physical_types[ + common_axis], "deg"), + "Data [km]", tuple(seq_axis1_lim_deg), + (min([min((cube.data * cube.unit).to(u.km).value) + for cube in seq_with_units[:, :, 0, 0].data]), + max([max((cube.data * cube.unit).to(u.km).value) + for cube in seq_with_units[:, :, 0, 0].data])))), + + (seq_with_units_and_some_uncertainty[:, :, 0, 0], {}, + (np.array([0.49998731, 0.99989848, 0.49998731, 0.99989848, + 0.49998731, 0.99989848, 0.49998731, 0.99989848]), + np.array([1, 2, 0.011, 0.022, 1, 2, 0.011, 0.022]), + "{0} [{1}]".format( + seq_with_units_and_some_uncertainty[:, :, 0, 0].cube_like_world_axis_physical_types[ + common_axis], "deg"), + "Data [km]", tuple(seq_axis1_lim_deg), + (min([min((cube.data * cube.unit).to(u.km).value) + for cube in seq_with_units[:, :, 0, 0].data]), + max([max((cube.data * cube.unit).to(u.km).value) + for cube in seq_with_units[:, :, 0, 0].data])))), + + (seq[:, :, 0, 0], {"axes_coordinates": "pix"}, + (seq[:, :, 0, 0].common_axis_extra_coords["pix"].value, + np.array([1, 2, 11, 22, 1, 2, 11, 22]), "pix [pix]", "Data [None]", + (min(seq[:, :, 0, 0].common_axis_extra_coords["pix"].value), + max(seq[:, :, 0, 0].common_axis_extra_coords["pix"].value)), + (min([cube.data.min() for cube in seq[:, :, 0, 0].data]), + max([cube.data.max() for cube in seq[:, :, 0, 0].data])))), + + (seq[:, :, 0, 0], + {"axes_coordinates": np.arange(10, 10+seq[:, :, 0, 0].cube_like_dimensions[0].value)}, + (np.arange(10, 10 + seq[:, :, 0, 0].cube_like_dimensions[0].value), + np.array([1, 2, 11, 22, 1, 2, 11, 22]), + "{0} [{1}]".format("", None), "Data [None]", + (10, 10 + seq[:, :, 0, 0].cube_like_dimensions[0].value - 1), + (min([cube.data.min() for cube in seq[:, :, 0, 0].data]), + max([cube.data.max() for cube in seq[:, :, 0, 0].data])))) + ]) +def test_sequence_plot_as_cube_1D_plot(test_input, test_kwargs, expected_values): + # Unpack expected values + expected_x_data, expected_y_data, expected_xlabel, expected_ylabel, \ + expected_xlim, expected_ylim = expected_values + # Run plot method + output = test_input.plot_as_cube(**test_kwargs) + # Check values are correct + # Check type of ouput plot object + assert isinstance(output, matplotlib.axes.Axes) + # Check x and y data are correct. + assert np.allclose(output.lines[0].get_xdata(), expected_x_data) + assert np.allclose(output.lines[0].get_ydata(), expected_y_data) + # Check x and y axis labels are correct. + assert output.axes.get_xlabel() == expected_xlabel + assert output.axes.get_ylabel() == expected_ylabel + # Check all data is contained within x and y axes limits. + output_xlim = output.axes.get_xlim() + assert output_xlim[0] <= expected_xlim[0] + assert output_xlim[1] >= expected_xlim[1] + output_ylim = output.axes.get_ylim() + assert output_ylim[0] <= expected_ylim[0] + assert output_ylim[1] >= expected_ylim[1] + + +def test_sequence_plot_as_cube_error(): + with pytest.raises(TypeError): + seq_no_common_axis.plot_as_cube() + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_values", [ + (seq[:, :, 0, 0], {}, + (seq_stack[:, :, 0, 0], + "custom:pos.helioprojective.lat [deg]", "meta.obs.sequence [None]", + tuple(seq_axis1_lim_deg + [0, len(seq.data)-1]))), + + (seq_with_units[:, :, 0, 0], {}, + (seq_stack_km[:, :, 0, 0], + "custom:pos.helioprojective.lat [deg]", "meta.obs.sequence [None]", + tuple(seq_axis1_lim_deg + [0, len(seq.data)-1]))), + + (seq[:, :, 0, 0], {"plot_axis_indices": [0, 1]}, + (seq_stack[:, :, 0, 0].transpose(), + "meta.obs.sequence [None]", "custom:pos.helioprojective.lat [deg]", + tuple([0, len(seq.data)-1] + seq_axis1_lim_deg))), + + (seq[:, :, 0, 0], {"axes_coordinates": ["pix", "distance"]}, + (seq_stack[:, :, 0, 0], + "pix [pix]", "distance [cm]", + (min(seq[0, :, 0, 0].extra_coords["pix"]["value"].value), + max(seq[0, :, 0, 0].extra_coords["pix"]["value"].value), + min(seq[:, :, 0, 0].sequence_axis_extra_coords["distance"].value), + max(seq[:, :, 0, 0].sequence_axis_extra_coords["distance"].value)))), + # This example shows weakness of current extra coord axis values on 2D plotting! + # Only the coordinates from the first cube are shown. + + (seq[:, :, 0, 0], {"axes_coordinates": [np.arange( + 10, 10+seq[:, :, 0, 0].dimensions[-1].value), "distance"], "axes_units": [None, u.m]}, + (seq_stack[:, :, 0, 0], + " [None]", "distance [m]", + (10, 10+seq[:, :, 0, 0].dimensions[-1].value-1, + min(seq[:, :, 0, 0].sequence_axis_extra_coords["distance"].to(u.m).value), + max(seq[:, :, 0, 0].sequence_axis_extra_coords["distance"].to(u.m).value)))), + + (seq[:, :, 0, 0], {"axes_coordinates": [np.arange( + 10, 10+seq[:, :, 0, 0].dimensions[-1].value)*u.deg, None], "axes_units": [u.arcsec, None]}, + (seq_stack[:, :, 0, 0], + " [arcsec]", "meta.obs.sequence [None]", + tuple(list( + (np.arange(10, 10+seq[:, :, 0, 0].dimensions[-1].value)*u.deg).to(u.arcsec).value) \ + + [0, len(seq.data)-1]))) + ]) +def test_sequence_plot_2D_image(test_input, test_kwargs, expected_values): + # Unpack expected values + expected_data, expected_xlabel, expected_ylabel, expected_extent = expected_values + # Run plot method + output = test_input.plot(**test_kwargs) + # Check values are correct + assert isinstance(output, matplotlib.axes.Axes) + np.testing.assert_array_equal(output.images[0].get_array(), expected_data) + assert output.xaxis.get_label_text() == expected_xlabel + assert output.yaxis.get_label_text() == expected_ylabel + assert np.allclose(output.images[0].get_extent(), expected_extent, rtol=1e-3) + # Also check x and y values????? + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_error", [ + (seq[:, :, 0, 0], {"axes_coordinates": [ + np.arange(10, 10+seq[:, :, 0, 0].dimensions[-1].value), None], + "axes_units": [u.m, None]}, ValueError), + + (seq[:, :, 0, 0], {"axes_coordinates": [ + None, np.arange(10, 10+seq[:, :, 0, 0].dimensions[0].value)], + "axes_units": [None, u.m]}, ValueError) + ]) +def test_sequence_plot_2D_image_errors(test_input, test_kwargs, expected_error): + with pytest.raises(expected_error): + output = test_input.plot(**test_kwargs) + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_values", [ + (seq[:, :, :, 0], {}, + (seq_concat[:, :, 0], + "em.wl [m]", "custom:pos.helioprojective.lat [deg]", + tuple(seq_axis2_lim_m + seq_axis1_lim_deg))), + + (seq_with_units[:, :, :, 0], {}, + (seq_concat_km[:, :, 0], + "em.wl [m]", "custom:pos.helioprojective.lat [deg]", + tuple(seq_axis2_lim_m + seq_axis1_lim_deg))), + + (seq[:, :, :, 0], {"plot_axis_indices": [0, 1], + "axes_coordinates": ["pix", "hi"]}, + (seq_concat[:, :, 0].transpose(), "pix [pix]", "hi [s]", + ((seq[:, :, :, 0].common_axis_extra_coords["pix"][0].value, + seq[:, :, :, 0].common_axis_extra_coords["pix"][-1].value, + seq[:, :, :, 0].data[0].extra_coords["hi"]["value"][0].value, + seq[:, :, :, 0].data[0].extra_coords["hi"]["value"][-1].value)))), + + (seq[:, :, :, 0], {"axes_coordinates": [ + np.arange(10, 10+seq[:, :, :, 0].cube_like_dimensions[-1].value) * u.m, + np.arange(10, 10+seq[:, :, :, 0].cube_like_dimensions[0].value) * u.m]}, + (seq_concat[:, :, 0], " [m]", " [m]", + (10, 10+seq[:, :, :, 0].cube_like_dimensions[-1].value-1, + 10, 10+seq[:, :, :, 0].cube_like_dimensions[0].value-1))), + + (seq[:, :, :, 0], {"axes_coordinates": [ + np.arange(10, 10+seq[:, :, :, 0].cube_like_dimensions[-1].value) * u.m, + np.arange(10, 10+seq[:, :, :, 0].cube_like_dimensions[0].value) * u.m], + "axes_units": ["cm", u.cm]}, + (seq_concat[:, :, 0], " [cm]", " [cm]", + (10*100, (10+seq[:, :, :, 0].cube_like_dimensions[-1].value-1)*100, + 10*100, (10+seq[:, :, :, 0].cube_like_dimensions[0].value-1)*100))) + ]) +def test_sequence_plot_as_cube_2D_image(test_input, test_kwargs, expected_values): + # Unpack expected values + expected_data, expected_xlabel, expected_ylabel, expected_extent = expected_values + # Run plot method + output = test_input.plot_as_cube(**test_kwargs) + # Check values are correct + assert isinstance(output, matplotlib.axes.Axes) + np.testing.assert_array_equal(output.images[0].get_array(), expected_data) + assert output.xaxis.get_label_text() == expected_xlabel + assert output.yaxis.get_label_text() == expected_ylabel + assert np.allclose(output.images[0].get_extent(), expected_extent, rtol=1e-3) + # Also check x and y values????? + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_error", [ + (seq[:, :, :, 0], {"axes_coordinates": [ + np.arange(10, 10+seq[:, :, :, 0].cube_like_dimensions[-1].value), None], + "axes_units": [u.m, None]}, ValueError), + + (seq[:, :, :, 0], {"axes_coordinates": [ + None, np.arange(10, 10+seq[:, :, :, 0].cube_like_dimensions[0].value)], + "axes_units": [None, u.m]}, ValueError) + ]) +def test_sequence_plot_as_cube_2D_image_errors(test_input, test_kwargs, expected_error): + with pytest.raises(expected_error): + output = test_input.plot_as_cube(**test_kwargs) + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_data", [ + (seq, {}, seq_stack.reshape(4, 1, 2, 3, 4)), + (seq_with_units, {}, seq_stack_km.reshape(4, 1, 2, 3, 4)) + ]) +def test_sequence_plot_ImageAnimator(test_input, test_kwargs, expected_data): + # Run plot method + output = test_input.plot(**test_kwargs) + # Check plot object properties are correct. + assert isinstance(output, ndcube.mixins.sequence_plotting.ImageAnimatorNDCubeSequence) + np.testing.assert_array_equal(output.data, expected_data) + + +@pytest.mark.parametrize("test_input, test_kwargs, expected_data", [ + (seq, {}, seq_concat.reshape(1, 8, 3, 4)), + (seq_with_units, {}, seq_concat_km.reshape(1, 8, 3, 4)) + ]) +def test_sequence_plot_as_cube_ImageAnimator(test_input, test_kwargs, expected_data): + # Run plot method + output = test_input.plot_as_cube(**test_kwargs) + # Check plot object properties are correct. + assert isinstance(output, ndcube.mixins.sequence_plotting.ImageAnimatorCubeLikeNDCubeSequence) + np.testing.assert_array_equal(output.data, expected_data) + + +@pytest.mark.parametrize("test_input, expected", [ + ((seq_with_unit0.data, None), (None, None)), + ((seq_with_unit0.data, u.km), (None, None)), + ((seq_with_units.data, None), ([u.km, u.m, u.km, u.m], u.km)), + ((seq_with_units.data, u.cm), ([u.km, u.m, u.km, u.m], u.cm))]) +def test_determine_sequence_units(test_input, expected): + output_seq_unit, output_unit = ndcube.mixins.sequence_plotting._determine_sequence_units( + test_input[0], unit=test_input[1]) + assert output_seq_unit == expected[0] + assert output_unit == expected[1] + + +def test_determine_sequence_units(): + with pytest.raises(ValueError): + output_seq_unit, output_unit = ndcube.mixins.sequence_plotting._determine_sequence_units( + seq.data, u.m) + + +@pytest.mark.parametrize("test_input, expected", [ + ((3, 1, "time", u.s), ([1], [None, 'time', None], [None, u.s, None])), + ((3, None, None, None), ([-1, -2], None, None))]) +def test_prep_axes_kwargs(test_input, expected): + output = ndcube.mixins.sequence_plotting._prep_axes_kwargs(*test_input) + for i in range(3): + assert output[i] == expected[i] + + +@pytest.mark.parametrize("test_input, expected_error", [ + ((3, [0, 1, 2], ["time", "pix"], u.s), ValueError), + ((3, 0, ["time", "pix"], u.s), ValueError), + ((3, 0, "time", [u.s, u.pix]), ValueError), + ((3, 0, 0, u.s), TypeError), + ((3, 0, "time", 0), TypeError)]) +def test_prep_axes_kwargs_errors(test_input, expected_error): + with pytest.raises(expected_error): + output = ndcube.mixins.sequence_plotting._prep_axes_kwargs(*test_input) diff --git a/ndcube/visualization/__init__.py b/ndcube/visualization/__init__.py deleted file mode 100644 index 4332d48cd..000000000 --- a/ndcube/visualization/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -'''Sunpycube Visualization''' - -from . import animation diff --git a/ndcube/visualization/animation.py b/ndcube/visualization/animation.py deleted file mode 100644 index 867363854..000000000 --- a/ndcube/visualization/animation.py +++ /dev/null @@ -1,181 +0,0 @@ -import numpy as np -from sunpy.visualization.imageanimator import ImageAnimatorWCS - -from ndcube import utils - - -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 seq._common_axis is not None: - raise ValueError("Common axis can't set set to use this class. " - "Use ImageAnimatorCommonAxisNDCubeSequence.") - 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