TESS WCS Transformation Implementation

This notebook implements a complete World Coordinate System (WCS) transformation for TESS images that perfectly matches astropy's functionality. The implementation handles both forward (pixel → RA/DEC) transformations with SIP (Simple Imaging Polynomial) distortion corrections.

The transformation follows these steps:

Forward Transformation (Pixel → RA/DEC)

1. Pixel coordinate normalization: Convert to detector coordinates relative to reference pixel
2. SIP distortion correction: Apply polynomial corrections to pixel coordinates
3. Linear transformation: Apply CD matrix to get intermediate world coordinates
4. Projection: Convert to celestial coordinates using TAN projection

import numpy as np
from astropy.io import fits
from astropy.wcs import WCS

# Load TESS data and WCS information
im1_file = "../data/tess/20_3_3/tess2020019135923-s0020-3-3-0165-s_ffic.fits"
im1_file = "../data/tess_ffi/s0020_3_3/tess2019359015923-s0020-3-3-0165-s_ffic.fits"
hdul = fits.open(im1_file)
tess_data = hdul[1].data
tess_header = hdul[1].header
data_shape = np.shape(tess_data)
tess_wcs = WCS(hdul[1].header)
hdul.close()

print(f"TESS image shape: {data_shape}")
print(f"Reference pixel: CRPIX = ({tess_header['CRPIX1']}, {tess_header['CRPIX2']})")
print(f"Reference coordinates: RA = {tess_header['CRVAL1']:.6f}°, DEC = {tess_header['CRVAL2']:.6f}°")

TESS image shape: (2078, 2136)
Reference pixel: CRPIX = (1045.0, 1001.0)
Reference coordinates: RA = 216.754373°, DEC = 76.392049°

# Create undistorted WCS for comparison (remove SIP coefficients)
import copy

header_no_distortion = copy.deepcopy(tess_header)

# Remove all SIP distortion coefficients
for keyword in list(tess_header.keys()):
    if keyword.startswith(('A_', 'B_', 'AP_', 'BP_')):
        del header_no_distortion[keyword]

# Change to standard TAN projection
header_no_distortion['CTYPE1'] = 'RA---TAN'
header_no_distortion['CTYPE2'] = 'DEC--TAN'

# Create undistorted WCS
wcs_no_distortion = WCS(header_no_distortion)

Mathematical Formulas

Forward Transformation: Pixel → RA/DEC

Step 1: Pixel Coordinate Normalization (1-based FITS convention) $$u = x - (\text{CRPIX}_1 - 1)$$ $$v = y - (\text{CRPIX}_2 - 1)$$

Step 2: SIP Distortion Correction $$f(u,v) = \sum_{i=0}^{m} \sum_{j=0}^{m-i} A_{i,j} \, u^i v^j$$ $$g(u,v) = \sum_{i=0}^{m} \sum_{j=0}^{m-i} B_{i,j} \, u^i v^j$$ $$u' = u + f(u,v)$$ $$v' = v + g(u,v)$$

Step 3: Linear Transformation (CD Matrix) $$\begin{pmatrix} \xi \\ \eta \end{pmatrix} = \begin{pmatrix} \text{CD}_{1,1} & \text{CD}_{1,2} \\ \text{CD}_{2,1} & \text{CD}_{2,2} \end{pmatrix} \begin{pmatrix} u' \\ v' \end{pmatrix}$$

Step 4: TAN Projection (Gnomonic) $$\rho = \sqrt{\xi^2 + \eta^2}$$ $$c = \arctan(\rho)$$ $$\delta = \arcsin\left(\cos(c) \sin(\delta_0) + \frac{\eta \sin(c) \cos(\delta_0)}{\rho}\right)$$ $$\alpha = \alpha_0 + \operatorname{atan2}\left( \xi \sin(c), \, \rho \cos(\delta_0) \cos(c) - \eta \sin(\delta_0) \sin(c) \right)$$

Implementation

Core Functions

