Merging d2im changes into trunk

git-svn-id: http://svn.stsci.edu/svn/ssb/stsci_python/stwcs/trunk@23819 fe389314-cf27-0410-b35b-8c050e845b92

1 files changed, 192 insertions, 114 deletions

diff --git a/lib/stwcs/updatewcs/det2im.py b/lib/stwcs/updatewcs/det2im.py
index 8235d83..6cf018a 100644
--- a/lib/stwcs/updatewcs/det2im.py
+++ b/lib/stwcs/updatewcs/det2im.py
@@ -1,141 +1,221 @@
 from __future__ import division # confidence high
 
-import time
 import pyfits
 from stsci.tools import fileutil
 import utils
+import numpy as np
 
-import logging
-logger = logging.getLogger('stwcs.updatewcs.Det2IM')
+import logging, time
+logger = logging.getLogger('stwcs.updatewcs.d2im')
 
 class DET2IMCorr(object):
     """
-    Stores a small correction to the detector coordinates as a d2imarr
-    extension in the science file.
+    Defines a Lookup table prior distortion correction as per WCS paper IV.
+    It uses a reference file defined by the D2IMFILE (suffix 'd2im') keyword
+    in the primary header.
 
     Notes
     -----
-    For the case of ACS/WFC every 68th column is wider than the rest.
-    To compensate for this a small correction needs to be applied to the
-    detector coordinates. We call this a detector to image transformation.
-    The so obtained image coordinates are the input to all other distortion
-    corrections. The correction is originally stored in an external
-    reference file pointed to by 'D2IMFILE' keyword in the primary header.
-    This class attaches the correction array as an extension to the science
-    file with extname = `d2imarr`.
+    - 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 'D2IMEXT' to the science extension header to store
+      the name of the reference file used to create the WCSDVARR extensions.
 
-    Other keywords used in this correction are:
-
-    `AXISCORR`: integer (1 or 2) - axis to which the detector to image
-                correction is applied
-
-    `D2IMEXT`:  string = name of reference file which was used to create
-                the lookup table extension
-
-    `D2IMERR`:  float, optional - maximum value of the correction
+    If WCSDVARR extensions exist and `D2IMFILE` is different from `D2IMEXT`,
+    a subsequent update will overwrite the existing extensions.
+    If WCSDVARR extensions were not found in the science file, they will be added.
 
     """
+
     def updateWCS(cls, fobj):
         """
         Parameters
         ----------
         fobj: pyfits object
-              Science file, for which a detector to image correction
-              is available
+                Science file, for which a distortion correction in a NPOLFILE is available
 
-        Notes
-        -----
-        Uses the file pointed to in the primary header keyword 'D2IMFILE'
-        to create an extension with a detector to image correction.
         """
-        logger.info("\n\tStarting Det2IM Correction: %s" % time.asctime())
+        logger.info("\n\tStarting DET2IM: %s" %time.asctime())
         try:
             assert isinstance(fobj, pyfits.HDUList)
         except AssertionError:
             logger.exception('\n\tInput must be a pyfits.HDUList object')
             raise
 
-        d2imfile = fileutil.osfn(fobj[0].header['D2IMFILE'])
-        axiscorr = cls.getAxisCorr(d2imfile)
-        d2imerr = pyfits.getdata(d2imfile, ext=1).max()
-        if axiscorr == None:
-            new_kw = {}
-        else:
-            new_kw = {'D2IMEXT': d2imfile, 'AXISCORR': axiscorr, 'D2IMERR': d2imerr}
-            cls.applyDet2ImCorr(fobj,axiscorr)
-        cls.updatehdr(fobj, new_kw)
+        cls.applyDet2ImCorr(fobj)
+        d2imfile = fobj[0].header['D2IMFILE']
+
+        new_kw = {'D2IMEXT': d2imfile}
+        return new_kw
 
     updateWCS = classmethod(updateWCS)
 
