Wcskey and alternate WCS used now in stwcs

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

4 files changed, 269 insertions, 101 deletions

diff --git a/lib/__init__.py b/lib/__init__.py
index c8d174d..4519077 100644
--- a/lib/__init__.py
+++ b/lib/__init__.py
@@ -7,6 +7,8 @@ import distortion
 import pywcs
 from pytools import fileutil
 
+__docformat__ = 'restructuredtext'
+
 DEGTORAD = fileutil.DEGTORAD
 RADTODEG = fileutil.RADTODEG
 
diff --git a/lib/utils.py b/lib/utils.py
index e57cc43..3d04195 100644
--- a/lib/utils.py
+++ b/lib/utils.py
@@ -2,16 +2,38 @@ from __future__ import division # confidence high
 
 from pytools import parseinput, fileutil
 import pyfits
-from wcsutil.mappings import basic_wcs
+import pywcs
+import numpy as np
+import string 
 
-def restoreWCS(fnames):
+def restoreWCS(fnames, wcskey, clobber=False):
     """
-    Given a list of fits file names restore the original basic WCS kw
-    and write out the files. This overwrites the original files.
+    Purpose
+    =======
+    Reads in a WCS defined with wcskey and saves it as the primary WCS.
+    If clobber is False, writes out new files whose names are the original 
+    names with an attached 3 character string representing _'WCSKEY'_. 
+    Otherwise overwrites the files. The WCS is restored from the 'SCI' 
+    extension but the primary WCS of all extensions are updated.
     
-    Affected keywords:
-    'CD1_1', 'CD1_2', 'CD2_1', 'CD2_2', 'CRVAL1','CRVAL2','CTYPE1', 'CTYPE2',
-    'CRPIX1', 'CRPIX2', 'CTYPE1', 'CTYPE2', 'ORIENTAT'
+    WCS keywords:
+    'CD1_1', 'CD1_2', 'CD2_1', 'CD2_2', 
+    'CRVAL*',
+    'CTYPE*',
+    'CRPIX*', 
+    'CDELT*',
+    'CUNIT*',
+    'ORIENTAT' - ?
+    'TDDALPHA', 'TDDBETA'
+    'A_x_x', B_x_x' - SIP coefficients
+    'CPERROR*', 'CPDIS*', 'DP*', 
+    
+    `fnames`: a python list of file names, a string of comma separated file names, 
+              an @file
+    `wcskey`: a charater
+              Used for one of 26 alternate WCS definitions.
+    `clobber`: boolean
+              A flag to define if the original files should be overwritten
     """
     files = parseinput.parseinput(fnames)[0]
     for f in files:
@@ -20,61 +42,99 @@ def restoreWCS(fnames):
             print "RestoreWCS works only with true fits files."
             return
         else:
-            fobj = pyfits.open(f, mode='update')
-            for ext in fobj:
+            if clobber:
+                print 'Overwriting original files\n'
+                fobj = pyfits.open(f, mode='update')
+                name = f
+            else:
+                fobj = pyfits.open(f)
+                name = (f.split('.fits')[0] + '_%s_' + '.fits') %wcskey
+            for e in range(len(fobj)):
                 try:
-                    extname = ext.header['EXTNAME'].lower()
+                    extname = fobj[e].header['EXTNAME'].lower()
                 except KeyError:
                     continue
