Merge pull request #669 from dirac-institute/parallax/vectorized_impl

Vectorized implementation of reflex motion of the Earth correction

src/kbmod/reprojection_utils.py

@@ -1,12 +1,22 @@
 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,
@@ -18,7 +28,8 @@ def correct_parallax(
     use_bounds=False,
 ):
     """Calculate the parallax corrected postions for a given object at a given
-    time, observation location on Earth, and user defined distance from the Sun.
+    time, observation location on Earth, and user defined distance from th
+    e Sun.
     By default, this function will use the geometric solution for objects beyond 1au.
     If the distance is less than 1au, the function will use the scipy minimizer
     to find the best geocentric distance.
@@ -217,6 +228,177 @@ def correct_parallax_geometrically(coord, obstime, point_on_earth, heliocentric_
     return answer, dist
 
 
+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.
+
+    This geometric solution is only applicable for objects beyond the 1au.
+    If the parallax correction failed, that is returned an unphysical result, the
+    coordinate distance is set to 0AU.
+
+    Given MJDs do not need to be the same, the returned parallax corrected coordinates
+    will be returned in the same order of the given coordinates and times.
+
+    Attributes
+    ----------
+    ra : `list[float]`
+        Right Ascension in decimal degrees
+    dec : `list[float]`
+        Declination in decimal degrees
+    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`, 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
+    unique_obstimes = np.unique(mjds)
+
+    # fetch ephemeris of planet earth
+    earth_ephems_x = np.zeros((len(unique_obstimes),))
+    earth_ephems_y = np.zeros((len(unique_obstimes),))
+    earth_ephems_z = np.zeros((len(unique_obstimes),))
+    # 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
+    location_ephems_y = earth_ephems_y + point_on_earth.y.to(u.AU).value
+    location_ephems_z = earth_ephems_z + point_on_earth.z.to(u.AU).value
+
+    # Move the given ICRS coordinates back to GCRS, this change of
+    # 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
+    mjd_times = Time(mjds, format="mjd")
+    los_earth_obj = SkyCoord(ra=ra * u.deg, dec=dec * u.deg, obstime=mjd_times).transform_to(
+        GCRS(obsgeoloc=point_on_earth_geocentric)
+    )
+    los_earth_obj_cart = los_earth_obj.cartesian
+
+    # Now that we have the ray elements, re-create the SkyCoords with but with an
+    # added distance placeholder, because that alters the underlying coordinate
+    # representation to SphericalRepresentation. Avoid data copying
+    los_earth_obj = SkyCoord(
+        ra=los_earth_obj.ra,
+        dec=los_earth_obj.dec,
+        distance=np.zeros((len(ra),)) * u.AU,
+        obstime=los_earth_obj.obstime,
+        obsgeoloc=los_earth_obj.obsgeoloc,
+        frame="gcrs",
+        copy=False,
+    )
+
+    # for each ephemeris calculate the geocentric distance
+    # that matches the given heliocentric distance guess and update
+    # our LOS coordinate with these new distances
+    for obstime, ext, eyt, ezt in zip(
+        unique_obstimes, location_ephems_x, location_ephems_y, location_ephems_z
+    ):
+        mjd_mask = mjds == obstime
+        vx = los_earth_obj_cart.x[mjd_mask].value
+        vy = los_earth_obj_cart.y[mjd_mask].value
+        vz = los_earth_obj_cart.z[mjd_mask].value
+
+        # Solve the quadratic equation for the ray leaving the earth and intersecting
+        # a sphere centered on the barycenter at (0, 0, 0) with radius = heliocentric_distance
+        a = vx * vx + vy * vy + vz * vz
+        b = 2 * vx * ext + 2 * vy * eyt + 2 * vz * ezt
+        c = ext * ext + eyt * eyt + ezt * ezt - heliocentric_distance * heliocentric_distance
+        disc = b * b - 4 * a * c
+
+        invalid_result_mask = disc < 0
+        disc[disc < 0] = 0
+
+        # Since the ray will be starting from within the sphere (we assume the
+        # heliocentric_distance is at least 1 au), one of the solutions should be positive
+        # and the other negative. We only use the positive one.
+        dist = (-b + np.sqrt(disc)) / (2 * a)
+        # we can't set the distance to negative in SkyCOord without a lot of wrangling
+        # the distance of 0 should correspond to UnitSphericalRepresentation
+        dist[invalid_result_mask] = 0.0
+        los_earth_obj.distance[mjd_mask] = dist * u.AU
+
+    # 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,18 @@
 
 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 +329,49 @@ 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()