Merge pull request #181 from B612-Asteroid-Institute/kk/variant_fix

Kk/variant fix

Author: KatKiker

pyproject.toml

@@ -10,7 +10,7 @@ authors = [
 ]
 description = "Core libraries for the ADAM platform"
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.11,<3.14"
 classifiers = [
   "Operating System :: OS Independent",
   "Development Status :: 4 - Beta",

src/adam_core/dynamics/tests/test_propagation.py

@@ -500,6 +500,7 @@ def test_propagate_2body_preserves_physical_parameters():
     np.testing.assert_allclose(have_H, expected_H)
     np.testing.assert_allclose(have_G, expected_G)
 
+
 @pytest.mark.profile
 def test_profile_propagate_2body_matrix(propagated_orbits, tmp_path):
     """Profile the propagate_2body function with different combinations of orbits and times.

src/adam_core/orbits/tests/test_variants.py

@@ -135,7 +135,9 @@ def test_VariantOrbits_collapse_by_object_id():
         orbit_id=["obj1", "obj1"],
         object_id=["obj1", "obj1"],
         variant_id=["0", "1"],
-        physical_parameters=PhysicalParameters.from_kwargs(H_v=[15.0, 15.0], G=[0.15, 0.15]),
+        physical_parameters=PhysicalParameters.from_kwargs(
+            H_v=[15.0, 15.0], G=[0.15, 0.15]
+        ),
         coordinates=CartesianCoordinates.from_kwargs(
             x=[1.0, 1.1],
             y=[1.0, 1.1],
@@ -156,7 +158,9 @@ def test_VariantOrbits_collapse_by_object_id():
         orbit_id=["obj1", "obj1"],
         object_id=["obj1", "obj1"],
         variant_id=["0", "1"],
-        physical_parameters=PhysicalParameters.from_kwargs(H_v=[15.0, 15.0], G=[0.15, 0.15]),
+        physical_parameters=PhysicalParameters.from_kwargs(
+            H_v=[15.0, 15.0], G=[0.15, 0.15]
+        ),
         coordinates=CartesianCoordinates.from_kwargs(
             x=[1.0, 1.1],
             y=[1.0, 1.1],

src/adam_core/propagator/propagator.py

@@ -798,7 +798,7 @@ def propagate_orbits(
         chunk_size: int = 100,
         max_processes: Optional[int] = 1,
         seed: Optional[int] = None,
-    ) -> Orbits:
+    ) -> Union[Orbits, VariantOrbits]:
         """
         Propagate each orbit in orbits to each time in times.
 
@@ -830,15 +830,21 @@ def propagate_orbits(
 
         Returns
         -------
-        propagated : `~adam_core.orbits.orbits.Orbits`
+        propagated : `~adam_core.orbits.orbits.Orbits` or `~adam_core.orbits.variants.VariantOrbits`
             Propagated orbits.
         """
+        if covariance is True and isinstance(orbits, VariantOrbits):
+            raise AssertionError("Covariance is not supported for VariantOrbits")
+
         if max_processes is None:
             max_processes = mp.cpu_count()
 
         if max_processes > 1:
             propagated_list: List[Orbits] = []
-            variants_list: List[VariantOrbits] = []
+            covariance_variants_list: List[VariantOrbits] = []
+            # When the input is VariantOrbits, do not treat them as covariance.
+            propagated_variants_input_list: List[VariantOrbits] = []
+            input_is_variants: Optional[bool] = None
 
             if RAY_INSTALLED is False:
                 raise ImportError(
@@ -856,13 +862,18 @@ def propagate_orbits(
                 times = ray.get(times_ref)
 
             if not isinstance(orbits, ObjectRef):
+                input_is_variants = isinstance(orbits, VariantOrbits)
                 orbits_ref = ray.put(orbits)
             else:
                 orbits_ref = orbits
                 # We need to dereference the orbits ObjectRef so we can
                 # check its length for chunking and determine
                 # if we need to propagate variants
                 orbits = ray.get(orbits_ref)
+                input_is_variants = isinstance(orbits, VariantOrbits)
+
+            if covariance is True and input_is_variants:
+                raise AssertionError("Covariance is not supported for VariantOrbits")
 
             # Create futures inputs
             futures_inputs = []
@@ -910,7 +921,10 @@ def propagate_orbits(
                     if isinstance(result, Orbits):
                         propagated_list.append(result)
                     elif isinstance(result, VariantOrbits):
-                        variants_list.append(result)
+                        if input_is_variants:
+                            propagated_variants_input_list.append(result)
+                        else:
+                            covariance_variants_list.append(result)
                     else:
                         raise ValueError(
                             f"Unexpected result type from propagation worker: {type(result)}"
@@ -923,22 +937,33 @@ def propagate_orbits(
                 if isinstance(result, Orbits):
                     propagated_list.append(result)
                 elif isinstance(result, VariantOrbits):
-                    variants_list.append(result)
+                    if input_is_variants:
+                        propagated_variants_input_list.append(result)
+                    else:
+                        covariance_variants_list.append(result)
                 else:
                     raise ValueError(
                         f"Unexpected result type from propagation worker: {type(result)}"
                     )
 
             # Concatenate propagated orbits
-            propagated = qv.concatenate(propagated_list)
-            if len(variants_list) > 0:
-                propagated_variants = qv.concatenate(variants_list)
-                # sort by variant_id and time
-                propagated_variants = propagated_variants.sort_by(
-                    ["variant_id", "coordinates.time.days", "coordinates.time.nanos"]
-                )
-            else:
+            if input_is_variants:
+                propagated = qv.concatenate(propagated_variants_input_list)
                 propagated_variants = None
+            else:
+                propagated = qv.concatenate(propagated_list)
+                if len(covariance_variants_list) > 0:
+                    propagated_variants = qv.concatenate(covariance_variants_list)
+                    # sort by variant_id and time
+                    propagated_variants = propagated_variants.sort_by(
+                        [
+                            "variant_id",
+                            "coordinates.time.days",
+                            "coordinates.time.nanos",
+                        ]
+                    )
+                else:
+                    propagated_variants = None
 
         else:
             propagated = self._propagate_orbits(orbits, times)

src/adam_core/propagator/tests/test_propagator.py

@@ -233,6 +233,22 @@ def test_propagator_multiple_workers_ray():
     assert len(have) == len(orbits) * len(times)
 
 
+@pytest.mark.skipif(not RAY_INSTALLED, reason="Ray not installed")
+def test_propagate_orbits_multiple_workers_ray_variant_orbits_input():
+    """
+    Regression test: VariantOrbits should be supported as an input to propagate_orbits
+    under the ray parallel dispatcher.
+    """
+    base = make_real_orbits(4)
+    variants = VariantOrbits.create(base, method="sigma-point")
+    times = Timestamp.from_iso8601(["2020-01-01T00:00:00", "2020-01-01T00:00:01"])
+
+    prop = MockPropagator()
+    have = prop.propagate_orbits(variants, times, max_processes=2)
+    assert isinstance(have, VariantOrbits)
+    assert len(have) == len(variants) * len(times)
+
+
 def test_propagate_different_origins():
     """
     Test that we are returning propagated orbits with their original origins