Fixes bug where the aberrated coordinates were misaligned with epheme…

…ris produced (#132)

src/adam_core/propagator/propagator.py

@@ -127,13 +127,24 @@ def _add_light_time(
 
                 # Calculate the topocentric distance
                 rho = np.linalg.norm(orbit_i.coordinates.r - observer_position)
+                # If rho becomes too large, we are probably simulating a close encounter
+                # and our propagation will break
+                if np.isnan(rho) or rho > 1e12:
+                    raise ValueError(
+                        "Distance from observer is NaN or too large and propagation will break."
+                    )
 
                 # Calculate the light travel time
                 lt = rho / C
 
                 # Calculate the change in light travel time since the previous iteration
                 dlt = np.abs(lt - lt_prev)
 
+                if np.isnan(lt) or lt > 1e12:
+                    raise ValueError(
+                        "Light travel time is NaN or too large and propagation will break."
+                    )
+
                 # Calculate the new epoch and propagate the initial orbit to that epoch
                 # Should be sufficient to use 2body propagation for this
                 orbit_i = propagate_2body(
@@ -160,7 +171,24 @@ def _generate_ephemeris(
         elif isinstance(orbits, VariantOrbits):
             ephemeris_total = VariantEphemeris.empty()
 
+        # Sort observers by time and code to ensure consistent ordering
+        # As further propagation will order by time as well
+        observers = observers.sort_by(
+            ["coordinates.time.days", "coordinates.time.nanos", "code"]
+        )
+
+        observers_barycentric = observers.set_column(
+            "coordinates",
+            transform_coordinates(
+                observers.coordinates,
+                CartesianCoordinates,
+                frame_out="ecliptic",
+                origin_out=OriginCodes.SOLAR_SYSTEM_BARYCENTER,
+            ),
+        )
+
         for orbit in orbits:
+            # Propagate orbits to sorted observer times
             propagated_orbits = self.propagate_orbits(orbit, observers.coordinates.time)
 
             # Transform both the orbits and observers to the barycenter if they are not already.
@@ -173,15 +201,7 @@ def _generate_ephemeris(
                     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(
@@ -248,6 +268,15 @@ def _generate_ephemeris(
 
             ephemeris_total = qv.concatenate([ephemeris_total, ephemeris])
 
+        ephemeris_total = ephemeris_total.sort_by(
+            [
+                "orbit_id",
+                "coordinates.time.days",
+                "coordinates.time.nanos",
+                "coordinates.origin.code",
+            ]
+        )
+
         return ephemeris_total
 
     def generate_ephemeris(

src/adam_core/propagator/tests/test_propagator.py

@@ -1,6 +1,9 @@
 import numpy as np
 import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
 import quivr as qv
+from adam_assist import ASSISTPropagator
 
 from ...coordinates.cartesian import CartesianCoordinates
 from ...coordinates.origin import Origin, OriginCodes
@@ -58,6 +61,8 @@ def _generate_ephemeris(self, orbits: Orbits, observers: Observers) -> Ephemeris
                 orbit_id=pa.array(np.full(len(coords), orbit.orbit_id[0].as_py())),
                 object_id=pa.array(np.full(len(coords), orbit.object_id[0].as_py())),
                 coordinates=coords.to_spherical(),
+                aberrated_coordinates=coords,
+                light_time=np.full(len(coords), 0.0),
             )
             ephemeris_list.append(ephemeris_i)
 
@@ -157,3 +162,129 @@ def test_propagate_different_origins():
         orbit_two_results.coordinates.origin.code.unique()[0].as_py()
         == "EARTH_MOON_BARYCENTER"
     )
+
+
+def test_light_time_distance_threshold():
+    """
+    Test that _add_light_time raises a ValueError when an object gets too far from the observer.
+    """
+    # Create a single orbit with very high velocity
+    orbit = Orbits.from_kwargs(
+        orbit_id=["1"],
+        object_id=["1"],
+        coordinates=CartesianCoordinates.from_kwargs(
+            x=[1],  # Start at origin
+            y=[1],
+            z=[1],
+            vx=[1e9],  # Very high velocity (will quickly exceed our limits)
+            vy=[0],
+            vz=[0],
+            time=Timestamp.from_mjd([60000], scale="tdb"),
+            frame="ecliptic",
+            origin=Origin.from_kwargs(code=["SOLAR_SYSTEM_BARYCENTER"]),
+        ),
+    )
+
+    # Create an observer at a fixed position
+    observer = Observers.from_kwargs(
+        code=["500"],  # Arbitrary observer code
+        coordinates=CartesianCoordinates.from_kwargs(
+            x=[0],
+            y=[0],
+            z=[0],
+            vx=[0],
+            vy=[0],
+            vz=[0],
+            time=Timestamp.from_mjd([60020], scale="tdb"),  # 1 day later
+            frame="ecliptic",
+            origin=Origin.from_kwargs(code=["SOLAR_SYSTEM_BARYCENTER"]),
+        ),
+    )
+
+    prop = MockPropagator()
+
+    # The object will have moved ~86400 * 1e6 AU in one day, which should trigger the threshold
+    with pytest.raises(
+        ValueError,
+        match="Distance from observer is NaN or too large and propagation will break.",
+    ):
+        prop._add_light_time(orbit, observer)
+
+
+@pytest.mark.parametrize("max_processes", [1, 4])
+def test_generate_ephemeris_unordered_observers(max_processes):
+    """
+    Test that ephemeris generation works correctly even when observers
+    are not ordered by time, verifying that physical positions don't get mixed up
+    between different objects.
+    """
+    # Create two distinct orbits - one near Earth and one near Mars
+    orbits = Orbits.from_kwargs(
+        orbit_id=["far", "near"],
+        object_id=["far", "near"],
+        coordinates=CartesianCoordinates.from_kwargs(
+            x=[3.0, 1],  # Pick something far from both observers (earth) and sun
+            y=[3.0, 1],
+            z=[3.0, 1],
+            vx=[0.1, 0.1],
+            vy=[0.0, 0.0],  # Different orbital velocities
+            vz=[0.0, 0.0],
+            time=Timestamp.from_mjd([59000, 59000], scale="tdb"),
+            origin=Origin.from_kwargs(code=["SUN", "SUN"]),
+            frame="ecliptic",
+        ),
+    )
+    # Create observers with deliberately unordered times
+    times = Timestamp.from_mjd([59005, 59004, 59005, 59004, 59006, 59006], scale="utc")
+    codes = ["500", "500", "X05", "X05", "X05", "500"]  # Mix of observatory codes
+    observers = Observers.from_codes(codes, times)
+
+    propagator = ASSISTPropagator()
+    ephemeris = propagator.generate_ephemeris(
+        orbits, observers, max_processes=max_processes
+    )
+
+    # Basic ordering checks
+    assert len(ephemeris) == 12  # 2 objects × 6 times
+
+    # Verify that coordinates.time - aberrated_coordinates.time is equal to light_time
+    time_difference_days, time_difference_nanos = ephemeris.coordinates.time.rescale(
+        "tdb"
+    ).difference(ephemeris.aberrated_coordinates.time)
+    fractional_days = pc.divide(time_difference_nanos, 86400 * 1e9)
+    time_difference = pc.add(time_difference_days, fractional_days)
+    np.testing.assert_allclose(
+        time_difference.to_numpy(zero_copy_only=False),
+        ephemeris.light_time.to_numpy(zero_copy_only=False),
+        atol=1e-6,
+    )
+
+    # Verify that the near-Earth object is consistently closer than the Mars object
+    far_ephem = ephemeris.select("orbit_id", "far")
+    near_ephem = ephemeris.select("orbit_id", "near")
+
+    # make sure the far object is always further than the near object
+    far_positions = np.linalg.norm(
+        far_ephem.aberrated_coordinates.values[:, :3], axis=1
+    )
+    near_positions = np.linalg.norm(
+        near_ephem.aberrated_coordinates.values[:, :3], axis=1
+    )
+    assert np.all(
+        far_positions > near_positions
+    ), "Far aberrated positions should be consistently further than near positions"
+
+    # Verify observer codes match the expected order after sorting for each object
+    sorted_observers = observers.sort_by(
+        ["coordinates.time.days", "coordinates.time.nanos", "code"]
+    )
+    for orbit_id in ["far", "near"]:
+        orbit_ephem = ephemeris.select("orbit_id", orbit_id)
+        np.testing.assert_array_equal(
+            orbit_ephem.coordinates.origin.code.to_numpy(zero_copy_only=False),
+            sorted_observers.code.to_numpy(zero_copy_only=False),
+        )
+
+    # Link back to observers to verify correct correspondence
+    linkage = ephemeris.link_to_observers(observers)
+    assert len(linkage.all_unique_values) == len(observers)