Source code for kcwidrp.primitives.SolveGeom

from keckdrpframework.primitives.base_primitive import BasePrimitive
from kcwidrp.primitives.kcwi_file_primitives import strip_fname
import os
import numpy as np
from skimage import transform as tf
import pickle


[docs]class SolveGeom(BasePrimitive): """Solve the overall geometry of the IFU""" 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 _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('polynomial', src, dst, order=3) invtf = tf.estimate_transform('polynomial', dst, src, order=3) # 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.name 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, "xsize": xsize, "ysize": ysize, "pxscl": pxscl, "slscl": slscl, "cbarsno": self.action.args.contbar_image_number, "cbarsfl": self.action.args.contbar_image, "arcno": self.action.args.arc_number, "arcfl": self.action.args.arc_image, "barsep": self.action.args.reference_bar_separation, "bar0": self.action.args.x0out, "waveall0": self.action.args.waveall0, "waveall1": self.action.args.waveall1, "wavegood0": self.action.args.wavegood0, "wavegood1": self.action.args.wavegood1, "wavemid": self.action.args.wavemid, "wavensall0": self.action.args.wavensall0, "wavensall1": self.action.args.wavensall1, "wavensgood0": self.action.args.wavensgood0, "wavensgood1": self.action.args.wavensgood1, "wavensmid": self.action.args.wavensmid, "dich_frac": dichroic_fraction, "dwout": dwout, "wave0out": self.action.args.wave0out, "wave1out": self.action.args.wave1out, "avwvsig": self.action.args.av_bar_sig, "sdwvsig": self.action.args.st_bar_sig, "xl0": xl0_out, "xl1": xl1_out, "tform": tform_list, "invtf": invtf_list } 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()