Impact Probability Fixes (#10)

Authored-by: Kathleen Kiker <[email protected]>

src/adam_core/propagator/adam_assist.py

@@ -226,7 +226,6 @@ def _propagate_orbits_inner(
                     object_id=orbits.object_id,
                 )
             elif isinstance(orbits, VariantOrbits):
-                # Retrieve the orbit id and weights from hash
                 # Retrieve the orbit id and weights from hash
                 particle_ids = [orbit_id_mapping[h] for h in orbit_id_hashes]
                 orbit_ids, variant_ids = zip(
@@ -352,8 +351,9 @@ def _detect_impacts(
         time_step_results: Union[None, OrbitType] = None
 
         # Step through each time, move the simulation forward and
-        # collect the results.
-        while past_integrator_time is False:
+        # collect the results. End if all orbits are removed from
+        # the simulation or the final integrator time is reached.
+        while past_integrator_time is False and len(orbits) > 0:
             sim.steps(1)
             # print(sim.dt_last_done)
             if sim.t >= final_integrator_time:
@@ -507,11 +507,22 @@ def _detect_impacts(
                 else:
                     results = qv.concatenate([results, impacting_orbits])
 
-        # Add the final positions of the particles to the results
-        if results is None:
-            results = time_step_results
-        else:
-            results = qv.concatenate([results, time_step_results])
+        # Add the final positions of the particles that are not already in the results
+        if time_step_results is not None:
+            if results is None:
+                results = time_step_results
+            else:
+                if isinstance(orbits, Orbits):
+                    still_in_simulation = pc.invert(
+                        pc.is_in(time_step_results.orbit_id, results.orbit_id)
+                    )
+                elif isinstance(orbits, VariantOrbits):
+                    still_in_simulation = pc.invert(
+                        pc.is_in(time_step_results.variant_id, results.variant_id)
+                    )
+                results = qv.concatenate(
+                    [results, time_step_results.apply_mask(still_in_simulation)]
+                )
 
         if earth_impacts is None:
             earth_impacts = EarthImpacts.from_kwargs(

tests/test_impacts.py

@@ -8,14 +8,18 @@
 )
 
 # Contains a likely impactor with ~60% chance of impact in 30 days
-IMPACTOR_FILE_PATH = "tests/data/I00007_orbit.parquet"
+IMPACTOR_FILE_PATH_60 = "tests/data/I00007_orbit.parquet"
+# Contains a likely impactor with 100% chance of impact in 30 days
+IMPACTOR_FILE_PATH_100 = "tests/data/I00008_orbit.parquet"
+# Contains a likely impactor with 0% chance of impact in 30 days
+IMPACTOR_FILE_PATH_0 = "tests/data/I00009_orbit.parquet"
 
 
 @pytest.mark.benchmark
 @pytest.mark.parametrize("processes", [1, 2])
 def test_calculate_impacts_benchmark(benchmark, processes):
     download_jpl_ephemeris_files()
-    impactor = Orbits.from_parquet(IMPACTOR_FILE_PATH)[0]
+    impactor = Orbits.from_parquet(IMPACTOR_FILE_PATH_60)[0]
     propagator = ASSISTPropagator()
     variants, impacts = benchmark(
         calculate_impacts,
@@ -28,3 +32,41 @@ 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.benchmark
+@pytest.mark.parametrize("processes", [1, 2])
+def test_calculate_impacts_benchmark(benchmark, processes):
+    download_jpl_ephemeris_files()
+    impactor = Orbits.from_parquet(IMPACTOR_FILE_PATH_100)[0]
+    propagator = ASSISTPropagator()
+    variants, impacts = benchmark(
+        calculate_impacts,
+        impactor,
+        60,
+        propagator,
+        num_samples=200,
+        processes=processes,
+        seed=42  # This allows us to predict exact number of impactors empirically
+    )
+    assert len(variants) == 200, "Should have 200 variants"
+    assert len(impacts) == 200, "Should have exactly 200 impactors"
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize("processes", [1, 2])
+def test_calculate_impacts_benchmark(benchmark, processes):
+    download_jpl_ephemeris_files()
+    impactor = Orbits.from_parquet(IMPACTOR_FILE_PATH_0)[0]
+    propagator = ASSISTPropagator()
+    variants, impacts = benchmark(
+        calculate_impacts,
+        impactor,
+        60,
+        propagator,
+        num_samples=200,
+        processes=processes,
+        seed=42  # This allows us to predict exact number of impactors empirically
+    )
+    assert len(variants) == 200, "Should have 200 variants"
+    assert len(impacts) == 0, "Should have exactly 0 impactors"