Per Warren's suggestion, moved the wcscorr and convertwcs modules out of betadrizzle and moved them into the stwcs.wcsutil package. Also cleaned up some imports and cleaned up some whitespace--might want to view this diff with ignore-whitespace.

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

2 files changed, 454 insertions, 0 deletions

diff --git a/wcsutil/convertwcs.py b/wcsutil/convertwcs.py
new file mode 100644
index 0000000..2a53d57
--- /dev/null
+++ b/wcsutil/convertwcs.py
@@ -0,0 +1,118 @@
+import pyfits
+try:
+    import stwcs
+    from stwcs import wcsutil
+except:
+    stwcs = None
+
+from pytools import fileutil
+
+OPUS_WCSKEYS = ['OCRVAL1','OCRVAL2','OCRPIX1','OCRPIX2',
+                'OCD1_1','OCD1_2','OCD2_1','OCD2_2',
+                'OCTYPE1','OCTYPE2']
+
+
+def archive_prefix_OPUS_WCS(fobj,extname='SCI'):
+    """ Identifies WCS keywords which were generated by OPUS and archived
+        using a prefix of 'O' for all 'SCI' extensions in the file
+
+        Parameters
+        ----------
+        fobj: string or pyfits.HDUList
+            Filename or pyfits object of a file
+
+    """
+    if stwcs is None:
+        print '====================='
+        print 'The STWCS package is needed to convert an old-style OPUS WCS to an alternate WCS'
+        print '====================='
+        raise ImportError
+
+
+    closefits = False
+    if isinstance(fobj,str):
+        # A filename was provided as input
+        fobj = pyfits.open(fobj,mode='update')
+        closefits=True
+
+    # Define the header
+    ext = ('sci',1)
+    hdr = fobj[ext].header
+
+    numextn = fileutil.countExtn(fobj)
+    extlist = []
+    for e in xrange(1,numextn+1):
+        extlist.append(('sci',e))
+
+    # Insure that the 'O' alternate WCS is present
+    if 'O' not in wcsutil.wcskeys(hdr):
+        # if not, archive the Primary WCS as the default OPUS WCS
+        wcsutil.archiveWCS(fobj,extlist, wcskey='O', wcsname='OPUS')
+
+    # find out how many SCI extensions are in the image
+    numextn = fileutil.countExtn(fobj,extname=extname)
+    if numextn == 0:
+        extname = 'PRIMARY'
+
+    # create HSTWCS object from PRIMARY WCS
+    wcsobj = wcsutil.HSTWCS(fobj,ext=ext,wcskey='O')
+    # get list of WCS keywords
+    wcskeys = wcsobj.wcs2header().keys()
+
+    # For each SCI extension...
+    for e in xrange(1,numextn+1):
+        # Now, look for any WCS keywords with a prefix of 'O'
+        for key in wcskeys:
+            okey = 'O'+key[:7]
+            hdr = fobj[(extname,e)].header
+            if hdr.has_key(okey):
+                # Update alternate WCS keyword with prefix-O OPUS keyword value
+                hdr[key] = hdr[okey]
+
+    if closefits:
+        fobj.close()
+
+def create_prefix_OPUS_WCS(fobj,extname='SCI'):
+    """ Creates alternate WCS with a prefix of 'O' for OPUS generated WCS values
+        to work with old MultiDrizzle.
+
+        Parameters
+        ----------
+        fobj: string or pyfits.HDUList
+            Filename or pyfits object of a file
+
+        Raises
+        ------
+        IOError:
+            if input PyFITS object was not opened in 'update' mode
+
+    """
+    # List of O-prefix keywords to create
+    owcskeys = OPUS_WCSKEYS
+
+    closefits = False
+    if isinstance(fobj,str):
+        # A filename was provided as input
+        fobj = pyfits.open(fobj,mode='update')
+        closefits=True
+    else:
+        # check to make sure this FITS obj has been opened in update mode
+        if fobj.fileinfo(0)['filemode'] != 'update':
+            print 'File not opened with "mode=update". Quitting...'
+            raise IOError
+
+    # check for existance of O-prefix WCS
+    if not fobj['sci',1].header.has_key(owcskeys[0]):
+
+        # find out how many SCI extensions are in the image
+        numextn = fileutil.countExtn(fobj,extname=extname)
+        if numextn == 0:
+            extname = ''
+        for extn in xrange(1,numextn+1):
+            hdr = fobj[(extname,extn)].header
+            for okey in owcskeys:
+                hdr.update(okey,hdr[okey[1:]+'O'])
+
+    # Close FITS image if we had to open it...
+    if closefits:
+        fobj.close()
diff --git a/wcsutil/wcscorr.py b/wcsutil/wcscorr.py
new file mode 100644
index 0000000..cb0497b
--- /dev/null
+++ b/wcsutil/wcscorr.py
@@ -0,0 +1,336 @@
+import os,copy
+import pyfits
+import numpy as np
+
+from pytools import fileutil
+import stwcs
+from stwcs.wcsutil import altwcs
+import convertwcs
+
+
+DEFAULT_WCS_KEYS = ['CRVAL1','CRVAL2','CRPIX1','CRPIX2',
+                    'CD1_1','CD1_2','CD2_1','CD2_2',
+                    'CTYPE1','CTYPE2']
+DEFAULT_PRI_KEYS = ['PA_V3']
+###
+### WCSEXT table related keyword archive functions
+###
+def init_wcscorr(input, force=False):
+    """ This function will initialize the WCSCORR table if it is not already present,
+        and look for WCS keywords with a prefix of 'O' as the original OPUS
+        generated WCS as the initial row for the table or use the current WCS
+        keywords as initial row if no 'O' prefix keywords are found.
+
+        This function will NOT overwrite any rows already present.
+
+        This function works on all SCI extensions at one time.
+    """
+
+    # TODO: Create some sort of decorator or (for Python2.5) context for
+    # opening a FITS file and closing it when done, if necessary
+    if not isinstance(input, pyfits.HDUList):
+        # input must be a filename, so open as PyFITS object
+        fimg = pyfits.open(input, mode='update')
+        need_to_close = True
+    else:
+        fimg = input
+        need_to_close = False
+
+    # Verify that a WCSCORR extension does not already exist...
+    for extn in fimg:
+        if extn.header.has_key('extname') and extn.header['extname'] == 'WCSCORR':
+            if not force:
+                return
+            else:
+                del fimg['WCSCORR']
+    # define the primary columns of the WCSEXT table with initial rows for each
+    # SCI extension for the original OPUS solution
+    numsci = fileutil.countExtn(fimg)
+
+    # create new table with more rows than needed initially to make it easier to
+    # add new rows later
+    wcsext = create_wcscorr(numrows=numsci,padding=numsci*4)
+    # Assign the correct EXTNAME value to this table extension
+    wcsext.header.update('TROWS',numsci*2,comment='Number of updated rows in table')
+    wcsext.header.update('EXTNAME','WCSCORR',comment='Table with WCS Update history')
+
+    used_wcskeys = None
+    wcs1 = stwcs.wcsutil.HSTWCS(fimg,ext=('SCI',1))
+    idc2header = True
+    if wcs1.idcscale is None:
+        idc2header = False
+    wcs_keywords = wcs1.wcs2header(idc2hdr=idc2header).keys()
+    # Now copy original OPUS values into table
+    for extver in xrange(1,numsci+1):
+        rowind = find_wcscorr_row(wcsext.data,{'WCS_ID':'OPUS','EXTVER':extver,'WCS_key':'O'})
+        # There should only EVER be a single row for each extension with OPUS values
+        rownum = np.where(rowind)[0][0]
+        #print 'Archiving OPUS WCS in row number ',rownum,' in WCSCORR table for SCI,',extver
+
+        hdr = fimg['SCI',extver].header
+        # define set of WCS keywords which need to be managed and copied to the table
+        if used_wcskeys is None:
+            used_wcskeys = altwcs.wcskeys(hdr)
+        # Check to see whether or not there is an OPUS alternate WCS present,
+        # if so, get its values directly, otherwise, archive the PRIMARY WCS
+        # as the OPUS values in the WCSCORR table
+        if 'O' in used_wcskeys:
+            wkey = 'O'
+        else:
+            wkey = " "
+        wcshdr = stwcs.wcsutil.HSTWCS(fimg,ext=('SCI',extver),wcskey=wkey).wcs2header(idc2hdr=idc2header)
+
+        for key in DEFAULT_PRI_KEYS:
+            prihdr_keys = []
+            if not hdr.has_key(key):
+                prihdr_keys.append(key)
+
+        if wcsext.data.field('CRVAL1')[rownum] != 0:
+            # If we find values for these keywords already in the table, do not
+            # overwrite them again
+            print 'WCS keywords already updated...'
+            break
+        for key in wcs_keywords:
+            if key in wcsext.data.names:
+                wcsext.data.field(key)[rownum] = wcshdr[(key+wkey)[:8]]
+        # Now get any keywords from PRIMARY header needed for WCS updates
+        for key in prihdr_keys:
+            wcsext.data.field(key)[rownum] = fimg[0].header[key]
+
+    # Now that we have archived the OPUS alternate WCS, remove it from the list
+    # of used_wcskeys
+    if "O" in used_wcskeys:
+        used_wcskeys.remove("O")
+
+    # Now copy remaining alternate WCSs into table
+    for extver in xrange(1,numsci+1):
+        hdr = fimg['SCI',extver].header
+        if wkey == ' ':
+            break
+        for uwkey in used_wcskeys:
+            wcshdr = stwcs.wcsutil.HSTWCS(fimg,ext=('SCI',extver),wcskey=uwkey).wcs2header()
+            if 'WCSNAME'+uwkey not in wcshdr:
+                idctab =fileutil.osfn(fimg[0].header['idctab'])
+                idcname = os.path.split(idctab)[-1]
+                idcname = idcname[:idcname.find('_')]
+                wcsid = 'IDC_'+idcname+'_'+fileutil.getDate()
+            else:
+                wcsid = wcshdr['WCSNAME'+uwkey]
+
+            # identify next empty row
+            rowind = find_wcscorr_row(wcsext.data,selections={'wcs_id':' '})
+            rows = np.where(rowind)
+            if len(rows[0]) > 0:
+                rownum = np.where(rowind)[0][0]
+            else:
+                print 'No available rows found for updating. '
+            #print 'Archiving current WCS row number ',rownum,' in WCSCORR table for SCI,',extver
+
+            # Update selection columns for this row with relevant values
+            wcsext.data.field('WCS_ID')[rownum] = wcsid
+            wcsext.data.field('EXTVER')[rownum] = extver
+            wcsext.data.field('WCS_key')[rownum] = uwkey
+
+            # Look for standard WCS keyword values
+            for key in wcs_keywords:
+                if key in wcsext.data.names:
+                    wcsext.data.field(key)[rownum] = wcshdr[key+uwkey]
+            # Now get any keywords from PRIMARY header needed for WCS updates
+            for key in prihdr_keys:
+                wcsext.data.field(key)[rownum] = fimg[0].header[key]
+
+    # Append this table to the image FITS file
+    fimg.append(wcsext)
+    # force an update now
+    # set the verify flag to 'warn' so that it will always succeed, but still
+    # tell the user if PyFITS detects any problems with the file as a whole
+    fimg.flush('warn')
+
+    if need_to_close:
+        fimg.close()
+
+def find_wcscorr_row(wcstab, selections):
+    """ Return an array of indices from the table (NOT HDU) 'wcstab' that matches the
+        selections specified by the user.
+
+        The row selection criteria must be specified as a dictionary with
+        column name as key and value(s) representing the valid desired row values.
+        For example, {'wcs_id':'OPUS','extver':2}.
+    """
+    mask = None
+    for i in selections:
+        bmask = (wcstab.field(i) == selections[i])
+        if mask is None:
+            mask = bmask.copy()
+        else:
+            mask = np.logical_and(mask,bmask)
+        del bmask
+    return mask
+
+def archive_wcs_file(image, wcs_id=None):
+    """ Update WCSCORR table with rows for each SCI extension to record the
+        newly updated WCS keyword values.
+    """
+
+    if not isinstance(image, pyfits.HDUList):
+        fimg = pyfits.open(image, mode='update')
+        close_image = True
+    else:
+        fimg = image
+        close_image = False
+    numsci = fileutil.countExtn(fimg)
+    for extn in range(1, numsci+1):
+        update_wcscorr(fimg, fimg['sci',extn].header, selections=wcs_id)
+    if close_image:
+        fimg.close()
+
+def update_wcscorr(fimg,hdr,selections=None):
+    """ Update WCSCORR table with a new row for this extension header. It
+    copies the current set of WCS keywords as a new row of the table based on
+    keyed WCSs as per Paper I Multiple WCS standard).
+    """
+    # Now update the table...
+    if selections is None:
+        # define the WCS ID for this update
+        wcs_key = altwcs.wcskeys(hdr)[-1]
+        selections = {'WCS_ID':'TWEAK_'+fileutil.getDate(),'EXTVER':hdr['extver'],'WCS_key':wcs_key}
+    if selections == 'IDC':
+        idcname = os.path.split(fileutil.osfn(fimg[0].header['idctab']))[1]
+        selections = {'WCS_ID':'IDC_'+idcname+'_'+fileutil.getDate(),'EXTVER':hdr['extver'],'WCS_key':wcs_key}
+    # create new table for hdr and populate it with the newly updated values
+    new_table = create_wcscorr()
+    if hdr.has_key('extname'):
+        extname = hdr['extname']
+    else:
+        extname = 'PRIMARY'
+    if hdr.has_key('extver'): extver = hdr['extver']
+    else: extver = 1
+    extn = (extname,extver)
+    wcshdr = stwcs.wcsutil.HSTWCS(fimg,ext=extn).wcs2header()
+    # ===> NOTE: need to add checks to insure that these IDs are unique
+    # Update selection column values
+    for key in selections:
+        new_table.data.field(key)[0] = selections[key]
+
+    for key in wcshdr:
+        if key in new_table.data.names:
+            new_table.data.field(key)[0] = hdr[key]
+
+    for key in DEFAULT_PRI_KEYS:
+        if key in new_table.data.names:
+            new_table.data.field(key)[0] = fimg[0].header[key]
+
+    # Now, we need to merge this into the existing table
+    old_table = fimg['WCSCORR']
+    rowind = find_wcscorr_row(old_table.data,{'wcs_id':' '})
+    old_nrows = np.where(rowind)[0][0]
+
+    # check to see if there is room for the new row
+    if (old_nrows + 1) > old_table.data.shape[0]:
+        # if not, create a new table with 'pad_rows' new empty rows
+        upd_table = pyfits.new_table(old_table.columns,nrows=old_table.data.shape[0]+pad_rows)
+    else:
+        upd_table = old_table
+    # Now, add
+    for name in old_table.columns.names:
+        upd_table.data.field(name)[old_nrows:old_nrows+1] = new_table.data.field(name)
+    upd_table.header.update('TROWS',old_nrows+1)
+
+    # replace old extension with newly updated table extension
+    fimg['WCSCORR'] = upd_table
+
+def restore_file_from_wcscorr(image, id='OPUS', wcskey=''):
+    """ Copies the values of the WCS from the WCSCORR based on ID specified by user.
+    The default will be to restore the original OPUS-derived values to the Primary WCS.
+    If wcskey is specified, the WCS with that key will be updated instead.
+    """
+
+    if not isinstance(image, pyfits.HDUList):
+        fimg = pyfits.open(image, mode='update')
+        close_image = True
+    else:
+        fimg = image
+        close_image = False
+    numsci = fileutil.countExtn(fimg)
+    wcs_table = fimg['WCSCORR']
+    orig_rows = (wcs_table.data.field('WCS_ID') == 'OPUS')
+    # create an HSTWCS object to figure out what WCS keywords need to be updated
+    wcsobj = stwcs.wcsutil.HSTWCS(fimg,ext=('sci',1))
+    wcshdr = wcsobj.wcs2header()
+    for extn in range(1,numsci+1):
+        # find corresponding row from table
+        ext_rows = (wcs_table.data.field('EXTVER') == extn)
+        erow = np.where(np.logical_and(ext_rows,orig_rows))[0][0]
+        for key in wcshdr:
+            if key in wcs_table.data.names: # insure that keyword is column in table
+                tkey = key
+
+                if 'orient' in key.lower():
+                    key = 'ORIENT'
+                if wcskey == '':
+                    skey = key
+                else:
+                    skey = key[:7]+wcskey
+                fimg['sci',extn].header.update(skey,wcs_table.data.field(tkey)[erow])
+        for key in DEFAULT_PRI_KEYS:
+            if key in wcs_table.data.names:
+                if wcskey == '':
+                    pkey = key
+                else:
+                    pkey = key[:7]+wcskey
+                fimg[0].header.update(pkey,wcs_table.data.field(key)[erow])
+
+    # close the image now that the update has been completed.
+    if close_image:
+        fimg.close()
+
+def create_wcscorr(descrip=False, numrows=1, padding=0):
+    """ Return the basic definitions for a WCSCORR table.
+    The dtype definitions for the string columns are set to the maximum allowed so
+    that all new elements will have the same max size which will be automatically
+    truncated to this limit upon updating (if needed).
+
+    The table is initialized with rows corresponding to the OPUS solution
+    for all the 'SCI' extensions.
+    """
+    trows = numrows + padding
+    # define initialized arrays as placeholders for column data
+    def_float64_zeros = np.array([0.0]*trows,dtype=np.float64)
+    def_float64_ones = def_float64_zeros + 1.0
+    def_range_1nrows = np.array(range(1,numrows+1),dtype=np.int16)
+    def_float_col = {'format':'D','array':def_float64_zeros.copy()}
+    def_float1_col = {'format':'D','array':def_float64_ones.copy()}
+    def_str40_col = {'format':'40A','array':np.array(['']*trows,dtype="S40")}
+    def_str24_col = {'format':'24A','array':np.array(['']*trows,dtype="S24")}
+    def_int32_col = {'format':'J','array':np.array([0]*trows,dtype=np.int32)}
+
+    # If more columns are needed, simply add their definitions to this list
+    col_names = [('CRVAL1',def_float_col),('CRVAL2',def_float_col),
+                ('CRPIX1',def_float_col),('CRPIX2',def_float_col),
+                ('CD1_1',def_float_col),('CD1_2',def_float_col),
+                ('CD2_1',def_float_col),('CD2_2',def_float_col),
+                ('CTYPE1',def_str24_col),('CTYPE2',def_str24_col),
+                ('ORIENTAT',def_float_col),('PA_V3',def_float_col),
+                ('Delta_RA',def_float_col),('Delta_Dec',def_float_col),
+                ('RMS_RA',def_float_col),('RMS_Dec',def_float_col),
+                ('Delta_Orientat',def_float_col),('Delta_Scale',def_float1_col),
+                ('NMatch',def_int32_col),('Catalog',def_str40_col)]
+
+    # Define selector columns
+    id_col = pyfits.Column(name='WCS_ID',format='40A',array=np.array(['OPUS']*numrows+['']*padding,dtype="S24"))
+    extver_col = pyfits.Column(name='EXTVER',format='I',array=np.array(range(1,numrows+1),dtype=np.int16))
+    wcskey_col = pyfits.Column(name='WCS_key',format='A',array=np.array(['O']*numrows+['']*padding,dtype="S"))
+    # create list of remaining columns to be added to table
+    col_list = [id_col,extver_col,wcskey_col] # start with selector columns
+    for c in col_names:
+        cdef = copy.deepcopy(c[1])
+        col_list.append(pyfits.Column(name=c[0],format=cdef['format'],array=cdef['array']))
+
+    if descrip:
+        col_list.append(pyfits.Column(name='Descrip',format='128A',array=np.array(['Original WCS computed by OPUS']*numrows,dtype="S128")))
+
+    # Now create the new table from the column definitions
+    newtab = pyfits.new_table(pyfits.ColDefs(col_list),nrows=trows)
+
+    return newtab
+