def evaluate_sip_polynomial(u, v, order, coeff_prefix, header):
    """
    Evaluate SIP polynomial correction
    
    Parameters:
    -----------
    u, v : float
        Pixel coordinates relative to reference pixel
    order : int
        Maximum polynomial order
    coeff_prefix : str
        'A' for f(u,v) or 'B' for g(u,v)
    header : dict
        FITS header containing SIP coefficients
        
    Returns:
    --------
    float : Polynomial correction value
    """
    correction = 0.0
    for i in range(order + 1):
        for j in range(order - i + 1):
            key = f"{coeff_prefix}_{i}_{j}"
            if key in header:
                term = header[key] * (u**i) * (v**j)
                correction += term
    return correction

def inverse_tan_projection(xi, eta, header):
    """
    Convert intermediate world coordinates to celestial coordinates using TAN projection
    
    Parameters:
    -----------
    xi, eta : float
        Intermediate world coordinates in degrees
    header : dict
        FITS header containing CRVAL1, CRVAL2
        
    Returns:
    --------
    tuple : (ra, dec) in degrees
    """
    # Convert to radians
    xi_rad = np.deg2rad(xi)
    eta_rad = np.deg2rad(eta)
    
    # Reference coordinates
    ra0 = np.deg2rad(header["CRVAL1"])
    dec0 = np.deg2rad(header["CRVAL2"])

    # Handle special case at reference point
    if np.allclose(xi, 0) and np.allclose(eta, 0):
        return (header["CRVAL1"], header["CRVAL2"])

    # TAN projection formulas
    rho = np.sqrt(xi_rad**2 + eta_rad**2)
    c = np.arctan(rho)
    
    sin_c = np.sin(c)
    cos_c = np.cos(c)
    sin_dec0 = np.sin(dec0)
    cos_dec0 = np.cos(dec0)

    # Calculate declination
    dec = np.arcsin(cos_c * sin_dec0 + (eta_rad * sin_c * cos_dec0) / rho)

    # Calculate right ascension
    y_term = xi_rad * sin_c
    x_term = rho * cos_dec0 * cos_c - eta_rad * sin_dec0 * sin_c
    ra = ra0 + np.arctan2(y_term, x_term)

    return (np.rad2deg(ra), np.rad2deg(dec))

def forward_tan_projection(ra, dec, header):
    """
    Convert celestial coordinates to intermediate world coordinates using TAN projection
    
    Parameters:
    -----------
    ra, dec : float
        Right ascension and declination in degrees
    header : dict
        FITS header containing CRVAL1, CRVAL2
        
    Returns:
    --------
    tuple : (xi, eta) intermediate world coordinates in degrees
    """
    # Convert to radians
    ra_rad = np.deg2rad(ra)
    dec_rad = np.deg2rad(dec)
    
    # Reference coordinates
    ra0 = np.deg2rad(header["CRVAL1"])
    dec0 = np.deg2rad(header["CRVAL2"])
    
    # Handle special case at reference point
    if np.allclose(ra, header["CRVAL1"]) and np.allclose(dec, header["CRVAL2"]):
        return (0.0, 0.0)
    
    # Forward TAN projection
    cos_dec = np.cos(dec_rad)
    sin_dec = np.sin(dec_rad)
    cos_dec0 = np.cos(dec0)
    sin_dec0 = np.sin(dec0)
    
    dra = ra_rad - ra0
    cos_dra = np.cos(dra)
    sin_dra = np.sin(dra)
    
    # Calculate angular distance
    cos_c = sin_dec0 * sin_dec + cos_dec0 * cos_dec * cos_dra
    
    if cos_c <= 0:
        raise ValueError("Point is more than 90 degrees from reference point")
    
    # Intermediate coordinates
    xi = cos_dec * sin_dra / cos_c
    eta = (cos_dec0 * sin_dec - sin_dec0 * cos_dec * cos_dra) / cos_c
    
    return (np.rad2deg(xi), np.rad2deg(eta))

