Implement impact_coordinates (#19)

B612-Asteroid-Institute · Feb 18, 2025 · f31bd91 · f31bd91
1 parent 7bf3409
commit f31bd91
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 10 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -19,7 +19,7 @@ classifiers = [
 ]
 requires-python = ">=3.11,<3.13"
 dependencies = [
-    "adam-core>=0.3.6",
+    "adam-core>=0.3.7",
     "naif-de440",
     "jpl-small-bodies-de441-n16",
     # Temporary until SPK support is merged and released upstream

diff --git a/src/adam_assist/propagator.py b/src/adam_assist/propagator.py
@@ -11,8 +11,13 @@
 import rebound
 from adam_core.constants import KM_P_AU
 from adam_core.constants import Constants as c
-from adam_core.coordinates import CartesianCoordinates, Origin, transform_coordinates
-from adam_core.coordinates.origin import OriginCodes
+from adam_core.coordinates import (
+    CartesianCoordinates,
+    Origin,
+    OriginCodes,
+    SphericalCoordinates,
+    transform_coordinates,
+)
 from adam_core.dynamics.impacts import EarthImpacts, ImpactMixin
 from adam_core.orbits import Orbits
 from adam_core.orbits.variants import VariantOrbits
@@ -467,21 +472,30 @@ def _detect_impacts(
             # If any are within our earth radius, we record the impact
             # and do bookkeeping to remove the particle from the simulation
             if np.any(within_radius):
-                distances = normalized_distance[within_radius]
                 impacting_orbits = time_step_results.apply_mask(within_radius)
 
                 if isinstance(orbits, VariantOrbits):
                     new_impacts = EarthImpacts.from_kwargs(
                         orbit_id=impacting_orbits.orbit_id,
-                        distance=distances,
-                        coordinates=impacting_orbits.coordinates,
                         variant_id=impacting_orbits.variant_id,
+                        coordinates=impacting_orbits.coordinates,
+                        impact_coordinates=transform_coordinates(
+                            impacting_orbits.coordinates,
+                            representation_out=SphericalCoordinates,
+                            origin_out=OriginCodes.EARTH,
+                            frame_out="itrf93",
+                        ),
                     )
                 elif isinstance(orbits, Orbits):
                     new_impacts = EarthImpacts.from_kwargs(
                         orbit_id=impacting_orbits.orbit_id,
-                        distance=distances,
                         coordinates=impacting_orbits.coordinates,
+                        impact_coordinates=transform_coordinates(
+                            impacting_orbits.coordinates,
+                            representation_out=SphericalCoordinates,
+                            origin_out=OriginCodes.EARTH,
+                            frame_out="itrf93",
+                        ),
                     )
                 if earth_impacts is None:
                     earth_impacts = new_impacts
@@ -522,7 +536,7 @@ def _detect_impacts(
         if earth_impacts is None:
             earth_impacts = EarthImpacts.from_kwargs(
                 orbit_id=[],
-                distance=[],
+                variant_id=[],
                 coordinates=CartesianCoordinates.from_kwargs(
                     x=[],
                     y=[],
@@ -536,7 +550,19 @@ def _detect_impacts(
                     ),
                     frame="equatorial",
                 ),
-                variant_id=[],
+                impact_coordinates=CartesianCoordinates.from_kwargs(
+                    x=[],
+                    y=[],
+                    z=[],
+                    vx=[],
+                    vy=[],
+                    vz=[],
+                    time=Timestamp.from_jd([], scale="tdb"),
+                    origin=Origin.from_kwargs(
+                        code=[],
+                    ),
+                    frame="itrf93",
+                ),
             )
 
         # Store the last simulation in a private variable for reference

diff --git a/tests/test_impacts.py b/tests/test_impacts.py
@@ -1,6 +1,10 @@
 import pytest
+import pyarrow.compute as pc
+
+from adam_core.constants import Constants as c
+from adam_core.constants import KM_P_AU
 from adam_core.dynamics.impacts import calculate_impacts
-from adam_core.orbits import Orbits
+from adam_core.orbits import Orbits, VariantOrbits
 from adam_core.orbits.query.horizons import query_horizons
 from adam_core.time import Timestamp
 
@@ -13,6 +17,8 @@
 # Contains a likely impactor with 0% chance of impact in 30 days
 IMPACTOR_FILE_PATH_0 = "tests/data/I00009_orbit.parquet"
 
+R_EARTH_KM = c.R_EARTH_EQUATORIAL * KM_P_AU
+
 
 @pytest.mark.benchmark
 @pytest.mark.parametrize("processes", [1, 2])
@@ -31,6 +37,24 @@ def test_calculate_impacts_benchmark(benchmark, processes):
     assert len(variants) == 200, "Should have 200 variants"
     assert len(impacts) == 138, "Should have exactly 138 impactors"
 
+@pytest.mark.parametrize("processes", [1, 2])
+def test_calculate_impacts(processes):
+    impactor = Orbits.from_parquet(IMPACTOR_FILE_PATH_60)[0]
+    propagator = ASSISTPropagator()
+
+    variants_orbits = VariantOrbits.create(impactor, method="monte-carlo", num_samples=200, seed=42)
+    variants, impacts = propagator.detect_impacts(
+        variants_orbits,
+        60,
+        max_processes=processes,
+    )
+    assert len(variants) == 200, "Should have 200 variants"
+    assert len(impacts) == 138, "Should have exactly 138 impactors"
+
+    assert impacts.impact_coordinates.frame == "itrf93"
+    assert pc.all(pc.equal(impacts.impact_coordinates.origin.code, "EARTH")).as_py()
+    assert pc.all(pc.less_equal(impacts.impact_coordinates.rho, R_EARTH_KM)).as_py()
+
 
 @pytest.mark.benchmark
 @pytest.mark.parametrize("processes", [1, 2])