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_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_get_item():
"""Spectrum class: testing __getitem__"""
# Set the shape and contents of the spectrum's data array.
shape = (5, )
data = np.arange(shape[0])
# Create a test spectrum with the above data array.
s = generate_spectrum(data=data,
shape=shape,
wave=WaveCoord(crval=1, cunit=u.angstrom))
s.primary_header['KEY'] = 'primary value'
s.data_header['KEY'] = 'data value'
# Select the whole spectrum.
r = s[:]
assert_array_equal(r.shape, s.shape)
assert_allclose(r.data, s.data)
assert r.primary_header['KEY'] == s.primary_header['KEY']
assert r.data_header['KEY'] == s.data_header['KEY']
assert isinstance(r, Spectrum)
assert r.wcs is None
assert r.wave.isEqual(s.wave)
# Select a sub-spectrum.
r = s[1:3]
assert_array_equal(r.shape, (2))
assert_allclose(r.data, s.data[1:3])
assert r.primary_header['KEY'] == s.primary_header['KEY']
assert r.data_header['KEY'] == s.data_header['KEY']
assert isinstance(r, Spectrum)
assert r.wave.isEqual(s.wave[1:3])
assert r.wcs is None
# Select a single pixel of the spectrum.
r = s[2]
assert np.isscalar(r)
assert_allclose(r, s.data[2])