Cap ASSIST dependencies for REBOUND 5 (#32)

Require ASSIST and REBOUND versions that remain compatible with the REBOUND 4 API used by ASSIST.

Also adds dependency metadata tests, satisfies mypy for pyarrow compute calls, and uses a serialized SBDB orbit fixture for the back-to-back propagation test while retaining broader Horizons/SBDB integration coverage.

Author: akoumjian

pyproject.toml

@@ -22,11 +22,11 @@ dependencies = [
     "adam-core>=0.5.4",
     "naif-de440",
     "jpl-small-bodies-de441-n16",
-    "assist==1.2.0",
+    "assist>=1.2.3,<1.3",
     "numpy",
     "ray",
     "spiceypy>=6.0.0",
-    "rebound>=4.4.10",
+    "rebound>=4.4.11,<5",
     "timezonefinder==8.0.0",
 ]
 

src/adam_assist/propagator.py

@@ -40,6 +40,13 @@
 # adam_core defines it in au but we need it in km
 EARTH_RADIUS_KM = c.R_EARTH_EQUATORIAL * KM_P_AU
 
+# pyarrow.compute type stubs are incomplete for the functions used here.
+pc_cast: Any = pc.cast
+pc_invert: Any = pc.invert  # type: ignore[attr-defined]
+pc_is_in: Any = pc.is_in  # type: ignore[attr-defined]
+pc_subtract: Any = pc.subtract  # type: ignore[attr-defined]
+pc_unique: Any = pc.unique  # type: ignore[attr-defined]
+
 
 def uint32_hash(s: str) -> c_uint32:
     sha256_result = hashlib.sha256(s.encode()).digest()
@@ -290,8 +297,8 @@ def _propagate_single_orbit_inner_optimized(
         )
 
         if is_variant:
-            orbit_ids_out = pa.repeat(pc.cast(orbit_id, pa.large_string()), N)
-            variant_ids_out = pa.repeat(pc.cast(variant_id, pa.large_string()), N)
+            orbit_ids_out = pa.repeat(pc_cast(orbit_id, pa.large_string()), N)
+            variant_ids_out = pa.repeat(pc_cast(variant_id, pa.large_string()), N)
             object_id_out = pa.repeat(orbit.object_id[0], N)
             weights_out = pa.repeat(orbit.weights[0], N)
             weights_cov_out = pa.repeat(orbit.weights_cov[0], N)
@@ -319,7 +326,7 @@ def _propagate_single_orbit_inner_optimized(
                 ),
             )
         else:
