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