Source code for ndcube.mixins.plotting

import numpy as np
import matplotlib.pyplot as plt

import astropy.units as u
from sunpy.visualization.imageanimator import ImageAnimatorWCS
import sunpy.visualization.wcsaxes_compat as wcsaxes_compat

__all__ = ['NDCubePlotMixin']


[docs]class NDCubePlotMixin: """ Add plotting functionality to a NDCube class. """
[docs] def plot(self, axes=None, image_axes=[-1, -2], unit_x_axis=None, unit_y_axis=None, axis_ranges=None, unit=None, origin=0, **kwargs): """ Plots an interactive visualization of this cube with a slider controlling the wavelength axis for data having dimensions greater than 2. Plots an x-y graph onto the current axes for 2D or 1D data. Keyword arguments are passed on to matplotlib. Parameters other than data and wcs are passed to ImageAnimatorWCS, which in turn passes them to imshow for data greater than 2D. Parameters ---------- image_axes: `list` The two axes that make the image. Like [-1,-2] this implies cube instance -1 dimension will be x-axis and -2 dimension will be y-axis. axes: `astropy.visualization.wcsaxes.core.WCSAxes` or None: The axes to plot onto. If None the current axes will be used. unit_x_axis: `astropy.units.Unit` The unit of x axis for 2D plots. unit_y_axis: `astropy.units.Unit` The unit of y axis for 2D plots. unit: `astropy.unit.Unit` The data is changed to the unit given or the cube.unit if not given, for 1D plots. 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. """ axis_data = ['x' for i in range(2)] axis_data[image_axes[1]] = 'y' if self.data.ndim >= 3: plot = self._plot_3D_cube(image_axes=image_axes, unit_x_axis=unit_x_axis, unit_y_axis=unit_y_axis, axis_ranges=axis_ranges, **kwargs) elif self.data.ndim is 2: plot = self._plot_2D_cube(axes=axes, image_axes=axis_data[::-1], **kwargs) elif self.data.ndim is 1: plot = self._plot_1D_cube(unit=unit, origin=origin) return plot
def _plot_3D_cube(self, image_axes=None, unit_x_axis=None, unit_y_axis=None, axis_ranges=None, **kwargs): """ Plots an interactive visualization of this cube using sliders to move through axes plot using in the image. Parameters other than data and wcs are passed to ImageAnimatorWCS, which in turn passes them to imshow. Parameters ---------- image_axes: `list` The two axes that make the image. Like [-1,-2] this implies cube instance -1 dimension will be x-axis and -2 dimension will be y-axis. unit_x_axis: `astropy.units.Unit` The unit of x axis. unit_y_axis: `astropy.units.Unit` The unit of y axis. 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. """ if not image_axes: image_axes = [-1, -2] i = ImageAnimatorWCS(self.data, wcs=self.wcs, image_axes=image_axes, unit_x_axis=unit_x_axis, unit_y_axis=unit_y_axis, axis_ranges=axis_ranges, **kwargs) return i def _plot_2D_cube(self, axes=None, image_axes=None, **kwargs): """ Plots a 2D image onto the current axes. Keyword arguments are passed on to matplotlib. Parameters ---------- axes: `astropy.visualization.wcsaxes.core.WCSAxes` or `None`: The axes to plot onto. If None the current axes will be used. image_axes: `list`. The first axis in WCS object will become the first axis of image_axes and second axis in WCS object will become the second axis of image_axes. Default: ['x', 'y'] """ if not image_axes: image_axes = ['x', 'y'] if axes is None: if self.wcs.naxis is not 2: missing_axis = self.missing_axis slice_list = [] index = 0 for i, bool_ in enumerate(missing_axis): if not bool_: slice_list.append(image_axes[index]) index += 1 else: slice_list.append(1) if index is not 2: raise ValueError("Dimensions of WCS and data don't match") axes = wcsaxes_compat.gca_wcs(self.wcs, slices=slice_list) plot = axes.imshow(self.data, **kwargs) return plot def _plot_1D_cube(self, unit=None, origin=0): """ Plots a graph. Keyword arguments are passed on to matplotlib. Parameters ---------- unit: `astropy.unit.Unit` The data is changed to the unit given or the cube.unit if not given. """ index_not_one = [] for i, _bool in enumerate(self.missing_axis): if not _bool: index_not_one.append(i) if unit is None: unit = self.wcs.wcs.cunit[index_not_one[0]] plot = plt.plot(self.pixel_to_world(*[u.Quantity(np.arange(self.data.shape[0]), unit=u.pix)])[0].to(unit), self.data) return plot