def orient(array, wcs, *extra_arrs):
# This is mostly lifted from astropy's spectral cube.
"""
Given a 3 or 4D cube and a WCS, swap around the axes so that the
axes are in the correct order: the first in Numpy notation, and the last
in WCS notation.
Parameters
----------
array : `~numpy.ndarray`
The input 3- or 4-d array with two position dimensions and one spectral
dimension.
wcs : `~sunpy.wcs.WCS`
The input 3- or 4-d WCS with two position dimensions and one spectral
dimension.
extra_arrs: one or more ndarrays, optional
Extra arrays to orient, corresponding to uncertainties and errors in
the data.
"""
if wcs.oriented: # If this wcs has already been oriented.
return (array, wcs) + extra_arrs
if array.ndim != 3 and array.ndim != 4:
raise ValueError("Input array must be 3- or 4-dimensional")
if not ((wcs.wcs.naxis == 3 and array.ndim == 3) or
(wcs.wcs.naxis == 4 and array.ndim == 4 and not wcs.was_augmented)
or (wcs.wcs.naxis == 4 and array.ndim == 3 and wcs.was_augmented)):
raise ValueError("WCS must have the same dimensions as the array")
axtypes = list(wcs.wcs.ctype)
if wcs.was_augmented:
array_order = select_order(axtypes[2::-1])
else:
array_order = select_order(axtypes)
result_array = array.transpose(array_order)
wcs_order = np.array(select_order(axtypes))[::-1]
result_wcs = wcs_util.reindex_wcs(wcs, wcs_order)
result_wcs.was_augmented = wcs.was_augmented
result_wcs.oriented = True
result_extras = [arr.transpose(array_order) for arr in extra_arrs]
return (result_array, result_wcs) + tuple(result_extras)
def convert_to_spectral_cube(self):
"""
Converts this cube into a SpectralCube. It will only work if the cube
has exactly three dimensions and one of those is a spectral axis.
"""
if self.data.ndim == 4:
raise cu.CubeError(4, "Too many dimensions: Can only convert a " +
"3D cube. Slice the cube before converting")
if 'WAVE' not in self.axes_wcs.wcs.ctype:
raise cu.CubeError(2, 'Spectral axis needed to create a spectrum')
axis = 0 if self.axes_wcs.wcs.ctype[-1] == 'WAVE' else 1
coordaxes = [1, 2] if axis == 0 else [0, 2] # Non-spectral axes
newwcs = wu.reindex_wcs(self.axes_wcs, np.arary(coordaxes))
time_or_x_size = self.data.shape[coordaxes[1]]
y_size = self.data.shape[coordaxes[0]]
spectra = np.empty((time_or_x_size, y_size), dtype=Spectrum)
for i in range(time_or_x_size):
for j in range(y_size):
spectra[i][j] = self.slice_to_spectrum(i, j)
return SpectralCube(spectra, newwcs, self.meta)
def slice_to_cube(self, axis, chunk, **kwargs):
"""
For a hypercube, return a 3-D cube that has been cut along the given
axis and with data corresponding to the given chunk.
Parameters
----------
axis: int
The axis to cut from the hypercube
chunk: int, astropy Quantity or tuple:
The data to take from the axis
"""
if self.data.ndim == 3:
raise cu.CubeError(4, 'Can only slice a hypercube into a cube')
item = [slice(None, None, None) for _ in range(4)]
if isinstance(chunk, tuple):
if cu.iter_isinstance(chunk, (u.Quantity, u.Quantity)):
pixel0 = cu.convert_point(chunk[0].value, chunk[0].unit,
self.axes_wcs, axis)
pixel1 = cu.convert_point(chunk[1].value, chunk[1].unit,
self.axes_wcs, axis)
item[axis] = slice(pixel0, pixel1, None)
elif cu.iter_isinstance((chunk, int, int)):
item[axis] = slice(chunk[0], chunk[1], None)
else:
raise cu.CubeError(5, "Parameters must be of the same type")
newdata = self.data[item].sum(axis)
else:
unit = chunk.unit if isinstance(chunk, u.Quantity) else None
pixel = cu.convert_point(chunk, unit, self.axes_wcs, axis)
item[axis] = pixel
newdata = self.data[item]
wcs_indices = [0, 1, 2, 3]
wcs_indices.remove(3 - axis)
newwcs = wu.reindex_wcs(self.axes_wcs, np.array(wcs_indices))
if axis == 2 or axis == 3:
newwcs = wu.add_celestial_axis(newwcs)
newwcs.was_augmented = True
cube = Cube(newdata, newwcs, meta=self.meta, **kwargs)
return cube