Merge pull request #52 from B612-Asteroid-Institute/jm/2body-ephemeris

Add 2body ephemeris generation and covariance mapping

REFERENCES.md

@@ -5,8 +5,10 @@
     Icarus, 166, 248–270. https://doi.org/10.1016/S0019-1035(03)00051-4
 * Curtis, H. D. (2014). Orbital Mechanics for Engineering Students. 3rd ed.,  
     Elsevier Ltd. ISBN-13: 978-0080977478
-* Danby, J. M. A. (1992). Fundamentals of Celestial Mechanics. 2nd ed.,
+* Danby, J. M. A. (1992). Fundamentals of Celestial Mechanics. 2nd ed.,  
     William-Bell, Inc. ISBN-13: 978-0943396200
     Notes: of particular interest is Danby's fantastic chapter on universal variables (6.9)
+* Urban, S. E; Seidelmann, P. K. (2013) Explanatory Supplement to the Astronomical Almanac. 3rd ed.,  
+    University Science Books. ISBN-13: 978-1891389856
 * Wan, E. A; Van Der Merwe, R. (2000). The unscented Kalman filter for nonlinear estimation.  
     Proceedings of the IEEE 2000 Adaptive Systems for Signal Processing, Communications, and Control Symposium, 153-158. https://doi.org/10.1109/ASSPCC.2000.882463

adam_core/coordinates/__init__.py

@@ -26,13 +26,18 @@
     _cartesian_to_cometary,
     _cartesian_to_keplerian,
     _cartesian_to_keplerian6,
+    _cartesian_to_spherical,
     _cometary_to_cartesian,
     _keplerian_to_cartesian_a,
     _keplerian_to_cartesian_p,
     _keplerian_to_cartesian_q,
+    _spherical_to_cartesian,
     cartesian_to_cometary,
     cartesian_to_keplerian,
+    cartesian_to_spherical,
     cometary_to_cartesian,
+    keplerian_to_cartesian,
+    spherical_to_cartesian,
     transform_coordinates,
 )
 from .variants import create_coordinate_variants

adam_core/coordinates/transform.py

@@ -1,14 +1,12 @@
 import logging
-from typing import Literal, Optional, Union
+from typing import Literal, Optional, TypeVar, Union
 
 import jax.numpy as jnp
 import numpy as np
 import pyarrow.compute as pc
 from jax import config, jit, lax, vmap
 
 from ..constants import Constants as c
-from ..dynamics.barker import solve_barker
-from ..dynamics.kepler import calc_mean_anomaly, solve_kepler
 from .cartesian import CartesianCoordinates
 from .cometary import CometaryCoordinates
 from .keplerian import KeplerianCoordinates
@@ -25,26 +23,16 @@
 
 logger = logging.getLogger(__name__)
 
