def test_checks(self):
cube = generate_cube(shape=(3, 2, 1))
cube.write(os.path.join(self.tmpdir, 'cube-tmp.fits'), savemask='nan')
clist = CubeList(self.cubenames[:1] + [cube.filename])
assert clist.check_dim() is False
assert clist.check_wcs() is False
cube = generate_cube(shape=self.shape, crval=12.)
cube.write(os.path.join(self.tmpdir, 'cube-tmp.fits'), savemask='nan')
clist = CubeList(self.cubenames[:1] + [cube.filename])
assert clist.check_dim() is True
assert clist.check_wcs() is False
def test_arithmetic_cubes():
image2 = generate_image(data=1, unit='2 ct')
cube1 = generate_cube(data=0.5, unit=u.Unit('2 ct'))
for op in (add, sub, mul, div):
assert_allclose(op(image2, cube1).data, op(image2.data, cube1.data))
assert_allclose(op(cube1, image2).data, op(cube1.data, image2.data))
def test_write(tmpdir):
"""Cube class: testing write"""
unit = u.Unit('1e-20 erg/s/cm2/Angstrom')
cube = generate_cube(data=1, wave=WaveCoord(crval=1, cunit=u.angstrom),
unit=unit)
cube.data[:, 0, 0] = ma.masked
cube.var = np.ones_like(cube.data)
testfile = str(tmpdir.join('cube.fits'))
cube.write(testfile)
hdu = fits.open(testfile)
assert_array_equal(hdu[1].data.shape, cube.shape)
assert_array_equal([h.name for h in hdu],
['PRIMARY', 'DATA', 'STAT', 'DQ'])
hdr = hdu[0].header
assert hdr['AUTHOR'] == 'MPDAF'
hdr = hdu[1].header
assert hdr['EXTNAME'] == 'DATA'
assert hdr['NAXIS'] == 3
assert u.Unit(hdr['BUNIT']) == unit
assert u.Unit(hdr['CUNIT3']) == u.angstrom
assert hdr['NAXIS1'] == cube.shape[2]
assert hdr['NAXIS2'] == cube.shape[1]
assert hdr['NAXIS3'] == cube.shape[0]
for key in ('CRPIX1', 'CRPIX2'):
assert hdr[key] == 1.0
hdu.close()
def test_aperture():
"""Cube class: testing spectrum extraction"""
cube = generate_cube(data=1, wave=WaveCoord(crval=1))
spe = cube.aperture(center=(2, 2.8), radius=1,
unit_center=None, unit_radius=None)
assert spe.shape[0] == 10
assert spe.get_start() == 1
def test_convolve():
shape = (3, 12, 25)
data = np.zeros(shape)
data[:, 7, 5] = 1.0
mask = np.zeros(shape, dtype=bool)
mask[:, 5, 3] = True
c = generate_cube(data=data, mask=mask, shape=shape,
wave=WaveCoord(crval=1, cunit=u.angstrom))
# Create a symmetric convolution kernel with an even number of elements
# along one dimension and and odd number along the other dimension.
# Make the kernel symmetric around (shape-1)//2. This requires that
# the final column be all zeros.
kern = np.array([[[0.1, 0.25, 0.1, 0.0],
[0.25, 0.50, 0.25, 0.0],
[0.1, 0.25, 0.1, 0.0]]])
# The image should consist of a copy of the convolution kernel, centered
# such that pixels (kern.shape-1)//2 is at pixel 7,5 of data.
expected_data = ma.array(data=np.zeros(shape), mask=mask)
expected_data.data[:, 6:9, 4:8] = kern
res = c.convolve(kern)
assert_masked_allclose(res.data, expected_data, atol=1e-15)
res = c.convolve(Image(data=kern))
assert_masked_allclose(res.data, expected_data, atol=1e-15)
res = c.fftconvolve(kern)
assert_masked_allclose(res.data, expected_data, atol=1e-15)
def test_subcube(mask):
"""Cube class: testing sub-cube extraction methods"""
cube1 = generate_cube(data=np.arange(10 * 6 * 5).reshape(10, 6, 5),
wave=WaveCoord(crval=1), mask=mask)
# Extract a sub-cube whose images are centered close to pixel
# (2.3, 2.8) of the cube and have a width and height of 2 pixels.
# The center of a 2x2 pixel region is at the shared corner between
# these pixels, and the closest corner to the requested center
# of 2.3,2.8 is 2.5,2.5. Thus the sub-images should be from pixels
# 2,3 along both the X and Y axes.
cube2 = cube1.subcube(center=(2.3, 2.8), size=2, lbda=(5, 8),
unit_center=None, unit_size=None, unit_wave=None)
assert_allclose(cube1.data[5:9, 2:4, 2:4], cube2.data)
# Test when subcube is on the edges
cube2 = cube1.subcube(center=(0.3, 0.8), size=4,
unit_center=None, unit_size=None)
assert_allclose(cube1.data[:, :3, :3], cube2.data[:, 1:, 1:])
# pixels inside the selected region are not masked
assert np.all(~cube2.mask[:, 1:, 1:])
# pixels outside the selected region are masked
assert np.all(cube2.mask[:, :, 0])
assert np.all(cube2.mask[:, 0, :])
# The following should select the same image area as above, followed by
# masking pixels in this area outside a circle of radius 1.
cube2 = cube1.subcube_circle_aperture(center=(2.3, 2.8), radius=1,
unit_center=None, unit_radius=None)
# masking the subcube should not mask the original cube
assert ma.count_masked(cube1[0].data) == 0
if cube2.mask is not ma.nomask:
assert bool(cube2.mask[0, 0, 0]) is True
assert_array_equal(cube2.get_start(), (1, 2, 2))
assert_array_equal(cube2.shape, (10, 2, 2))
def test_arithmetic():
"""Cube class: tests arithmetic functions"""
cube1 = generate_cube(uwave=u.nm)
image1 = generate_image(wcs=cube1.wcs, unit=2 * u.ct)
spectrum1 = generate_spectrum(data=2.3, cdelt=30.0, crval=5)
cube2 = image1 + cube1
for op in (add, sub, mul, div):
cube3 = op(cube1, cube2)
assert_almost_equal(cube3.data, op(cube1.data, cube2.data))
# with spectrum
sp1 = spectrum1.data[:, np.newaxis, np.newaxis]
for op in (add, sub, mul, div):
cube3 = op(cube1, spectrum1)
assert_almost_equal(cube3.data, op(cube1.data, sp1))
# with image
im1 = image1.data.data[np.newaxis, :, :] * image1.unit
for op in (add, sub, mul, div):
cube3 = op(cube1, image1)
assert_almost_equal(
cube3._data,
op(cube1._data * cube1.unit, im1).to(cube3.unit).value)
cube2 = cube1 / 25.3
assert_almost_equal(cube2.data, cube1.data / 25.3)
def test_iter_spe():
"""Cube class: tests Spectrum iterator"""
cube1 = generate_cube(data=0.)
for (spe, (p, q)) in iter_spe(cube1, True):
spe[:] = spe + p + q
y, x = np.mgrid[:cube1.shape[1], :cube1.shape[2]]
assert_array_equal(*np.broadcast_arrays(cube1.data.data, y + x))
def test_truncate():
"""Cube class: testing truncation"""
cube1 = generate_cube(data=2, wave=WaveCoord(crval=1))
coord = [2, 0, 1, 5, 1, 3]
cube2 = cube1.truncate(coord, unit_wcs=cube1.wcs.unit,
unit_wave=cube1.wave.unit)
assert_array_equal(cube2.shape, (4, 2, 3))
assert_array_equal(cube2.get_start(), (2, 0, 1))
assert_array_equal(cube2.get_end(), (5, 1, 3))
def test_mean():
"""Cube class: testing mean method"""
# Specify the dimensions of the test cube.
shape = (2, 3, 4)
# Create 3D data, variance and mask arrays.
d = np.arange(np.asarray(shape).prod(), dtype=float).reshape(shape)
v = d / 10.0
m = d % 5 == 0
# Create a cube with the above values.
cube = generate_cube(data=d, var=v, mask=m)
# Create masked array versions of the data and variance arrays.
data = ma.array(d, mask=m)
var = ma.array(v, mask=m)
# Test the mean over each of the supported axes.
for axis in (None, 0, (1, 2)):
# Test the two weighting options. The loop-variable tuple
# specifies the weights argument of cube.mean(), and a masked
# array of the weights that are expected to be used.
for weights, w in [
(None, np.ones(shape, dtype=float)), # Unweighted
(np.sqrt(d), np.sqrt(d))
]: # Weighted
# Compute the mean.
mean = cube.mean(axis=axis, weights=weights)
# Mask the array of the expected weights.
w = ma.array(w, mask=m)
# Compute the expected values of the mean and its variance,
# using a different implementation to cube.mean().
expected_data = np.sum(data * w, axis=axis) / np.sum(w, axis=axis)
expected_var = (np.sum(var * w**2, axis=axis) /
np.sum(w, axis=axis)**2)
# In the case any of the following tests fail, provide an
# error message that indicates which arguments were passed
# to cube.mean()
errmsg = "Failure of cube.mean(axis=%s, weights=%s)" % (axis,
weights)
# See if the mean pixels and variances of these pixels have the
# expected values.
if axis is None:
assert_allclose(mean, expected_data, err_msg=errmsg)
else:
assert_allclose(mean.data, expected_data, err_msg=errmsg)
assert_allclose(mean.var, expected_var, err_msg=errmsg)
def test_arithmetic_errors(cube):
cube = generate_cube()
image1 = generate_image(wcs=cube.wcs)
image1.wcs.set_crval1(10)
with pytest.raises(ValueError):
cube2 = image1 + cube
spectrum1 = generate_spectrum(wave=cube.wave)
spectrum1.wave.set_crval(25)
with pytest.raises(ValueError):
cube2 = spectrum1 + cube
spectrum1 = generate_spectrum(wave=cube.wave)
spectrum1.wave.shape = 12
with pytest.raises(ValueError):
cube2 = spectrum1 + cube
def test_sum():
"""Cube class: testing sum method"""
cube1 = generate_cube(data=1, wave=WaveCoord(crval=1))
ind = np.arange(10)
refsum = ind.sum()
cube1.data = (ind[:, np.newaxis, np.newaxis] *
np.ones((6, 5))[np.newaxis, :, :])
assert cube1.sum() == 6 * 5 * refsum
assert_array_equal(cube1.sum(axis=0).data, np.full((6, 5), refsum, float))
weights = np.ones(shape=(10, 6, 5))
assert cube1.sum(weights=weights) == 6 * 5 * refsum
weights = np.ones(shape=(10, 6, 5)) * 2
assert cube1.sum(weights=weights) == 6 * 5 * refsum
assert_array_equal(cube1.sum(axis=(1, 2)).data, ind * 6 * 5)
def test_arithmetic_variance(cube):
cube = generate_cube()
image1 = generate_image(wcs=cube.wcs, var=None)
cube2 = image1 + cube
assert_almost_equal(cube.var, cube2.var)
cube2 = image1 * cube
assert_almost_equal(cube2.var, cube.var * image1.data * image1.data)
cube.var = None
image1 = generate_image(wcs=cube.wcs)
cube2 = image1 + cube
assert_almost_equal(cube2.var, np.tile(image1.var, (cube2.shape[0], 1, 1)))
cube2 = image1 * cube
assert_almost_equal(cube2.var, image1.var * cube.data * cube.data)
def test_min():
"""Cube class: testing min method"""
cube1 = generate_cube(data=1., wave=WaveCoord(crval=1))
ind = np.arange(10)
minimum = np.amin(ind)
cube1.data = (ind[:, np.newaxis, np.newaxis] *
np.ones((6, 5))[np.newaxis, :, :])
m = cube1.min()
assert m == minimum
m = cube1.min(axis=0)
assert m[3, 3] == minimum
m = cube1.min(axis=(1, 2))
assert_array_equal(m.data, ind)
with pytest.raises(ValueError):
m = cube1.min(axis=-1)
def test_multiprocess():
"""Cube class: tests multiprocess"""
data_value = 2.2
cube1 = generate_cube(data=data_value)
# Test image processing using an Image method.
spe = cube1.loop_ima_multiprocessing(Image.get_rot, cpu=2, verbose=False)
assert_allclose(spe.data, cube1.get_rot())
# Test image processing using a normal function.
spe = cube1.loop_ima_multiprocessing(_multiproc_func, cpu=2, verbose=False)
assert_allclose(spe.data, data_value * 10.0)
# Test spectrum processing using a Spectrum method.
im = cube1.loop_spe_multiprocessing(Spectrum.mean, cpu=2, verbose=False)
assert_allclose(im[2, 3], cube1[:, 2, 3].mean())
# Test spectrum processing using a normal function.
im = cube1.loop_spe_multiprocessing(_multiproc_func, cpu=2, verbose=False)
assert_allclose(im.data, data_value * 10.0)
def setUpClass(cls):
cls.tmpdir = tempfile.mkdtemp()
print('\n>>> Create cubes in', cls.tmpdir)
cls.cubenames = []
cls.arr = np.arange(np.prod(cls.shape)).reshape(cls.shape)
cls.scaledcube = np.mean(
[(cls.arr * val + cls.offsetlist[k]) * cls.scalelist[k]
for k, val in enumerate(cls.cubevals)],
axis=0)
cls.combined_cube = np.mean([(cls.arr * i) for i in cls.cubevals],
axis=0)
for i in cls.cubevals:
cube = generate_cube(data=cls.arr * i, shape=cls.shape)
cube.primary_header['CUBEIDX'] = i
cube.primary_header['OBJECT'] = 'OBJECT %d' % i
cube.primary_header['EXPTIME'] = 100
filename = os.path.join(cls.tmpdir, 'cube-%d.fits' % i)
cube.write(filename, savemask='nan')
cls.cubenames.append(filename)
cls.expmap = np.full(cls.shape, cls.ncubes, dtype=int)
def test_rebin():
"""Cube class: testing rebin methods"""
# Create spectral and spatial world coordinates that make each
# pixel equal its index.
wcs = WCS(crval=(0, 0), crpix=(1, 1), cdelt=(1.0, 1.0))
wave = WaveCoord(crval=0.0, crpix=1.0, cdelt=1.0)
# Create a cube with even valued dimensions, filled with ones.
data = ma.ones((4, 6, 8)) # Start with all pixels 1.0
data.reshape(4 * 6 * 8)[::2] = 0.0 # Set every second pixel to 0.0
data.mask = data < -1 # Unmask all pixels.
cube1 = generate_cube(data=data.data, mask=data.mask, wcs=wcs, wave=wave)
# Rebin each of the axes such as to leave only 2 pixels along each
# dimension.
factor = (2, 3, 4)
cube2 = cube1.rebin(factor=factor)
# Compute the expected output cube, given that the input cube is a
# repeating pattern of 1,0, and we divided the x-axis by a
# multiple of 2, the output pixels should all be 0.5.
expected = ma.ones((2, 2, 2)) * 0.5
expected.mask = expected < 0 # All pixels unmasked.
assert_masked_allclose(cube2.data, expected)
# Do the same experiment but with the zero valued pixels all masked.
data = ma.ones((4, 6, 8)) # Start with all pixels 1.0
data.reshape(4 * 6 * 8)[::2] = 0.0 # Set every second pixel to 0.0
data.mask = data < 0.1 # Mask the pixels that are 0.0
cube1 = generate_cube(data=data.data, mask=data.mask, wcs=wcs, wave=wave)
# Rebin each of the axes such as to leave only 2 pixels along each
# dimension.
factor = np.array([2, 3, 4])
cube2 = cube1.rebin(factor=factor)
# Compute the expected output cube. The zero valued pixels are all
# masked, leaving just pixels with values of 1, so the mean that is
# recorded in each output pixel should be 1.0.
expected = ma.ones((2, 2, 2)) * 1.0
expected.mask = expected < 0 # All output pixels should be unmasked.
assert_masked_allclose(cube2.data, expected)
# Check that the world coordinates are correct. We averaged
# factor[] pixels whose coordinates were equal to their pixel
# indexes, so the coordinates of the first pixel of the rebinned
# cube should be the mean of the first factor indexes along each
# dimension. The sum from k=0 to factor-1 is
# ((factor-1)*factor)/2, and dividing this by the number of pixels
# gives (factor-1)/2.
assert_allclose(np.asarray(cube2.get_start()), (factor - 1) / 2.0)
# Create a cube that has a larger number of pixels along the
# y and x axes of the images, so that we can divide those axes
# by a number whose remainder is large enough to test selection
# of the truncated part of the cube.
shape = np.array([4, 17, 15])
data = ma.ones(shape) # Start with all pixels 1.0
data.reshape(shape.prod())[::2] = 0.0 # Set every second pixel to 0.0
data.mask = data < -1 # Unmask all pixels.
cube1 = generate_cube(data=data.data, mask=data.mask, wcs=wcs, wave=wave)
# Choose the rebinning factors such that there is a significant
# remainder after dividing the final two dimensions of the cube by
# the specified factor. We don't do this for the first axis because
# we want the interleaved pattern of 0s and 1s to remain.
factor = np.array([2, 7, 9])
cube2 = cube1.rebin(factor=factor, margin='origin')
# Compute the expected output cube. Given that the input cube is a
# repeating pattern of 1,0, and we divided the x-axis by a
# multiple of 2, the output pixels should all be 0.5.
expected_shape = cube1.shape // factor
expected = ma.ones(expected_shape) * 0.5
expected.mask = expected < 0 # All pixels unmasked.
assert_masked_allclose(cube2.data, expected)
# We chose a margin value of 'origin', so that the outer corner of
# pixel 0,0,0 of the input cube would also be the outer corner of
# pixel 0,0,0 of the rebinned cube. The output world coordinates
# can be calculated as described for the previous test.
assert_allclose(np.asarray(cube2.get_start()), (factor - 1) / 2.0)
# Do the same test, but with margin='center'.
cube2 = cube1.rebin(factor=factor, margin='center')
# Compute the expected output cube. The values should be the
# same as the previous test.
expected_shape = cube1.shape // factor
expected = ma.ones(expected_shape) * 0.5
expected.mask = expected < 0 # All pixels unmasked.
assert_masked_allclose(cube2.data, expected)
# We chose a margin value of 'center'. We need to know which
# pixels should have contributed to pixel 0,0,0. First calculate
# how many pixels remain after dividing the shape by the reduction
# factors. This is the number of extra pixels that should have been
# discarded. Divide this by 2 to determine the number of pixels that
# are removed from the start of each axis.
cut = np.mod(shape, factor).astype(int) // 2
# The world coordinates of the input cube were equal to its pixel
# indexes, and the world coordinates of a pixel of the output cube
# is thus the mean of the indexes of the pixels that were combined
# to make that pixel. In this case, we combine pixels
# data[cut:cut+factor] along each axis. This can be calculated as
# the sum of pixel indexes from 0 to cut+factor, minus the sum of
# pixel indexes from 0 to cut, with the result divided by the number
# of pixels that were averaged. Again we make use of the series sum,
# sum[n=0..N] = (n*(n-1))/2.
tmp = cut + factor
assert_allclose(np.asarray(cube2.get_start()),
((tmp * (tmp - 1)) / 2.0 -
(cut * (cut - 1)) / 2.0) / factor)
def test_bandpass_image():
"""Cube class: testing bandpass_image"""
shape = (7, 2, 2)
# Create a rectangular shaped bandpass response whose ends are half
# way into pixels.
wavelengths = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
sensitivities = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
# Specify a ramp for the values of the pixels in the cube versus
# wavelength.
spectral_values = np.arange(shape[0], dtype=float)
# Specify a ramp for the variances of the pixels in the cube
# versus wavelength.
spectral_vars = np.arange(shape[0], dtype=float) * 0.5
# Calculate the expected weights versus wavelength for each
# spectral pixels. The weight of each pixel is supposed to be
# integral of the sensitivity function over the width of the pixel.
#
# | 0 | 1 | 2 | 3 | 4 | 5 | 6 | Pixel indexes
# _______________________
# _______| |_________ Sensitivities
#
# 0.0 0.5 1.0 1.0 1.0 0.5 0.0 Weights
# 0.0 1.0 2.0 3.0 4.0 5.0 6.0 Pixel values vs wavelength
# 0.0 0.5 2.0 3.0 4.0 2.5 0.0 Pixel values * weights
weights = np.array([0.0, 0.5, 1.0, 1.0, 1.0, 0.5, 0.0])
# Compute the expected weighted mean of the spectral pixel values,
# assuming that no pixels are unmasked.
unmasked_mean = (weights * spectral_values).sum() / weights.sum()
# Compute the expected weighted mean if pixel 1 is masked.
masked_pixel = 1
masked_mean = (((weights * spectral_values).sum() -
weights[masked_pixel] * spectral_values[masked_pixel]) /
(weights.sum() - weights[masked_pixel]))
# Compute the expected variances of the unmasked and masked means.
unmasked_var = (weights**2 * spectral_vars).sum() / weights.sum()**2
masked_var = (((weights**2 * spectral_vars).sum() -
weights[masked_pixel]**2 * spectral_vars[masked_pixel]) /
(weights.sum() - weights[masked_pixel])**2)
# Create the data array of the cube, giving all map pixels the
# same data and variance spectrums.
data = spectral_values[:, np.newaxis, np.newaxis] * np.ones(shape)
var = spectral_vars[:, np.newaxis, np.newaxis] * np.ones(shape)
# Create a mask with all pixels unmasked.
mask = np.zeros(shape)
# Mask spectral pixel 'masked_pixel' of map index 1,1.
mask[masked_pixel, 1, 1] = True
# Also mask all pixels of map pixel 0,0.
mask[:, 0, 0] = True
# Create a test cube with the above data and mask arrays.
c = generate_cube(shape=shape, data=data, mask=mask, var=var,
wave=WaveCoord(crval=0.0, cdelt=1.0, crpix=1.0,
cunit=u.angstrom))
# Extract an image that has the above bandpass response.
im = c.bandpass_image(wavelengths, sensitivities)
# Only the map pixel in which all spectral pixels are masked should
# be masked in the output, so just map pixel [0,0] should be masked.
expected_mask = np.array([[True, False],
[False, False]], dtype=bool)
# What do we expect?
expected_data = ma.array(
data=[[unmasked_mean, unmasked_mean], [unmasked_mean, masked_mean]],
mask=expected_mask)
expected_var = ma.array(
data=[[unmasked_var, unmasked_var], [unmasked_var, masked_var]],
mask=expected_mask)
# Are the results consistent with the predicted values?
assert_masked_allclose(im.data, expected_data)
assert_masked_allclose(im.var, expected_var)
def test_get_item():
"""Cube class: testing __getitem__"""
# Set the shape and contents of the cube's data array.
shape = (3, 4, 5)
data = np.arange(shape[0] * shape[1] * shape[2])\
.reshape(shape[0], shape[1], shape[2])
# Create a test cube with the above data array.
c = generate_cube(data=data, shape=shape,
wave=WaveCoord(crval=1, cunit=u.angstrom))
c.primary_header['KEY'] = 'primary value'
c.data_header['KEY'] = 'data value'
# Select the whole cube.
for r in [c[:, :, :], c[:, :], c[:]]:
assert_array_equal(r.shape, c.shape)
assert_allclose(r.data, c.data)
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Cube)
assert r.wcs.isEqual(c.wcs)
assert r.wave.isEqual(c.wave)
# Select a subimage that only has one pixel along the y axis.
for r in [c[1, 2, :], c[1, 2]]:
assert_array_equal(r.shape, (1, c.shape[2]))
assert_allclose(r.data.ravel(), c.data[1, 2, :].ravel())
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Image)
assert r.wcs.isEqual(c.wcs[2, :])
assert r.wave is None
# Select a subcube that only has one pixel along the y axis.
for r in [c[:, 2, :], c[:, 2]]:
assert_array_equal(r.shape, (c.shape[0], 1, c.shape[2]))
assert_allclose(r.data.ravel(), c.data[:, 2, :].ravel())
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Cube)
assert r.wcs.isEqual(c.wcs[2, :])
assert r.wave.isEqual(c.wave)
# Select a subimage that only has one pixel along the x axis.
r = c[1, :, 2]
assert_array_equal(r.shape, (c.shape[1], 1))
assert_allclose(r.data.ravel(), c.data[1, :, 2].ravel())
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Image)
assert r.wcs.isEqual(c.wcs[:, 2])
assert r.wave is None
# Select a subcube that only has one pixel along the x axis.
r = c[:, :, 2]
assert_array_equal(r.shape, (c.shape[0], c.shape[1], 1))
assert_allclose(r.data.ravel(), c.data[:, :, 2].ravel())
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Cube)
assert r.wcs.isEqual(c.wcs[:, 2])
assert r.wave.isEqual(c.wave)
# Select a sub-cube using a non-scalar slice along each axis.
r = c[:, :2, :2]
assert_array_equal(r.shape, (c.shape[0], 2, 2))
assert_allclose(r.data.ravel(), c.data[:, :2, :2].ravel())
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Cube)
assert r.wcs.isEqual(c.wcs[:2, :2])
assert r.wave.isEqual(c.wave)
# Select the image of a single spectral pixel.
for r in [c[1, :, :], c[1, :], c[1]]:
assert_array_equal(r.shape, (c.shape[1], c.shape[2]))
assert_allclose(r.data.ravel(), c.data[1, :, :].ravel())
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Image)
assert r.wcs.isEqual(c.wcs)
assert r.wave is None
# Select the spectrum of a single spatial pixel.
r = c[:, 2, 2]
assert_array_equal(r.shape, (c.shape[0], ))
assert_allclose(r.data.ravel(), c.data[:, 2, 2].ravel())
assert r.primary_header['KEY'] == c.primary_header['KEY']
assert r.data_header['KEY'] == c.data_header['KEY']
assert isinstance(r, Spectrum)
assert r.wave.isEqual(c.wave)
assert r.wcs is None
# Select a single pixel.
r = c[2, 2, 2]
assert np.isscalar(r)
assert_allclose(r, c.data[2, 2, 2])
def test_mask(cube):
"""Cube class: testing mask functionalities"""
# A region of half-width=1 and half-height=1 should have a size of
# 2x2 pixels. A 2x2 region of pixels has a center at the shared
# corner of the 4 pixels, and the closest corner to the requested
# center of 2.1,1.8 is 2.5,1.5, so we expect the square of unmasked pixels
# to be pixels 2,3 along the Y axis, and pixels 1,2 along the X axis.
cube.mask_region((2.1, 1.8), (1, 1), lmin=2, lmax=5, inside=True,
unit_center=None, unit_radius=None, unit_wave=None)
# The expected mask for the images between lmin and lmax.
expected_mask = np.array([[False, False, False, False, False],
[False, False, False, False, False],
[False, True, True, False, False],
[False, True, True, False, False],
[False, False, False, False, False],
[False, False, False, False, False]], dtype=bool)
assert_array_equal(np.any(cube.mask[:2, :, :], axis=0),
np.zeros(cube.shape[1:]))
assert_array_equal(np.all(cube._mask[2:5, :, :], axis=0), expected_mask)
assert_array_equal(np.any(cube.mask[5:, :, :], axis=0),
np.zeros(cube.shape[1:]))
cube.unmask()
# Do the same experiment, but this time mask pixels outside the region
# instead of within the region.
cube.mask_region((2.1, 1.8), (1, 1), lmin=2, lmax=5, inside=False,
unit_center=None, unit_radius=None, unit_wave=None)
assert_array_equal(np.any(cube.mask[:2, :, :], axis=0),
np.ones(cube.shape[1:]))
assert_array_equal(np.all(cube._mask[2:5, :, :], axis=0), ~expected_mask)
assert_array_equal(np.any(cube.mask[5:, :, :], axis=0),
np.ones(cube.shape[1:]))
cube.unmask()
# Do the same experiment, but this time with the center and size
# of the region specified using the equivalent world coordinates.
wcs = WCS(deg=True)
wave = WaveCoord(cunit=u.angstrom)
cube = Cube(data=cube.data, wave=wave, wcs=wcs, copy=False)
cube.mask_region(wcs.pix2sky([2.1, 1.8]), (3600, 3600), lmin=2, lmax=5,
inside=False)
assert_array_equal(np.any(cube.mask[:2, :, :], axis=0),
np.ones(cube.shape[1:]))
assert_array_equal(np.all(cube._mask[2:5, :, :], axis=0), ~expected_mask)
assert_array_equal(np.any(cube.mask[5:, :, :], axis=0),
np.ones(cube.shape[1:]))
cube.unmask()
# Mask around a region of half-width and half-height 1.1 pixels,
# specified in arcseconds, centered close to pixel 2.4,3.8. This
# ideally corresponds to a region of 2.2x2.2 pixels. The closest
# possible size is 2x2 pixels. A region of 2x2 pixels has its
# center at the shared corner of these 4 pixels, and the nearest
# corner to the desired central index of (2.4,3.8) is (2.5,3.5).
# So all of the image should be masked, except for a 2x2 area of
# pixel indexes 2,3 along the Y axis and pixel indexes 3,4 along
# the X axis.
cube.mask_region(wcs.pix2sky([2.4, 3.8]), 1.1 * 3600.0, inside=False,
lmin=2, lmax=5)
# The expected mask for the images between lmin and lmax.
expected_mask = np.array([[True, True, True, True, True],
[True, True, True, True, True],
[True, True, True, False, False],
[True, True, True, False, False],
[True, True, True, True, True],
[True, True, True, True, True]], dtype=bool)
assert_array_equal(np.any(cube.mask[:2, :, :], axis=0),
np.ones(cube.shape[1:]))
assert_array_equal(np.all(cube._mask[2:5, :, :], axis=0), expected_mask)
assert_array_equal(np.any(cube.mask[5:, :, :], axis=0),
np.ones(cube.shape[1:]))
cube.unmask()
# Mask outside an elliptical region centered at pixel 3.5,3.5.
# The boolean expected_mask array given below was a verified
# output of mask_ellipse() for the specified ellipse parameters.
cube = generate_cube(shape=(10, 8, 8), wcs=wcs, var=None)
cube.mask_ellipse([3.5, 3.5], (2.5, 3.5), 45.0, unit_radius=None,
unit_center=None, inside=False, lmin=2, lmax=5)
expected_mask = np.array([
[True, True, True, True, True, True, True, True],
[True, True, True, False, False, False, True, True],
[True, True, False, False, False, False, False, True],
[True, False, False, False, False, False, False, True],
[True, False, False, False, False, False, False, True],
[True, False, False, False, False, False, True, True],
[True, True, False, False, False, True, True, True],
[True, True, True, True, True, True, True, True]],
dtype=bool)
assert_array_equal(np.any(cube.mask[:2, :, :], axis=0),
np.ones(cube.shape[1:]))
assert_array_equal(np.all(cube._mask[2:5, :, :], axis=0), expected_mask)
assert_array_equal(np.any(cube.mask[5:, :, :], axis=0),
np.ones(cube.shape[1:]))
# Check that we can select the same mask via the output of np.where()
# passed to mask_selection().
ksel = np.where(cube.data.mask)
cube.unmask()
cube.mask_selection(ksel)
assert_array_equal(np.any(cube.mask[:2, :, :], axis=0),
np.ones(cube.shape[1:]))
assert_array_equal(np.all(cube._mask[2:5, :, :], axis=0), expected_mask)
assert_array_equal(np.any(cube.mask[5:, :, :], axis=0),
np.ones(cube.shape[1:]))
# The cube was generated without any variance information.
# Check that mask_variance() raises an error due to this.
with pytest.raises(ValueError):
cube.mask_variance(0.1)
# Add an array of variances to the cube and check that mask_variance()
# masks all pixels that have variances greater than 0.1.
cube.unmask()
cube.var = np.random.randn(*cube.shape)
mask = cube.var > 0.1
cube.mask_variance(0.1)
assert_array_equal(cube.data.mask, mask)