def pixel_to_world(x, y, header):
    """
    Transform pixel coordinates to world coordinates (RA/DEC)
    
    This is the COMPLETE implementation that perfectly matches astropy.
    
    Parameters:
    -----------
    x, y : float
        Pixel coordinates (0-based)
    header : dict
        FITS header containing WCS information
        
    Returns:
    --------
    tuple : ((ra_undist, dec_undist), (ra_dist, dec_dist))
            Undistorted and distorted coordinates in degrees
    """
    # Step 1: Convert to detector coordinates (1-based FITS convention)
    u = x - (header['CRPIX1'] - 1)
    v = y - (header['CRPIX2'] - 1)

    # Step 2: Linear transformation for undistorted coordinates
    xi_undist = header['CD1_1'] * u + header['CD1_2'] * v
    eta_undist = header['CD2_1'] * u + header['CD2_2'] * v

    # Step 3: SIP distortion correction (apply to PIXEL coordinates first)
    f_uv = evaluate_sip_polynomial(u, v, int(header['A_ORDER']), 'A', header)
    g_uv = evaluate_sip_polynomial(u, v, int(header['B_ORDER']), 'B', header)
    
    # Correct pixel coordinates
    u_corrected = u + f_uv
    v_corrected = v + g_uv
    
    # Then apply CD transformation to corrected pixels
    xi_dist = header['CD1_1'] * u_corrected + header['CD1_2'] * v_corrected
    eta_dist = header['CD2_1'] * u_corrected + header['CD2_2'] * v_corrected

    # Step 4: Convert to celestial coordinates
    ra_undistorted, dec_undistorted = inverse_tan_projection(xi_undist, eta_undist, header)
    ra_distorted, dec_distorted = inverse_tan_projection(xi_dist, eta_dist, header)
    
    return (ra_undistorted, dec_undistorted), (ra_distorted, dec_distorted)

x, y = 0, 0
(ra_undist, dec_undist), (ra_dist, dec_dist) = pixel_to_world(x, y, tess_header)
astropy_result = tess_wcs.all_pix2world(x, y, 0)
astropy_result_nodist = wcs_no_distortion.all_pix2world(x, y, 0)

print(f"Pixel ({x}, {y}) corresponds to:")
print(f"No distortion 1. Astropy  2. This result:")
print(f"RA={ra_undist:.6f}°, DEC={dec_undist:.6f}°")
print(f"RA={astropy_result_nodist[0]:.6f}°, DEC={astropy_result_nodist[1]:.6f}°")
print(f"Distorted 1. Astropy 2. This result:")
print(f"RA={astropy_result[0]:.6f}°, DEC={astropy_result[1]:.6f}°")
print(f"RA={ra_dist:.6f}°, DEC={dec_dist:.6f}°")

Pixel (0, 0) corresponds to:
No distortion 1. Astropy  2. This result:
RA=225.164999°, DEC=68.687529°
RA=225.164999°, DEC=68.687529°
Distorted 1. Astropy 2. This result:
RA=224.892441°, DEC=69.066314°
RA=224.892441°, DEC=69.066314°

---

Reverse Transformation: RA/DEC → Pixel

Now we implement the reverse transformation to convert celestial coordinates (RA/DEC) back to pixel coordinates. This is the inverse of the forward transformation we just implemented.

The reverse transformation follows these steps in reverse order:

Reverse Transformation (RA/DEC → Pixel)

1. Projection: Convert celestial coordinates to intermediate world coordinates using inverse TAN projection
2. Linear transformation: Apply inverse CD matrix to get distortion-corrected pixel offsets
3. Inverse SIP distortion correction: Apply inverse polynomial corrections using AP/BP coefficients
4. Pixel coordinate denormalization: Add back the reference pixel offset to get final pixel coordinates

Mathematical Formulas for Reverse Transformation

Reverse Transformation: RA/DEC → Pixel

Step 1: Forward TAN Projection (Celestial to Intermediate)

Given celestial coordinates $(\alpha, \delta)$ and reference point $(\alpha_0, \delta_0)$:

$$\Delta\alpha = \alpha - \alpha_0$$ $$\cos c = \sin(\delta_0) \sin(\delta) + \cos(\delta_0) \cos(\delta) \cos(\Delta\alpha)$$ $$\xi = \frac{\cos(\delta) \sin(\Delta\alpha)}{\cos c}$$ $$\eta = \frac{\cos(\delta_0) \sin(\delta) - \sin(\delta_0) \cos(\delta) \cos(\Delta\alpha)}{\cos c}$$

Step 2: Apply Inverse CD Matrix

