1 files changed, 95 insertions, 0 deletions
diff --git a/updatewcs/alignwcs.py b/updatewcs/alignwcs.py
new file mode 100755
index 0000000..db96ef6
--- /dev/null
+++ b/updatewcs/alignwcs.py
@@ -0,0 +1,95 @@
+#!/usr/bin/env python
+
+from pydrizzle import process_input
+import pywcs, pyfits, numpy as np
+from updatewcs.distortion import utils
+from pytools import fileutil as fu
+from updatewcs import wcsutil
+
+
+def align_wcs(input, shiftfile=None, writeasn=False, asnname=None):
+    """
+    Purpose
+    =======
+    Correct the WCS of a fits file, so that multidrizzle aligns the images.
+    To view the format of the shift file:
+    >>>from pytools.asnutil import ShiftFile  
+    >>>print ShiftFile.__doc__
+
+    Example
+    =======
+    >>>import alignwcs
+    >>>alignwcs.align_wcs('*flt.fits', shiftfile='shifts.txt')
+    It works also on the command line:
+    %./alignwcs.py *flt.fits shifts.txt
+    
+    Dependencies 
+    ============
+    `pytools`
+    `pyfits`
+    `pywcs`
+    `numpy`
+    `pydrizzle`
+
+    :Parameters:
+    `input`: a python list of file names or a string (wild card characters allowed)
+             input files may be in fits, geis or waiver fits format
+    `shiftfile`: a shift file, as the output of tweakshifts
+    `writeasn`: boolean
+                defines whether to write out the asn table updated with shift information
+    `asnname`: string
+                name for the output asn table
+    """
+    if shiftfile == None:
+        print 'A shift file is required but was not provided\n'
+        return
+    
+    asndict,b,c = process_input.process_input(input, shiftfile=shiftfile, updatewcs=False)
+    if writeasn:
+        if asnname == None:
+            asnname = c.split('.fits')[0] + '_shifts' + '.fits'
+        a.write(asnname)
+    outwcs = get_outwcs(asndict)
+    apply_shifts(asndict, outwcs)
+
+def apply_shifts(asndict, outwcs):
+    for mem in asndict['members']:
+        filename = fu.buildRootname(mem)
+        xsh = asndict['members'][mem]['xshift']
+        ysh = asndict['members'][mem]['yshift']
+        shifts = np.array([[xsh, ysh]])
+        f = pyfits.open(filename)
+        for extn in f:
+            if extn.header.has_key('extname') and extn.header['extname'].lower() == 'sci':
+                extver = extn.header['extver']
+                owcs = pywcs.WCS(extn.header, f)
+                crvals = np.array([owcs.wcs.crval])
+                px = outwcs.wcs.s2p(crvals, 1)['pixcrd'] + shifts
+                ncrvals = outwcs.all_pix2sky(px, 1)
+                pyfits.setval(filename, 'CRVAL1', value=ncrvals[0,0], ext=extver)
+                pyfits.setval(filename, 'CRVAL2', value=ncrvals[0,1], ext=extver)
+                print 'Updated %s with shifts ' % filename, shifts
+        f.close()
+
+
+def get_outwcs(asndict):
+    wcslist = []
+    for mem in asndict['members']:
+        filename = fu.buildRootname(mem)
+        f = pyfits.open(filename)
+        hdr0 = f[0].header
+        for extn in f:
+            if extn.header.has_key('extname') and extn.header['extname'].lower() == 'sci':
+                owcs = wcsutil.HSTWCS(hdr0, extn.header,f)
+                wcslist.append(owcs)
+        f.close()
+    outwcs = utils.output_wcs(wcslist)
+    return outwcs
+
+if __name__ == '__main__':
+    import sys
+    args = sys.argv[1:]
+    input = args[:-1]
+    shifts = args[-1]
+    
+    align_wcs(input, shiftfile=shifts)