-    def getAxisCorr(cls, refname):
-        try:
-            direction = pyfits.getval(refname, 'EXTNAME', ext=1)
-            if direction == 'DX': return 1
-            elif direction == 'DY': return 2
-            else:
-                logger.warning('\n\tDET2IM correction expects the reference file to have \
-                an EXTNAME keyword of value "DX"  or "DY", EXTNAMe %s detected' % direction)
-                return None
-        except KeyError:
-            logger.exception("\n\tD2IMFILE %s is missing EXTNAME keyword. Unable to determine axis \
-            to which to apply the correction." % refname)
-            direction = None
-        return direction
-    getAxisCorr = classmethod(getAxisCorr)
-
-    def applyDet2ImCorr(cls,fobj, axiscorr):
-        binned = utils.getBinning(fobj)
-        hdu = cls.createDgeoHDU(fobj, axiscorr, binned)
-        d2imarr_ind = cls.getD2imIndex(fobj)
-        if d2imarr_ind:
-            fobj[d2imarr_ind] = hdu
-        else:
-            fobj.append(hdu)
+    def applyDet2ImCorr(cls, fobj):
+        """
+        For each science extension in a pyfits file object:
+            - create a WCSDVARR extension
+            - update science header
+            - add/update D2IMEXT keyword
+        """
+        d2imfile = fileutil.osfn(fobj[0].header['D2IMFILE'])
+        # Map D2IMARR EXTVER numbers to FITS extension numbers
+        wcsdvarr_ind = cls.getWCSIndex(fobj)
+        d2im_num_ext = 1
+        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
+                dx, dy = cls.getData(d2imfile, ccdchip)
+                # Determine EXTVER for the D2IMARR extension from the D2I file (EXTNAME, EXTVER) kw.
+                # This is used to populate DPj.EXTVER kw
+                for ename in zip(['DX', 'DY'], [dx, dy]):
+                    if ename[1] is not None:
+                        error_val = ename[1].max()
+                        cls.addSciExtKw(header, wdvarr_ver=d2im_num_ext, d2im_extname=ename[0], error_val=error_val)
+                        hdu = cls.createD2ImHDU(header, d2imfile=d2imfile, \
+                            wdvarr_ver=d2im_num_ext, d2im_extname=ename[0], data=ename[1],ccdchip=ccdchip, binned=binned)
+                        if wcsdvarr_ind and d2im_num_ext in wcsdvarr_ind:
+                            fobj[wcsdvarr_ind[d2im_num_ext]] = hdu
+                        else:
+                            fobj.append(hdu)
+                        d2im_num_ext = d2im_num_ext + 1
     applyDet2ImCorr = classmethod(applyDet2ImCorr)
 
-    def getD2imIndex(cls,fobj):
-        index = None
+    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 == 'D2IMARR':
-                index = e
-        return index
-    getD2imIndex = classmethod(getD2imIndex)
+                wcsd[fobj[e].header['EXTVER']] = e
+        logger.debug("A map of D2IMARR extensions %s" % wcsd)
+        return wcsd
 
-    def createDgeoHDU(cls, fobj, axiscorr, binned=1):
-        d2imfile = fileutil.osfn(fobj[0].header['D2IMFILE'])
-        d2im_data = pyfits.getdata(d2imfile, ext=1)
-        sci_hdr = fobj['sci',1].header
-        d2im_hdr = cls.createDet2ImHdr(fobj, binned)
-        hdu = pyfits.ImageHDU(header=d2im_hdr, data=d2im_data)
+    getWCSIndex = classmethod(getWCSIndex)
 