$$\begin{pmatrix} u' \\ v' \end{pmatrix} = \begin{pmatrix} \text{CD}_{1,1} & \text{CD}_{1,2} \\ \text{CD}_{2,1} & \text{CD}_{2,2} \end{pmatrix}^{-1} \begin{pmatrix} \xi \\ \eta \end{pmatrix}$$

Step 3: Inverse SIP Distortion Correction

If inverse coefficients (AP, BP) are available: $$f^{-1}(u',v') = \sum_{i=0}^{m} \sum_{j=0}^{m-i} \text{AP}_{i,j} \, {u'}^i {v'}^j$$ $$g^{-1}(u',v') = \sum_{i=0}^{m} \sum_{j=0}^{m-i} \text{BP}_{i,j} \, {u'}^i {v'}^j$$ $$u = u' + f^{-1}(u',v')$$ $$v = v' + g^{-1}(u',v')$$

If inverse coefficients are not available, use iterative Newton-Raphson method to solve: $$u' = u + f(u,v)$$ $$v' = v + g(u,v)$$

Step 4: Convert to Final Pixel Coordinates (FITS 1-based convention)

$$x = u + (\text{CRPIX}_1 - 1)$$ $$y = v + (\text{CRPIX}_2 - 1)$$

Implementation of Reverse Transformation

Helper Functions for RA/DEC → Pixel

def evaluate_inverse_sip_polynomial(u_prime, v_prime, order, coeff_prefix, header):
    """
    Evaluate inverse SIP polynomial correction
    
    This function applies the inverse distortion correction using AP or BP coefficients.
    These coefficients are pre-computed to directly map from distorted to undistorted coordinates.
    
    Parameters:
    -----------
    u_prime, v_prime : float
        Distortion-corrected pixel coordinates relative to reference pixel
    order : int
        Maximum polynomial order
    coeff_prefix : str
        'AP' for inverse f(u,v) or 'BP' for inverse g(u,v)
    header : dict
        FITS header containing inverse SIP coefficients
        
    Returns:
    --------
    float : Inverse polynomial correction value
    """
    correction = 0.0
    for i in range(order + 1):
        for j in range(order - i + 1):
            key = f"{coeff_prefix}_{i}_{j}"
            if key in header:
                term = header[key] * (u_prime**i) * (v_prime**j)
                correction += term
    return correction

def invert_sip_distortion_iterative(u_prime, v_prime, header, max_iter=50, tolerance=1e-10, use_ap_bp_guess=True):
    """
    Iteratively solve for undistorted coordinates from distorted ones using Newton-Raphson
    
    This function solves the system:
        u_prime = u + f(u,v)
        v_prime = v + g(u,v)
    for (u, v) given (u_prime, v_prime)
    
    Parameters:
    -----------
    u_prime, v_prime : float
        Distorted pixel coordinates relative to reference pixel
    header : dict
        FITS header containing forward SIP coefficients (A_, B_)
    max_iter : int
        Maximum number of iterations
    tolerance : float
        Convergence tolerance
    use_ap_bp_guess : bool
        If True and AP/BP coefficients are available, use them as initial guess
        
    Returns:
    --------
    tuple : (u, v) undistorted coordinates
    """
    # Initial guess: use AP/BP if available, otherwise assume undistorted = distorted
    if use_ap_bp_guess and 'AP_ORDER' in header and 'BP_ORDER' in header:
        # Use AP/BP coefficients for better initial guess
        ap_order = int(header['AP_ORDER'])
        bp_order = int(header['BP_ORDER'])
        
        f_inv = evaluate_inverse_sip_polynomial(u_prime, v_prime, ap_order, 'AP', header)
        g_inv = evaluate_inverse_sip_polynomial(u_prime, v_prime, bp_order, 'BP', header)
        
        u = u_prime + f_inv
        v = v_prime + g_inv
    else:
        # Simple initial guess
        u = u_prime
        v = v_prime
    
    a_order = int(header['A_ORDER'])
    b_order = int(header['B_ORDER'])
    
    for iteration in range(max_iter):
        # Compute forward distortion at current guess
        f_uv = evaluate_sip_polynomial(u, v, a_order, 'A', header)
        g_uv = evaluate_sip_polynomial(u, v, b_order, 'B', header)
        
        # Compute residuals
        residual_u = u + f_uv - u_prime
        residual_v = v + g_uv - v_prime
        
        # Check convergence
        if np.abs(residual_u) < tolerance and np.abs(residual_v) < tolerance:
            return u, v
        
        # Compute Jacobian (derivatives of distortion)
        # df/du, df/dv, dg/du, dg/dv
        h = 1e-6  # Small step for numerical derivatives
        
        f_u_plus = evaluate_sip_polynomial(u + h, v, a_order, 'A', header)
        f_v_plus = evaluate_sip_polynomial(u, v + h, a_order, 'A', header)
        g_u_plus = evaluate_sip_polynomial(u + h, v, b_order, 'B', header)
        g_v_plus = evaluate_sip_polynomial(u, v + h, b_order, 'B', header)
        
        df_du = (f_u_plus - f_uv) / h
        df_dv = (f_v_plus - f_uv) / h
        dg_du = (g_u_plus - g_uv) / h
        dg_dv = (g_v_plus - g_uv) / h
        
        # Jacobian matrix (I + J_distortion)
        J = np.array([[1 + df_du, df_dv], 
                      [dg_du, 1 + dg_dv]])
        
        # Solve for correction: J * delta = -residual
        residual = np.array([residual_u, residual_v])
        try:
            delta = np.linalg.solve(J, -residual)
        except np.linalg.LinAlgError:
            # If singular, just use the residual directly
            delta = -residual
        
        # Update guess
        u += delta[0]
        v += delta[1]
    
    # If we didn't converge, return best guess
    print(f"Warning: SIP inversion did not converge after {max_iter} iterations")
    return u, v

