from keckdrpframework.primitives.base_primitive import BasePrimitive
from kcwidrp.primitives.kcwi_file_primitives import strip_fname
import os
import numpy as np
from kcwidrp.core import geometric as tf
import pickle
[docs]class SolveGeom(BasePrimitive):
"""
Solve the overall geometry of the IFU.
Takes individual bar arc spectra and generates a fit for each slice, which
contains five bar arc spectra each. Given that there are only five points
in the 'x' or spatial direction and many more in the 'y' or wavelength
direction, an asymmetric polynomial is fit with order 2 in the x and order
4 in the y directions. The wavelength coverage of the observation is
recorded in parameters with the range that includes all data being in
the waveall0 and waveall1 parameters, and the range that includes only
good data in wavegood0 and wavegood1 parameters. The middle of the
wavelength range is recorded in the wavemid parameter.
Forward and inverse transforms, along with all the parameters for the
geometric fit are written out as a python pickled dictionary in a
\*_geom.pkl file.
"""
def __init__(self, action, context):
BasePrimitive.__init__(self, action, context)
self.action.args.geometry_file = None
self.action.args.x0out = None
self.action.args.wave0out = None
self.action.args.wave1out = None
self.action.args.wavegood0 = None
self.action.args.wavegood1 = None
self.action.args.waveall0 = None
self.action.args.waveall1 = None
self.action.args.wavemid = None
self.logger = context.pipeline_logger
def _pre_condition(self):
self.logger.info("Checking for master arc")
if 'MARC' in self.action.args.ccddata.header['IMTYPE']:
return True
else:
return False
def _perform(self):
self.logger.info("Solving overall geometry")
# Get some geometry constraints
goody0 = 0
goody1 = max(self.action.args.xsvals)
# N&S limits
goodnsy0 = self.action.args.shufrows
goodnsy1 = goodnsy0 + self.action.args.shufrows
# Calculate wavelength ranges
y0wvs = []
y1wvs = []
# N&S ranges
y0nswvs = []
y1nswvs = []
# Get wavelength extremes for each bar
for fcfs in self.action.args.fincoeff:
y0wvs.append(float(np.polyval(fcfs, goody0)))
y1wvs.append(float(np.polyval(fcfs, goody1)))
y0nswvs.append(float(np.polyval(fcfs, goodnsy0)))
y1nswvs.append(float(np.polyval(fcfs, goodnsy1)))
# Now get ensemble extremes
y0max = max(y0wvs)
y0min = min(y0wvs)
y1max = max(y1wvs)
y1min = min(y1wvs)
y0nsmax = max(y0nswvs)
y0nsmin = min(y0nswvs)
y1nsmax = max(y1nswvs)
y1nsmin = min(y1nswvs)
# Cube trimming wavelengths
trimw0 = y0min
trimw1 = y1max
# Check for negative dispersion
if trimw0 > trimw1:
trimw0 = y1min
trimw1 = y0max
# Calculate output wavelengths
dwout = self.action.args.dwout
ndels = int((trimw0 - self.config.instrument.WAVEFID) / dwout)
self.action.args.wave0out = \
self.config.instrument.WAVEFID + float(ndels) * dwout
ndels = int((trimw1 - self.config.instrument.WAVEFID) / dwout)
self.action.args.wave1out = \
self.config.instrument.WAVEFID + float(ndels) * dwout
self.logger.info("WAVE RANGE: %.2f - %.2f" %
(self.action.args.wave0out, self.action.args.wave1out))
# Calculate wavelength limits
self.action.args.wavegood0 = min([y0max, y1max])
self.action.args.wavegood1 = max([y0min, y1min])
self.action.args.waveall0 = min([y0min, y1min])
self.action.args.waveall1 = max([y0max, y1max])
self.action.args.wavemid = np.average([self.action.args.wavegood0,
self.action.args.wavegood1,
self.action.args.waveall0,
self.action.args.waveall1])
self.action.args.wavensgood0 = min([y0nsmax, y1nsmax])
self.action.args.wavensgood1 = max([y0nsmin, y1nsmin])
self.action.args.wavensall0 = min([y0nsmin, y1nsmin])
self.action.args.wavensall1 = max([y0nsmax, y1nsmax])
self.action.args.wavensmid = np.average([self.action.args.wavensgood0,
self.action.args.wavensgood1,
self.action.args.wavensall0,
self.action.args.wavensall1])
self.logger.info("WAVE GOOD: %.2f - %.2f" %
(self.action.args.wavegood0,
self.action.args.wavegood1))
self.logger.info("WAVE ALL: %.2f - %.2f" %
(self.action.args.waveall0, self.action.args.waveall1))
self.logger.info("WAVE MID: %.2f" % self.action.args.wavemid)
# Start setting up slice transforms
self.action.args.x0out = \
int(self.action.args.reference_bar_separation / 2.) + 1
self.refoutx = np.arange(0, 5) * \
self.action.args.reference_bar_separation + self.action.args.x0out
# Variables for output control points
srcw = []
# Loop over source control points
for ixy, xy in enumerate(self.action.args.source_control_points):
# Calculate y wavelength
yw = float(np.polyval(
self.action.args.fincoeff[self.action.args.bar_id[ixy]], xy[1]))
# Convert to output pixels
yw = (yw - self.action.args.wave0out) / dwout
srcw.append([xy[0], yw])
# Use extremes to define output size
ysize = int((self.action.args.waveall1 - self.action.args.wave0out)
/ dwout)
xsize = int(5. * self.action.args.reference_bar_separation) + 1
self.logger.info("Output slices will be %d x %d px" % (xsize, ysize))
# Now loop over slices and get relevant control points for each slice
# Output variables
xl0_out = []
xl1_out = []
tform_list = []
invtf_list = []
# Loop over 24 slices
for isl in range(0, 24):
# Get control points
xw = []
yw = []
xi = []
yi = []
# Loop over all control points
for ixy, xy in enumerate(srcw):
# Only use the ones for this slice
if self.action.args.slice_id[ixy] == isl:
# Index in to reference output x array
ib = self.action.args.bar_id[ixy] % 5
# Geometrically corrected control points
xw.append(self.refoutx[ib])
yw.append(xy[1])
# Input control points
xi.append(self.action.args.destination_control_points[
ixy][0])
yi.append(self.action.args.destination_control_points[
ixy][1])
# get image limits
xl0 = int(min(xi) - self.action.args.reference_bar_separation)
if xl0 < 0:
xl0 = 0
xl1 = int(max(xi) + self.action.args.reference_bar_separation)
if xl1 > (self.action.args.ccddata.data.shape[0] - 1):
xl1 = self.action.args.ccddata.data.shape[0] - 1
# Store for output
xl0_out.append(xl0)
xl1_out.append(xl1)
self.logger.info("Slice %d arc image x limits: %d - %d" %
(isl, xl0, xl1))
# adjust control points
xit = [x - float(xl0) for x in xi]
# fit transform
dst = np.column_stack((xit, yi))
src = np.column_stack((xw, yw))
self.logger.info("Fitting wavelength and spatial control points")
tform = tf.estimate_transform('asympolynomial', src, dst,
order=(2, 4))
invtf = tf.estimate_transform('asympolynomial', dst, src,
order=(2, 4))
# Store for output
tform_list.append(tform)
invtf_list.append(invtf)
# Pixel scales
pxscl = self.config.instrument.PIXSCALE * self.action.args.xbinsize
ifunum = self.action.args.ifunum
if ifunum == 2:
slscl = self.config.instrument.SLICESCALE / 2.0
elif ifunum == 3:
slscl = self.config.instrument.SLICESCALE / 4.0
else:
slscl = self.config.instrument.SLICESCALE
# Dichroic fraction
try:
dichroic_fraction = self.action.args.dichroic_fraction
except AttributeError:
dichroic_fraction = 1.
# Package geometry data
ofname = self.action.args.ccddata.header['OFNAME']
self.action.args.geometry_file = os.path.join(
self.config.instrument.output_directory,
strip_fname(ofname) + '_geom.pkl')
if os.path.exists(self.action.args.geometry_file):
self.logger.error("Geometry file already exists: %s" %
self.action.args.geometry_file)
else:
geom = {
"geom_file": self.action.args.geometry_file, # output save file for geom struct
"xsize": xsize, "ysize": ysize, # xy size of image
"pxscl": pxscl, "slscl": slscl, # degrees per unbinned spatial and slice pixel
"cbarsno": self.action.args.contbar_image_number, # cont bars image number
"cbarsfl": self.action.args.contbar_image, # cont bars image name
"arcno": self.action.args.arc_number, # arc image number
"arcfl": self.action.args.arc_image, # arc image name
"barsep": self.action.args.reference_bar_separation, # reference delta x between bars (pix)
"bar0": self.action.args.x0out, # output spatial zeropoint (unbinned pix)
"waveall0": self.action.args.waveall0, # low wavelength that includes all data
"waveall1": self.action.args.waveall1, # high wavelength that includes all data
"wavegood0": self.action.args.wavegood0, # low wavelength that includes good data
"wavegood1": self.action.args.wavegood1, # high wavelength that includes good data
"wavemid": self.action.args.wavemid, # wavelength of the middle of the range
"wavensall0": self.action.args.wavensall0, # low wavelength that includes all masked data
"wavensall1": self.action.args.wavensall1, # high wavelength that includes all masked data
"wavensgood0": self.action.args.wavensgood0, # low wavelength that includes good masked data
"wavensgood1": self.action.args.wavensgood1, # high wavelength that includes good masked data
"wavensmid": self.action.args.wavensmid, # wavelength of the middle of the masked range
"dich_frac": dichroic_fraction, # fraction of wavelength range impacted by dichroic
"dwout": dwout, # output disperson (Ang/pix)
"wave0out": self.action.args.wave0out, # output wavelength zeropoint
"wave1out": self.action.args.wave1out, # output ending wavelength
"avwvsig": self.action.args.av_bar_sig, # average bar wavelength sigma (Ang)
"sdwvsig": self.action.args.st_bar_sig, # standard deviation of bar sigmas (Ang)
"xl0": xl0_out, "xl1": xl1_out, # slice spatial limits
"tform": tform_list, "invtf": invtf_list # transform and inverse transforms for each slice
}
with open(self.action.args.geometry_file, 'wb') as ofile:
pickle.dump(geom, ofile)
self.logger.info("Geometry written to: %s" %
self.action.args.geometry_file)
log_string = SolveGeom.__module__
self.action.args.ccddata.header['HISTORY'] = log_string
self.logger.info(log_string)
return self.action.args
# END: class SolveGeom()