def test_center_gaussian():
"""
Test the centering algorithm with Gaussian weighting
"""
ngau = 10
npix = 200
# Make a Gaussian comb with some random subpixel offsets
xi = np.arange(ngau, npix, npix // ngau, dtype=float)
xt = xi + np.round(np.random.normal(scale=1., size=xi.size), decimals=1)
img = np.zeros((1, npix), dtype=float)
x = np.arange(npix)
for c in xt:
img[0, :] += (erf((x - c + 0.5) / np.sqrt(2) / 2) - erf(
(x - c - 0.5) / np.sqrt(2) / 2)) / 2.
xr, xe, bad = moment.moment1d(img,
xi,
4.,
row=0,
weighting='gaussian',
order=1)
assert np.mean(np.absolute(xi-xt)) - np.mean(np.absolute(xr-xt)) > 0, \
'Recentering did not improve the mean difference'
assert np.std(xi-xt) - np.std(xr-xt) > 0, \
'Recentering did not improve the standard deviation in the difference'
def test_bounded():
c = [45, 50, 55]
img = np.zeros((len(c), 100), dtype=float)
x = np.arange(100)
sig = 5.
for i, _c in enumerate(c):
img[i, :] = (erf((x - c[i] + 0.5) / np.sqrt(2) / sig) - erf(
(x - c[i] - 0.5) / np.sqrt(2) / sig)) / 2.
# Should find first moment is less than 49, meaning that the moment
# is set to the input and flagged
mu, mue, flag = moment.moment1d(img,
50,
40.,
row=0,
order=[0, 1, 2],
bounds=([0.9, -1., 4], [1.1, 1., 6]))
assert mu[1] == 50. and flag[1]
def test_basics():
c = [45, 50, 55]
img = np.zeros((len(c), 100), dtype=float)
x = np.arange(100)
sig = 5.
for i, _c in enumerate(c):
img[i, :] = (erf((x - c[i] + 0.5) / np.sqrt(2) / sig) - erf(
(x - c[i] - 0.5) / np.sqrt(2) / sig)) / 2.
# Calculate all moments at one column and row
mu, mue, flag = moment.moment1d(img, 50, 40., row=0, order=[0, 1, 2])
assert np.allclose(mu, np.array([0.99858297, 45.02314924, 4.97367636]))
# Calculate all moments at one column for all rows
assert np.allclose(
moment.moment1d(img, 50, 40., order=[0, 1, 2])[0],
np.array([[0.99858297, 0.99993125, 0.99858297],
[45.02314924, 50., 54.97685076],
[4.97367636, 5.00545947, 4.97367636]]))
# Calculate zeroth moments in all rows centered at column 50
assert np.allclose(
moment.moment1d(img, 50, 40., order=0)[0],
np.array([0.99858297, 0.99993125, 0.99858297]))
# Calculate zeroth moments in all rows for three column positions
assert np.allclose(
moment.moment1d(img, [45, 50, 55], 40., order=0, mesh=True)[0],
np.array([[0.99993125, 0.99858297, 0.97670951],
[0.99858297, 0.99993125, 0.99858297],
[0.97670951, 0.99858297, 0.99993125]]))
# Calculate zeroth moments in each row with one column center per row
assert np.allclose(
moment.moment1d(img, [45, 50, 55], 40., row=[0, 1, 2], order=0)[0],
np.array([0.99993125, 0.99993125, 0.99993125]))
# Calculate the first moment in a column for all rows
assert np.allclose(
moment.moment1d(img, 50, 40., row=[0, 1, 2], order=1)[0],
np.array([45.02314924, 50., 54.97685076]))
# Or pick a column unique to each row
assert np.allclose(
moment.moment1d(img, [43, 52, 57], 40., row=[0, 1, 2], order=1)[0],
np.array([44.99688181, 50.00311819, 55.00311819]))
def test_extract():
"""
Test a simple zeroth moment flux extraction
"""
ngau = 10
npix = 2000
sig = 10.
delt = sig
# Make a Gaussian comb with some random subpixel offsets
xi = np.arange(npix // ngau // 2, npix, npix // ngau, dtype=float)
xt = xi + np.round(np.random.normal(scale=1., size=xi.size), decimals=1)
img = np.zeros((1, npix), dtype=float)
x = np.arange(npix)
for c in xt:
img[0, :] += (erf((x - c + 0.5) / np.sqrt(2) / sig) - erf(
(x - c - 0.5) / np.sqrt(2) / sig)) / 2.
xr, xe, bad = moment.moment1d(img, xt, delt * 2, row=0, order=0)
truth = (erf(delt / np.sqrt(2) / sig) - erf(-delt / np.sqrt(2) / sig)) / 2.
assert np.mean(np.absolute(xr - truth)) < 1e-3, 'Extraction inaccurate'
def test_width():
"""
Test the measurement of the second moment
"""
ngau = 10
npix = 2000
sig = 10.
delt = 3 * sig
# Make a Gaussian comb with some random subpixel offsets
xi = np.arange(npix // ngau // 2, npix, npix // ngau, dtype=float)
xt = xi + np.round(np.random.normal(scale=1., size=xi.size), decimals=1)
img = np.zeros((1, npix), dtype=float)
x = np.arange(npix)
for c in xt:
# img[0,:] += np.exp(-np.square((x-c)/sig)/2)
img[0, :] += (erf((x - c + 0.5) / np.sqrt(2) / sig) - erf(
(x - c - 0.5) / np.sqrt(2) / sig)) / 2.
xr, xe, bad = moment.moment1d(img, xt, delt * 2, row=0, order=2)
assert np.absolute(
np.mean(xr / sig) -
1) < 0.02, 'Second moment should be good to better than 2%'
def rectify_image(img,
col,
bpm=None,
ocol=None,
max_ocol=None,
extract_width=None,
mask_threshold=0.5):
r"""
Rectify the image by shuffling flux along columns using the provided
column mapping.
The image recification is one dimensional, treating each image row
independently. It can be done either by a direct resampling of the
image columns using the provided mapping of output to input column
location (see `col` and :class:`Resample`) or by an extraction along
the provided column locations (see `extract_width`). The latter is
generally faster; however, when resampling each row, the flux is
explicitly conserved (see the `conserve` argument of
:class:`Resample`).
Args:
img (`numpy.ndarray`_):
The 2D image to rectify. Shape is :math:`(N_{\rm row},
N_{\rm col})`.
col (`numpy.ndarray`_):
The array mapping each output column to its location in
the input image. That is, e.g., `col[:,0]` provides the
column coordinate in `img` that should be rectified to
column 0 in the output image. Shape is :math:`(N_{\rm
row}, N_{\rm map})`.
bpm (`numpy.ndarray`_, optional):
Boolean bad-pixel mask for pixels to ignore in input
image. If None, no pixels are masked in the
rectification. If provided, shape must match `img`.
ocol (`numpy.ndarray`_, optional):
The column in the output image for each column in `col`.
If None, assume::
ocol = numpy.arange(col.shape[1])
These coordinates can fall off the output image (i.e.,
:math:`<0` or :math:`\geq N_{\rm out,col}`), but those
columns are removed from the output).
max_ocol (:obj:`int`, optional):
The last viable column *index* to include in the output
image; ie., for an image with `ncol` columns, this should
be `ncol-1`. If None, assume `max(ocol)`.
extract_width (:obj:`float`, optional):
The width of the extraction aperture to use for the image
rectification. If None, the image recification is performed
using :class:`Resample` along each row.
mask_threshold (:obj:`float`, optional):
Either due to `bpm` or the bounds of the provided `img`,
pixels in the rectified image may not be fully covered by
valid pixels in `img`. Pixels in the output image with
less than this fractional coverage of an input pixel are
flagged in the output.
Returns:
Two `numpy.ndarray`_ objects are returned both with shape
`(nrow,max_ocol+1)`, the rectified image and its boolean
bad-pixel mask.
"""
# Check the input
if img.ndim != 2:
raise ValueError('Input image must be 2D.')
if bpm is not None and bpm.shape != img.shape:
raise ValueError('Image bad-pixel mask must match image shape.')
_img = numpy.ma.MaskedArray(img, mask=bpm)
nrow, ncol = _img.shape
if col.ndim != 2:
raise ValueError('Column mapping array must be 2D.')
if col.shape[0] != nrow:
raise ValueError(
'Number of rows in column mapping array must match image to rectify.'
)
_ocol = numpy.arange(
col.shape[1]) if ocol is None else numpy.atleast_1d(ocol)
if _ocol.ndim != 1:
raise ValueError('Output column indices must be provided as a vector.')
if _ocol.size != col.shape[1]:
raise ValueError(
'Output column indices must match columns in column mapping array.'
)
_max_ocol = numpy.amax(_ocol) if max_ocol is None else max_ocol
# Use an aperture extraction to rectify the image
if extract_width is not None:
# Select viable columns
indx = (_ocol >= 0) & (_ocol <= _max_ocol)
# Initialize the output image as all masked
out_img = numpy.ma.masked_all((nrow, _max_ocol + 1), dtype=float)
# Perform the extraction
out_img[:, _ocol[indx]] = moment.moment1d(_img.data,
col[:, indx],
extract_width,
bpm=_img.mask)[0]
# Determine what fraction of the extraction fell off the image
coo = col[:,indx,None] + numpy.arange(numpy.ceil(extract_width)).astype(int)[None,None,:] \
- extract_width/2
in_image = (coo >= 0) | (coo < ncol)
out_bpm = numpy.sum(in_image, axis=2) / extract_width < mask_threshold
# Return the filled numpy.ndarray and boolean mask
return out_img.filled(0.0) / extract_width, out_img.mask | out_bpm
# Directly resample the image
#
# `col` provides the column in the input image that should
# resampled to a given position in the output image: the value of
# the flux at img[:,col[:,0]] should be rectified to `outimg[:,0]`.
# To run the resampling algorithm we need to invert this. That is,
# instead of having a function that provides the output column as a
# function of the input column, we want a function that provies the
# input column as a function of the output column.
# Instantiate the output image
out_img = numpy.zeros((nrow, _max_ocol + 1), dtype=float)
out_bpm = numpy.zeros((nrow, _max_ocol + 1), dtype=bool)
icol = numpy.arange(ncol)
for i in range(nrow):
# Get the coordinate vector of the output column
_icol = interpolate.interp1d(col[i, :],
_ocol,
copy=False,
bounds_error=False,
fill_value='extrapolate',
assume_sorted=True)(icol)
# Resample it
r = Resample(_img[i, :],
x=_icol,
newRange=[0, _max_ocol],
newpix=_max_ocol + 1,
newLog=False,
conserve=True)
# Save the resampled data
out_img[i, :] = r.outy
# Flag pixels
out_bpm[i, :] = r.outf < mask_threshold
return out_img, out_bpm
def get_wave_grid(self, **kwargs_wave):
"""
Routine to create a wavelength grid for 2d coadds using all of the wavelengths of the extracted objects. Calls
coadd1d.get_wave_grid.
Args:
**kwargs_wave (dict):
Optional argumments for coadd1d.get_wve_grid function
Returns:
tuple: Returns the following:
- wave_grid (np.ndarray): New wavelength grid, not
masked
- wave_grid_mid (np.ndarray): New wavelength grid
evaluated at the centers of the wavelength bins, that
is this grid is simply offset from wave_grid by
dsamp/2.0, in either linear space or log10 depending
on whether linear or (log10 or velocity) was
requested. For iref or concatenate the linear
wavelength sampling will be calculated.
- dsamp (float): The pixel sampling for wavelength grid
created.
"""
nobjs_tot = int(np.array([len(spec) for spec in self.stack_dict['specobjs_list']]).sum())
# TODO: Do we need this flag since we can determine whether or not we have specobjs from nobjs_tot?
# This all seems a bit hacky
if self.par['coadd2d']['use_slits4wvgrid'] or nobjs_tot==0:
nslits_tot = np.sum([slits.nslits for slits in self.stack_dict['slits_list']])
waves = np.zeros((self.nspec, nslits_tot*3))
gpm = np.zeros_like(waves, dtype=bool)
box_radius = 3.
indx = 0
# Loop on the exposures
for waveimg, slitmask, slits in zip(self.stack_dict['waveimg_stack'],
self.stack_dict['slitmask_stack'],
self.stack_dict['slits_list']):
slits_left, slits_righ, _ = slits.select_edges()
row = np.arange(slits_left.shape[0])
# Loop on the slits
for kk, spat_id in enumerate(slits.spat_id):
mask = slitmask == spat_id
# Create apertures at 5%, 50%, and 95% of the slit width to cover full range of wavelengths
# on this slit
trace_spat = slits_left[:, kk][:,np.newaxis] + np.outer((slits_righ[:,kk] - slits_left[:,kk]),[0.05,0.5,0.95])
box_denom = moment1d(waveimg * mask > 0.0, trace_spat, 2 * box_radius, row=row)[0]
wave_box = moment1d(waveimg * mask, trace_spat, 2 * box_radius,
row=row)[0] / (box_denom + (box_denom == 0.0))
waves[:, indx:indx+3] = wave_box
# TODO -- This looks a bit risky
gpm[:, indx: indx+3] = wave_box > 0.
indx += 3
else:
waves = np.zeros((self.nspec, nobjs_tot))
gpm = np.zeros_like(waves, dtype=bool)
indx = 0
for spec_this in self.stack_dict['specobjs_list']:
for spec in spec_this:
waves[:, indx] = spec.OPT_WAVE
# TODO -- OPT_MASK is likely to become a bpm with int values
gpm[:, indx] = spec.OPT_MASK
indx += 1
wave_grid, wave_grid_mid, dsamp = coadd.get_wave_grid(waves, masks=gpm, **kwargs_wave)
return wave_grid, wave_grid_mid, dsamp
def spec_flexure_correct(self, mode="local", sobjs=None):
""" Correct for spectral flexure
Spectra are modified in place (wavelengths are shifted)
Args:
mode (str):
"local" - Use sobjs to determine flexure correction
"global" - Use waveimg and global_sky to determine flexure correction at the centre of the slit
sobjs (:class:`pypeit.specobjs.SpecObjs`, None):
Spectrally extracted objects
"""
if self.par['flexure']['spec_method'] == 'skip':
msgs.info('Skipping flexure correction.')
return
# Perform some checks
if mode == "local" and sobjs is None:
msgs.error("No spectral extractions provided for flexure, using slit center instead")
elif mode not in ["local", "global"]:
msgs.error("mode must be 'global' or 'local'. Assuming 'global'.")
# Prepare a list of slit spectra, if required.
if mode == "global":
gd_slits = np.logical_not(self.extract_bpm)
# TODO :: Need to think about spatial flexure - is the appropriate spatial flexure already included in trace_spat via left/right slits?
trace_spat = 0.5 * (self.slits_left + self.slits_right)
trace_spec = np.arange(self.slits.nspec)
slit_specs = []
for ss in range(self.slits.nslits):
if not gd_slits[ss]:
slit_specs.append(None)
continue
slit_spat = self.slits.spat_id[ss]
thismask = (self.slitmask == slit_spat)
box_denom = moment1d(self.waveimg * thismask > 0.0, trace_spat[:, ss], 2, row=trace_spec)[0]
wghts = (box_denom + (box_denom == 0.0))
slit_sky = moment1d(self.global_sky * thismask, trace_spat[:, ss], 2, row=trace_spec)[0] / wghts
# Denom is computed in case the trace goes off the edge of the image
slit_wave = moment1d(self.waveimg * thismask, trace_spat[:, ss], 2, row=trace_spec)[0] / wghts
# TODO :: Need to remove this XSpectrum1D dependency - it is required in: flexure.spec_flex_shift
slit_specs.append(xspectrum1d.XSpectrum1D.from_tuple((slit_wave, slit_sky)))
# Measure flexure
# If mode == global: specobjs = None and slitspecs != None
flex_list = flexure.spec_flexure_slit(self.slits, self.slits.slitord_id, self.extract_bpm,
self.par['flexure']['spectrum'],
method=self.par['flexure']['spec_method'],
mxshft=self.par['flexure']['spec_maxshift'],
specobjs=sobjs, slit_specs=slit_specs)
# Store the slit shifts that were applied to each slit
# These corrections are later needed so the specobjs metadata contains the total spectral flexure
self.slitshift = np.zeros(self.slits.nslits)
for islit in range(self.slits.nslits):
if (not gd_slits[islit]) or len(flex_list[islit]['shift']) == 0:
continue
self.slitshift[islit] = flex_list[islit]['shift'][0]
# Apply flexure to the new wavelength solution
msgs.info("Regenerating wavelength image")
self.waveimg = self.wv_calib.build_waveimg(self.tilts, self.slits,
spat_flexure=self.spat_flexure_shift,
spec_flexure=self.slitshift)
elif mode == "local":
# Measure flexure:
# If mode == local: specobjs != None and slitspecs = None
flex_list = flexure.spec_flexure_slit(self.slits, self.slits.slitord_id, self.extract_bpm,
self.par['flexure']['spectrum'],
method=self.par['flexure']['spec_method'],
mxshft=self.par['flexure']['spec_maxshift'],
specobjs=sobjs, slit_specs=None)
# Apply flexure to objects
for islit in range(self.slits.nslits):
i_slitord = self.slits.slitord_id[islit]
indx = sobjs.slitorder_indices(i_slitord)
this_specobjs = sobjs[indx]
this_flex_dict = flex_list[islit]
# Loop through objects
cntr = 0
for ss, sobj in enumerate(this_specobjs):
if sobj is None or sobj['BOX_WAVE'] is None: # Nothing extracted; only the trace exists
continue
# Interpolate
new_sky = sobj.apply_spectral_flexure(this_flex_dict['shift'][cntr],
this_flex_dict['sky_spec'][cntr])
flex_list[islit]['sky_spec'][cntr] = new_sky.copy()
cntr += 1
# Save QA
basename = f'{self.basename}_{mode}_{self.spectrograph.get_det_name(self.det)}'
out_dir = os.path.join(self.par['rdx']['redux_path'], 'QA')
flexure.spec_flexure_qa(self.slits.slitord_id, self.extract_bpm, basename, flex_list,
specobjs=sobjs, out_dir=out_dir)