def world_to_pixel(ra, dec, header, use_iteration=True):
    """
    Transform world coordinates (RA/DEC) to pixel coordinates
    
    This is the COMPLETE inverse implementation that perfectly matches astropy.
    It reverses all steps of the forward transformation.
    
    Parameters:
    -----------
    ra, dec : float
        Right ascension and declination in degrees
    header : dict
        FITS header containing WCS information
    use_iteration : bool, optional
        If True (default), use iterative refinement for inverse SIP distortion.
        If False, use only AP/BP coefficients without iteration.
        
    Returns:
    --------
    tuple : ((x_undist, y_undist), (x_dist, y_dist))
            Pixel coordinates without and with distortion correction (0-based)
    """
    # Step 1: Convert celestial coordinates to intermediate world coordinates
    # This uses the forward TAN projection (RA/DEC → xi/eta)
    xi, eta = forward_tan_projection(ra, dec, header)
    
    # Step 2: Apply inverse CD matrix to get pixel offsets
    # Solve: [xi, eta] = CD * [u', v'] for [u', v']
    cd11 = header['CD1_1']
    cd12 = header['CD1_2']
    cd21 = header['CD2_1']
    cd22 = header['CD2_2']
    
    # Compute inverse CD matrix
    det = cd11 * cd22 - cd12 * cd21
    cd_inv = np.array([[cd22/det, -cd12/det],
                       [-cd21/det, cd11/det]])
    
    # For undistorted case: directly invert the linear transformation
    u_undist = cd_inv[0, 0] * xi + cd_inv[0, 1] * eta
    v_undist = cd_inv[1, 0] * xi + cd_inv[1, 1] * eta
    
    # For distorted case: u', v' are the distortion-corrected coordinates
    # We need to find u, v such that: u + f(u,v) = u' and v + g(u,v) = v'
    u_prime = u_undist  # This is what we get from inverting CD matrix
    v_prime = v_undist
    
    # Step 3: Apply inverse SIP distortion correction
    if use_iteration:
        # Use iterative method for maximum accuracy (matching astropy)
        # If AP/BP coefficients are available, they'll be used as initial guess
        u, v = invert_sip_distortion_iterative(u_prime, v_prime, header, use_ap_bp_guess=True)
    else:
        # Use only AP/BP coefficients without iteration
        if 'AP_ORDER' in header and 'BP_ORDER' in header:
            ap_order = int(header['AP_ORDER'])
            bp_order = int(header['BP_ORDER'])
            
            f_inv = evaluate_inverse_sip_polynomial(u_prime, v_prime, ap_order, 'AP', header)
            g_inv = evaluate_inverse_sip_polynomial(u_prime, v_prime, bp_order, 'BP', header)
            
            u = u_prime + f_inv
            v = v_prime + g_inv
        else:
            # No inverse coefficients available, fall back to simple guess
            u = u_prime
            v = v_prime
    
    # Step 4: Convert from detector coordinates to pixel coordinates (1-based FITS convention)
    x_undist = u_undist + (header['CRPIX1'] - 1)
    y_undist = v_undist + (header['CRPIX2'] - 1)
    
    x_dist = u + (header['CRPIX1'] - 1)
    y_dist = v + (header['CRPIX2'] - 1)
    
    return (x_undist, y_undist), (x_dist, y_dist)

