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_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_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_arithmetic_images(image):
image2 = generate_image(data=1, unit=u.Unit('2 ct'))
for op in (add, sub, mul, div):
image3 = op(image, image2)
assert_allclose(
image3._data,
op(image._data * image.unit, image2._data * image2.unit)
.to(image3.unit).value)
def test_get_item():
"""Image class: testing __getitem__"""
# Set the shape and contents of the image's data array.
shape = (4, 5)
data = np.arange(shape[0] * shape[1]).reshape(shape[0], shape[1])
# Create a test image with the above data array.
im = generate_image(data=data, shape=shape)
im.primary_header['KEY'] = 'primary value'
im.data_header['KEY'] = 'data value'
# Select the whole image.
for r in [im[:, :], im[:]]:
assert_array_equal(r.shape, im.shape)
assert_allclose(r.data, im.data)
assert r.primary_header['KEY'] == im.primary_header['KEY']
assert r.data_header['KEY'] == im.data_header['KEY']
assert isinstance(r, Image)
assert r.wcs.isEqual(im.wcs)
assert r.wave is None
# Select a subimage that only has one pixel along the y axis.
for r in [im[2, :], im[2]]:
assert_array_equal(r.shape, (1, im.shape[1]))
assert_allclose(r.data.ravel(), im.data[2, :].ravel())
assert r.primary_header['KEY'] == im.primary_header['KEY']
assert r.data_header['KEY'] == im.data_header['KEY']
assert isinstance(r, Image)
assert r.wcs.isEqual(im.wcs[2, :])
assert r.wave is None
# Select a subimage that only has one pixel along the x axis.
r = im[:, 2]
assert_array_equal(r.shape, (im.shape[0], 1))
assert_allclose(r.data.ravel(), im.data[:, 2].ravel())
assert r.primary_header['KEY'] == im.primary_header['KEY']
assert r.data_header['KEY'] == im.data_header['KEY']
assert isinstance(r, Image)
assert r.wcs.isEqual(im.wcs[:, 2])
assert r.wave is None
# Select a sub-image using a non-scalar slice along each axis.
r = im[1:2, 2:4]
assert_array_equal(r.shape, (1, 2))
assert_allclose(r.data.ravel(), im.data[1:2, 2:4].ravel())
assert r.primary_header['KEY'] == im.primary_header['KEY']
assert r.data_header['KEY'] == im.data_header['KEY']
assert isinstance(r, Image)
assert r.wcs.isEqual(im.wcs[1:2, 2:4])
assert r.wave is None
# Select a single pixel.
r = im[2, 3]
assert np.isscalar(r)
assert_allclose(r, im.data[2, 3])
def test_prepare_data():
image = generate_image(data=2.0)
image[1, 1] = np.ma.masked
image[3:5, 2:4] = np.ma.masked
data = image._prepare_data()
assert not np.ma.is_masked(data)
assert np.allclose(data, 2.0)
data = image._prepare_data(interp='linear')
assert not np.ma.is_masked(data)
assert np.allclose(data, 2.0)
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_crop():
"""Image class: testing crop method"""
# Create an image whose pixels are all masked.
image1 = generate_image(shape=(9, 7), data=2.0, var=0.5, mask=True)
# Create a masked array of unmasked values to be assigned to the
# part of the image, with just a diamond shaped area of pixels
# unmasked.
diamond = np.ma.array(data=[[6.0, 2.0, 9.0],
[1.0, 4.0, 8.0],
[0.0, 5.0, 3.0]],
mask=[[True, False, True],
[False, False, False],
[True, False, True]])
# Assign the above array to part of the image to clear the mask of
# an irregular rectangular area of pixels there.
image1.data[2:5, 1:4] = diamond
# The following should crop all but the rectangular area that was
# assigned above.
image1.crop()
# Check that the masked data array is as expected.
assert_masked_allclose(image1.data, diamond)
# The cropped variance array should look like the following array.
expected_var = np.ma.array(data=[[0.5, 0.5, 0.5],
[0.5, 0.5, 0.5],
[0.5, 0.5, 0.5]], mask=diamond.mask)
# Check that the masked variance array is as expected.
assert_masked_allclose(image1.var, expected_var)
# Check the WCS characteristics of the cropped image.
assert_image_equal(image1, shape=(3, 3), start=(2, 1), end=(4, 3))
assert image1.get_rot() == 0