-__all__ = [
-    "transform_coordinates",
-    "_cartesian_to_keplerian",
-    "_cartesian_to_keplerian6",
-    "cartesian_to_keplerian",
-    "_keplerian_to_cartesian_a",
-    "_keplerian_to_cartesian_p",
-    "_keplerian_to_cartesian_q",
-    "_cartesian_to_cometary",
-    "cartesian_to_cometary",
-    "_cometary_to_cartesian",
-    "cometary_to_cartesian",
-]
-
-CoordinateType = Union[
-    CartesianCoordinates,
-    KeplerianCoordinates,
-    CometaryCoordinates,
-    SphericalCoordinates,
-]
+
+CoordinateType = TypeVar(
+    "CoordinateType",
+    bound=Union[
+        CartesianCoordinates,
+        KeplerianCoordinates,
+        CometaryCoordinates,
+        SphericalCoordinates,
+    ],
+)
 CoordinatesClasses = (
     CartesianCoordinates,
     KeplerianCoordinates,
@@ -355,6 +343,8 @@ def _cartesian_to_keplerian(
     [1] Bate, R. R; Mueller, D. D; White, J. E. (1971). Fundamentals of Astrodynamics. 1st ed.,
         Dover Publications, Inc. ISBN-13: 978-0486600611
     """
+    from ..dynamics.kepler import calc_mean_anomaly
+
     coords_keplerian = jnp.zeros(13, dtype=jnp.float64)
     r = coords_cartesian[0:3]
     v = coords_cartesian[3:6]
@@ -622,6 +612,9 @@ def _keplerian_to_cartesian_p(
         vy : y-velocity in units of au per day.
         vz : z-velocity in units of au per day.
     """
+    from ..dynamics.barker import solve_barker
+    from ..dynamics.kepler import solve_kepler
+
     coords_cartesian = jnp.zeros(6, dtype=jnp.float64)
 
     p = coords_keplerian[0]

adam_core/dynamics/__init__.py

@@ -1,6 +1,8 @@
 # flake8: noqa: F401
+from .aberrations import add_light_time, add_stellar_aberration
 from .barker import solve_barker
 from .chi import calc_chi
+from .ephemeris import generate_ephemeris_2body
 from .kepler import solve_kepler
 from .lagrange import apply_lagrange_coefficients, calc_lagrange_coefficients
 from .propagation import propagate_2body

adam_core/dynamics/aberrations.py

@@ -0,0 +1,208 @@
+from typing import Tuple
+
+import jax.numpy as jnp
+from jax import config, jit, lax, vmap
+
+from ..constants import Constants as c
+from .propagation import _propagate_2body
+
+config.update("jax_enable_x64", True)
+
+MU = c.MU
+C = c.C
+
+
+@jit
+def _add_light_time(
+    orbit: jnp.ndarray,
+    t0: float,
+    observer_position: jnp.ndarray,
+    lt_tol: float = 1e-10,
+    mu: float = MU,
+    max_iter: int = 1000,
+    tol: float = 1e-15,
+) -> Tuple[jnp.ndarray, jnp.float64]:
+    """
+    When generating ephemeris, orbits need to be backwards propagated to the time
+    at which the light emitted or relflected from the object towards the observer.
+
+    Light time correction must be added to orbits in expressed in an inertial frame (ie, orbits
+    must be barycentric).
+
+    Parameters
+    ----------
+    orbit : `~jax.numpy.ndarray` (6)
+        Barycentric orbit in cartesian elements to correct for light time delay.
+    t0 : float
+        Epoch at which orbits are defined.
+    observer_positions : `~jax.numpy.ndarray` (3)
+        Location of the observer in barycentric cartesian elements at the time of observation.
+    lt_tol : float, optional
+        Calculate aberration to within this value in time (units of days.)
+    mu : float, optional
+        Gravitational parameter (GM) of the attracting body in units of
+        AU**3 / d**2.
+    max_iter : int, optional
+        Maximum number of iterations over which to converge for propagation.
+    tol : float, optional
+        Numerical tolerance to which to compute universal anomaly during propagation using the Newtown-Raphson
+        method.
+
+    Returns
+    -------
+    corrected_orbit : `~jax.numpy.ndarray` (6)
+        Orbit adjusted for light travel time.
+    lt : float
+        Light time correction (t0 - corrected_t0).
+    """
+    dlt = 1e30
+    lt = 1e30
+
+    @jit
+    def _iterate_light_time(p):
+
+        orbit_i = p[0]
+        t0 = p[1]
+        lt0 = p[2]
+        dlt = p[3]
+
+        # Calculate topocentric distance
+        rho = jnp.linalg.norm(orbit_i[:3] - observer_position)
+
+        # Calculate initial guess of light time
+        lt = rho / C
+
+        # Calculate difference between previous light time correction
+        # and current guess
+        dlt = jnp.abs(lt - lt0)
+
+        # Propagate backwards to new epoch
+        t1 = t0 - lt
+        orbit_propagated = _propagate_2body(
+            orbit, t0, t1, mu=mu, max_iter=max_iter, tol=tol
+        )
+
+        p[0] = orbit_propagated
+        p[1] = t1
+        p[2] = lt
+        p[3] = dlt
+        return p
+
+    @jit
+    def _while_condition(p):
+        dlt = p[-1]
+        return dlt > lt_tol
+
+    p = [orbit, t0, lt, dlt]
+    p = lax.while_loop(_while_condition, _iterate_light_time, p)
+
+    orbit_aberrated = p[0]
+    t0_aberrated = p[1]  # noqa: F841
+    lt = p[2]
+    return orbit_aberrated, lt
+
+
+# Vectorization Map: _add_light_time
+_add_light_time_vmap = jit(
+    vmap(_add_light_time, in_axes=(0, 0, 0, None, None, None, None), out_axes=(0, 0))
+)
+
+
+@jit
+def add_light_time(
+    orbits: jnp.ndarray,
+    t0: jnp.ndarray,
+    observer_positions: jnp.ndarray,
+    lt_tol: float = 1e-10,
+    mu: float = MU,
+    max_iter: int = 1000,
+    tol: float = 1e-15,
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    When generating ephemeris, orbits need to be backwards propagated to the time
+    at which the light emitted or relflected from the object towards the observer.
+
+    Light time correction must be added to orbits in expressed in an inertial frame (ie, orbits
+    must be barycentric).
+
+    Parameters
+    ----------
+    orbits : `~jax.numpy.ndarray` (N, 6)
+        Barycentric orbits in cartesian elements to correct for light time delay.
+    t0 : `~jax.numpy.ndarray` (N)
+        Epoch at which orbits are defined.
+    observer_positions : `~jax.numpy.ndarray` (N, 3)
+        Location of the observer in barycentric cartesian elements at the time of observation.
+    lt_tol : float, optional
+        Calculate aberration to within this value in time (units of days.)
+    mu : float, optional
+        Gravitational parameter (GM) of the attracting body in units of
+        AU**3 / d**2.
+    max_iter : int, optional
+        Maximum number of iterations over which to converge for propagation.
+    tol : float, optional
+        Numerical tolerance to which to compute universal anomaly during propagation using the Newtown-Raphson
+        method.
+
+    Returns
+    -------
+    corrected_orbits : `~jax.numpy.ndarray` (N, 6)
+        Orbits adjusted for light travel time.
+    lt : `~jax.numpy.ndarray` (N)
+        Light time correction (t0 - corrected_t0).
+    """
+    orbits_aberrated, lts = _add_light_time_vmap(
+        orbits, t0, observer_positions, lt_tol, mu, max_iter, tol
+    )
+    return orbits_aberrated, lts
+
+
+@jit
+def add_stellar_aberration(
+    orbits: jnp.ndarray, observer_states: jnp.ndarray
+) -> jnp.ndarray:
+    """
+    The motion of the observer in an inertial frame will cause an object
+    to appear in a different location than its true geometric location. This
+    aberration is typically applied after light time corrections have been added.
+
+    The velocity of the input orbits are unmodified only the position
+    vector is modified with stellar aberration.
+
+    Parameters
+    ----------
+    orbits : `~jax.numpy.ndarray` (N, 6)
+        Orbits in barycentric cartesian elements.
+    observer_states : `~jax.numpy.ndarray` (N, 6)
+        Observer states in barycentric cartesian elements.
+
+    Returns
+    -------
+    rho_aberrated : `~jax.numpy.ndarray` (N, 3)
+        The topocentric position vector for each orbit with
+        added stellar aberration.
+
+    References
+    ----------
+    [1] Urban, S. E; Seidelmann, P. K. (2013) Explanatory Supplement to the Astronomical Almanac. 3rd ed.,
+        University Science Books. ISBN-13: 978-1891389856
+    """
+    topo_states = orbits - observer_states
+    rho_aberrated = jnp.zeros((len(topo_states), 3), dtype=jnp.float64)
+    rho_aberrated = rho_aberrated.at[:].set(topo_states[:, :3])
+
+    v_obs = observer_states[:, 3:]
+    gamma = v_obs / C
+    beta_inv = jnp.sqrt(1 - jnp.linalg.norm(gamma, axis=1, keepdims=True) ** 2)
+    delta = jnp.linalg.norm(topo_states[:, :3], axis=1, keepdims=True)
+
+    # Equation 7.40 in Urban & Seidelmann (2013) [1]
+    rho = topo_states[:, :3] / delta
+    rho_dot_gamma = jnp.sum(rho * gamma, axis=1, keepdims=True)
+    rho_aberrated = rho_aberrated.at[:].set(
+        (beta_inv * rho + gamma + rho_dot_gamma * gamma / (1 + beta_inv))
+        / (1 + rho_dot_gamma)
+    )
+    rho_aberrated = rho_aberrated.at[:].multiply(delta)
+
+    return rho_aberrated