-            orbit_ids_out = pa.repeat(pc.cast(orbit_id, pa.large_string()), N)
+            orbit_ids_out = pa.repeat(pc_cast(orbit_id, pa.large_string()), N)
             object_id_out = pa.repeat(orbit.object_id[0], N)
             physical_parameters_out = orbit.physical_parameters.take(
                 np.zeros(N, dtype=np.int64)
@@ -408,7 +415,7 @@ def _propagate_orbits_inner(
 
         # Prepare the times as jd - jd_ref
         integrator_times = times.rescale("tdb").jd()
-        integrator_times = pc.subtract(integrator_times, ephem.jd_ref)
+        integrator_times = pc_subtract(integrator_times, ephem.jd_ref)
         integrator_times = integrator_times.to_numpy()
 
         # Unified accumulation for both Orbits and VariantOrbits
@@ -532,7 +539,7 @@ def _detect_collisions(
         conditions: CollisionConditions,
     ) -> Tuple[VariantOrbits, CollisionEvent]:
         # Assert that the time for each orbit definition is the same for the simulator to work
-        assert len(pc.unique(orbits.coordinates.time.mjd())) == 1
+        assert len(pc_unique(orbits.coordinates.time.mjd())) == 1
 
         # The coordinate frame is the equatorial International Celestial Reference Frame (ICRF).
         # This is also the native coordinate system for the JPL binary files.
@@ -807,12 +814,12 @@ def _detect_collisions(
                             # For some reason, it fails if we let rebound convert the hash to c_uint32
 
                         if isinstance(orbits, VariantOrbits):
-                            keep_mask = pc.invert(
-                                pc.is_in(orbits.variant_id, colliding_orbits.variant_id)
+                            keep_mask = pc_invert(
+                                pc_is_in(orbits.variant_id, colliding_orbits.variant_id)
                             )
                         else:
-                            keep_mask = pc.invert(
-                                pc.is_in(orbits.orbit_id, colliding_orbits.orbit_id)
+                            keep_mask = pc_invert(
+                                pc_is_in(orbits.orbit_id, colliding_orbits.orbit_id)
                             )
 
                         orbits = orbits.apply_mask(keep_mask)
@@ -836,12 +843,12 @@ def _detect_collisions(
                 results_list.append(time_step_results)
             else:
                 if isinstance(orbits, Orbits):
-                    still_in_simulation = pc.invert(
-                        pc.is_in(time_step_results.orbit_id, colliders.orbit_id)
+                    still_in_simulation = pc_invert(
+                        pc_is_in(time_step_results.orbit_id, colliders.orbit_id)
                     )
                 elif isinstance(orbits, VariantOrbits):
-                    still_in_simulation = pc.invert(
-                        pc.is_in(time_step_results.variant_id, colliders.variant_id)
+                    still_in_simulation = pc_invert(
+                        pc_is_in(time_step_results.variant_id, colliders.variant_id)
                     )
                 addl = time_step_results.apply_mask(still_in_simulation)
                 results_list.append(addl)

tests/data/2013_RR165_orbit.parquet

@@ -0,0 +1,17 @@
+import tomllib
+from pathlib import Path
+
+
+def _project_dependencies() -> list[str]:
+    pyproject_path = Path(__file__).resolve().parents[1] / "pyproject.toml"
+    with pyproject_path.open("rb") as pyproject_file:
+        pyproject = tomllib.load(pyproject_file)
+    return list(pyproject["project"]["dependencies"])
+
+
+def test_assist_dependency_uses_rebound_header_compatible_release() -> None:
+    assert "assist>=1.2.3,<1.3" in _project_dependencies()
+
+
+def test_rebound_dependency_stays_on_assist_supported_major_version() -> None:
+    assert "rebound>=4.4.11,<5" in _project_dependencies()

tests/test_dependency_metadata.py

@@ -5,7 +5,6 @@
 from adam_core.coordinates import CartesianCoordinates, Origin
 from adam_core.coordinates.residuals import Residuals
 from adam_core.orbits import Orbits
-from adam_core.orbits.query import query_sbdb
 from adam_core.orbits.query.horizons import query_horizons
 from adam_core.time import Timestamp
 from astropy import units as u
@@ -14,6 +13,7 @@
 
 DEFAULT_POSITION_TOLERANCE = (50 * u.m).to(u.au).value
 DEFAULT_VELOCITY_TOLERANCE = (1 * u.mm / u.s).to(u.au / u.day).value
+BACK_TO_BACK_ORBIT_FILE_PATH = "tests/data/2013_RR165_orbit.parquet"
 
 
 OBJECTS = {
@@ -35,7 +35,7 @@
     },
     "2000 PH5": {
         # Accomodate 2 km uncertainty
-        "position": (2 * u.km).to(u.au).value, 
+        "position": (2 * u.km).to(u.au).value,
         "velocity": DEFAULT_VELOCITY_TOLERANCE,
     },
     "1977 HB": {
@@ -131,6 +131,7 @@
     # "A802 FA",
 }
 
+
 @pytest.mark.parametrize("object_id", list(OBJECTS.keys()))
 def test_propagate(object_id):
     """
@@ -153,8 +154,8 @@ def test_propagate(object_id):
     )
 
     ephem_times_difference = pc.subtract(
-        assist_propagated_orbits.coordinates.time.mjd(), 
-        horizons_propagated_orbits.coordinates.time.mjd()
+        assist_propagated_orbits.coordinates.time.mjd(),
+        horizons_propagated_orbits.coordinates.time.mjd(),
     )
     np.testing.assert_array_less(
         np.abs(ephem_times_difference.to_numpy(zero_copy_only=False)),
@@ -177,15 +178,19 @@ def test_propagate(object_id):
     pos_tol = OBJECTS.get(object_id).get("position")
     vel_tol = OBJECTS.get(object_id).get("velocity")
 
-    np.testing.assert_array_less(absolute_position, pos_tol, f"Failed position for {object_id}")
-    np.testing.assert_array_less(absolute_velocity, vel_tol, f"Failed velocity for {object_id}")
+    np.testing.assert_array_less(
+        absolute_position, pos_tol, f"Failed position for {object_id}"
+    )
+    np.testing.assert_array_less(
+        absolute_velocity, vel_tol, f"Failed velocity for {object_id}"
+    )
 
 
 def test_propagate_different_input_times(mocker):
     """
     Ensure that we can pass in vectors with different epochs
     """
-    
+
     prop = ASSISTPropagator()
     watched_propagate_orbits_inner = mocker.spy(prop, "_propagate_orbits_inner")
     orbits = Orbits.from_kwargs(
@@ -204,23 +209,29 @@ def test_propagate_different_input_times(mocker):
         ),
     )
 
-    propagated_orbits = prop.propagate_orbits(orbits, Timestamp.from_mjd([60005, 60006], scale="tdb"))
+    propagated_orbits = prop.propagate_orbits(
+        orbits, Timestamp.from_mjd([60005, 60006], scale="tdb")
+    )
 
-    assert watched_propagate_orbits_inner.call_count == 2, "Inner function should be called once for each unique input epoch"
+    assert (
+        watched_propagate_orbits_inner.call_count == 2
+    ), "Inner function should be called once for each unique input epoch"
 
     assert len(propagated_orbits.coordinates.time.unique()) == 2
-    assert len(propagated_orbits) == 8, "Should have input orbits x epochs number of results"
+    assert (
+        len(propagated_orbits) == 8
+    ), "Should have input orbits x epochs number of results"
 
 
 def test_back_to_back_propagations():
     """
-    Ensure that back-to-back multiprocessed propagations work. This test should 
+    Ensure that back-to-back multiprocessed propagations work. This test should
     fail at the moment since the ray remote cannot initialize a propagator object with an already
     defined simulation.
 
     """
     prop = ASSISTPropagator()
-    orbits = query_sbdb(["2013 RR165"])
+    orbits = Orbits.from_parquet(BACK_TO_BACK_ORBIT_FILE_PATH)
 
     time = Timestamp.from_mjd([60000], scale="tdb")
     first_prop = prop.propagate_orbits(orbits, time, max_processes=1)