From 2e7bffbf87c2c3c70da13f19a2009cf7fa0038c9 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Thu, 13 Jul 2023 14:20:12 -0700
Subject: [PATCH 01/10] Add add_light_time and add_stellar_aberration

---
 REFERENCES.md                     |   4 +-
 adam_core/dynamics/__init__.py    |   1 +
 adam_core/dynamics/aberrations.py | 207 ++++++++++++++++++++++++++++++
 3 files changed, 211 insertions(+), 1 deletion(-)
 create mode 100644 adam_core/dynamics/aberrations.py

diff --git a/REFERENCES.md b/REFERENCES.md
index 0fa27e7e..94e79487 100644
--- a/REFERENCES.md
+++ b/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
diff --git a/adam_core/dynamics/__init__.py b/adam_core/dynamics/__init__.py
index f1bfc6d7..7e9afbbf 100644
--- a/adam_core/dynamics/__init__.py
+++ b/adam_core/dynamics/__init__.py
@@ -1,4 +1,5 @@
 # flake8: noqa: F401
+from .aberrations import add_light_time, add_stellar_aberration
 from .barker import solve_barker
 from .chi import calc_chi
 from .kepler import solve_kepler
diff --git a/adam_core/dynamics/aberrations.py b/adam_core/dynamics/aberrations.py
new file mode 100644
index 00000000..15504d5e
--- /dev/null
+++ b/adam_core/dynamics/aberrations.py
@@ -0,0 +1,207 @@
+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:]
+    beta = v_obs / C
+    gamma_inv = jnp.sqrt(1 - jnp.linalg.norm(beta, 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_aberrated = rho_aberrated.at[:].set(
+        (gamma_inv * rho + beta + rho * beta * beta / (1 + gamma_inv))
+        / (1 + rho * beta)
+    )
+    rho_aberrated = rho_aberrated.at[:].multiply(delta)
+
+    return rho_aberrated

From ca718b68a90b5a07d52af0a3839d21f5b09140d4 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Mon, 31 Jul 2023 10:28:15 -0700
Subject: [PATCH 02/10] Add _generate_ephemeris_2body and
 _generate_ephemeris_2body_vmap

---
 adam_core/dynamics/ephemeris.py | 99 +++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)
 create mode 100644 adam_core/dynamics/ephemeris.py

diff --git a/adam_core/dynamics/ephemeris.py b/adam_core/dynamics/ephemeris.py
new file mode 100644
index 00000000..4dffd310
--- /dev/null
+++ b/adam_core/dynamics/ephemeris.py
@@ -0,0 +1,99 @@
+from typing import Tuple
+
+import jax.numpy as jnp
+import numpy as np
+from jax import jit, vmap
+
+from ..constants import Constants as c
+from ..coordinates.transform import _cartesian_to_spherical
+from .aberrations import _add_light_time, add_stellar_aberration
+
+MU = c.MU
+
+
+@jit
+def _generate_ephemeris_2body(
+    propagated_orbit: np.ndarray,
+    observation_time: float,
+    observer_coordinates: jnp.ndarray,
+    lt_tol: float = 1e-10,
+    mu: float = MU,
+    max_iter: int = 100,
+    tol: float = 1e-15,
+) -> Tuple[jnp.ndarray, jnp.float64]:
+    """
+    Given a propagated orbit, generate its on-sky ephemeris as viewed from the observer.
+    This function calculates the light time delay between the propagated orbit and the observer,
+    and then propagates the orbit backward by that amount to when the light from object was actually
+    emitted towards the observer.
+
+    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 is known as
+    stellar aberration. Stellar aberration is will also be applied after
+    light time correction has been added.
+
+    The velocity of the input orbits are unmodified only the position
+    vector is modified with stellar aberration.
+
+    Parameters
+    ----------
+    propagated_orbit : `~jax.numpy.ndarray` (6)
+        Barycentric Cartesian orbit propagated to the given time.
+    observation_time : float
+        Epoch at which orbit and observer coordinates are defined.
+    observer_coordinates : `~jax.numpy.ndarray` (3)
+        Barycentric Cartesian observer coordinates.
+    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
+    -------
+    ephemeris_spherical : `~jax.numpy.ndarray` (6)
+        Topocentric Spherical ephemeris.
+    lt : float
+        Light time correction (t0 - corrected_t0).
+    """
+    # Add light time correction
+    propagated_orbits_aberrated, light_time = _add_light_time(
+        propagated_orbit,
+        observation_time,
+        observer_coordinates[0:3],
+        lt_tol=lt_tol,
+        mu=mu,
+        max_iter=max_iter,
+        tol=tol,
+    )
+
+    # Calculate topocentric coordinates
+    topocentric_coordinates = propagated_orbits_aberrated - observer_coordinates
+
+    # Apply stellar aberration to topocentric coordinates
+    topocentric_coordinates = topocentric_coordinates.at[0:3].set(
+        add_stellar_aberration(
+            propagated_orbits_aberrated.reshape(1, -1),
+            observer_coordinates.reshape(1, -1),
+        )[0]
+    )
+
+    # Convert to spherical coordinates
+    ephemeris_spherical = _cartesian_to_spherical(topocentric_coordinates)
+
+    return ephemeris_spherical, light_time
+
+
+# Vectorization Map: _generate_ephemeris_2body
+_generate_ephemeris_2body_vmap = jit(
+    vmap(
+        _generate_ephemeris_2body,
+        in_axes=(0, 0, 0, None, None, None, None),
+        out_axes=(0, 0),
+    )
+)

From f2cc8d3d8f067a78fb2ee37793aa7da89c33b9d6 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Mon, 21 Aug 2023 16:46:03 -0700
Subject: [PATCH 03/10] Add light_time to the Ephemeris class

---
 adam_core/orbits/ephemeris.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/adam_core/orbits/ephemeris.py b/adam_core/orbits/ephemeris.py
index d92bc156..705ddf92 100644
--- a/adam_core/orbits/ephemeris.py
+++ b/adam_core/orbits/ephemeris.py
@@ -1,13 +1,13 @@
-from quivr import StringColumn, Table
+import quivr as qv
 
 from ..coordinates.cartesian import CartesianCoordinates
 from ..coordinates.spherical import SphericalCoordinates
 
 
-class Ephemeris(Table):
+class Ephemeris(qv.Table):
 
-    orbit_id = StringColumn()
-    object_id = StringColumn(nullable=True)
+    orbit_id = qv.StringColumn()
+    object_id = qv.StringColumn(nullable=True)
     coordinates = SphericalCoordinates.as_column()
 
     # The coordinates as observed by the observer will be the result of
@@ -16,4 +16,5 @@ class Ephemeris(Table):
     # will be different from its actual geometric coordinates at the time of observation.
     # Aberrated coordinates are coordinates that account for the light travel time
     # from the time of emission/reflection to the time of observation
+    light_time = qv.Float64Column(nullable=True)
     aberrated_coordinates = CartesianCoordinates.as_column(nullable=True)

From 09f6122911b397ef90b80ec4e375d93af76aafe8 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Mon, 21 Aug 2023 17:10:03 -0700
Subject: [PATCH 04/10] Add generate_ephemeris_2body

---
 adam_core/coordinates/transform.py |  17 +--
 adam_core/dynamics/__init__.py     |   1 +
 adam_core/dynamics/ephemeris.py    | 162 ++++++++++++++++++++++++++++-
 adam_core/dynamics/propagation.py  |   2 +-
 4 files changed, 173 insertions(+), 9 deletions(-)

diff --git a/adam_core/coordinates/transform.py b/adam_core/coordinates/transform.py
index dd6d730f..40d93aff 100644
--- a/adam_core/coordinates/transform.py
+++ b/adam_core/coordinates/transform.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Literal, Optional, Union
+from typing import Literal, Optional, TypeVar, Union
 
 import jax.numpy as jnp
 import numpy as np
@@ -39,12 +39,15 @@
     "cometary_to_cartesian",
 ]
 