-                if extname in ['sci', 'err', 'sdq']:
-                    hdr = ext.header
-                    backup = get_archive(hdr)
-                    if not backup:
-                        #print 'No archived keywords found.\n'
-                        continue
-                    else:
-                        for kw in basic_wcs:
-                            nkw = ('O'+kw)[:7]
-                            if backup.has_key(kw):
-                                hdr.update(kw, hdr[nkw])
-                    tddalpha = hdr.get('TDDALPHA', None)
-                    tddbeta = hdr.get('TDDBETA', None)
-                    if tddalpha or tddbeta:
-                        hdr.update('TDDALPHA', 0.0)
-                        hdr.update('TDDBETA', 0.0)
+                #Restore always from a 'SCI' extension but write it out to 'ERR' and 'DQ'
+                if extname == 'sci':
+                    sciver = fobj[e].header['extver']
+                    try:
+                        nwcs = pywcs.WCS(fobj[e].header, fobj=fobj, key=wcskey)
+                    except:
+                        print 'utils.restoreWCS: Could not read WCS with key %s in file %s,  \
+                        extension %d' % (wcskey, f, e)
+                        return #raise
+                    hwcs = nwcs.to_header()
+                    if nwcs.wcs.has_pc():
+                        for c in ['1_1', '1_2', '2_1', '2_2']:
+                            del hwcs['CD'+c+wcskey]
+                    elif nwcs.wcs.has_cd():
+                        for c in ['1_1', '1_2', '2_1', '2_2']:
+                            hwcs.update(key='CD'+c+wcskey, value=hwcs['PC'+c+wcskey])
+                            del hwcs['PC'+c]
+                    for k in hwcs.keys():
+                        key = k[:-1]
+                        if key in fobj[e].header.keys():
+                            fobj[e].header.update(key=key, value = hwcs[k])
+                        else:
+                            continue
+                    if wcskey == 'O':
+                        fobj[e].header['TDDALPHA'] = 0.0
+                        fobj[e].header['TDDBETA'] = 0.0
+                    if fobj[e].header.has_key('ORIENTAT'):
+                        cd12 = 'CD1_2%s' % wcskey
+                        cd22 = 'CD2_2%s' % wcskey
+                        norient = np.rad2deg(np.arctan2(hwcs[cd12],hwcs[cd22]))
+                        fobj[e].header.update(key='ORIENTAT', value=norient)
+                elif extname in ['err', 'dq', 'sdq']:
+                    cextver = fobj[e].header['extver']
+                    if cextver == sciver:
+                        for k in hwcs.keys():
+                            key = k[:-1]
+                            fobj[e].header.update(key=key, value = hwcs[k])
+                        if fobj[e].header.has_key('ORIENTAT'):
+                            cd12 = 'CD1_2%s' % wcskey
+                            cd22 = 'CD2_2%s' % wcskey
+                            norient = np.rad2deg(np.arctan2(hwcs[cd12],hwcs[cd22]))
+                            fobj[e].header.update(key='ORIENTAT', value=norient)
+                else:
+                    continue
+                
+            if not clobber:
+                fobj.writeto(name)
             fobj.close()
 
-def get_archive(header):
-    """
-    Returns a dictionary with the archived basic WCS keywords.
-    """    
-
-    backup = {}
-    for k in basic_wcs:
-        try:
-            nkw = ('O'+k)[:7]
-            backup[k] = header[nkw]
-        except KeyError:
-            pass
-    return backup
-    
-def write_archive(header):
+def archiveWCS(fname, ext, wcskey, wcsname=" "):
     """
-    Archives original WCS kw before recalculating them.
+    Copy the primary WCS to an alternate WCS 
+    with wcskey and name WCSNAME.
     """
-    backup_kw = get_archive(header)
-    if backup_kw != {}:
-        #print "Archive already exists\n"
+    f = pyfits.open(fname, mode='update')
+    w = pywcs.WCS(f[ext].header, fobj=f)
+    assert len(wcskey) == 1
+    if wcskey == " ": 
+        print "Please provide a valid wcskey for this WCS."
+        print 'Use "utils.next_wcskey" to obtain a valid wcskey.'
+        print 'Use utils.restoreWCS to write to the primary WCS.'
         return
