Add generic Ephemeris implementation to Propagator with light time co…

…rrection

src/adam_core/_version.py

@@ -0,0 +1 @@
+__version__ = '0.2.1.dev1+g7f7fcd5.d20240725'

src/adam_core/orbits/query/horizons.py

@@ -157,10 +157,11 @@ def _get_horizons_ephemeris(
         as seen from the observer location at the given times.
     """
     dfs = []
+    mjd_utc = times.rescale("utc").mjd().to_numpy(zero_copy_only=False)
     for i, obj_id in enumerate(object_ids):
         obj = Horizons(
             id=obj_id,
-            epochs=times.rescale("utc").mjd().to_numpy(zero_copy_only=False),
+            epochs=mjd_utc,
             location=location,
             id_type=id_type,
         )
@@ -171,7 +172,7 @@ def _get_horizons_ephemeris(
             cache=False,
         ).to_pandas()
         ephemeris.insert(0, "orbit_id", f"{i:05d}")
-        ephemeris.insert(2, "mjd_utc", times.utc.mjd)
+        ephemeris.insert(2, "mjd_utc", mjd_utc)
         ephemeris.insert(3, "observatory_code", location)
 
         dfs.append(ephemeris)

src/adam_core/propagator/propagator.py

@@ -6,12 +6,16 @@
 import numpy.typing as npt
 import quivr as qv
 
-from adam_core.ray_cluster import initialize_use_ray
-
+from ..coordinates.cartesian import CartesianCoordinates
+from ..coordinates.origin import Origin, OriginCodes
+from ..coordinates.spherical import SphericalCoordinates
+from ..coordinates.transform import transform_coordinates
+from ..dynamics.aberrations import C
 from ..observers.observers import Observers
 from ..orbits.ephemeris import Ephemeris
 from ..orbits.orbits import Orbits
 from ..orbits.variants import VariantEphemeris, VariantOrbits
+from ..ray_cluster import initialize_use_ray
 from ..time import Timestamp
 from .utils import _iterate_chunks
 
@@ -89,17 +93,155 @@ class EphemerisMixin:
     Subclasses should implement the _generate_ephemeris method.
     """
 
-    @abstractmethod
+    def _add_light_time(
+        self, 
+        orbits,
+        observers,
+        lt_tol: float = 1e-10,
+    ):
+        """
+        Accounts for light time correction in the ephemeris generation.
+
+        This differs from the _add_light_time method in the aberrations module
+        in that is uses the underlying propagator to propagate the orbits as opposed
+        to the _propagate_2body function.
+        """
+        orbits_aberrated = orbits.empty()
+        lts = np.zeros(len(orbits))
+        for i, (orbit, observer) in enumerate(zip(orbits, observers)):
+            lt0 = 0
+            dlt = 1
+            iterations = 0
+            while dlt > lt_tol:
+                iterations += 1
+                if iterations > 3:
+                    break
+                observer_position = observer.coordinates.values[0,:3]
+                orbit_i = orbit
+                t0 = orbit_i.coordinates.time.mjd()[0].as_py()
+
+                rho = np.linalg.norm(orbit_i.coordinates.values[0,:3] - observer_position)
+
+                lt = rho / C
+
+                dlt = np.abs(lt - lt0)
+
+                t1 = t0 - lt
+                t1 = Timestamp.from_mjd([t1], scale="tdb")
+                orbit_propagated = self.propagate_orbits(orbit, t1)
+
+                orbit_i = orbit_propagated
+                t0 = t1
+                lt0 = lt
+                dlt = dlt
+            orbits_aberrated = qv.concatenate([orbits_aberrated, orbit])
+            lts[i] = lt
+
+        return orbits_aberrated, lts
+
     def _generate_ephemeris(
-        self, orbits: EphemerisType, observers: ObserverType
+        self, orbits: OrbitType, observers: ObserverType, lt_tol: float = 1e-10
     ) -> EphemerisType:
+
         """
-        Generate ephemerides for the given orbits as observed by
-        the observers.
-
-        THIS FUNCTION SHOULD BE DEFINED BY THE USER.
+        A generic ephemeris implementation, which can be used or overridden by subclasses.
         """
-        pass
+
+        if isinstance(orbits, Orbits):
+            ephemeris_total = Ephemeris.empty()
+        elif isinstance(orbits, VariantOrbits):
+            ephemeris_total = VariantEphemeris.empty()
+
+        for orbit in orbits:
+            propagated_orbits = self.propagate_orbits(
+                orbit, observers.coordinates.time
+            )
+
+            # 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,
+                ),
+            )
+
+            num_orbits = len(propagated_orbits_barycentric.orbit_id.unique())
+
+            observer_codes = np.tile(
+                observers.code.to_numpy(zero_copy_only=False), num_orbits
+            )
+
+            propagated_orbits_aberrated, light_time = self._add_light_time(
+                propagated_orbits_barycentric,
+                observers_barycentric,
+                lt_tol=lt_tol,
+            )
+
+            topocentric_coordinates = CartesianCoordinates.from_kwargs(
+                x=propagated_orbits_aberrated.coordinates.values[:, 0]
+                - observers_barycentric.coordinates.values[:, 0],
+                y=propagated_orbits_aberrated.coordinates.values[:, 1]
+                - observers_barycentric.coordinates.values[:, 1],
+                z=propagated_orbits_aberrated.coordinates.values[:, 2]
+                - observers_barycentric.coordinates.values[:, 2],
+                vx=propagated_orbits_aberrated.coordinates.values[:, 3]
+                - observers_barycentric.coordinates.values[:, 3],
+                vy=propagated_orbits_aberrated.coordinates.values[:, 4]
+                - observers_barycentric.coordinates.values[:, 4],
+                vz=propagated_orbits_aberrated.coordinates.values[:, 5]
+                - observers_barycentric.coordinates.values[:, 5],
+                covariance=None,
+                time=propagated_orbits_aberrated.coordinates.time,
+                origin=Origin.from_kwargs(code=observer_codes),
+                frame="ecliptic",
+            )
+
+            spherical_coordinates = SphericalCoordinates.from_cartesian(
+                topocentric_coordinates
+            )
+            light_time = np.array(light_time)
+
+            spherical_coordinates = transform_coordinates(
+                spherical_coordinates, SphericalCoordinates, frame_out="equatorial"
+            )
+
+            if isinstance(orbits, Orbits):
+
+                ephemeris = Ephemeris.from_kwargs(
+                    orbit_id=propagated_orbits_barycentric.orbit_id,
+                    object_id=propagated_orbits_barycentric.object_id,
+                    coordinates=spherical_coordinates,
+                    light_time=light_time,
+                    aberrated_coordinates=propagated_orbits_aberrated.coordinates,
+                )
+
+            elif isinstance(orbits, VariantOrbits):
+                weights = orbits.weights
+                weights_cov = orbits.weights_cov
+
+                ephemeris = VariantEphemeris.from_kwargs(
+                    orbit_id=propagated_orbits_barycentric.orbit_id,
+                    object_id=propagated_orbits_barycentric.object_id,
+                    coordinates=spherical_coordinates,
+                    weights=weights,
+                    weights_cov=weights_cov,
+                )
+
+            ephemeris_total = qv.concatenate([ephemeris_total, ephemeris])
+
+        return ephemeris_total
 
     def generate_ephemeris(
         self,
@@ -261,7 +403,7 @@ def generate_ephemeris(
         )
 
 
-class Propagator(ABC):
+class Propagator(ABC, EphemerisMixin):
     """
     Abstract class for propagating orbits and related functions.