diff options
Diffstat (limited to 'lib/stwcs/updatewcs/npol.py')
| -rw-r--r-- | lib/stwcs/updatewcs/npol.py | 319 |
1 files changed, 319 insertions, 0 deletions
diff --git a/lib/stwcs/updatewcs/npol.py b/lib/stwcs/updatewcs/npol.py new file mode 100644 index 0000000..db928bf --- /dev/null +++ b/lib/stwcs/updatewcs/npol.py @@ -0,0 +1,319 @@ +from __future__ import division # confidence high + +import pyfits +from stsci.tools import fileutil +import utils +import numpy as np + +import logging, time +logger = logging.getLogger('stwcs.updatewcs.npol') + +class NPOLCorr(object): + """ + Defines a Lookup table prior distortion correction as per WCS paper IV. + It uses a reference file defined by the NPOLFILE (suffix 'NPL') keyword + in the primary header. + + Notes + ----- + - Using extensions in the reference file create a WCSDVARR extensions + and add them to the science file. + - Add record-valued keywords to the science extension header to describe + the lookup tables. + - Add a keyword 'NPOLEXT' to the science extension header to store + the name of the reference file used to create the WCSDVARR extensions. + + If WCSDVARR extensions exist and `NPOLFILE` is different from `NPOLEXT`, + a subsequent update will overwrite the existing extensions. + If WCSDVARR extensions were not found in the science file, they will be added. + + It is assumed that the NPL reference files were created to work with IDC tables + but will be applied with SIP coefficients. A transformation is applied to correct + for the fact that the lookup tables will be applied before the first order coefficients + which are in the CD matrix when the SIP convention is used. + """ + + def updateWCS(cls, fobj): + """ + Parameters + ---------- + fobj: pyfits object + Science file, for which a distortion correction in a NPOLFILE is available + + """ + logger.info("\n\tStarting CompSIP: %s" %time.asctime()) + try: + assert isinstance(fobj, pyfits.HDUList) + except AssertionError: + logger.exception('\n\tInput must be a pyfits.HDUList object') + raise + + cls.applyNPOLCorr(fobj) + nplfile = fobj[0].header['NPOLFILE'] + + new_kw = {'NPOLEXT': nplfile} + return new_kw + + updateWCS = classmethod(updateWCS) + + def applyNPOLCorr(cls, fobj): + """ + For each science extension in a pyfits file object: + - create a WCSDVARR extension + - update science header + - add/update NPOLEXT keyword + """ + nplfile = fileutil.osfn(fobj[0].header['NPOLFILE']) + # Map WCSDVARR EXTVER numbers to extension numbers + wcsdvarr_ind = cls.getWCSIndex(fobj) + for ext in fobj: + try: + extname = ext.header['EXTNAME'].lower() + except KeyError: + continue + if extname == 'sci': + extversion = ext.header['EXTVER'] + ccdchip = cls.get_ccdchip(fobj, extname='SCI', extver=extversion) + binned = utils.getBinning(fobj, extversion) + header = ext.header + # get the data arrays from the reference file and transform them for use with SIP + dx,dy = cls.getData(nplfile, ccdchip) + idccoeffs = cls.getIDCCoeffs(header) + + if idccoeffs != None: + dx, dy = cls.transformData(dx,dy, idccoeffs) + + # Determine EXTVER for the WCSDVARR extension from the NPL file (EXTNAME, EXTVER) kw. + # This is used to populate DPj.EXTVER kw + wcsdvarr_x_version = 2 * extversion -1 + wcsdvarr_y_version = 2 * extversion + + for ename in zip(['DX', 'DY'], [wcsdvarr_x_version,wcsdvarr_y_version],[dx, dy]): + cls.addSciExtKw(header, wdvarr_ver=ename[1], npol_extname=ename[0]) + hdu = cls.createNpolHDU(header, npolfile=nplfile, \ + wdvarr_ver=ename[1], npl_extname=ename[0], data=ename[2],ccdchip=ccdchip, binned=binned) + if wcsdvarr_ind: + fobj[wcsdvarr_ind[ename[1]]] = hdu + else: + fobj.append(hdu) + + + applyNPOLCorr = classmethod(applyNPOLCorr) + + def getWCSIndex(cls, fobj): + + """ + If fobj has WCSDVARR extensions: + returns a mapping of their EXTVER kw to file object extension numbers + if fobj does not have WCSDVARR extensions: + an empty dictionary is returned + """ + wcsd = {} + for e in range(len(fobj)): + try: + ename = fobj[e].header['EXTNAME'] + except KeyError: + continue + if ename == 'WCSDVARR': + wcsd[fobj[e].header['EXTVER']] = e + logger.debug("A map of WSCDVARR externsions %s" % wcsd) + return wcsd + + getWCSIndex = classmethod(getWCSIndex) + + def addSciExtKw(cls, hdr, wdvarr_ver=None, npol_extname=None): + """ + Adds kw to sci extension to define WCSDVARR lookup table extensions + + """ + if npol_extname =='DX': + j=1 + else: + j=2 + + cperror = 'CPERROR%s' %j + cpdis = 'CPDIS%s' %j + dpext = 'DP%s.' %j + 'EXTVER' + dpnaxes = 'DP%s.' %j +'NAXES' + dpaxis1 = 'DP%s.' %j+'AXIS.1' + dpaxis2 = 'DP%s.' %j+'AXIS.2' + keys = [cperror, cpdis, dpext, dpnaxes, dpaxis1, dpaxis2] + values = {cperror: 0.0, cpdis: 'Lookup', dpext: wdvarr_ver, dpnaxes: 2, + dpaxis1: 1, dpaxis2: 2} + + comments = {cperror: 'Maximum error of NPOL correction for axis %s' % j, + cpdis: 'Prior distortion funcion type', + dpext: 'Version number of WCSDVARR extension containing lookup distortion table', + dpnaxes: 'Number of independent variables in distortion function', + dpaxis1: 'Axis number of the jth independent variable in a distortion function', + dpaxis2: 'Axis number of the jth independent variable in a distortion function' + } + + for key in keys: + hdr.update(key=key, value=values[key], comment=comments[key], before='HISTORY') + + addSciExtKw = classmethod(addSciExtKw) + + def getData(cls,nplfile, ccdchip): + """ + Get the data arrays from the reference NPOL files + Make sure 'CCDCHIP' in the npolfile matches "CCDCHIP' in the science file. + """ + npl = pyfits.open(nplfile) + for ext in npl: + nplextname = ext.header.get('EXTNAME',"") + nplccdchip = ext.header.get('CCDCHIP',1) + if nplextname == 'DX' and nplccdchip == ccdchip: + xdata = ext.data.copy() + continue + elif nplextname == 'DY' and nplccdchip == ccdchip: + ydata = ext.data.copy() + continue + else: + continue + npl.close() + return xdata, ydata + getData = classmethod(getData) + + def transformData(cls, dx, dy, coeffs): + """ + Transform the NPOL data arrays for use with SIP + """ + ndx, ndy = np.dot(coeffs, [dx.ravel(), dy.ravel()]) + ndx.shape = dx.shape + ndy.shape=dy.shape + return ndx, ndy + + transformData = classmethod(transformData) + + def getIDCCoeffs(cls, header): + """ + Return a matrix of the scaled first order IDC coefficients. + """ + try: + ocx10 = header['OCX10'] + ocx11 = header['OCX11'] + ocy10 = header['OCY10'] + ocy11 = header['OCY11'] + coeffs = np.array([[ocx11, ocx10], [ocy11,ocy10]], dtype=np.float32) + except KeyError: + logger.exception('\n\tFirst order IDCTAB coefficients are not available. \n\ + Cannot convert SIP to IDC coefficients.') + return None + try: + idcscale = header['IDCSCALE'] + except KeyError: + logger.exception("IDCSCALE not found in header - setting it to 1.") + idcscale = 1 + + return np.linalg.inv(coeffs/idcscale) + + getIDCCoeffs = classmethod(getIDCCoeffs) + + def createNpolHDU(cls, sciheader, npolfile=None, wdvarr_ver=1, npl_extname=None,data = None, ccdchip=1, binned=1): + """ + Creates an HDU to be added to the file object. + """ + hdr = cls.createNpolHdr(sciheader, npolfile=npolfile, wdvarr_ver=wdvarr_ver, npl_extname=npl_extname, ccdchip=ccdchip, binned=binned) + hdu=pyfits.ImageHDU(header=hdr, data=data) + return hdu + + createNpolHDU = classmethod(createNpolHDU) + + def createNpolHdr(cls, sciheader, npolfile, wdvarr_ver, npl_extname, ccdchip, binned): + """ + Creates a header for the WCSDVARR extension based on the NPOL reference file + and sci extension header. The goal is to always work in image coordinates + (also for subarrays and binned images. The WCS for the WCSDVARR extension + i ssuch that a full size npol table is created and then shifted or scaled + if the science image is a subarray or binned image. + """ + npl = pyfits.open(npolfile) + for ext in npl: + try: + nplextname = ext.header['EXTNAME'] + nplextver = ext.header['EXTVER'] + except KeyError: + continue + nplccdchip = cls.get_ccdchip(npl, extname=nplextname, extver=nplextver) + if nplextname == npl_extname and nplccdchip == ccdchip: + npol_header = ext.header + break + else: + continue + npl.close() + + naxis = pyfits.getval(npolfile, ext=1, key='NAXIS') + ccdchip = nplextname #npol_header['CCDCHIP'] + + kw = { 'NAXIS': 'Size of the axis', + 'CDELT': 'Coordinate increment along axis', + 'CRPIX': 'Coordinate system reference pixel', + 'CRVAL': 'Coordinate system value at reference pixel', + } + + kw_comm1 = {} + kw_val1 = {} + for key in kw.keys(): + for i in range(1, naxis+1): + si = str(i) + kw_comm1[key+si] = kw[key] + + for i in range(1, naxis+1): + si = str(i) + kw_val1['NAXIS'+si] = npol_header.get('NAXIS'+si) + kw_val1['CDELT'+si] = npol_header.get('CDELT'+si, 1.0) + kw_val1['CRPIX'+si] = npol_header.get('CRPIX'+si, 0.0) + kw_val1['CRVAL'+si] = npol_header.get('CRVAL'+si, 0.0) + + kw_comm0 = {'XTENSION': 'Image extension', + 'BITPIX': 'IEEE floating point', + 'NAXIS': 'Number of axes', + 'EXTNAME': 'WCS distortion array', + 'EXTVER': 'Distortion array version number', + 'PCOUNT': 'Special data area of size 0', + 'GCOUNT': 'One data group', + } + + kw_val0 = { 'XTENSION': 'IMAGE', + 'BITPIX': -32, + 'NAXIS': naxis, + 'EXTNAME': 'WCSDVARR', + 'EXTVER': wdvarr_ver, + 'PCOUNT': 0, + 'GCOUNT': 1, + 'CCDCHIP': ccdchip, + } + + cdl = pyfits.CardList() + for key in kw_comm0.keys(): + cdl.append(pyfits.Card(key=key, value=kw_val0[key], comment=kw_comm0[key])) + for key in kw_comm1.keys(): + cdl.append(pyfits.Card(key=key, value=kw_val1[key], comment=kw_comm1[key])) + + + hdr = pyfits.Header(cards=cdl) + + return hdr + + createNpolHdr = classmethod(createNpolHdr) + + def get_ccdchip(cls, fobj, extname, extver): + """ + Given a science file or npol file determine CCDCHIP + """ + ccdchip = 1 + if fobj[0].header['INSTRUME'] == 'ACS' and fobj[0].header['DETECTOR'] == 'WFC': + ccdchip = fobj[extname, extver].header['CCDCHIP'] + elif fobj[0].header['INSTRUME'] == 'WFC3' and fobj[0].header['DETECTOR'] == 'UVIS': + ccdchip = fobj[extname, extver].header['CCDCHIP'] + elif fobj[0].header['INSTRUME'] == 'WFPC2': + ccdchip = fobj[extname, extver].header['DETECTOR'] + elif fobj[0].header['INSTRUME'] == 'STIS': + ccdchip = fobj[extname, extver].header['DETECTOR'] + elif fobj[0].header['INSTRUME'] == 'NICMOS': + ccdchip = fobj[extname, extver].header['CAMERA'] + return ccdchip + + get_ccdchip = classmethod(get_ccdchip) +
\ No newline at end of file |