/
test_arithmetic.py
107 lines (68 loc) · 3.15 KB
/
test_arithmetic.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
import astropy.units as u
import numpy as np
import pytest
from ..spectra.spectrum1d import Spectrum1D
from .spectral_examples import simulated_spectra
def test_spectral_axes():
flux1 = (np.random.sample(49) * 100).astype(int)
flux2 = (np.random.sample(49) * 100).astype(int)
flux3 = flux1 + flux2
spec1 = Spectrum1D(spectral_axis=np.arange(1, 50) * u.nm,
flux=flux1 * u.Jy)
spec2 = Spectrum1D(spectral_axis=np.arange(1, 50) * u.nm,
flux=flux2 * u.Jy)
spec3 = spec1 + spec2
assert np.allclose(spec3.flux.value, flux3)
def test_add_basic_spectra(simulated_spectra):
# Get the numpy array of data
flux1 = simulated_spectra.s1_um_mJy_e1_flux
flux2 = simulated_spectra.s1_um_mJy_e2_flux
flux3 = flux1 + flux2
# Calculate using the spectrum1d/nddata code
spec3 = simulated_spectra.s1_um_mJy_e1 + simulated_spectra.s1_um_mJy_e2
assert np.allclose(spec3.flux.value, flux3)
def test_add_diff_flux_prefix(simulated_spectra):
# Get the numpy array of data
# this assumes output will be in mJy units
flux1 = simulated_spectra.s1_AA_mJy_e3_flux
flux2 = simulated_spectra.s1_AA_nJy_e4_flux
flux3 = flux1 + (flux2 / 1000000)
# Calculate using the spectrum1d/nddata code
spec3 = simulated_spectra.s1_AA_mJy_e3 + simulated_spectra.s1_AA_nJy_e4
assert np.allclose(spec3.flux.value, flux3)
def test_subtract_basic_spectra(simulated_spectra):
# Get the numpy array of data
flux1 = simulated_spectra.s1_um_mJy_e1_flux
flux2 = simulated_spectra.s1_um_mJy_e2_flux
flux3 = flux2 - flux1
# Calculate using the spectrum1d/nddata code
spec3 = simulated_spectra.s1_um_mJy_e2 - simulated_spectra.s1_um_mJy_e1
assert np.allclose(spec3.flux.value, flux3)
def test_divide_basic_spectra(simulated_spectra):
# Get the numpy array of data
flux1 = simulated_spectra.s1_um_mJy_e1_flux
flux2 = simulated_spectra.s1_um_mJy_e2_flux
flux3 = flux1 / flux2
# Calculate using the spectrum1d/nddata code
spec3 = simulated_spectra.s1_um_mJy_e1 / simulated_spectra.s1_um_mJy_e2
assert np.allclose(spec3.flux.value, flux3)
def test_multiplication_basic_spectra(simulated_spectra):
# Get the numpy array of data
flux1 = simulated_spectra.s1_um_mJy_e1_flux
flux2 = simulated_spectra.s1_um_mJy_e2_flux
flux3 = flux1 * flux2
# Calculate using the spectrum1d/nddata code
spec3 = simulated_spectra.s1_um_mJy_e1 * simulated_spectra.s1_um_mJy_e2
assert np.allclose(spec3.flux.value, flux3)
@pytest.mark.xfail(raises=ValueError)
def test_add_diff_spectral_axis(simulated_spectra):
# Calculate using the spectrum1d/nddata code
spec3 = simulated_spectra.s1_um_mJy_e1 + simulated_spectra.s1_AA_mJy_e3
def test_masks(simulated_spectra):
masked_spec = simulated_spectra.s1_um_mJy_e1_masked
masked_sum = masked_spec + masked_spec
assert np.all(masked_sum.mask == simulated_spectra.s1_um_mJy_e1_masked.mask)
masked_sum.mask[:50] = True
masked_diff = masked_sum - masked_spec
assert u.allclose(masked_diff.flux, masked_spec.flux)
assert np.all(masked_diff.mask == masked_sum.mask | masked_spec.mask)