def test_profile_fit_with_noise_with_blaze():
nx, ny = 401, 403
input_profile, trace_centers, input_traces = make_simple_traces(nx,
ny,
blaze=True)
read_noise = 10.0
# Filter out the numerical noise
input_profile[input_profile < 1e-15] = 0.0
input_profile += np.random.poisson(input_profile)
input_profile += np.random.normal(0.0,
read_noise,
size=input_profile.shape)
uncertainty = np.sqrt(input_profile + read_noise**2.0)
image = NRESObservationFrame([
CCDData(data=input_profile.copy(),
uncertainty=uncertainty,
meta={'OBJECTS': 'tung&tung&none'})
], 'foo.fits')
image.traces = input_traces
input_context = context.Context({})
stage = ProfileFitter(input_context)
image = stage.do_stage(image)
for i in range(1, np.max(image.traces) + 1):
this_trace = get_trace_region(image.traces == i)
scale_factor = np.max(input_profile[this_trace]) / np.max(
image.profile[this_trace] * image.blaze[i - 1]['blaze'])
np.testing.assert_allclose(input_profile[this_trace],
image.profile[this_trace] *
image.blaze[i - 1]['blaze'] * scale_factor,
rtol=1.5e-2,
atol=1.0)
def do_stage(self, image):
image.add_or_update(
ArrayData(np.zeros_like(image.data), meta={}, name='PROFILE'))
blazes = np.zeros((image.num_traces, image.data.shape[1]), dtype=float)
blazes_errors = np.zeros((image.num_traces, image.data.shape[1]),
dtype=float)
trace_ids = np.arange(1, image.num_traces + 1, dtype=int)
for i, trace_id in enumerate(trace_ids):
# Extract the pixels from the spectrum extension in that order
this_trace = get_trace_region(image.traces == trace_id)
# Evaluate the best fit spline along the center of the trace, normalize, save as BLAZE
# do a flux weighted mean in the y direction
fluxes = image.data[this_trace]
flux_errors = image.uncertainty[this_trace]
blaze = (fluxes * fluxes).sum(axis=0) / fluxes.sum(axis=0)
# Uncertainty propagation is a bear
blaze_errors = np.sqrt(
np.sum((np.sqrt(2.0) * fluxes * flux_errors)**2.0, axis=0))
flux_sum_error = np.sqrt((flux_errors * flux_errors).sum(axis=0))
blaze_errors = blaze * np.sqrt(
(blaze_errors / (fluxes * fluxes).sum(axis=0))**2.0 +
(flux_sum_error / fluxes.sum(axis=0))**2.0)
# Normalize so the sum of the blaze is 1
blaze_sum = blaze.sum()
blaze_sum_error = np.sqrt((blaze * blaze).sum())
blaze /= blaze_sum
blaze_errors = blaze * np.sqrt((blaze_errors /
(blaze * blaze_sum))**2 +
(blaze_sum_error / blaze_sum)**2.0)
# Evaluate the best fit spline/BLAZE save in an empty array, normalize (PROFILE extension)
image.profile[this_trace] = fluxes / blaze
image.profile[this_trace] /= image.profile[this_trace].sum(axis=0)
image.profile[this_trace][fluxes < 0.0] = 0.0
# Use the 4 bit to represent negative values in the profile
image.mask[this_trace][fluxes < 0.0] |= 4
# Divide the original image by best fit spline (Pixel-to-pixel sensitivity)
# TODO: propagate the division of blaze and profile to the uncertainties of the image!!!
image.data[this_trace] /= blaze
image.data[this_trace] /= image.profile[this_trace]
# Normalize the pixel-to-pixel sensitivity by the median
image.data[this_trace] /= np.median(this_trace)
# get the horizontal (x) extent of the trace.
x_extent = slice(np.min(this_trace[1]), np.max(this_trace[1]) + 1)
blazes[i, x_extent] = blaze
blazes_errors[i, x_extent] = blaze_errors
image.add_or_update(
DataTable(Table({
'id': trace_ids,
'blaze': blazes,
'blaze_error': blazes_errors
}),
name='BLAZE'))
return image
def do_stage(self, image: NRESObservationFrame):
if image.profile is None:
logger.error('Profile missing. Rejecting image.', image=image)
return None
image.add_or_update(
ArrayData(np.zeros_like(image.data, dtype=float), name='WEIGHTS'))
trace_ids = np.arange(1, image.num_traces + 1)
for trace_id in trace_ids:
yx = get_trace_region(np.isclose(image.traces, trace_id))
image.weights[yx] = self.weights(image.profile[yx],
image.uncertainty[yx]**2,
image.mask[yx])
return image
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_get_trace_region():
nx, ny = 1001, 1003
test_data = np.zeros((ny, nx))
x0 = 498
trace = np.round(2e-4 * (np.arange(nx) - x0)**2.0 + 0.01 *
(np.arange(nx) - x0) + 502.0).astype(int)
counter = 1
for i in np.arange(nx, dtype=int):
for j in range(-7, 8):
test_data[trace[i] + j, i] = counter
counter += 1
traces = np.zeros((ny, nx), dtype=bool)
traces[test_data > 0] = True
trace_region = trace_utils.get_trace_region(traces)
counter = 1
expected_trace = np.zeros((15, nx), dtype=int)
for i in np.arange(nx, dtype=int):
for j in range(15):
expected_trace[j, i] = counter
counter += 1
np.testing.assert_equal(test_data[trace_region], expected_trace)
def test_get_region():
traces = np.zeros((4, 4))
traces[[2, 3], :] = 1
trace_yx_pos = trace_utils.get_trace_region(traces == 1)
assert np.allclose(traces[trace_yx_pos], 1)