From ae44bf49bdbe976c714c60c98bd5f54392864509 Mon Sep 17 00:00:00 2001 From: Joachim Moeyens Date: Tue, 17 Oct 2023 09:09:28 -0700 Subject: [PATCH] Return Ephemeris from generate_ephemeris functions (#81) * Return Ephemeris (not qv.MultikeyLinkage) from ephemeris generation function * Add Ephemeris.link_to_observers --- adam_core/dynamics/ephemeris.py | 21 +---- adam_core/dynamics/tests/test_ephemeris.py | 3 - adam_core/orbits/ephemeris.py | 80 +++++++++++++++++++ adam_core/orbits/tests/test_ephemeris.py | 90 ++++++++++++++++++++++ adam_core/propagator/propagator.py | 19 +++-- adam_core/propagator/pyoorb.py | 41 ++-------- 6 files changed, 192 insertions(+), 62 deletions(-) create mode 100644 adam_core/orbits/tests/test_ephemeris.py diff --git a/adam_core/dynamics/ephemeris.py b/adam_core/dynamics/ephemeris.py index 8203f5dd..710e35f7 100644 --- a/adam_core/dynamics/ephemeris.py +++ b/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 diff --git a/adam_core/dynamics/tests/test_ephemeris.py b/adam_core/dynamics/tests/test_ephemeris.py index 43657cb4..152e4d79 100644 --- a/adam_core/dynamics/tests/test_ephemeris.py +++ b/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] diff --git a/adam_core/orbits/ephemeris.py b/adam_core/orbits/ephemeris.py index 705ddf92..8c3dc2d7 100644 --- a/adam_core/orbits/ephemeris.py +++ b/adam_core/orbits/ephemeris.py @@ -1,7 +1,11 @@ +import warnings + +import pyarrow as pa import quivr as qv from ..coordinates.cartesian import CartesianCoordinates from ..coordinates.spherical import SphericalCoordinates +from ..observers.observers import Observers class Ephemeris(qv.Table): @@ -18,3 +22,79 @@ 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, + } + linkage = qv.MultiKeyLinkage(self, observers, left_keys, right_keys) + + # Check to make sure we have the correct number of linkages by calculating the number + # of unique observers in the table and checking to see if that matches the number of + # unique keys in the linkage + observers_table = pa.table( + [ + observers.code, + observers.coordinates.time.days, + observers.coordinates.time.nanos, + ], + names=["observatory_code", "days", "nanos"], + ) + unique_observers = observers_table.group_by( + ["observatory_code", "days", "nanos"] + ).aggregate([]) + + expected_length = len(unique_observers) + actual_length = len(linkage.all_unique_values) + if expected_length != actual_length: + warnings.warn( + "The number of unique keys in the linkage does not match the number" + " of unique observers. Linkage precision may be too low." + f"Expected {expected_length} unique keys, got {actual_length}.", + UserWarning, + ) + return linkage diff --git a/adam_core/orbits/tests/test_ephemeris.py b/adam_core/orbits/tests/test_ephemeris.py new file mode 100644 index 00000000..3f899794 --- /dev/null +++ b/adam_core/orbits/tests/test_ephemeris.py @@ -0,0 +1,90 @@ +import pyarrow as pa +import pytest +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 + with pytest.warns(UserWarning): + 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 + with pytest.warns(UserWarning): + 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 diff --git a/adam_core/propagator/propagator.py b/adam_core/propagator/propagator.py index 0bd7b36a..51c81c9f 100644 --- a/adam_core/propagator/propagator.py +++ b/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", + ] + ) diff --git a/adam_core/propagator/pyoorb.py b/adam_core/propagator/pyoorb.py index 24316851..e0fe1b68 100644 --- a/adam_core/propagator/pyoorb.py +++ b/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)