Cleanup code for vectorized parallax.

src/kbmod/reprojection_utils.py

@@ -1,12 +1,28 @@
 import astropy.units as u
 import numpy as np
 from astropy import units as u
-from astropy.coordinates import SkyCoord, GCRS, ICRS, get_body_barycentric
+from astropy.coordinates import (
+    SkyCoord,
+    GCRS,
+    ICRS,
+    solar_system_ephemeris,
+    get_body_barycentric
+)
 from astropy.time import Time
 from astropy.wcs.utils import fit_wcs_from_points
 from scipy.optimize import minimize
 
 
+__all__ = [
+    "correct_parallax",
+    "invert_correct_parallax",
+    "fit_barycentric_wcs",
+    "transform_wcses_to_ebd",
+    "correct_parallax_geometrically_vectorized",
+    "invert_correct_parallax_vectorized"
+]
+
+
 def correct_parallax(
     coord,
     obstime,
@@ -218,7 +234,8 @@ def correct_parallax_geometrically(coord, obstime, point_on_earth, heliocentric_
     return answer, dist
 
 
-def correct_parallax_geometrically_vectorized(ra, dec, mjds, point_on_earth, heliocentric_distance, return_geo_dists=True):
+def correct_parallax_geometrically_vectorized(ra, dec, mjds, heliocentric_distance,
+                                              point_on_earth=None, return_geo_dists=True):
     """Calculate the parallax corrected postions for a given object at a given time,
     position on Earth, and a hypothetical distance from the Sun.
 
@@ -231,23 +248,28 @@ def correct_parallax_geometrically_vectorized(ra, dec, mjds, point_on_earth, hel
 
     Attributes
     ----------
-    ra : `float`
-        Right Ascencion in decimal degrees
-    dec : `float`
+    ra : `list[float]`
+        Right Ascension in decimal degrees
+    dec : `list[float]`
         Declination in decimal degrees
-    mjds : `float`
+    mjds : `list[float]`
         MJD timestamps of the times the ``ra`` and ``dec`` were recorded.
     heliocentric_distance : `float`
         The guess distance to the object from the Sun in AU.
+    point_on_earth : `EarthLocation` or `None`, optional
+        Observation is returned from the geocenter by default. Provide an
+        EarthLocation if you want to also account for the position of the
+        observatory. If not provided, assumed to be geocenter.
     return_geo_dists : `bool`, default: `True`
         Return calculated geocentric distances (in AU).
 
     Returns
     ----------
     parallax_corrected: `astropy.coordinates.SkyCoord`
         `SkyCoord` vector of parallax corrected coordinates.
-    geocentric_distances: `list`
-        A list of calculated geocentric distances.
+    geocentric_distances: `list`, optional
+        A list of calculated geocentric distances. Returned for
+        compatibility.
     """
     # Processing has to be batched over same times since Earth
     # position changes at each time stamp
@@ -257,13 +279,15 @@ def correct_parallax_geometrically_vectorized(ra, dec, mjds, point_on_earth, hel
     earth_ephems_x = np.zeros((len(unique_obstimes), ))
     earth_ephems_y = np.zeros((len(unique_obstimes), ))
     earth_ephems_z = np.zeros((len(unique_obstimes), ))
-    for i, ot in enumerate(Time(unique_obstimes, format="mjd")):
-        # figure out what coordinate this is pointing to - center of mass, or?
-        # TODO: redo the calls to make use of JPL ephems with all the obstime at once
-        earth_coords = get_body_barycentric("earth", ot)
-        earth_ephems_x[i] = earth_coords.x.to(u.AU).value
-        earth_ephems_y[i] = earth_coords.y.to(u.AU).value
-        earth_ephems_z[i] = earth_coords.z.to(u.AU).value
+    # figure out what coordinate this is pointing to - center of mass, or?
+    # TODO: redo the calls to make use of JPL ephems with all the obstime at once
+    # because there could be a lot of timestamps here
+    with solar_system_ephemeris.set('de432s'):
+        for i, ot in enumerate(Time(unique_obstimes, format="mjd")):
+            earth_coords = get_body_barycentric("earth", ot)
+            earth_ephems_x[i] = earth_coords.x.to(u.AU).value
+            earth_ephems_y[i] = earth_coords.y.to(u.AU).value
+            earth_ephems_z[i] = earth_coords.z.to(u.AU).value
 
     # calculate the ephemeris of a particular location on earth, f.e. CTIO
     location_ephems_x = earth_ephems_x + point_on_earth.x.to(u.AU).value
@@ -274,7 +298,7 @@ def correct_parallax_geometrically_vectorized(ra, dec, mjds, point_on_earth, hel
     # coordinate origins changes the RA, DEC so that its line of sight
     # now contains the object, which isn't true for ICRS coordinate
     # Copy the cartesian elements of the LOS ray for use later
-    point_on_earth_geocentric=point_on_earth.geocentric*u.m
+    point_on_earth_geocentric = point_on_earth.geocentric*u.m
     mjd_times = Time(mjds, format="mjd")
     los_earth_obj = SkyCoord(
         ra=ra*u.deg,
@@ -326,15 +350,64 @@ def correct_parallax_geometrically_vectorized(ra, dec, mjds, point_on_earth, hel
         # the distance of 0 should correspond to UnitSphericalRepresentation
         dist[invalid_result_mask] = 0.0
         los_earth_obj.distance[mjd_mask] = dist*u.AU
-        break
 
     # finally, transform the coordinates with the new distance guesses
     # back to ICRS, which will now include the correct parallax correction
+    # Don't return geo dists if the user doesn't want them
     if return_geo_dists:
         return los_earth_obj.transform_to(ICRS()), los_earth_obj.distance
     return los_earth_obj.transform_to(ICRS())
 
 
+def invert_correct_parallax_vectorized(coords, obstimes, point_on_earth=None):
+    """Converts a given ICRS coordinate with distance into an ICRS coordinate,
+    without accounting for the reflex correction.
+
+    ICRS coordinates corrected for reflex motion of the Earth are ICRS coordinates
+    that are aware of the finite distance to the object, and thus able to account
+    for it during transformation.
+    ICRS coordinates that are reported by observers on Earth, assume that distances
+    to all objects are ~infinity, Thus, the "true" ICRS coordinate of an object
+    with a finite distance is converted without the appropriate correction for the
+    parallax. Note the loose use of "true ICRS coordinate" as it is the "wrong"
+    ICRS coordinate for all other objects.
+    This function takes an ICRS coordinate capable of accounting for the parallax
+    due to the finite distance to the object, and returns an ICRS coordinate as it
+    would be reported by an observer from Earth at a given time, if they did not
+    know how their distance to the object.
+
+    Parameters
+    ----------
+    coords : `SkyCoord`
+        True coordinates.
+    obstimes : `Time` or `list[float]`
+        Timestamps of observations of the object.
+    point_on_earth : `EarthLocation` or `None`, optional
+        Observation is returned from the geocenter by default. Provide an
+        EarthLocation if you want to also account for the position of the
+        observatory.
+
+    Returns
+    -------
+    icrs : `SkyCoord`
+        ICRS coordinate of the object as it would be observed from Earth at
+        the given timestamp(s) without knowing the distance to the object.
+    """
+    obstimes = Time(obstimes, format="mjd") if not isinstance(obstimes, Time) else obstimes
+    obsgeoloc = point_on_earth.geocentric*u.m
+    los = coords.transform_to(GCRS(obstime=obstimes,  obsgeoloc=obsgeoloc))
+    los_earth_obj = SkyCoord(
+        ra=los.ra,
+        dec=los.dec,
+        obsgeoloc=los.obsgeoloc,
+        obstime=los.obstime,
+        frame="gcrs",
+        copy=False,
+    )
+    return los_earth_obj.icrs
+
+
+
 def invert_correct_parallax(coord, obstime, point_on_earth, geocentric_distance, heliocentric_distance):
     """Calculate the original ICRS coordinates of a point in EBD space, i.e. a result from `correct_parallax`.
 

tests/test_reprojection_utils.py

@@ -2,15 +2,23 @@
 
 import numpy as np
 import numpy.testing as npt
-from astropy.coordinates import EarthLocation, SkyCoord, solar_system_ephemeris
+import astropy.units as u
 from astropy.time import Time
 from astropy.wcs import WCS
+from astropy.coordinates import (
+    EarthLocation,
+    SkyCoord,
+    solar_system_ephemeris,
+    GCRS
+)
 
 from kbmod.reprojection_utils import (
     correct_parallax,
     invert_correct_parallax,
     fit_barycentric_wcs,
     transform_wcses_to_ebd,
+    correct_parallax_geometrically_vectorized,
+    invert_correct_parallax_vectorized
 )
 
 
@@ -326,6 +334,42 @@ def test_parallax_with_method_and_no_bounds(self):
         assert type(corrected_coord1) is SkyCoord
         assert type(corrected_coord2) is SkyCoord
 
+    def test_equinox_vectorized_parallax_correction(self):
+        # Chosen so that at equinox the position of the objects
+        # is ~ lon=0, lat=0 in ecliptic coordinates, when Earth
+        # and Sun have ecliptic lat ~0. This makes ICRS of the obj
+        # ~ra=90, dec=23.4 in ICRS by definition
+        t = Time("2023-03-20T16:00:00", format="isot", scale="utc")
+        true_ra = 90*u.degree
+        true_dec = 23.43952556*u.degree
+        true_distance = 50*u.au
+        truth = SkyCoord(true_ra, true_dec, distance=true_distance, frame="icrs")
+
+        with solar_system_ephemeris.set("de432s"):
+            ctio = EarthLocation.of_site("ctio")
+        earth_truth = truth.transform_to(GCRS(obsgeoloc=ctio.geocentric*u.m, obstime=t))
+
+        # finally, synthesize the observation as it would have been seen
+        # from the earth at the time, without any knowledge that the
+        # object has a finite distance
+        obs = SkyCoord(earth_truth.ra, earth_truth.dec, obstime=t, obsgeoloc=ctio.geocentric*u.m, frame="gcrs").icrs
+
+        # Now forward solve the problem and confirm the numbers match
+        obstimes = [obs.obstime.mjd, ]
+        corr = correct_parallax_geometrically_vectorized(
+            [obs.ra.deg, ],
+            [obs.dec.deg, ],
+            obstimes,
+            true_distance.value,
+            ctio,
+            return_geo_dists=False
+        )
+        self.assertLessEqual(corr.separation(truth).arcsecond, 1e-4)
+
+        inverted = invert_correct_parallax_vectorized(corr, obstimes, ctio)
+        self.assertLessEqual(inverted.separation(obs).arcsecond, 1e-4)
+
+
 
 if __name__ == "__main__":
     unittest.main()