-    else:
-        cmt = 'archived value'
-        for kw in basic_wcs:
-            nkw = 'O'+kw
-            try:
-                header.update(nkw[:8], header[kw], comment=cmt)
-            except KeyError:  #ORIENTAT is not always present in headers
-                pass
-        header.update('WCSCDATE', fileutil.getLTime(), comment='Local time the WCS kw were archived')
+    if wcskey not in available_wcskeys(f[ext].header):
+        print 'wcskey %s is already used in this header.' % wcskey
+        print 'Use "utils.next_wcskey" to obtain a valid wcskey'
+    hwcs = w.to_header()
+    wkey = 'WCSNAME' + wcskey
+    f[ext].header.update(key=wkey, value=wcsname)
+    if w.wcs.has_pc():
+        for c in ['CD1_1', 'CD1_2', 'CD2_1', 'CD2_2']:
+            del hwcs[c]
+    elif w.wcs.has_cd():
+        for c in ['PC1_1', 'PC1_2', 'PC2_1', 'PC2_2']:
+            del hwcs[c]
+    for k in hwcs.keys():
+        key = k+wcskey
+        f[ext].header.update(key=key, value = hwcs[k])
+    norient = np.rad2deg(np.arctan2(hwcs['CD1_2'],hwcs['CD2_2']))
+    okey = 'ORIENT%s' % wcskey
+    f[ext].header.update(key=okey, value=norient) 
+    f.close()
     
 
 def diff_angles(a,b):
@@ -102,3 +162,62 @@ def getBinning(fobj, extver=1):
     else:
         binned = fobj['SCI', extver].header.get('BINAXIS',1)
     return binned
+
+def wcsnames(header):
+    """
+    Purpose
+    =======
+    Return a dictionary of wcskey: WCSNAME pairs
+    """
+    names = header["WCSNAME*"]
+    d = {}
+    for c in names:
+        d[c.key[-1]] = c.value
+    return d
+    
+
+def wcskeys(header):
+    """
+    Purpose
+    =======
+    Returns a list of characters used in the header for alternate 
+    WCS description via WCSNAME keyword
+    
+    `header`: pyfits.Header
+    """
+    names = header["WCSNAME*"]
+    return [key.split('WCSNAME')[1] for key in names.keys()]
+
+def available_wcskeys(header):
+    """
+    Purpose
+    =======
+    Returns a list of characters which are not used in the header
+    with WCSNAME keyword. Any of them can be used to save a new
+    WCS.
+    
+    `header`: pyfits.Header
+    """
+    all_keys = list(string.ascii_uppercase)
+    used_keys = wcskeys(header)
+    try: 
+        used_keys.remove("")
+    except ValueError: 
+        pass
+    [all_keys.remove(key) for key in used_keys]
+    return all_keys
+
+def next_wcskey(header):
+    """
+    Purpose
+    =======
+    Returns next available character to be used for an alternate WCS
+    
+    `header`: pyfits.Header    
+    """
+    allkeys = available_wcskeys(header)
+    if allkeys != []:
+        return allkeys[0]
+    else:
+        return None
+    
\ No newline at end of file
diff --git a/updatewcs/__init__.py b/updatewcs/__init__.py
index b1d88d4..95119f1 100644
--- a/updatewcs/__init__.py
+++ b/updatewcs/__init__.py
@@ -66,7 +66,7 @@ def updatewcs(input, vacorr=True, tddcorr=True, dgeocorr=True, d2imcorr=True, ch
             tddcorr=tddcorr,dgeocorr=dgeocorr, d2imcorr=d2imcorr)
         
         #restore the original WCS keywords
-        utils.restoreWCS(f)
+        utils.restoreWCS(f, wcskey='O', clobber=True)
         makecorr(f, acorr)
     return files
 
@@ -88,7 +88,7 @@ def makecorr(fname, allowed_corr):
     nrefchip, nrefext = getNrefchip(f)
     ref_wcs = HSTWCS(fobj=f, ext=nrefext)
     ref_wcs.readModel(update=True,header=f[nrefext].header)
-    utils.write_archive(f[nrefext].header)
+    ref_wcs.copyWCS(header=f[nrefext].header, wcskey='O', wcsname='OPUS', clobber=True)
     
     if 'DET2IMCorr' in allowed_corr:
         det2im.DET2IMCorr.updateWCS(f)
@@ -96,28 +96,58 @@ def makecorr(fname, allowed_corr):
     for i in range(len(f))[1:]:
         # Perhaps all ext headers should be corrected (to be consistent)
         extn = f[i]
