Merge pull request #6 from B612-Asteroid-Institute/jm/object-id-fix

Small fixes
B612-Asteroid-Institute · May 24, 2024 · 6961c73 · 6961c73
2 parents cf8b386 + 14bee2d
commit 6961c73
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/adam_core/propagator/adam_assist.py b/src/adam_core/propagator/adam_assist.py
@@ -12,6 +12,7 @@
 import quivr as qv
 import rebound
 import urllib3
+from adam_core.constants import KM_P_AU, Constants
 from adam_core.coordinates import CartesianCoordinates, Origin, transform_coordinates
 from adam_core.coordinates.origin import OriginCodes
 from adam_core.dynamics.impacts import EarthImpacts, ImpactMixin
@@ -30,7 +31,9 @@
 
 DATA_DIR = os.getenv("ASSIST_DATA_DIR", "~/.adam_assist_data")
 
-EARTH_RADIUS_KM = 6371.0
+# Use the Earth's equatorial radius as used in DE4XX ephemerides
+# adam_core defines it in au but we need it in km
+EARTH_RADIUS_KM = Constants.R_EARTH * KM_P_AU
 
 
 def download_jpl_ephemeris_files(data_dir: str = DATA_DIR) -> None:
@@ -194,6 +197,7 @@ def _propagate_orbits(self, orbits: OrbitType, times: TimestampType) -> OrbitTyp
                         frame="equatorial",
                     ),
                     orbit_id=orbit_ids,
+                    object_id=orbits.object_id,
                 )
             elif isinstance(orbits, VariantOrbits):
                 # Retrieve the orbit id and weights from hash
@@ -368,6 +372,7 @@ def _detect_impacts(
                         frame="equatorial",
                     ),
                     orbit_id=orbit_ids,
+                    object_id=orbits.object_id,
                 )
             elif isinstance(orbits, VariantOrbits):
                 # Retrieve the orbit id and weights from hash
@@ -442,7 +447,8 @@ def _detect_impacts(
             diff = time_step_results.coordinates.values - earth_geo.values
 
             # Calculate the distance in KM
-            normalized_distance = np.linalg.norm(diff[:, :3], axis=1) * 149597870.691
+            # We use the IAU definition of the astronomical unit (149_597_870.7 km)
+            normalized_distance = np.linalg.norm(diff[:, :3], axis=1) * KM_P_AU
 
             # Calculate which particles are within an Earth radius
             within_radius = normalized_distance < EARTH_RADIUS_KM