def test_unpack_nddata(testdata, testvar, testmask):
nd = NDAstroData(testdata, mask=testmask)
nd.variance = testvar
out_data, out_mask, out_var = NDStacker.none(nd)
assert_allclose(out_data, testdata)
assert_allclose(out_var, testvar)
assert_allclose(out_mask, testmask)
def wrapper(sequence, scale=None, zero=None, *args, **kwargs):
nddata_list = list(sequence)
if scale is None:
scale = [1.0] * len(nddata_list)
if zero is None:
zero = [0.0] * len(nddata_list)
# Coerce all data to 32-bit floats. FITS data on disk is big-endian
# and preserving that datatype will cause problems with Cython
# stacking if the compiler is little-endian.
dtype = np.float32
data = np.empty((len(nddata_list),)+nddata_list[0].data.shape, dtype=dtype)
for i, (ndd, s, z) in enumerate(zip(nddata_list, scale, zero)):
data[i] = ndd.data * s + z
if any(ndd.mask is None for ndd in nddata_list):
mask = None
else:
mask = np.empty_like(data, dtype=DQ.datatype)
for i, ndd in enumerate(nddata_list):
mask[i] = ndd.mask
if any(ndd.variance is None for ndd in nddata_list):
variance = None
else:
variance = np.empty_like(data)
for i, (ndd, s, z) in enumerate(zip(nddata_list, scale, zero)):
variance[i] = ndd.variance * s*s
out_data, out_mask, out_var = fn(data=data, mask=mask,
variance=variance, *args, **kwargs)
# Can't instantiate NDAstroData with variance
ret_value = NDAstroData(out_data, mask=out_mask)
if out_var is not None:
ret_value.variance = out_var
return ret_value
def test_varclip():
# Confirm rejection of high pixel and correct output DQ
data = np.array([1., 1., 2., 2., 2., 100.]).reshape(6, 1)
ndd = NDAstroData(data)
ndd.mask = np.zeros_like(data, dtype=DQ.datatype)
ndd.mask[5, 0] = DQ.saturated
ndd.variance = np.ones_like(data)
stackit = NDStacker(combine="mean", reject="varclip")
result = stackit(ndd)
np.testing.assert_array_almost_equal(result.data, [1.6])
np.testing.assert_array_equal(result.mask, [0])
def test_varclip():
# Confirm rejection of high pixel and correct output DQ
data = np.array([1., 1., 2., 2., 2., 100.]).reshape(6, 1)
ndd = NDAstroData(data)
ndd.mask = np.zeros_like(data, dtype=DQ.datatype)
ndd.mask[5, 0] = DQ.saturated
ndd.variance = np.ones_like(data)
stackit = NDStacker(combine="mean", reject="varclip")
result = stackit(ndd)
assert_allclose(result.data, 1.6) # 100 is rejected
assert_allclose(result.mask, 0)
data = np.array([1., 1., 2., 2., 2., 100.]).reshape(6, 1)
ndd = NDAstroData(data)
ndd.variance = np.ones_like(data)
ndd.variance[5] = 400
stackit = NDStacker(combine="mean", reject="varclip", lsigma=3, hsigma=3)
result = stackit(ndd)
assert_allclose(result.data, 1.6) # 100 is rejected
stackit = NDStacker(combine="mean", reject="varclip", lsigma=5, hsigma=5)
result = stackit(ndd)
assert_allclose(result.data, 18) # 100 is not rejected
def standard_extraction(self, data, mask, var, aper_lower, aper_upper):
"""Uniform extraction across an aperture of width pixels"""
slitlength = data.shape[0]
all_x1 = self._center_pixels + aper_lower
all_x2 = self._center_pixels + aper_upper
ext = NDAstroData(data, mask=mask)
ext.variance = var
results = [sum1d(ext[:,i], x1, x2) for i, (x1, x2) in enumerate(zip(all_x1, all_x2))]
self.data[:] = [result.data for result in results]
if mask is not None:
self.mask[:] = [result.mask for result in results]
if var is not None:
self.var[:] = [result.variance for result in results]
def test_sum1d():
big_value = 10
data = np.ones((10, ))
ndd = NDAstroData(data)
ndd.mask = np.zeros_like(data, dtype=DQ.datatype)
ndd.mask[1] = 1
ndd.mask[2] = 2
ndd.data[4] = big_value
ndd.variance = np.ones_like(data)
x1 = -0.5
for x2 in np.arange(0., 4.5, 0.5):
result = sum1d(ndd, x1, x2, proportional_variance=True)
if x2 > 3.5:
np.testing.assert_almost_equal(4 + big_value * (x2 - 3.5),
result.data)
else:
np.testing.assert_almost_equal(x2 - x1, result.data)
np.testing.assert_almost_equal(x2 - x1, result.variance)
if x2 > 1.5:
assert result.mask == 3
elif x2 > 0.5:
assert result.mask == 1
else:
assert result.mask == 0
def extract(self,
ext,
width=None,
aper_lower=None,
aper_upper=None,
method='standard',
dispaxis=None,
viewer=None):
"""
Extract a 1D spectrum by following the model trace and extracting in
an aperture of a given number of pixels.
Parameters
----------
ext: single-slice AD object/ndarray
spectral image from which spectrum is to be extracted
width: float/None
full width of extraction aperture (in pixels)
aper_lower: float/None
lower extraction limit (pixels)
aper_upper: float/None
upper extraction limit (pixels)
method: str (standard|optimal)
type of extraction
dispaxis: int/None
dispersion axis (python sense)
viewer: Viewer/None
Viewer object on which to display extraction apertures
Returns
-------
NDAstroData: 1D spectrum
"""
if width is not None:
aper_upper = 0.5 * width
aper_lower = -aper_upper
if aper_lower is None:
aper_lower = self.aper_lower
if aper_upper is None:
aper_upper = self.aper_upper
if dispaxis is None:
dispaxis = 2 - ext.dispersion_axis() # python sense
slitlength = ext.shape[1 - dispaxis]
npix = ext.shape[dispaxis]
direction = "row" if dispaxis == 0 else "column"
self.check_domain(npix)
if aper_lower > aper_upper:
log.warning("Aperture lower limit is greater than upper limit.")
aper_lower, aper_upper = aper_upper, aper_lower
if aper_lower > 0:
log.warning("Aperture lower limit is greater than zero.")
if aper_upper < 0:
log.warning("Aperture upper limit is less than zero.")
# make data look like it's dispersed horizontally
# (this is best for optimal extraction, but not standard)
try:
mask = ext.mask
except AttributeError: # ext is just an ndarray
data = ext if dispaxis == 1 else ext.T
mask = None
var = None
else:
data = ext.data if dispaxis == 1 else ext.data.T
if dispaxis == 0 and mask is not None:
mask = mask.T
var = ext.variance
if dispaxis == 0 and var is not None:
var = var.T
# Avoid having to recalculate them
self._center_pixels = self.model(np.arange(npix))
all_x1 = self._center_pixels + aper_lower
all_x2 = self._center_pixels + aper_upper
if viewer is not None:
# Display extraction edges on viewer, every 10 pixels (for speed)
pixels = np.arange(npix)
edge_coords = np.array([pixels, all_x1]).T
viewer.polygon(edge_coords[::10],
closed=False,
xfirst=(dispaxis == 1),
origin=0)
edge_coords = np.array([pixels, all_x2]).T
viewer.polygon(edge_coords[::10],
closed=False,
xfirst=(dispaxis == 1),
origin=0)
# Remember how pixel coordinates are defined!
off_low = np.where(all_x1 < -0.5)[0]
if len(off_low):
log.warning(
"Aperture extends off {} of image between {}s {} and {}".
format(("left", "bottom")[dispaxis], direction, min(off_low),
max(off_low)))
off_high = np.where(all_x2 >= slitlength - 0.5)[0]
if len(off_high):
log.warning(
"Aperture extends off {} of image between {}s {} and {}".
format(("right", "top")[dispaxis], direction, min(off_high),
max(off_high)))
# Create the outputs here so the extraction function has access to them
self.data = np.zeros((npix, ), dtype=np.float32)
self.mask = None if mask is None else np.zeros_like(self.data,
dtype=DQ.datatype)
self.var = None if var is None else np.zeros_like(self.data)
extraction_func = getattr(self, "{}_extraction".format(method))
extraction_func(data, mask, var, aper_lower, aper_upper)
del self._center_pixels
self._last_extraction = (aper_lower, aper_upper)
ndd = NDAstroData(self.data, mask=self.mask)
ndd.variance = self.var
try:
ndd.meta['header'] = ext.hdr.copy()
except AttributeError: # we only had an ndarray
pass
return ndd