-CoordinateType = Union[
-    CartesianCoordinates,
-    KeplerianCoordinates,
-    CometaryCoordinates,
-    SphericalCoordinates,
-]
+CoordinateType = TypeVar(
+    "CoordinateType",
+    bound=Union[
+        CartesianCoordinates,
+        KeplerianCoordinates,
+        CometaryCoordinates,
+        SphericalCoordinates,
+    ],
+)
 CoordinatesClasses = (
     CartesianCoordinates,
     KeplerianCoordinates,
diff --git a/adam_core/dynamics/__init__.py b/adam_core/dynamics/__init__.py
index 7e9afbbf..3e53667c 100644
--- a/adam_core/dynamics/__init__.py
+++ b/adam_core/dynamics/__init__.py
@@ -2,6 +2,7 @@
 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
diff --git a/adam_core/dynamics/ephemeris.py b/adam_core/dynamics/ephemeris.py
index 4dffd310..7eb432e6 100644
--- a/adam_core/dynamics/ephemeris.py
+++ b/adam_core/dynamics/ephemeris.py
@@ -2,10 +2,21 @@
 
 import jax.numpy as jnp
 import numpy as np
+import quivr as qv
 from jax import jit, vmap
 
 from ..constants import Constants as c
-from ..coordinates.transform import _cartesian_to_spherical
+from ..coordinates.cartesian import CartesianCoordinates
+from ..coordinates.covariances import (
+    CoordinateCovariances,
+    transform_covariances_jacobian,
+)
+from ..coordinates.origin import Origin, OriginCodes
+from ..coordinates.spherical import SphericalCoordinates
+from ..coordinates.transform import _cartesian_to_spherical, transform_coordinates
+from ..observers.observers import Observers
+from ..orbits.ephemeris import Ephemeris
+from ..orbits.orbits import Orbits
 from .aberrations import _add_light_time, add_stellar_aberration
 
 MU = c.MU
@@ -97,3 +108,152 @@ def _generate_ephemeris_2body(
         out_axes=(0, 0),
     )
 )