-        return hdu
+    def addSciExtKw(cls, hdr, wdvarr_ver=None, d2im_extname=None, error_val=0.0):
+        """
+        Adds kw to sci extension to define WCSDVARR lookup table extensions
 
-    createDgeoHDU = classmethod(createDgeoHDU)
+        """
+        if d2im_extname =='DX':
+            j=1
+        else:
+            j=2
+
+        d2imerror = 'D2IMERR%s' %j
+        d2imdis = 'D2IMDIS%s' %j
+        d2imext = 'D2IM%s.' %j + 'EXTVER'
+        d2imnaxes = 'D2IM%s.' %j +'NAXES'
+        d2imaxis1 = 'D2IM%s.' %j+'AXIS.1'
+        d2imaxis2 = 'D2IM%s.' %j+'AXIS.2'
+        keys = [d2imerror, d2imdis, d2imext, d2imnaxes, d2imaxis1, d2imaxis2]
+        values = {d2imerror: error_val,
+                  d2imdis: 'Lookup',
+                  d2imext: wdvarr_ver,
+                  d2imnaxes: 2,
+                  d2imaxis1: 1,
+                  d2imaxis2: 2}
+
+        comments = {d2imerror: 'Maximum error of NPOL correction for axis %s' % j,
+                    d2imdis: 'Detector to image correction type',
+                    d2imext: 'Version number of WCSDVARR extension containing d2im lookup table',
+                    d2imnaxes: 'Number of independent variables in d2im function',
+                    d2imaxis1: 'Axis number of the jth independent variable in a d2im function',
+                    d2imaxis2: 'Axis number of the jth independent variable in a d2im function'
+                    }
+        # Look for HISTORY keywords. If present, insert new keywords before them
+        before_key = 'HISTORY'
+        if before_key not in hdr:
+            before_key = None
+
+        for key in keys:
+            hdr.update(key=key, value=values[key], comment=comments[key], before=before_key)
+
+    addSciExtKw = classmethod(addSciExtKw)
+
+    def getData(cls,d2imfile, ccdchip):
+        """
+        Get the data arrays from the reference D2I files
+        Make sure 'CCDCHIP' in the npolfile matches "CCDCHIP' in the science file.
+        """
+        xdata, ydata = (None, None)
+        d2im = pyfits.open(d2imfile)
+        for ext in d2im:
+            d2imextname  = ext.header.get('EXTNAME',"")
+            d2imccdchip  = ext.header.get('CCDCHIP',1)
+            if d2imextname == 'DX' and d2imccdchip == ccdchip:
+                xdata = ext.data.copy()
+                continue
+            elif d2imextname == 'DY' and d2imccdchip == ccdchip:
+                ydata = ext.data.copy()
+                continue
+            else:
+                continue
+        d2im.close()
+        return xdata, ydata
+    getData = classmethod(getData)
 
-    def createDet2ImHdr(cls, fobj, binned=1):
+    def createD2ImHDU(cls, sciheader, d2imfile=None, wdvarr_ver=1, d2im_extname=None,data = None, ccdchip=1, binned=1):
         """
-        Creates a header for the D2IMARR extension based on the
-        reference file recorded in D2IMFILE keyword in the primary header.
-        fobj - the science  file
+        Creates an HDU to be added to the file object.
+        """
+        hdr = cls.createD2ImHdr(sciheader, d2imfile=d2imfile, wdvarr_ver=wdvarr_ver, d2im_extname=d2im_extname, ccdchip=ccdchip, binned=binned)
+        hdu=pyfits.ImageHDU(header=hdr, data=data)
+        return hdu
 