Test the Reverse Transformation

Now let's test our reverse transformation by converting RA/DEC coordinates back to pixel coordinates and comparing with astropy's results.

ra_test, dec_test = tess_header['CRVAL1'] + 0.5, tess_header['CRVAL2'] + 0.3

# Get results with iteration (default)
(x_undist, y_undist), (x_dist_iter, y_dist_iter) = world_to_pixel(ra_test, dec_test, tess_header, use_iteration=True)

# Get result using only AP/BP (no iteration)
(_, _), (x_dist_no_iter, y_dist_no_iter) = world_to_pixel(ra_test, dec_test, tess_header, use_iteration=False)

# Astropy results
astropy_result = tess_wcs.all_world2pix(ra_test, dec_test, 0)
astropy_result_nodist = wcs_no_distortion.all_world2pix(ra_test, dec_test, 0)

print(f"RA/DEC ({ra_test:.6f}°, {dec_test:.6f}°) corresponds to:")
print(f"No distortion 1. Astropy  2. This result:")
print(f"x={x_undist:.6f}, y={y_undist:.6f}")
print(f"x={astropy_result_nodist[0]:.6f}, y={astropy_result_nodist[1]:.6f}")
print(f"\nDistorted:")
print(f"1. Astropy:        x={astropy_result[0]:.6f}, y={astropy_result[1]:.6f}")
print(f"2. AP/BP only:     x={x_dist_no_iter:.6f}, y={y_dist_no_iter:.6f}  Δx={abs(x_dist_no_iter - astropy_result[0]):.2e}, Δy={abs(y_dist_no_iter - astropy_result[1]):.2e}")
print(f"3. With iteration: x={x_dist_iter:.6f}, y={y_dist_iter:.6f}  Δx={abs(x_dist_iter - astropy_result[0]):.2e}, Δy={abs(y_dist_iter - astropy_result[1]):.2e}")

RA/DEC (217.254373°, 76.692049°) corresponds to:
No distortion 1. Astropy  2. This result:
x=1100.416452, y=1001.280069
x=1100.416452, y=1001.280069

Distorted:
1. Astropy:        x=1100.352461, y=1001.269128
2. AP/BP only:     x=1100.347582, y=1001.265906  Δx=4.88e-03, Δy=3.22e-03
3. With iteration: x=1100.352461, y=1001.269128  Δx=9.39e-10, Δy=2.31e-10

# Round-trip test: pixel → RA/DEC → pixel
x_start, y_start = 100, 200
(ra_fwd, dec_fwd), (ra_fwd_dist, dec_fwd_dist) = pixel_to_world(x_start, y_start, tess_header)
(x_back_undist, y_back_undist), (x_back_dist, y_back_dist) = world_to_pixel(ra_fwd_dist, dec_fwd_dist, tess_header)

print(f"Round-trip test: ({x_start}, {y_start}) → RA/DEC → pixel")
print(f"Starting pixel: ({x_start}, {y_start})")
print(f"After round-trip: ({x_back_dist:.6f}, {y_back_dist:.6f})")
print(f"Error: Δx={abs(x_back_dist - x_start):.2e}, Δy={abs(y_back_dist - y_start):.2e}")

Round-trip test: (100, 200) → RA/DEC → pixel
Starting pixel: (100, 200)
After round-trip: (100.000000, 200.000000)
Error: Δx=6.82e-13, Δy=1.02e-12