Return Ephemeris from generate_ephemeris functions

adam_core/dynamics/ephemeris.py

@@ -2,7 +2,6 @@
 
 import jax.numpy as jnp
 import numpy as np
-import quivr as qv
 from jax import jit, lax, vmap
 
 from ..coordinates.cartesian import CartesianCoordinates
@@ -129,7 +128,7 @@ def generate_ephemeris_2body(
     max_iter: int = 1000,
     tol: float = 1e-15,
     stellar_aberration: bool = False,
-) -> qv.MultiKeyLinkage[Ephemeris, Observers]:
+) -> Ephemeris:
     """
     Generate on-sky ephemerides for each propagated orbit as viewed by the observers.
     This function calculates the light time delay between the propagated orbit and the observer,
@@ -174,7 +173,7 @@ def generate_ephemeris_2body(
 
     Returns
     -------
-    ephemeris : `~quivr.linkage.MultikeyLinkage` (N)
+    ephemeris : `~adam_core.orbits.ephemeris.Ephemeris` (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.
@@ -263,23 +262,9 @@ def generate_ephemeris_2body(
         spherical_coordinates, SphericalCoordinates, frame_out="equatorial"
     )
 
-    ephems = Ephemeris.from_kwargs(
+    return 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

adam_core/dynamics/tests/test_ephemeris.py

@@ -93,9 +93,6 @@ def test_generate_ephemeris_2body(object_id, propagated_orbits, ephemeris):
 
     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]

adam_core/orbits/ephemeris.py

@@ -2,6 +2,7 @@
 
 from ..coordinates.cartesian import CartesianCoordinates
 from ..coordinates.spherical import SphericalCoordinates
+from ..observers.observers import Observers
 
 
 class Ephemeris(qv.Table):
@@ -18,3 +19,53 @@ class Ephemeris(qv.Table):
     # from the time of emission/reflection to the time of observation
     light_time = qv.Float64Column(nullable=True)
     aberrated_coordinates = CartesianCoordinates.as_column(nullable=True)
+
+    def link_to_observers(
+        self, observers: Observers, precision="ns"
+    ) -> qv.MultiKeyLinkage["Ephemeris", Observers]:
+        """
+        Link these ephemerides back to the observers that generated them. This is useful if
+        you want or need to use the observer's position as part of any computation for
+        any given set of ephemerides.
+
+        Not all propagators will return ephemerides exactly at the time of the input observers.
+        As an example, PYOORB stores times as a single MJD, when converting from two integers to
+        this singular float there will be a loss of precision. To mitigate this, the user may
+        optionally define the precision to which would like to link back to observers. Times
+        for both the ephemerides and observers will be rounded to this precision before linking.
+
+        Parameters
+        ----------
+        observers : `~adam_core.observers.observers.Observers` (N)
+            Observers that generated the ephemerides.
+        precision : str, optional
+            Precision to which to link back to observers, by default "ns".
+
+        Returns
+        -------
+        `~qv.MultiKeyLinkage[
+                `~adam_core.orbits.ephemeris.Ephemeris`,
+                `~adam_core.observers.observersObservers
+        ]`
+            Linkage between ephemerides and observers.
+        """
+        if self.coordinates.time.scale != observers.coordinates.time.scale:
+            observers = observers.set_column(
+                "coordinates.time",
+                observers.coordinates.time.rescale(self.coordinates.time.scale),
+            )
+
+        rounded = self.coordinates.time.rounded(precision)
+        observers_rounded = observers.coordinates.time.rounded(precision)
+
+        left_keys = {
+            "days": rounded.days,
+            "nanos": rounded.nanos,
+            "observatory_code": self.coordinates.origin.code,
+        }
+        right_keys = {
+            "days": observers_rounded.days,
+            "nanos": observers_rounded.nanos,
+            "observatory_code": observers.code,
+        }
+        return qv.MultiKeyLinkage(self, observers, left_keys, right_keys)

adam_core/orbits/tests/test_ephemeris.py

@@ -0,0 +1,86 @@
+import pyarrow as pa
+import quivr as qv
+
+from ...coordinates.origin import Origin
+from ...coordinates.spherical import SphericalCoordinates
+from ...observers.observers import Observers
+from ...time import Timestamp
+from ..ephemeris import Ephemeris
+
+
+def test_Ephemeris_link_to_observers():
+    # Test that we can link ephemerides to observers with different
+    # precisions
+    observer_01 = Observers.from_code(
+        "X05",
+        Timestamp.from_kwargs(
+            days=[59000, 59000, 59000],
+            nanos=[9, 999, 999_999],
+            scale="utc",
+        ),
+    )
+    observer_02 = Observers.from_code(
+        "500",
+        Timestamp.from_kwargs(
+            days=[59000, 59000, 59000],
+            nanos=[9, 999, 999_999],
+            scale="utc",
+        ),
+    )
+    observers = qv.concatenate([observer_01, observer_02])
+
+    ephemeris = Ephemeris.from_kwargs(
+        orbit_id=pa.array(["00000" for i in range(6)]),
+        object_id=pa.array(["00000" for i in range(6)]),
+        coordinates=SphericalCoordinates.from_kwargs(
+            time=Timestamp.from_kwargs(
+                days=[59000, 59000, 59000, 59000, 59000, 59000],
+                nanos=[9, 999, 999_999, 9, 999, 999_999],
+                scale="utc",
+            ),
+            lon=[0, 0, 0, 0, 0, 0],
+            lat=[0, 0, 0, 0, 0, 0],
+            frame="equatorial",
+            origin=Origin.from_kwargs(
+                code=pa.array(["X05", "X05", "X05", "500", "500", "500"])
+            ),
+        ),
+    )
+
+    linkage = ephemeris.link_to_observers(observers, precision="ns")
+    assert len(linkage.all_unique_values) == 6
+    e1, o1 = linkage.select((59000, 9, "X05"))
+    assert len(e1) == len(o1) == 1
+    e2, o2 = linkage.select((59000, 999, "X05"))
+    assert len(e2) == len(o2) == 1
+    e3, o3 = linkage.select((59000, 999_999, "X05"))
+    assert len(e3) == len(o3) == 1
+    e4, o4 = linkage.select((59000, 9, "500"))
+    assert len(e4) == len(o4) == 1
+    e5, o5 = linkage.select((59000, 999, "500"))
+    assert len(e5) == len(o5) == 1
+    e6, o6 = linkage.select((59000, 999_999, "500"))
+    assert len(e6) == len(o6) == 1
+
+    # Reduce precision to microseconds
+    linkage = ephemeris.link_to_observers(observers, precision="us")
+
+    # First two times should be grouped together
+    # Last 2 times should be round-down to the previous microsecond
+    assert len(linkage.all_unique_values) == 4
+    e1, o1 = linkage.select((59000, 0, "X05"))
+    assert len(e1) == len(o1) == 2
+    e2, o2 = linkage.select((59000, 0, "500"))
+    assert len(e2) == len(o2) == 2
+    e3, o3 = linkage.select((59000, 999_000, "X05"))
+    assert len(e3) == len(o3) == 1
+    e4, o4 = linkage.select((59000, 999_000, "500"))
+    assert len(e4) == len(o4) == 1
+
+    # Reduce precision to milliseconds
+    linkage = ephemeris.link_to_observers(observers, precision="ms")
+    assert len(linkage.all_unique_values) == 2
+    e1, o1 = linkage.select((59000, 0, "X05"))
+    assert len(e1) == len(o1) == 3
+    e2, o2 = linkage.select((59000, 0, "500"))
+    assert len(e2) == len(o2) == 3

adam_core/propagator/propagator.py

@@ -26,7 +26,7 @@ def propagation_worker(
 
 def ephemeris_worker(
     orbits: Orbits, observers: Observers, propagator: "Propagator"
-) -> qv.MultiKeyLinkage[Ephemeris, Observers]:
+) -> Ephemeris:
     ephemeris = propagator._generate_ephemeris(orbits, observers)
     return ephemeris
 
@@ -165,7 +165,7 @@ def generate_ephemeris(
         observers: Observers,
         chunk_size: int = 100,
         max_processes: Optional[int] = 1,
-    ) -> qv.MultiKeyLinkage[Ephemeris, Observers]:
+    ) -> Ephemeris:
         """
         Generate ephemerides for each orbit in orbits as observed by each observer
         in observers.
@@ -186,7 +186,7 @@ def generate_ephemeris(
 
         Returns
         -------
-        ephemeris : List[`~quivr.linkage.MultiKeyLinkage`] (M)
+        ephemeris : `~adam_core.orbits.ephemeris.Ephemeris` (M)
             Predicted ephemerides for each orbit observed by each
             observer.
         """
@@ -204,7 +204,16 @@ def generate_ephemeris(
                 for future in concurrent.futures.as_completed(futures):
                     ephemeris_list.append(future.result())
 
-            return qv.combine_multilinkages(ephemeris_list)
+            ephemeris = qv.concatenate(ephemeris_list)
 
         else:
-            return self._generate_ephemeris(orbits, observers)
+            ephemeris = self._generate_ephemeris(orbits, observers)
+
+        return ephemeris.sort_by(
+            [
+                "orbit_id",
+                "coordinates.time.days",
+                "coordinates.time.nanos",
+                "coordinates.origin.code",
+            ]
+        )

adam_core/propagator/pyoorb.py

@@ -319,22 +319,20 @@ def _propagate_orbits(self, orbits: OrbitType, times: Timestamp) -> OrbitType:
 
         return propagated_orbits
 
-    def _generate_ephemeris(
-        self, orbits: Orbits, observers: Observers
-    ) -> qv.MultiKeyLinkage[Ephemeris, Observers]:
+    def _generate_ephemeris(self, orbits: OrbitType, observers: Observers) -> Ephemeris:
         """
         Generate ephemerides for orbits as viewed from observers using PYOORB.
 
         Parameters
         ----------
-        orbits : `~adam_core.orbits.orbits.Orbits` (N)
-            Orbits to generate ephemerides for.
+        orbits : {`~adam_core.orbits.orbits.Orbits`, `~adam_core.orbits.orbits.VariantOrbits`} (N)
+            Orbits to propagate.
         observers : `~adam_core.observers.observers.Observers` (M)
             Observers to generate ephemerides for.
 
         Returns
         -------
-        ephemeris : `~quivr.linkage.MultikeyLinkage` (M)
+        ephemeris : `~adam_core.orbits.ephemeris.Ephemeris` (M)
             Ephemerides for each orbit as viewed from each observer.
         """
         # Convert orbits into PYOORB format
@@ -480,33 +478,4 @@ def _generate_ephemeris(
             )
             ephemeris_list.append(ephemeris)
 
-        # Concatenate ephemeris into single table
-        ephemeris = qv.concatenate(ephemeris_list)
-
-        # We use mjd as the time linking key because openorb uses
-        # mjd. If we try to use coordinates.time.jd1 and jd2 directly,
-        # we will get loss-of-precision issues, and the keys won't
-        # align.
-        linkages = qv.MultiKeyLinkage(
-            left_table=ephemeris,
-            right_table=observers,
-            left_keys={
-                "code": ephemeris.coordinates.origin.code,
-                "mjd": ephemeris.coordinates.time.mjd(),
-            },
-            right_keys={
-                "code": observers_utc.code,
-                "mjd": observers_utc.coordinates.time.mjd(),
-            },
-        )
-
-        # Check to make sure we have the correct number of linkages
-        expected_length = len(observers)
-        actual_length = len(linkages.all_unique_values)
-        if expected_length != actual_length:
-            raise ValueError(
-                "Ephemerides were not correctly linked to observers. "
-                f"Expected {expected_length} unique keys, got {actual_length}."
-            )
-
-        return linkages
+        return qv.concatenate(ephemeris_list)