+    createD2ImHDU = classmethod(createD2ImHDU)
+
+    def createD2ImHdr(cls, sciheader, d2imfile, wdvarr_ver, d2im_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.
         """
-        d2imfile = fileutil.osfn(fobj[0].header['D2IMFILE'])
-        axiscorr = cls.getAxisCorr(d2imfile)
-        sci_hdr = fobj[1].header
         d2im = pyfits.open(d2imfile)
-        data_shape = d2im[1].shape
-        naxis = d2im[1].header['NAXIS']
         d2im_phdr = d2im[0].header
+        for ext in d2im:
+            try:
+                d2imextname = ext.header['EXTNAME']
+                d2imextver = ext.header['EXTVER']
+            except KeyError:
+                continue
+            d2imccdchip = cls.get_ccdchip(d2im, extname=d2imextname, extver=d2imextver)
+            if d2imextname == d2im_extname and d2imccdchip == ccdchip:
+                d2im_header = ext.header
+                break
+            else:
+                continue
         d2im.close()
 
+        naxis = d2im[1].header['NAXIS']
+        ccdchip = d2imextname #npol_header['CCDCHIP']
+
         kw = { 'NAXIS': 'Size of the axis',
-                'CRPIX': 'Coordinate system reference pixel',
-                'CRVAL': 'Coordinate system value at reference pixel',
-                'CDELT': 'Coordinate increment along axis'}
+               'CDELT': 'Coordinate increment along axis',
+               'CRPIX': 'Coordinate system reference pixel',
+               'CRVAL': 'Coordinate system value at reference pixel',
+               }
 
         kw_comm1 = {}
         kw_val1 = {}
@@ -146,12 +226,10 @@ class DET2IMCorr(object):
 
         for i in range(1, naxis+1):
             si = str(i)
-            kw_val1['NAXIS'+si] = data_shape[i-1]
-            kw_val1['CRPIX'+si] = data_shape[i-1]/2.
-            kw_val1['CDELT'+si] = 1./binned
-            kw_val1['CRVAL'+si] = (sci_hdr.get('NAXIS'+si, 1)/2. + \
-                                        sci_hdr.get('LTV'+si, 0.)) / binned
-
+            kw_val1['NAXIS'+si] = d2im_header.get('NAXIS'+si)
+            kw_val1['CDELT'+si] = d2im_header.get('CDELT'+si, 1.0) / binned
+            kw_val1['CRPIX'+si] = d2im_header.get('CRPIX'+si, 0.0)
+            kw_val1['CRVAL'+si] = (d2im_header.get('CRVAL'+si, 0.0) + sciheader.get('LTV'+str(i), 0)) / binned
 
         kw_comm0 = {'XTENSION': 'Image extension',
                     'BITPIX': 'IEEE floating point',
@@ -160,19 +238,17 @@ class DET2IMCorr(object):
                     'EXTVER': 'Distortion array version number',
                     'PCOUNT': 'Special data area of size 0',
                     'GCOUNT': 'One data group',
-                    'AXISCORR': 'Direction in which the det2im correction is applied'}
+                    }
 
         kw_val0 = { 'XTENSION': 'IMAGE',
                     'BITPIX': -32,
                     'NAXIS': naxis,
                     'EXTNAME': 'D2IMARR',
-                    'EXTVER':  1,
+                    'EXTVER':  wdvarr_ver,
                     'PCOUNT': 0,
                     'GCOUNT': 1,
-                    'AXISCORR': axiscorr
+                    'CCDCHIP': ccdchip,
                 }
-
-
         cdl = []
         for key in kw_comm0.keys():
             cdl.append((key, kw_val0[key], kw_comm0[key]))
@@ -193,24 +269,26 @@ class DET2IMCorr(object):
                 cdl.append(card)
 
         hdr = pyfits.Header(cards=cdl)
+
         return hdr
 
-    createDet2ImHdr = classmethod(createDet2ImHdr)
+    createD2ImHdr = classmethod(createD2ImHdr)
 
-    def updatehdr(cls, fobj, kwdict):
+    def get_ccdchip(cls, fobj, extname, extver):
         """
-        Update extension headers to keep record of the files used for the
-        detector to image correction.
+        Given a science file or npol file determine CCDCHIP
         """
-        for ext in fobj:
-            try:
-                extname = ext.header['EXTNAME'].lower()
-            except KeyError:
-                continue
-            if extname == 'sci':
-                for kw in kwdict:
-                    ext.header.update(kw, kwdict[kw])
-            else:
-                continue
-    updatehdr = classmethod(updatehdr)
-
+        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)