-        if extn.header.has_key('extname') and extn.header['extname'].lower() == 'sci':
-            ref_wcs.restore(f[nrefext].header)
-            hdr = extn.header
-            utils.write_archive(hdr)
-            ext_wcs = HSTWCS(fobj=f, ext=i)
-            ext_wcs.readModel(update=True,header=hdr)
-            for c in allowed_corr:
-                if c != 'DGEOCorr' and c != 'DET2IMCorr':
-                    corr_klass = corrections.__getattribute__(c)
-                    kw2update = corr_klass.updateWCS(ext_wcs, ref_wcs)
-                    for kw in kw2update:
-                        hdr.update(kw, kw2update[kw])
         
+        if extn.header.has_key('extname'):
+            extname = extn.header['extname'].lower()
+            if  extname == 'sci':
+                
+                sciextver = extn.header['extver']
+                ref_wcs.restore(f[nrefext].header, wcskey="O")
+    
+                hdr = extn.header
+                ext_wcs = HSTWCS(fobj=f, ext=i)
+                ext_wcs.copyWCS(header=hdr, wcskey='O', wcsname='OPUS', clobber=True)
+                ext_wcs.readModel(update=True,header=hdr)
+                for c in allowed_corr:
+                    if c != 'DGEOCorr' and c != 'DET2IMCorr':
+                        corr_klass = corrections.__getattribute__(c)
+                        kw2update = corr_klass.updateWCS(ext_wcs, ref_wcs)
+                        for kw in kw2update:
+                            hdr.update(kw, kw2update[kw])
+                    
+                
+                idcname = os.path.split(fileutil.osfn(ext_wcs.idctab))[1]
+                wname = ''.join(['IDC_',idcname.split('_idc.fits')[0]])
+                wkey = getKey(hdr, wname)
+                ext_wcs.copyWCS(header=hdr, wcskey=wkey, wcsname=wname, clobber=True)
+            elif extname in ['err', 'dq', 'sdq']:
+                cextver = extn.header['extver']
+                if cextver == sciextver:
+                    ext_wcs.copyWCS(header=extn.header, wcskey=" ", wcsname=" ")
+            else:
+                cextver = extn.header['extver']
+                continue
+    
     if 'DGEOCorr' in allowed_corr:
         kw2update = dgeo.DGEOCorr.updateWCS(f)
         for kw in kw2update:
-            f[1].header.update(kw, kw2update[kw])
-    
-    
-            
+            f[1].header.update(kw, kw2update[kw])       
             
     f.close()
+
+def getKey(header, wcsname):
+    """
+    If WCSNAME is found in header, return its key, else return 
+    the next available key. This is used to update a specific WCS
+    repeatedly and not generate new keys every time.
+    """
+    wkey = utils.next_wcskey(header)
+    names = utils.wcsnames(header)
+    for item in names.items():
+        if item[1] == wcsname:
+            wkey = item[0]
+    return wkey
+
 def getNrefchip(fobj):
     """
     This handles the fact that WFPC2 subarray observations
diff --git a/wcsutil/__init__.py b/wcsutil/__init__.py
index 7e87fa6..5bdbeb9 100644
--- a/wcsutil/__init__.py
+++ b/wcsutil/__init__.py
@@ -5,6 +5,7 @@ from pywcs import WCS
 import pyfits
 import instruments
 from stwcs.distortion import models, coeff_converter
+from stwcs import utils
 import numpy as np
 from pytools import fileutil
 from pytools.fileutil import DEGTORAD, RADTODEG
@@ -22,11 +23,11 @@ class HSTWCS(WCS):
     Purpose
     =======
     Create a WCS object based on the instrument.
-    It has all basic WCS kw as attribbutes (set by pywcs).
+    It has all basic WCS kw as attributes (set by pywcs).
     It also uses the primary and extension header to define
     instrument specific attributes.
     """
