Improve code for computing VA scale and offsets

CHANGES.rst

@@ -2,7 +2,12 @@
 ===================
 
 
+lib
+---
 
+- Update code in ``set_velocity_aberration.py`` functions based on Colin Cox
+  suggestions (simplify DVA scale computation and improve offset
+  computation). [#5666]
 
 
 0.18.3 (2021-01-25)
@@ -59,7 +64,7 @@ datamodels
 
 - Updated schemas for new keywords CROWDFLD, PRIDTYPE, PRIDTPTS, PATTNPTS, SMGRDPAT,
   changed name of SUBPXPNS to SUBPXPTS, and new allowed values for PATTTYPE. [#5618]
-  
+
 flat_field
 ----------
 
@@ -100,7 +105,7 @@ datamodels
 extract_1d
 ----------
 
-- For IFU data (NIRSpec and MIRI) the extraction radius is now a varying size 
+- For IFU data (NIRSpec and MIRI) the extraction radius is now a varying size
   based on wavelength. The apcorr correction is a function of wavelength and
   radius size. Fixes a bug in units conversion for applying the apcorr correction.
   The units are now correctly converted from arcseconds to pixels. Added an

jwst/lib/set_velocity_aberration.py

@@ -18,16 +18,18 @@
 in the header other than what is required by the standard.
 '''
 
-import astropy.io.fits as fits
 import logging
+import numpy as np
+import astropy.io.fits as fits
+from gwcs.geometry import SphericalToCartesian
+from scipy.constants import speed_of_light
 import math
 
-
 # Configure logging
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
-SPEED_OF_LIGHT = 299792.458     # km / s
+SPEED_OF_LIGHT = speed_of_light / 1000     # km / s
 d_to_r = math.pi / 180.
 
 
@@ -55,40 +57,16 @@ def aberration_scale(velocity_x, velocity_y, velocity_z,
         this value to obtain the image scale corrected for the "aberration
         of starlight" due to the velocity of JWST with respect to the Sun.
     """
-
-    speed = math.sqrt(velocity_x**2 + velocity_y**2 + velocity_z**2)
-    if speed == 0.0:
-        logger.warning('Speed is zero. Forcing scale to 1.0')
-        return 1.0
-
-    beta = speed / SPEED_OF_LIGHT
-    gamma = 1. / math.sqrt(1. - beta**2)
-
-    # [targ_x, targ_y, targ_z] is a unit vector.
-    r_xy = math.cos(targ_dec * d_to_r)          # radial distance in xy-plane
-    targ_x = r_xy * math.cos(targ_ra * d_to_r)
-    targ_y = r_xy * math.sin(targ_ra * d_to_r)
-    targ_z = math.sin(targ_dec * d_to_r)
-
-    dot_prod = (velocity_x * targ_x +
-                velocity_y * targ_y +
-                velocity_z * targ_z)
-    cos_theta = dot_prod / speed
-    # This sin_theta is only valid over the range [0, pi], but so is the
-    # angle between the velocity vector and the direction toward the target.
-    sin_theta = math.sqrt(1. - cos_theta**2)
-
-    tan_theta_p = sin_theta / (gamma * (cos_theta + beta))
-    theta_p = math.atan(tan_theta_p)
-
-    scale_factor = (gamma * (cos_theta + beta)**2 /
-                    (math.cos(theta_p)**2 * (1. + beta * cos_theta)))
-
+    beta = np.array([velocity_x, velocity_y, velocity_z]) / SPEED_OF_LIGHT
+    beta2 = np.dot(beta, beta)
+    u = SphericalToCartesian()(targ_ra, targ_dec)
+    beta_u =  np.dot(beta, u)
+    igamma = np.sqrt(1.0 - beta2)  # inverse of usual gamma
+    scale_factor = (1.0 + beta_u) /  igamma
     return scale_factor
 
 
-def aberration_offset(velocity_x, velocity_y, velocity_z,
-                      targ_ra, targ_dec):
+def aberration_offset(velocity_x, velocity_y, velocity_z, targ_ra, targ_dec):
     """Compute the RA/Dec offsets due to velocity aberration.
 
     Parameters
@@ -108,23 +86,20 @@ def aberration_offset(velocity_x, velocity_y, velocity_z,
     -------
     delta_ra, delta_dec: float
         The offset to be added to the input RA/Dec, in units of radians.
-    """
-
-    xdot = velocity_x / SPEED_OF_LIGHT
-    ydot = velocity_y / SPEED_OF_LIGHT
-    zdot = velocity_z / SPEED_OF_LIGHT
 
-    sin_alpha = math.sin(targ_ra * d_to_r)
-    cos_alpha = math.cos(targ_ra * d_to_r)
-    sin_delta = math.sin(targ_dec * d_to_r)
-    cos_delta = math.cos(targ_dec * d_to_r)
-
-    delta_ra = (-xdot * sin_alpha + ydot * cos_alpha) / cos_delta
-    delta_dec = (-xdot * cos_alpha * sin_delta -
-                 ydot * sin_alpha * sin_delta +
-                 zdot * cos_delta)
-
-    return delta_ra, delta_dec
+    """
+    # Algorithm from Colin Cox notebook
+    beta = np.array([velocity_x, velocity_y, velocity_z]) / SPEED_OF_LIGHT
+    u = np.array(SphericalToCartesian()(targ_ra, targ_dec))
+    beta2 = np.dot(beta, beta)
+    if beta2 == 0.0:
+        return 0.0, 0.0
+    igamma = np.sqrt(1.0 - beta2)  # inverse of usual gamma
+    beta_u = np.dot(beta, u)
+    u_corr = (igamma * u + beta * (1.0 + (1.0 - igamma) * beta_u / beta2)) / (1.0 + beta_u)
+    ra_corr = np.arctan2(u_corr[1], u_corr[0])
+    dec_corr = np.arcsin(u_corr[2])
+    return ra_corr - np.deg2rad(targ_ra), dec_corr -  np.deg2rad(targ_dec)
 
 
 def add_dva(filename):