+
+
+def generate_ephemeris_2body(
+    propagated_orbits: Orbits,
+    observers: Observers,
+    lt_tol: float = 1e-10,
+    mu: float = MU,
+    max_iter: int = 1000,
+    tol: float = 1e-15,
+) -> qv.MultiKeyLinkage[Ephemeris, Observers]:
+    """
+    Generate on-sky ephemerides for each propagated orbit as viewed by the observers.
+    This function calculates the light time delay between the propagated orbits and the observers,
+    and then propagates the orbits backward by that amount to when the light from each object was actually
+    emitted towards the observer.
+
+    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 is known as
+    stellar aberration. Stellar aberration is will also be applied after
+    light time correction has been added.
+
+    The velocity of the input orbits are unmodified only the position
+    vector is modified with stellar aberration.
+
+    Parameters
+    ----------
+    propagated_orbits : `~adam_core.orbits.orbits.Orbits` (N)
+        Propagated orbits.
+    observers : `~adam_core.observers.observers.Observers` (N)
+        Observers for which to generate ephemerides. Orbits should already have been
+        propagated to the same times as the observers.
+    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
+    -------
+    ephemeris : `~quivr.linkage.MultikeyLinkage` (N)
+        Topocentric ephemerides for each propagated orbit as observed by the given observers.
+    """
+    # Transform both the orbits and observers to the barycenter if they are not already.
+    propagated_orbits_barycentric = propagated_orbits.set_column(
+        "coordinates",
+        transform_coordinates(
+            propagated_orbits.coordinates,
+            CartesianCoordinates,
+            frame_out="ecliptic",
+            origin_out=OriginCodes.SOLAR_SYSTEM_BARYCENTER,
+        ),
+    )
+    observers_barycentric = observers.set_column(
+        "coordinates",
+        transform_coordinates(
+            observers.coordinates,
+            CartesianCoordinates,
+            frame_out="ecliptic",
+            origin_out=OriginCodes.SOLAR_SYSTEM_BARYCENTER,
+        ),
+    )
+
+    # Stack the observer coordinates and codes for each orbit in the propagated orbits
+    num_orbits = len(propagated_orbits_barycentric.orbit_id.unique())
+    observer_coordinates = np.tile(
+        observers_barycentric.coordinates.values, (num_orbits, 1)
+    )
+    observer_codes = np.tile(observers.code.to_numpy(zero_copy_only=False), num_orbits)
+
+    times = propagated_orbits.coordinates.time.to_astropy()
+    ephemeris_spherical, light_time = _generate_ephemeris_2body_vmap(
+        propagated_orbits_barycentric.coordinates.values,
+        times.mjd,
+        observer_coordinates,
+        lt_tol,
+        mu,
+        max_iter,
+        tol,
+    )
+    ephemeris_spherical = np.array(ephemeris_spherical)
+    light_time = np.array(light_time)
+
+    if not propagated_orbits.coordinates.covariance.is_all_nan():
+
+        cartesian_covariances = propagated_orbits.coordinates.covariance.to_matrix()
+        covariances_spherical = transform_covariances_jacobian(
+            propagated_orbits.coordinates.values,
+            cartesian_covariances,
+            _generate_ephemeris_2body,
+            in_axes=(0, 0, 0, None, None, None, None),
+            out_axes=(0, 0),
+            observation_times=times.utc.mjd,
+            observer_coordinates=observer_coordinates,
+            lt_tol=lt_tol,
+            mu=mu,
+            max_iter=max_iter,
+            tol=tol,
+        )
+        covariances_spherical = CoordinateCovariances.from_matrix(
+            np.array(covariances_spherical)
+        )
+
+    else:
+        covariances_spherical = None
+
+    spherical_coordinates = SphericalCoordinates.from_kwargs(
+        time=propagated_orbits.coordinates.time,
+        rho=ephemeris_spherical[:, 0],
+        lon=ephemeris_spherical[:, 1],
+        lat=ephemeris_spherical[:, 2],
+        vrho=ephemeris_spherical[:, 3],
+        vlon=ephemeris_spherical[:, 4],
+        vlat=ephemeris_spherical[:, 5],
+        covariance=covariances_spherical,
+        origin=Origin.from_kwargs(code=observer_codes),
+        frame="ecliptic",
+    )
+
+    # Rotate the spherical coordinates from the ecliptic frame
+    # to the equatorial frame
+    spherical_coordinates = transform_coordinates(
+        spherical_coordinates, SphericalCoordinates, frame_out="equatorial"
+    )
+
+    ephems = Ephemeris.from_kwargs(
+        orbit_id=propagated_orbits_barycentric.orbit_id,
+        object_id=propagated_orbits_barycentric.object_id,
+        coordinates=spherical_coordinates,
+        light_time=light_time,
+    )
+
+    linkages = qv.MultiKeyLinkage(
+        left_table=ephems,
+        right_table=observers,
+        left_keys={
+            "code": ephems.coordinates.origin.code,
+            "mjd": ephems.coordinates.time.mjd(),
+        },
+        right_keys={
+            "code": observers.code,
+            "mjd": observers.coordinates.time.mjd(),
+        },
+    )
+    return linkages
diff --git a/adam_core/dynamics/propagation.py b/adam_core/dynamics/propagation.py
index 50afdee2..49a407b7 100644
--- a/adam_core/dynamics/propagation.py
+++ b/adam_core/dynamics/propagation.py
@@ -149,7 +149,7 @@ def propagate_2body(
     else:
         cartesian_covariances = None
 
-    origin_code = np.empty(n_orbits * n_times, dtype=str)
+    origin_code = np.empty(n_orbits * n_times, dtype="U3")
     origin_code.fill("SUN")
 
     orbits_propagated = Orbits.from_kwargs(

From 159680afd5af1a8ecc4e4a5ad888c6fb5890b7b8 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Fri, 25 Aug 2023 11:48:31 -0700
Subject: [PATCH 05/10] Bug fix: Correctly define dot product in for-loop less
 stellar aberration function

For reference, the original code was written in numba and looped over individual topocentric states:
https://github.com/moeyensj/thor/commit/29f28f5bd144be08e27d8710d2e5db0167f8abe8
---
 adam_core/dynamics/aberrations.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/adam_core/dynamics/aberrations.py b/adam_core/dynamics/aberrations.py
index 15504d5e..d6b271bc 100644
--- a/adam_core/dynamics/aberrations.py
+++ b/adam_core/dynamics/aberrations.py
@@ -192,15 +192,16 @@ def add_stellar_aberration(
     rho_aberrated = rho_aberrated.at[:].set(topo_states[:, :3])
 
     v_obs = observer_states[:, 3:]
-    beta = v_obs / C
-    gamma_inv = jnp.sqrt(1 - jnp.linalg.norm(beta, axis=1, keepdims=True) ** 2)
+    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(
-        (gamma_inv * rho + beta + rho * beta * beta / (1 + gamma_inv))
-        / (1 + rho * beta)
+        (beta_inv * rho + gamma + rho_dot_gamma * gamma / (1 + beta_inv))
+        / (1 + rho_dot_gamma)
     )
     rho_aberrated = rho_aberrated.at[:].multiply(delta)
 

From 0d0a1e29e5839c9deaab2884b6920e9d01f2a9a8 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Fri, 25 Aug 2023 11:59:31 -0700
Subject: [PATCH 06/10] Make stellar aberration optional

---
 adam_core/dynamics/ephemeris.py | 68 +++++++++++++++++++++++----------
 1 file changed, 48 insertions(+), 20 deletions(-)

diff --git a/adam_core/dynamics/ephemeris.py b/adam_core/dynamics/ephemeris.py
index 7eb432e6..41ab5e93 100644
--- a/adam_core/dynamics/ephemeris.py
+++ b/adam_core/dynamics/ephemeris.py
@@ -3,7 +3,7 @@
 import jax.numpy as jnp
 import numpy as np
 import quivr as qv
-from jax import jit, vmap
+from jax import jit, lax, vmap
 
 from ..constants import Constants as c
 from ..coordinates.cartesian import CartesianCoordinates
@@ -31,21 +31,29 @@ def _generate_ephemeris_2body(
     mu: float = MU,
     max_iter: int = 100,
     tol: float = 1e-15,
+    stellar_aberration: bool = False,
 ) -> Tuple[jnp.ndarray, jnp.float64]:
     """
     Given a propagated orbit, generate its on-sky ephemeris as viewed from the observer.
     This function calculates the light time delay between the propagated orbit and the observer,
     and then propagates the orbit backward by that amount to when the light from object was actually
-    emitted towards the observer.
+    emitted towards the observer ("astrometric coordinates").
 
     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 is known as
-    stellar aberration. Stellar aberration is will also be applied after
-    light time correction has been added.
+    to appear in a different location than its true location, this is known as
+    stellar aberration (often referred to in combination with other aberrations as "apparent
+    coordinates"). Stellar aberration can optionally be applied after
+    light time correction has been added but it should not be necessary when comparing to ephemerides
+    of solar system small bodies extracted from astrometric catalogs. The stars to which the
+    catalog is calibrated undergo the same aberration as the moving objects as seen from the observer.
 
-    The velocity of the input orbits are unmodified only the position
+    If stellar aberration is applied then the velocity of the input orbits are unmodified, only the position
     vector is modified with stellar aberration.
 
+    For more details on aberrations see:
+        https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/FORTRAN/req/abcorr.html
+        https://ssd.jpl.nasa.gov/horizons/manual.html#defs
+
     Parameters
     ----------
     propagated_orbit : `~jax.numpy.ndarray` (6)
@@ -64,6 +72,8 @@ def _generate_ephemeris_2body(
     tol : float, optional
         Numerical tolerance to which to compute universal anomaly during propagation using the Newtown-Raphson
         method.
+    stellar_aberration : bool, optional
+        Apply stellar aberration to the ephemerides.
 
     Returns
     -------
@@ -87,11 +97,16 @@ def _generate_ephemeris_2body(
     topocentric_coordinates = propagated_orbits_aberrated - observer_coordinates
 
     # Apply stellar aberration to topocentric coordinates
-    topocentric_coordinates = topocentric_coordinates.at[0:3].set(
-        add_stellar_aberration(
-            propagated_orbits_aberrated.reshape(1, -1),
-            observer_coordinates.reshape(1, -1),
-        )[0]
+    topocentric_coordinates = lax.cond(
+        stellar_aberration,
+        lambda topocentric_coords: topocentric_coords.at[0:3].set(
+            add_stellar_aberration(
+                propagated_orbits_aberrated.reshape(1, -1),
+                observer_coordinates.reshape(1, -1),
+            )[0],
+        ),
+        lambda topocentric_coords: topocentric_coords,
+        topocentric_coordinates,
     )
 
     # Convert to spherical coordinates
@@ -104,7 +119,7 @@ def _generate_ephemeris_2body(
 _generate_ephemeris_2body_vmap = jit(
     vmap(
         _generate_ephemeris_2body,
-        in_axes=(0, 0, 0, None, None, None, None),
+        in_axes=(0, 0, 0, None, None, None, None, None),
         out_axes=(0, 0),
     )
 )
@@ -117,21 +132,30 @@ def generate_ephemeris_2body(
     mu: float = MU,
     max_iter: int = 1000,
     tol: float = 1e-15,
+    stellar_aberration: bool = False,
 ) -> qv.MultiKeyLinkage[Ephemeris, Observers]:
     """
     Generate on-sky ephemerides for each propagated orbit as viewed by the observers.
-    This function calculates the light time delay between the propagated orbits and the observers,
-    and then propagates the orbits backward by that amount to when the light from each object was actually
-    emitted towards the observer.
+    This function calculates the light time delay between the propagated orbit and the observer,
+    and then propagates the orbit backward by that amount to when the light from object was actually
+    emitted towards the observer ("astrometric coordinates").
 
     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 is known as
-    stellar aberration. Stellar aberration is will also be applied after
-    light time correction has been added.
+    to appear in a different location than its true location, this is known as
+    stellar aberration (often referred to in combination with other aberrations as "apparent
+    coordinates"). Stellar aberration can optionally be applied after
+    light time correction has been added but it should not be necessary when comparing to ephemerides
+    of solar system small bodies extracted from astrometric catalogs. The stars to which the
+    catalog is calibrated undergo the same aberration as the moving objects as seen from the observer.
 
-    The velocity of the input orbits are unmodified only the position
+    If stellar aberration is applied then the velocity of the input orbits are unmodified, only the position
     vector is modified with stellar aberration.
 
+    For more details on aberrations see:
+        https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/FORTRAN/req/abcorr.html
+        https://ssd.jpl.nasa.gov/horizons/manual.html#defs
+
+
     Parameters
     ----------
     propagated_orbits : `~adam_core.orbits.orbits.Orbits` (N)
@@ -149,6 +173,8 @@ def generate_ephemeris_2body(
     tol : float, optional
         Numerical tolerance to which to compute universal anomaly during propagation using the Newtown-Raphson
         method.
+    stellar_aberration : bool, optional
+        Apply stellar aberration to the ephemerides.
 
     Returns
     -------
@@ -191,6 +217,7 @@ def generate_ephemeris_2body(
         mu,
         max_iter,
         tol,
+        stellar_aberration,
     )
     ephemeris_spherical = np.array(ephemeris_spherical)
     light_time = np.array(light_time)
@@ -202,7 +229,7 @@ def generate_ephemeris_2body(
             propagated_orbits.coordinates.values,
             cartesian_covariances,
             _generate_ephemeris_2body,
-            in_axes=(0, 0, 0, None, None, None, None),
+            in_axes=(0, 0, 0, None, None, None, None, None),
             out_axes=(0, 0),
             observation_times=times.utc.mjd,
             observer_coordinates=observer_coordinates,
@@ -210,6 +237,7 @@ def generate_ephemeris_2body(
             mu=mu,
             max_iter=max_iter,
             tol=tol,
+            stellar_aberration=stellar_aberration,
         )
         covariances_spherical = CoordinateCovariances.from_matrix(
             np.array(covariances_spherical)

From cc782282a7301c26b77042964166a41badaa688c Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Fri, 25 Aug 2023 16:21:48 -0700
Subject: [PATCH 07/10] Add generate_ephemeris_2body test

---
 adam_core/dynamics/tests/conftest.py       | 28 +++++++++
 adam_core/dynamics/tests/test_ephemeris.py | 69 ++++++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 adam_core/dynamics/tests/test_ephemeris.py

diff --git a/adam_core/dynamics/tests/conftest.py b/adam_core/dynamics/tests/conftest.py
index 138c8de3..64174ff9 100644
--- a/adam_core/dynamics/tests/conftest.py
+++ b/adam_core/dynamics/tests/conftest.py
@@ -3,6 +3,8 @@
 import pandas as pd
 import pytest
 
+from adam_core.orbits import Orbits
+
 
 @pytest.fixture
 def orbital_elements():
@@ -13,3 +15,29 @@ def orbital_elements():
         orbital_elements_file, index_col=False, float_precision="round_trip"
     )
     return df
+
+
+@pytest.fixture
+def propagated_orbits():
+    propagated_orbits_file = files("adam_core.utils.helpers.data").joinpath(
+        "propagated_orbits.csv"
+    )
+    df = pd.read_csv(
+        propagated_orbits_file,
+        index_col=False,
+        float_precision="round_trip",
+        dtype={"orbit_id": str},
+    )
+    return Orbits.from_dataframe(df, frame="ecliptic")
+
+
+@pytest.fixture
+def ephemeris():
+    ephemeris_file = files("adam_core.utils.helpers.data").joinpath("ephemeris.csv")
+    df = pd.read_csv(
+        ephemeris_file,
+        index_col=False,
+        float_precision="round_trip",
+        dtype={"orbit_id": str},
+    )
+    return df
diff --git a/adam_core/dynamics/tests/test_ephemeris.py b/adam_core/dynamics/tests/test_ephemeris.py
new file mode 100644
index 00000000..c726bb9d
--- /dev/null
+++ b/adam_core/dynamics/tests/test_ephemeris.py
@@ -0,0 +1,69 @@
+import numpy as np
+from astropy import units as u
+from astropy.time import Time
+
+from ...observers import Observers
+from ..ephemeris import generate_ephemeris_2body
+
+
+def test_generate_ephemeris_2body(propagated_orbits, ephemeris):
+    # For our catalog of test orbits, generate ephemeris using Horizons generated state vectors
+    # and compare the results to the Horizons generated ephemeris
+    orbit_ids = propagated_orbits.orbit_id.unique().to_numpy(zero_copy_only=False)
+    for orbit_id in orbit_ids:
+
+        propagated_orbit = propagated_orbits.select("orbit_id", orbit_id)
+        ephemeris_orbit = ephemeris[ephemeris["orbit_id"].isin([orbit_id])]
+
+        observers = Observers.from_code(
+            "X05",
+            Time(ephemeris_orbit["mjd_utc"].values, format="mjd", scale="utc"),
+        )
+
+        ephemeris_orbit_2body = generate_ephemeris_2body(propagated_orbit, observers)
+
+        # Extract only the ephemeris table
+        ephemeris_orbit_2body = ephemeris_orbit_2body.left_table
+
+        # Calculate the offset in light travel time
+        light_time_diff = np.abs(
+            ephemeris_orbit_2body.light_time.to_numpy()
+            - (ephemeris_orbit["lighttime"].values * 60 / 86400)
+        ) * (u.d).to(u.microsecond)
+
+        # Assert that the light time is less than 100 microseconds
+        np.testing.assert_array_less(light_time_diff, 100)
+
+        # Calculate the difference in RA and Dec
+        angular_diff = np.abs(
+            ephemeris_orbit_2body.coordinates.values[:, 1:3]
+            - ephemeris_orbit[["RA", "DEC"]].values
+        ) * (u.deg).to(u.milliarcsecond)
+
+        # Topocentric difference
+        range_diff = np.abs(
+            ephemeris_orbit_2body.coordinates.values[:, 0]
+            - ephemeris_orbit["delta"].values
+        ) * (u.au).to(u.m)
+
+        if orbit_id == "00014":
+            # Assert that the position is less than 50 milliarcsecond
+            # Orbit 14 is 6522 Aci (1991 NQ) an inner main-belt asteroid
+            np.testing.assert_array_less(angular_diff, 50)
+
+            # Assert that the range is less than 5000 m
+            np.testing.assert_array_less(range_diff, 5000)
+        elif orbit_id == "00019":
+            # Assert that the position is less than 10 milliarcsecond
+            # Orbit 19 is 911 Agamemnon (A919 FB), a Jupiter Trojan, which is interesting
+            # since Orbit 20 is also a Jupiter Trojan.
+            np.testing.assert_array_less(angular_diff, 10)
+
+            # Assert that the range is less than 12000 m
+            np.testing.assert_array_less(range_diff, 12000)
+        else:
+            # Assert that the position is less than 1 milliarcsecond
+            np.testing.assert_array_less(angular_diff, 1)
+
+            # Assert that the range is less than 200 m
+            np.testing.assert_array_less(range_diff, 200)

From 0baec63149d8b437069d011402d0c410d04f15e7 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Wed, 13 Sep 2023 15:30:51 -0700
Subject: [PATCH 08/10] Update ephemeris test tolerances to account for newly
 generated test data. Performance is excellent!

---
 adam_core/dynamics/tests/test_ephemeris.py | 75 ++++++++++++++++------
 1 file changed, 57 insertions(+), 18 deletions(-)

diff --git a/adam_core/dynamics/tests/test_ephemeris.py b/adam_core/dynamics/tests/test_ephemeris.py
index c726bb9d..baaf8124 100644
--- a/adam_core/dynamics/tests/test_ephemeris.py
+++ b/adam_core/dynamics/tests/test_ephemeris.py
@@ -1,4 +1,5 @@
 import numpy as np
+import quivr as qv
 from astropy import units as u
 from astropy.time import Time
 
@@ -15,10 +16,20 @@ def test_generate_ephemeris_2body(propagated_orbits, ephemeris):
         propagated_orbit = propagated_orbits.select("orbit_id", orbit_id)
         ephemeris_orbit = ephemeris[ephemeris["orbit_id"].isin([orbit_id])]
 
-        observers = Observers.from_code(
-            "X05",
-            Time(ephemeris_orbit["mjd_utc"].values, format="mjd", scale="utc"),
-        )
+        observers_list = []
+        for observatory_code in ephemeris_orbit["observatory_code"].unique():
+            observatory_mask = ephemeris_orbit["observatory_code"] == observatory_code
+            observer_i = Observers.from_code(
+                observatory_code,
+                Time(
+                    ephemeris_orbit[observatory_mask]["mjd_utc"].values,
+                    format="mjd",
+                    scale="utc",
+                ),
+            )
+            observers_list.append(observer_i)
+
+        observers = qv.concatenate(observers_list)
 
         ephemeris_orbit_2body = generate_ephemeris_2body(propagated_orbit, observers)
 
@@ -31,8 +42,8 @@ def test_generate_ephemeris_2body(propagated_orbits, ephemeris):
             - (ephemeris_orbit["lighttime"].values * 60 / 86400)
         ) * (u.d).to(u.microsecond)
 
-        # Assert that the light time is less than 100 microseconds
-        np.testing.assert_array_less(light_time_diff, 100)
+        # Assert that the light time is less than 10 microseconds
+        np.testing.assert_array_less(light_time_diff, 10)
 
         # Calculate the difference in RA and Dec
         angular_diff = np.abs(
@@ -46,24 +57,52 @@ def test_generate_ephemeris_2body(propagated_orbits, ephemeris):
             - ephemeris_orbit["delta"].values
         ) * (u.au).to(u.m)
 
+        # Orbit 00014 is 6522 Aci (1991 NQ) an inner main-belt asteroid
+        # with a diameter of about 5 km. What's interesting here is that of the 90
+        # ephemeris positions only 2 show a range difference of more than 15 m
+        # (many are less than 10 m). Their values are 92.17 and 177.24 m.
+        # These are the same observations that have offsets in RA and Dec greater
+        # than 1 milliarcsecond. The maximum offset for those observations not
+        # corresponding to the ones with a discrepant range offset is
+        # only 0.013 milliarcseconds!
         if orbit_id == "00014":
-            # Assert that the position is less than 50 milliarcsecond
-            # Orbit 14 is 6522 Aci (1991 NQ) an inner main-belt asteroid
-            np.testing.assert_array_less(angular_diff, 50)
+            # Assert that the position is less than 2 milliarcsecond
+            np.testing.assert_array_less(angular_diff, 2)
+
+            # Assert that the range is less than 200 m
+            np.testing.assert_array_less(range_diff, 200)
 
-            # Assert that the range is less than 5000 m
-            np.testing.assert_array_less(range_diff, 5000)
+        # Orbit 00019 is 1143 Odysseus (1930 BH), a Jupiter Trojan with a
+        # diameter of about 115 km. What's interesting here is that of the 90
+        # ephemeris positions only 2 show a range difference of more than 10 m.
+        # Their values are 231.79 and 468.26 m.  These are the same observations
+        # that have offsets in RA and Dec greater than 0.5 milliarcsecond.
+        # However, the maximum offset for those observations not corresponding to the
+        # ones with a discrepant range offset is only 0.004 milliarcseconds!
         elif orbit_id == "00019":
-            # Assert that the position is less than 10 milliarcsecond
-            # Orbit 19 is 911 Agamemnon (A919 FB), a Jupiter Trojan, which is interesting
-            # since Orbit 20 is also a Jupiter Trojan.
-            np.testing.assert_array_less(angular_diff, 10)
+            # Assert that the position is less than 1 milliarcsecond
+            np.testing.assert_array_less(angular_diff, 1)
 
-            # Assert that the range is less than 12000 m
-            np.testing.assert_array_less(range_diff, 12000)
-        else:
+            # Assert that the range is less than 500 m
+            np.testing.assert_array_less(range_diff, 500)
+
+        # Orbit 00000 is 594913 'Aylo'chaxnim (2020 AV2) (an orbit completely
+        # interior to Venus' orbit). It would not surprise me that there are
+        # some dynamics that we are not modeling fully for this orbit in terms
+        # of light travel time and/or the ephemeris generation.
+        elif orbit_id == "00000":
             # Assert that the position is less than 1 milliarcsecond
             np.testing.assert_array_less(angular_diff, 1)
 
             # Assert that the range is less than 200 m
             np.testing.assert_array_less(range_diff, 200)
+
+        # These tolerance apply to the rest of the orbits (25/28) and
+        # show excellent results. RA, DEC to within 0.1 milliarcsecond
+        # and range to within 20 m.
+        else:
+            # Assert that the position is less than 0.1 milliarcsecond
+            np.testing.assert_array_less(angular_diff, 0.1)
+
+            # Assert that the range is less than 20 m
+            np.testing.assert_array_less(range_diff, 20)

From ceee549a52e748388a7070070889866cca1fede5 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Mon, 25 Sep 2023 15:21:13 -0700
Subject: [PATCH 09/10] Use deferred imports for dynamics functions used in
 transform.py

---
 adam_core/coordinates/__init__.py  |  5 +++++
 adam_core/coordinates/transform.py | 20 +++++---------------
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/adam_core/coordinates/__init__.py b/adam_core/coordinates/__init__.py
index ebdb3cbf..62c36089 100644
--- a/adam_core/coordinates/__init__.py
+++ b/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
diff --git a/adam_core/coordinates/transform.py b/adam_core/coordinates/transform.py
index 40d93aff..e2a7d0ab 100644
--- a/adam_core/coordinates/transform.py
+++ b/adam_core/coordinates/transform.py
@@ -7,8 +7,6 @@
 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,19 +23,6 @@
 
 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 = TypeVar(
     "CoordinateType",
@@ -358,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]
@@ -625,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]

From 85dde909c49ab109dfd355a2c03a1da64749a620 Mon Sep 17 00:00:00 2001
From: Joachim Moeyens <moeyensj@gmail.com>
Date: Wed, 4 Oct 2023 08:50:11 -0700
Subject: [PATCH 10/10] Parameterize the generate_ephemeris_2body test and
 update tolerance for new data

---
 adam_core/dynamics/tests/conftest.py       |  11 +-
 adam_core/dynamics/tests/test_ephemeris.py | 217 +++++++++++----------
 2 files changed, 121 insertions(+), 107 deletions(-)

diff --git a/adam_core/dynamics/tests/conftest.py b/adam_core/dynamics/tests/conftest.py
index 64174ff9..2581e52f 100644
--- a/adam_core/dynamics/tests/conftest.py
+++ b/adam_core/dynamics/tests/conftest.py
@@ -19,16 +19,9 @@ def orbital_elements():
 
 @pytest.fixture
 def propagated_orbits():
-    propagated_orbits_file = files("adam_core.utils.helpers.data").joinpath(
-        "propagated_orbits.csv"
+    return Orbits.from_parquet(
+        files("adam_core.utils.helpers.data").joinpath("propagated_orbits.parquet")
     )
-    df = pd.read_csv(
-        propagated_orbits_file,
-        index_col=False,
-        float_precision="round_trip",
-        dtype={"orbit_id": str},
-    )
-    return Orbits.from_dataframe(df, frame="ecliptic")
 
 
 @pytest.fixture
diff --git a/adam_core/dynamics/tests/test_ephemeris.py b/adam_core/dynamics/tests/test_ephemeris.py
index baaf8124..62be0523 100644
--- a/adam_core/dynamics/tests/test_ephemeris.py
+++ b/adam_core/dynamics/tests/test_ephemeris.py
@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 import quivr as qv
 from astropy import units as u
 from astropy.time import Time
@@ -6,103 +7,123 @@
 from ...observers import Observers
 from ..ephemeris import generate_ephemeris_2body
 
-
-def test_generate_ephemeris_2body(propagated_orbits, ephemeris):
+OBJECT_IDS = [
+    "594913 'Aylo'chaxnim (2020 AV2)",
+    "163693 Atira (2003 CP20)",
+    "(2010 TK7)",
+    "3753 Cruithne (1986 TO)",
+    "54509 YORP (2000 PH5)",
+    "2063 Bacchus (1977 HB)",
+    "1221 Amor (1932 EA1)",
+    "433 Eros (A898 PA)",
+    "3908 Nyx (1980 PA)",
+    "434 Hungaria (A898 RB)",
+    "1876 Napolitania (1970 BA)",
+    "2001 Einstein (1973 EB)",
+    "2 Pallas (A802 FA)",
+    "6 Hebe (A847 NA)",
+    "6522 Aci (1991 NQ)",
+    "10297 Lynnejones (1988 RJ13)",
+    "17032 Edlu (1999 FM9)",
+    "202930 Ivezic (1998 SG172)",
+    "911 Agamemnon (A919 FB)",
+    "1143 Odysseus (1930 BH)",
+    "1172 Aneas (1930 UA)",
+    "3317 Paris (1984 KF)",
+    "5145 Pholus (1992 AD)",
+    "5335 Damocles (1991 DA)",
+    "15760 Albion (1992 QB1)",
+    "15788 (1993 SB)",
+    "15789 (1993 SC)",
+    "1I/'Oumuamua (A/2017 U1)",
+]
+
+TOLERANCES = {
+    # range (m), angular difference (mas), light_time difference (microseconds)
+    # Default range is about consistent with our ability to get
+    # observatory state vector positions compared to Horizons
+    "default": (20, 0.1, 10),
+    # 594913 'Aylo'chaxnim (2020 AV2) : Vatira
+    # Range difference for first 30 observations between
+    # 32.8 m and 80.4 m, for last 60 observations range is
+    # between 0.13 m and 26.3 m
+    "594913 'Aylo'chaxnim (2020 AV2)": (100, 1, 10),
+    # 3753 Cruithne (1986 TO) : NEO (ATEN)
+    # Range anywhere between 0.01 m and 40.5 m throughout
+    # the 90 observations
+    "3753 Cruithne (1986 TO)": (50, 1, 10),
+    # 6522 Aci (1991 NQ) : Inner Main Belt asteroid
+    # 89/90 RA, Decs have a mean offset of 0.007 mas
+    # 1/90 RA, Decs has an offset of 1.24 mas
+    # 88/90 ranges have a mean offset of 10.6 m
+    # 2/90 ranges are 89.5, 174.3 m
+    "6522 Aci (1991 NQ)": (200, 2, 10),
+    # 1143 Odysseus (1930 BH) : Jupiter Trojan
+    # 89/90 light times have a mean offset of 0.14 microseconds
+    # 1/90 light times has an offset of 38.4 microseconds
+    # 89/90 RA, Decs have a mean offset of 0.0056 mas
+    # 1/90 RA, Decs has an offset of 8.9 mas
+    # 89/90 ranges have a mean offset of 6.9 m
+    # 1/90 ranges has an offset of 11575 m? ...
+    "1143 Odysseus (1930 BH)": (11600, 10, 50),
+}
+
+
+@pytest.mark.parametrize("object_id", OBJECT_IDS)
+def test_generate_ephemeris_2body(object_id, propagated_orbits, ephemeris):
     # For our catalog of test orbits, generate ephemeris using Horizons generated state vectors
     # and compare the results to the Horizons generated ephemeris
-    orbit_ids = propagated_orbits.orbit_id.unique().to_numpy(zero_copy_only=False)
-    for orbit_id in orbit_ids:
-
-        propagated_orbit = propagated_orbits.select("orbit_id", orbit_id)
-        ephemeris_orbit = ephemeris[ephemeris["orbit_id"].isin([orbit_id])]
-
-        observers_list = []
-        for observatory_code in ephemeris_orbit["observatory_code"].unique():
-            observatory_mask = ephemeris_orbit["observatory_code"] == observatory_code
-            observer_i = Observers.from_code(
-                observatory_code,
-                Time(
-                    ephemeris_orbit[observatory_mask]["mjd_utc"].values,
-                    format="mjd",
-                    scale="utc",
-                ),
-            )
-            observers_list.append(observer_i)
-
-        observers = qv.concatenate(observers_list)
-
-        ephemeris_orbit_2body = generate_ephemeris_2body(propagated_orbit, observers)
-
-        # Extract only the ephemeris table
-        ephemeris_orbit_2body = ephemeris_orbit_2body.left_table
-
-        # Calculate the offset in light travel time
-        light_time_diff = np.abs(
-            ephemeris_orbit_2body.light_time.to_numpy()
-            - (ephemeris_orbit["lighttime"].values * 60 / 86400)
-        ) * (u.d).to(u.microsecond)
-
-        # Assert that the light time is less than 10 microseconds
-        np.testing.assert_array_less(light_time_diff, 10)
-
-        # Calculate the difference in RA and Dec
-        angular_diff = np.abs(
-            ephemeris_orbit_2body.coordinates.values[:, 1:3]
-            - ephemeris_orbit[["RA", "DEC"]].values
-        ) * (u.deg).to(u.milliarcsecond)
-
-        # Topocentric difference
-        range_diff = np.abs(
-            ephemeris_orbit_2body.coordinates.values[:, 0]
-            - ephemeris_orbit["delta"].values
-        ) * (u.au).to(u.m)
-
-        # Orbit 00014 is 6522 Aci (1991 NQ) an inner main-belt asteroid
-        # with a diameter of about 5 km. What's interesting here is that of the 90
-        # ephemeris positions only 2 show a range difference of more than 15 m
-        # (many are less than 10 m). Their values are 92.17 and 177.24 m.
-        # These are the same observations that have offsets in RA and Dec greater
-        # than 1 milliarcsecond. The maximum offset for those observations not
-        # corresponding to the ones with a discrepant range offset is
-        # only 0.013 milliarcseconds!
-        if orbit_id == "00014":
-            # Assert that the position is less than 2 milliarcsecond
-            np.testing.assert_array_less(angular_diff, 2)
-
-            # Assert that the range is less than 200 m
-            np.testing.assert_array_less(range_diff, 200)
-
-        # Orbit 00019 is 1143 Odysseus (1930 BH), a Jupiter Trojan with a
-        # diameter of about 115 km. What's interesting here is that of the 90
-        # ephemeris positions only 2 show a range difference of more than 10 m.
-        # Their values are 231.79 and 468.26 m.  These are the same observations
-        # that have offsets in RA and Dec greater than 0.5 milliarcsecond.
-        # However, the maximum offset for those observations not corresponding to the
-        # ones with a discrepant range offset is only 0.004 milliarcseconds!
-        elif orbit_id == "00019":
-            # Assert that the position is less than 1 milliarcsecond
-            np.testing.assert_array_less(angular_diff, 1)
-
-            # Assert that the range is less than 500 m
-            np.testing.assert_array_less(range_diff, 500)
-
-        # Orbit 00000 is 594913 'Aylo'chaxnim (2020 AV2) (an orbit completely
-        # interior to Venus' orbit). It would not surprise me that there are
-        # some dynamics that we are not modeling fully for this orbit in terms
-        # of light travel time and/or the ephemeris generation.
-        elif orbit_id == "00000":
-            # Assert that the position is less than 1 milliarcsecond
-            np.testing.assert_array_less(angular_diff, 1)
-
-            # Assert that the range is less than 200 m
-            np.testing.assert_array_less(range_diff, 200)
-
-        # These tolerance apply to the rest of the orbits (25/28) and
-        # show excellent results. RA, DEC to within 0.1 milliarcsecond
-        # and range to within 20 m.
-        else:
-            # Assert that the position is less than 0.1 milliarcsecond
-            np.testing.assert_array_less(angular_diff, 0.1)
-
-            # Assert that the range is less than 20 m
-            np.testing.assert_array_less(range_diff, 20)
+    propagated_orbit = propagated_orbits.select("object_id", object_id)
+    ephemeris_orbit = ephemeris[ephemeris["targetname"].isin([object_id])]
+
+    observers_list = []
+    for observatory_code in ephemeris_orbit["observatory_code"].unique():
+        observatory_mask = ephemeris_orbit["observatory_code"] == observatory_code
+        observer_i = Observers.from_code(
+            observatory_code,
+            Time(
+                ephemeris_orbit[observatory_mask]["mjd_utc"].values,
+                format="mjd",
+                scale="utc",
+            ),
+        )
+        observers_list.append(observer_i)
+
+    observers = qv.concatenate(observers_list)
+
+    ephemeris_orbit_2body = generate_ephemeris_2body(propagated_orbit, observers)
+
+    # Extract only the ephemeris table
+    ephemeris_orbit_2body = ephemeris_orbit_2body.left_table
+
+    # Get the tolerances for this orbit
+    if object_id in TOLERANCES:
+        range_tolerance, angular_tolerance, light_time_tolerance = TOLERANCES[object_id]
+    else:
+        range_tolerance, angular_tolerance, light_time_tolerance = TOLERANCES["default"]
+
+    # Calculate the offset in light travel time
+    light_time_diff = np.abs(
+        ephemeris_orbit_2body.light_time.to_numpy()
+        - (ephemeris_orbit["lighttime"].values * 60 / 86400)
+    ) * (u.d).to(u.microsecond)
+
+    # Assert that the light time is less than the tolerance
+    np.testing.assert_array_less(light_time_diff, light_time_tolerance)
+
+    # Calculate the difference in RA and Dec
+    angular_diff = np.abs(
+        ephemeris_orbit_2body.coordinates.values[:, 1:3]
+        - ephemeris_orbit[["RA", "DEC"]].values
+    ) * (u.deg).to(u.milliarcsecond)
+
+    # Topocentric difference
+    range_diff = np.abs(
+        ephemeris_orbit_2body.coordinates.values[:, 0] - ephemeris_orbit["delta"].values
+    ) * (u.au).to(u.m)
+
+    # Assert that the position is less than the tolerance
+    np.testing.assert_array_less(angular_diff, angular_tolerance)
+
+    # Assert that the range is less than the tolerance
+    np.testing.assert_array_less(range_diff, range_tolerance)