def test_do_stage(mock_normalize):
# tests that the normalized spectrum are assigned correctly.
expected = 1
# make a single order spectrum
flux = np.arange(1, 101) * expected
x = np.arange(1, 101)
spectrum = Table({
'wavelength': [x, 2 * x],
'flux': [flux, 2 * flux],
'blaze': [flux, 2 * flux],
'blaze_error': [flux, 2 * flux],
'uncertainty': [flux, 2 * flux],
'fiber': [0, 0],
'order': [1, 2]
})
image = NRESObservationFrame(
[CCDData(np.zeros(
(1, 1)), meta={'OBJECTS': 'tung&tung&none'})], 'test.fits')
image.spectrum = Spectrum1D(spectrum)
# make it so that ContinuumNormalizer.normalize just returns ones.
mock_normalize.return_value = (expected * np.ones_like(flux),
expected * np.ones_like(flux))
# Run the normalizer code
stage = ContinuumNormalizer(SimpleNamespace(db_address='foo'))
image = stage.do_stage(image)
for order in [1, 2]:
assert np.allclose(image.spectrum[0, order]['normflux'], expected)
assert np.allclose(image.spectrum[0, order]['normuncertainty'],
expected)
def do_stage(self, image: NRESObservationFrame):
if image.weights is None:
logger.error('Extraction weights are missing. Rejecting image.',
image=image)
return None
if image.wavelengths is None:
logger.error('Wavelengths are missing. Rejecting image.',
image=image)
return None
# consider adding a method to image, i.e. image.extracted_spectrum_shape.
flux = np.zeros((image.num_traces, image.data.shape[1]), dtype=float)
wavelength = np.zeros_like(flux, dtype=float)
variance = np.zeros_like(flux, dtype=float)
mask = np.zeros_like(flux, dtype=np.uint8)
trace_ids = np.arange(1, image.num_traces + 1, dtype=int)
for i, trace_id in enumerate(trace_ids):
this_trace = get_trace_region(np.isclose(image.traces, trace_id))
# get the horizontal (x) extent of the trace. Consider making this a get_extent function.
x_extent = slice(np.min(this_trace[1]), np.max(this_trace[1]) + 1)
flux[i, x_extent] = self.extract_order(image.data[this_trace],
image.weights[this_trace])
variance[i, x_extent] = self.extract_order(
image.uncertainty[this_trace]**2, image.weights[this_trace]**2)
# TODO: some fix that allows for actual wavelength extraction:
# wavelength[i, x_extent] = self.extract_order(image.wavelengths[this_trace],
# weights=1/image.wavelengths[this_trace].shape[0])
wavelength[i, x_extent] = np.median(image.wavelengths[this_trace],
axis=0)
mask[i, x_extent] = image.weights[this_trace].sum(axis=0) == 0.0
image.spectrum = Spectrum1D({
'id': trace_ids,
'order': image.fibers['order'],
'fiber': image.fibers['fiber'],
'wavelength': wavelength,
'flux': flux,
'uncertainty': np.sqrt(variance),
'blaze': image.blaze['blaze'],
'blaze_error': image.blaze['blaze_error'],
'mask': mask
})
return image
def test_do_stage_does_not_fit_non_science_fiber():
# make a single order spectrum
flux = np.arange(1, 101) * 2
x = np.arange(1, 101)
spectrum = Table({
'wavelength': [x],
'flux': [flux],
'blaze': [flux],
'blaze_error': [flux],
'uncertainty': [flux],
'fiber': [1],
'order': [1]
})
image = NRESObservationFrame(
[CCDData(np.zeros(
(1, 1)), meta={'OBJECTS': 'tung&tung&none'})], 'test.fits')
image.spectrum = Spectrum1D(spectrum)
# Run the normalizer code
stage = ContinuumNormalizer(SimpleNamespace(db_address='foo'))
image = stage.do_stage(image)
assert 'normflux' not in image.spectrum._table.colnames
assert 'normuncertainty' not in image.spectrum._table.colnames
def test_rv(mock_loader, mock_db):
# parameters that define the test data
num_orders = 5
lam_per_order = 70
res = 0.1
# Make fake data
test_wavelengths = np.arange(4500.0,
4500 + num_orders * lam_per_order + 100, res)
flux = np.ones(test_wavelengths.shape) * 1000
read_noise = 10.0
# add in fake lines
for i in range(1000):
central_wavelength = np.random.uniform() * 2000.0 + 4500.0
width = 0.1 # angstroms
sigma = width / 2 / np.sqrt(2.0 * np.log(2.0))
flux += gaussian(test_wavelengths, central_wavelength,
sigma) * 10000.0 * np.random.uniform()
noisy_flux = np.random.poisson(flux).astype(float)
noisy_flux += np.random.normal(0.0, read_noise, size=len(flux))
uncertainty = np.sqrt(flux + read_noise**2.0)
class MockLoader:
def load(self, *args):
return {'wavelength': test_wavelengths, 'flux': flux}
mock_loader.return_value = MockLoader()
true_v = 1.205 * units.km / units.s
# Make the fake image
# Set the site to the north pole, and the ra and dec to the ecliptic pole. The time we chose is just
# to minimize the rv correction
header = fits.Header({
'DATE-OBS': '2020-09-12T00:00:00.000000',
'RA': '18:00:00',
'DEC': '+66:33:32.8178',
'OBJECTS': 'foo&none&none',
'EXPTIME': 1200.0
})
site_info = {'longitude': 0.0, 'latitude': 90.0, 'elevation': 0.0}
# Each NRES order is ~70 Angstroms wide
spectrum = []
dx = int(lam_per_order / res)
for i in range(5):
order = slice(i * dx, (i + 1) * dx, 1)
row = {
'wavelength':
test_wavelengths[order] * (1.0 + true_v / constants.c),
'normflux': noisy_flux[order],
'normuncertainty': uncertainty[order],
'fiber': 0,
'order': i
}
spectrum.append(row)
image = NRESObservationFrame([CCDData(np.zeros((1, 1)), meta=header)],
'test.fits')
image.spectrum = Spectrum1D(spectrum)
image.instrument = SimpleNamespace()
image.instrument.site = 'npt'
# Classification just can't be None so that the stage does not abort.
image.classification = SimpleNamespace(**{
'T_effective': 5000.0,
'log_g': 0.0,
'metallicity': 0.0,
'alpha': 0.0
})
mock_db.return_value = SimpleNamespace(**site_info)
# Run the RV code
stage = RVCalculator(
SimpleNamespace(db_address='foo',
MIN_ORDER_TO_CORRELATE=0,
MAX_ORDER_TO_CORRELATE=num_orders - 1))
image = stage.do_stage(image)
# Assert that the true_v + rv_correction is what is in the header within 5 m/s
assert np.abs(
true_v.to('m / s').value + image.meta['BARYCORR'] -
image.meta['RV']) < 5.0