-    def __init__(self, fobj='DEFAULT', ext=None, minerr=0.0):
+    def __init__(self, fobj='DEFAULT', ext=None, minerr=0.0, wcskey=" "):
         """
         :Parameters:
         `fobj`: string or PyFITS HDUList object or None
@@ -45,13 +46,14 @@ class HSTWCS(WCS):
 
         self.inst_kw = ins_spec_kw
         self.minerr = minerr
+        self.wcskey = wcskey
         
         if fobj != 'DEFAULT':
             filename, hdr0, ehdr, phdu = self.parseInput(f=fobj, ext=ext)
             self.filename = filename
             self.instrument = hdr0['INSTRUME']
             
-            WCS.__init__(self, ehdr, fobj=phdu, minerr=minerr)
+            WCS.__init__(self, ehdr, fobj=phdu, minerr=self.minerr, key=self.wcskey)
             # If input was a pyfits HDUList object, it's the user's
             # responsibility to close it, otherwise, it's closed here.
             if not isinstance(fobj, pyfits.HDUList):
@@ -65,7 +67,7 @@ class HSTWCS(WCS):
         else:
             # create a default HSTWCS object 
             self.instrument = 'DEFAULT'
-            WCS.__init__(self, minerr=minerr)
+            WCS.__init__(self, minerr=self.minerr, key=self.wcskey)
             self.wcs.cd = np.array([[1.0, 0.0], [0.0, 1.0]], np.double)
             self.wcs.crval = np.zeros((self.naxis,), np.double)
             self.wcs.crpix = np.zeros((self.naxis,), np.double)
@@ -260,37 +262,52 @@ class HSTWCS(WCS):
                 self._updatehdr(header)
     
     
-    def restore(self, header=None):
+    def restore(self, header, wcskey):
         """
-        Restore a WCS archive in memory and update the WCS object.
-        Restored are the basic WCS kw as well as pscale and orient.
+        If WCS with wcskey exists - read it in and update the primary WCS.
+        If not, keep the original primary WCS.
         """
         from pywcs import Wcsprm
-        backup = {}
-        if header == None:
-            print 'Need a valid header in order to restore archive\n'
-            return
-
-        for k in basic_wcs:
-            try:
-                nkw = ('O'+k)[:7]
-                backup[k] = header[nkw]
-            except KeyError:
-                pass
-        if backup == {}:
-            print 'No archive was found\n'
-            return
-        cdl=pyfits.CardList()
-        for item in backup.items():
-            card = pyfits.Card(key=item[0], value=item[1])
-            cdl.append(card)
-
-        h = pyfits.Header(cdl)
-        wprm = Wcsprm(repr(h.ascard)) 
+        try:
+            wprm = Wcsprm(repr(header.ascard), key=wcskey)
+        except KeyError:
+            print "Could not restore WCS with key %s" % wcskey
+            return 
         self.wcs = wprm
         self.setPscale()
         self.setOrient()
-
+    
+    def copyWCS(self, header=None, wcskey=None, wcsname=" ", clobber=True):
+        """
+        Copy the current primary WCS as an alternate WCS 
+        using wcskey and WCSNAME. If wcskey = " ", it overwrites the 
+        primary WCS.
+        """
+        assert isinstance(header, pyfits.Header)
+        assert len(wcskey) == 1
+        #if wcskey == " ": wcskey = None
+        if header and not wcskey:
+            print "Please provide a valid wcskey for this WCS"
+        if wcskey not in utils.available_wcskeys(header) and not clobber:
+            print 'wcskey %s is already used in this header.' % wcskey
+            print 'Use "utils.next_wcskey" to obtain a valid wcskey'
+        hwcs = self.to_header()
+        wkey = 'WCSNAME' + wcskey
+        header.update(key=wkey, value=wcsname)
+        if self.wcs.has_pc():
+            for c in ['CD1_1', 'CD1_2', 'CD2_1', 'CD2_2']:
+                del hwcs[c]
+        elif self.wcs.has_cd():
+            for c in ['1_1', '1_2', '2_1', '2_2']:
+                hwcs.update(key='CD'+c, value=hwcs['PC'+c])
+                del hwcs['PC'+c]
+        for k in hwcs.keys():
+            key = k+wcskey
+            header.update(key=key, value = hwcs[k])
+        norient = np.rad2deg(np.arctan2(hwcs['CD1_2'],hwcs['CD2_2']))
+        okey = 'ORIENT%s' % wcskey
+        header.update(key=okey, value=norient) 
+        
     def _updatehdr(self, ext_hdr):
         #kw2add : OCX10, OCX11, OCY10, OCY11
         # record the model in the header for use by pydrizzle