Merge pull request #9 from nden/refactor-and-tests

restructure and add stwcs tests

1 files changed, 273 insertions, 0 deletions

diff --git a/stwcs/updatewcs/makewcs.py b/stwcs/updatewcs/makewcs.py
new file mode 100644
index 0000000..06c6f9c
--- /dev/null
+++ b/stwcs/updatewcs/makewcs.py
@@ -0,0 +1,273 @@
+from __future__ import absolute_import, division # confidence high
+
+import datetime
+
+import numpy as np
+import logging, time
+from math import sin, sqrt, pow, cos, asin, atan2,pi
+
+from stsci.tools import fileutil
+from . import utils
+
+logger = logging.getLogger(__name__)
+
+class MakeWCS(object):
+    """
+    Recompute basic WCS keywords based on PA_V3 and distortion model.
+
+    Notes
+    -----
+    - Compute the reference chip WCS:
+
+        -- CRVAL: transform the model XREF/YREF to the sky
+        -- PA_V3 is calculated at the target position and adjusted
+           for each chip orientation
+        -- CD: PA_V3 and the model scale are used to cnstruct a CD matrix
+
+    - Compute the second chip WCS:
+
+        -- CRVAL: - the distance between the zero points of the two
+                  chip models on the sky
+        -- CD matrix: first order coefficients are added to the components
+            of this distance and transfered on the sky. The difference
+            between CRVAL and these vectors is the new CD matrix for each chip.
+        -- CRPIX: chip's model zero point in pixel space (XREF/YREF)
+
+    - Time dependent distortion correction is applied to both chips' V2/V3 values.
+
+    """
+    tdd_xyref = {1: [2048, 3072], 2:[2048, 1024]}
+    def updateWCS(cls, ext_wcs, ref_wcs):
+        """
+        recomputes the basic WCS kw
+        """
+        logger.info("\n\tStarting MakeWCS: %s" % time.asctime())
+        if not ext_wcs.idcmodel:
+            logger.info("IDC model not found, turning off Makewcs")
+            return {}
+        ltvoff, offshift = cls.getOffsets(ext_wcs)
+
+        v23_corr = cls.zero_point_corr(ext_wcs)
+        rv23_corr = cls.zero_point_corr(ref_wcs)
+
+        cls.uprefwcs(ext_wcs, ref_wcs, rv23_corr, ltvoff, offshift)
+        cls.upextwcs(ext_wcs, ref_wcs, v23_corr, rv23_corr, ltvoff, offshift)
+
+        kw2update = {'CD1_1': ext_wcs.wcs.cd[0,0],
+                    'CD1_2': ext_wcs.wcs.cd[0,1],
+                    'CD2_1': ext_wcs.wcs.cd[1,0],
+                    'CD2_2': ext_wcs.wcs.cd[1,1],
+                    'CRVAL1': ext_wcs.wcs.crval[0],
+                    'CRVAL2': ext_wcs.wcs.crval[1],
+                    'CRPIX1': ext_wcs.wcs.crpix[0],
+                    'CRPIX2': ext_wcs.wcs.crpix[1],
+                    'IDCTAB': ext_wcs.idctab,
+                    'OCX10' : ext_wcs.idcmodel.cx[1,0],
+                    'OCX11' : ext_wcs.idcmodel.cx[1,1],
+                    'OCY10' : ext_wcs.idcmodel.cy[1,0],
+                    'OCY11' : ext_wcs.idcmodel.cy[1,1]
+            }
+        return kw2update
+
+    updateWCS = classmethod(updateWCS)
+
+    def upextwcs(cls, ext_wcs, ref_wcs, v23_corr, rv23_corr, loff, offsh):
+        """
+        updates an extension wcs
+        """
+        ltvoffx, ltvoffy = loff[0], loff[1]
+        offshiftx, offshifty = offsh[0], offsh[1]
+        ltv1 = ext_wcs.ltv1
+        ltv2 = ext_wcs.ltv2
+        if ltv1 != 0. or ltv2 != 0.:
+            offsetx = ext_wcs.wcs.crpix[0] - ltv1 - ext_wcs.idcmodel.refpix['XREF']
+            offsety = ext_wcs.wcs.crpix[1] - ltv2 - ext_wcs.idcmodel.refpix['YREF']
+            ext_wcs.idcmodel.shift(offsetx,offsety)
+
+        fx, fy = ext_wcs.idcmodel.cx, ext_wcs.idcmodel.cy
+
+        ext_wcs.setPscale()
+        tddscale = (ref_wcs.pscale/fx[1,1])
+        v2 = ext_wcs.idcmodel.refpix['V2REF'] + v23_corr[0,0] * tddscale
+        v3 = ext_wcs.idcmodel.refpix['V3REF'] - v23_corr[1,0] * tddscale
+        v2ref = ref_wcs.idcmodel.refpix['V2REF'] + rv23_corr[0,0] * tddscale
+        v3ref = ref_wcs.idcmodel.refpix['V3REF'] - rv23_corr[1,0] * tddscale
+
+        R_scale = ref_wcs.idcmodel.refpix['PSCALE']/3600.0
+        off = sqrt((v2-v2ref)**2 + (v3-v3ref)**2)/(R_scale*3600.0)
+
+        if v3 == v3ref:
+            theta=0.0
+        else:
+            theta = atan2(ext_wcs.parity[0][0]*(v2-v2ref), ext_wcs.parity[1][1]*(v3-v3ref))
+
+        if ref_wcs.idcmodel.refpix['THETA']: theta += ref_wcs.idcmodel.refpix['THETA']*pi/180.0
+
+        dX=(off*sin(theta)) + offshiftx
+        dY=(off*cos(theta)) + offshifty
+
+        px = np.array([[dX,dY]])
+        newcrval = ref_wcs.wcs.p2s(px, 1)['world'][0]
+        newcrpix = np.array([ext_wcs.idcmodel.refpix['XREF'] + ltvoffx,
+                                ext_wcs.idcmodel.refpix['YREF'] + ltvoffy])
+        ext_wcs.wcs.crval = newcrval
+        ext_wcs.wcs.crpix = newcrpix
+        ext_wcs.wcs.set()
+
+        # Create a small vector, in reference image pixel scale
+        delmat = np.array([[fx[1,1], fy[1,1]], \
+                              [fx[1,0], fy[1,0]]]) / R_scale/3600.
+
+        # Account for subarray offset
+        # Angle of chip relative to chip
+        if ext_wcs.idcmodel.refpix['THETA']:
+            dtheta = ext_wcs.idcmodel.refpix['THETA'] - ref_wcs.idcmodel.refpix['THETA']
+        else:
+            dtheta = 0.0
+
+        rrmat = fileutil.buildRotMatrix(dtheta)
+        # Rotate the vectors
+        dxy = np.dot(delmat, rrmat)
+        wc = ref_wcs.wcs.p2s((px + dxy), 1)['world']
+
+        # Calculate the new CDs and convert to degrees
+        cd11 = utils.diff_angles(wc[0,0],newcrval[0])*cos(newcrval[1]*pi/180.0)
+        cd12 = utils.diff_angles(wc[1,0],newcrval[0])*cos(newcrval[1]*pi/180.0)
+        cd21 = utils.diff_angles(wc[0,1],newcrval[1])
+        cd22 = utils.diff_angles(wc[1,1],newcrval[1])
+        cd = np.array([[cd11, cd12], [cd21, cd22]])
+        ext_wcs.wcs.cd = cd
+        ext_wcs.wcs.set()
+
+    upextwcs = classmethod(upextwcs)
+
+    def uprefwcs(cls, ext_wcs, ref_wcs, rv23_corr_tdd, loff, offsh):
+        """
+        Updates the reference chip
+        """
+        ltvoffx, ltvoffy = loff[0], loff[1]
+        ltv1 = ref_wcs.ltv1
+        ltv2 = ref_wcs.ltv2
+        if ref_wcs.ltv1 != 0. or ref_wcs.ltv2 != 0.:
+            offsetx = ref_wcs.wcs.crpix[0] - ltv1 - ref_wcs.idcmodel.refpix['XREF']
+            offsety = ref_wcs.wcs.crpix[1] - ltv2 - ref_wcs.idcmodel.refpix['YREF']
+            ref_wcs.idcmodel.shift(offsetx,offsety)
+
+        rfx, rfy = ref_wcs.idcmodel.cx, ref_wcs.idcmodel.cy
+
+        offshift = offsh
+        dec = ref_wcs.wcs.crval[1]
+        tddscale = (ref_wcs.pscale/ref_wcs.idcmodel.cx[1,1])
+        rv23 = [ref_wcs.idcmodel.refpix['V2REF'] + (rv23_corr_tdd[0,0] *tddscale),
+            ref_wcs.idcmodel.refpix['V3REF'] - (rv23_corr_tdd[1,0] * tddscale)]
+        # Get an approximate reference position on the sky
+        rref = np.array([[ref_wcs.idcmodel.refpix['XREF']+ltvoffx ,
+                            ref_wcs.idcmodel.refpix['YREF']+ltvoffy]])
+
+        crval = ref_wcs.wcs.p2s(rref, 1)['world'][0]
+        # Convert the PA_V3 orientation to the orientation at the aperture
+        # This is for the reference chip only - we use this for the
+        # reference tangent plane definition
+        # It has the same orientation as the reference chip
+        pv = troll(ext_wcs.pav3,dec,rv23[0], rv23[1])
+        # Add the chip rotation angle
+        if ref_wcs.idcmodel.refpix['THETA']:
+            pv += ref_wcs.idcmodel.refpix['THETA']
+
+
+        # Set values for the rest of the reference WCS
+        ref_wcs.wcs.crval = crval
+        ref_wcs.wcs.crpix = np.array([0.0,0.0])+offsh
+        parity = ref_wcs.parity
+        R_scale = ref_wcs.idcmodel.refpix['PSCALE']/3600.0
+        cd11 = parity[0][0] *  cos(pv*pi/180.0)*R_scale
+        cd12 = parity[0][0] * -sin(pv*pi/180.0)*R_scale
+        cd21 = parity[1][1] *  sin(pv*pi/180.0)*R_scale
+        cd22 = parity[1][1] *  cos(pv*pi/180.0)*R_scale
+
+        rcd = np.array([[cd11, cd12], [cd21, cd22]])
+        ref_wcs.wcs.cd = rcd
+        ref_wcs.wcs.set()
+
+    uprefwcs = classmethod(uprefwcs)
+
+    def zero_point_corr(cls,hwcs):
+        if hwcs.idcmodel.refpix['skew_coeffs'] is not None and 'TDD_CY_BETA' in hwcs.idcmodel.refpix['skew_coeffs']:
+            v23_corr = np.array([[0.],[0.]])
+            return v23_corr
+        else:
+            try:
+                alpha = hwcs.idcmodel.refpix['TDDALPHA']
+                beta = hwcs.idcmodel.refpix['TDDBETA']
+            except KeyError:
+                alpha = 0.0
+                beta = 0.0
+                v23_corr = np.array([[0.],[0.]])
+                logger.debug("\n\tTDD Zero point correction for chip %s defaulted to: %s" % (hwcs.chip, v23_corr))
+                return v23_corr
+
+        tdd = np.array([[beta, alpha], [alpha, -beta]])
+        mrotp = fileutil.buildRotMatrix(2.234529)/2048.
+        xy0 = np.array([[cls.tdd_xyref[hwcs.chip][0]-2048.], [cls.tdd_xyref[hwcs.chip][1]-2048.]])
+        v23_corr = np.dot(mrotp,np.dot(tdd,xy0)) * 0.05
+        logger.debug("\n\tTDD Zero point correction for chip %s: %s" % (hwcs.chip, v23_corr))
+        return v23_corr
+
+    zero_point_corr = classmethod(zero_point_corr)
+
+    def getOffsets(cls, ext_wcs):
+        ltv1 = ext_wcs.ltv1
+        ltv2 = ext_wcs.ltv2
+
+        offsetx = ext_wcs.wcs.crpix[0] - ltv1 - ext_wcs.idcmodel.refpix['XREF']
+        offsety = ext_wcs.wcs.crpix[1] - ltv2 - ext_wcs.idcmodel.refpix['YREF']
+
+        shiftx = ext_wcs.idcmodel.refpix['XREF'] + ltv1
+        shifty = ext_wcs.idcmodel.refpix['YREF'] + ltv2
+        if ltv1 != 0. or ltv2 != 0.:
+            ltvoffx = ltv1 + offsetx
+            ltvoffy = ltv2 + offsety
+            offshiftx = offsetx + shiftx
+            offshifty = offsety + shifty
+        else:
+            ltvoffx = 0.
+            ltvoffy = 0.
+            offshiftx = 0.
+            offshifty = 0.
+
+        ltvoff = np.array([ltvoffx, ltvoffy])
+        offshift = np.array([offshiftx, offshifty])
+        return ltvoff, offshift
+
+    getOffsets = classmethod(getOffsets)
+
+
+def troll(roll, dec, v2, v3):
+    """ Computes the roll angle at the target position based on:
+        the roll angle at the V1 axis(roll),
+        the dec of the target(dec), and
+        the V2/V3 position of the aperture (v2,v3) in arcseconds.
+
+        Based on algorithm provided by Colin Cox and used in
+        Generic Conversion at STScI.
+    """
+    # Convert all angles to radians
+    _roll = np.deg2rad(roll)
+    _dec = np.deg2rad(dec)
+    _v2 = np.deg2rad(v2 / 3600.)
+    _v3 = np.deg2rad(v3 / 3600.)
+
+    # compute components
+    sin_rho = sqrt((pow(sin(_v2),2)+pow(sin(_v3),2)) - (pow(sin(_v2),2)*pow(sin(_v3),2)))
+    rho = asin(sin_rho)
+    beta = asin(sin(_v3)/sin_rho)
+    if _v2 < 0: beta = pi - beta
+    gamma = asin(sin(_v2)/sin_rho)
+    if _v3 < 0: gamma = pi - gamma
+    A = pi/2. + _roll - beta
+    B = atan2( sin(A)*cos(_dec), (sin(_dec)*sin_rho - cos(_dec)*cos(rho)*cos(A)))
+
+    # compute final value
+    troll = np.rad2deg(pi - (gamma+B))
+
+    return troll