Skip to content

Commit

Permalink
Use new serialization method for Propagator in adam-core
Browse files Browse the repository at this point in the history
* Uses new serialization in adam-core

---------

Co-authored-by: Joachim Moeyens <moeyensj@gmail.com>
  • Loading branch information
akoumjian and moeyensj authored Jan 24, 2025
1 parent 134c18c commit 6c2ddd3
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 3 deletions.
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ classifiers = [
"Programming Language :: Python :: 3",
"Topic :: Scientific/Engineering :: Astronomy"
]
requires-python = ">=3.10,<4.0"
requires-python = ">=3.11,<3.13"
dependencies = [
"adam-core>=0.2.5",
"naif-de440",
Expand Down Expand Up @@ -57,7 +57,7 @@ dev = [
"pytest-benchmark",
"black",
"isort",
"ipython"
"ipython",
]

[tool.pdm.scripts]
Expand Down
8 changes: 8 additions & 0 deletions src/adam_assist/propagator.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,14 @@ def __init__(
self.adaptive_mode = adaptive_mode
self.epsilon = epsilon

def __getstate__(self):
state = self.__dict__.copy()
state.pop("_last_simulation", None)
return state

def __setstate__(self, state):
self.__dict__.update(state)

def _propagate_orbits(self, orbits: OrbitType, times: TimestampType) -> OrbitType:
"""
Propagate the orbits to the specified times.
Expand Down
22 changes: 21 additions & 1 deletion tests/test_propagate.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
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
Expand Down Expand Up @@ -208,4 +209,23 @@ def test_propagate_different_input_times(mocker):
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
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"])

time = Timestamp.from_mjd([60000], scale="tdb")
first_prop = prop.propagate_orbits(orbits, time, max_processes=1)

# Propagator has to be pickleable, which uses __getstate__ and __setstate__
# This doesn't work if _last_simulation is in the state
first_dict = prop.__getstate__()
second_prop = ASSISTPropagator(**first_dict)

0 comments on commit 6c2ddd3

Please sign in to comment.