diff --git a/README.md b/README.md index 04bd804..ff2a443 100644 --- a/README.md +++ b/README.md @@ -24,14 +24,6 @@ pip install adam-assist ``` -You will need to download the JPL ephemeris files. It is recommended to use the built in utility for this. - -```python -from adam_core.propagator.adam_assist import download_jpl_ephemeris_files - -download_jpl_ephemeris_files() -``` - ## Usage ### Propagating Orbits diff --git a/pyproject.toml b/pyproject.toml index 563e2cd..168236a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,11 +17,13 @@ classifiers = [ "Programming Language :: Python :: 3", "Topic :: Scientific/Engineering :: Astronomy" ] -requires-python = ">=3.11,<4.0" +requires-python = ">=3.10,<4.0" dependencies = [ - "adam-core>=0.2.3", - "assist", + "adam-core>=0.2.5", "naif-de440", + "jpl-small-bodies-de441-n16", + # Temporary until SPK support is merged and released upstream + "assist-adam-fork==1.1.9a2", "numpy", "ray", "spiceypy>=6.0.0" @@ -39,9 +41,9 @@ write_template = "__version__ = '{}'" [project.urls] -"Documentation" = "https://github.com/unknown/adam-assist#readme" -"Issues" = "https://github.com/unknown/adam-assist/issues" -"Source" = "https://github.com/unknown/adam-assist" +"Documentation" = "https://github.com/B612-Asteroid-Institute/adam-assist#readme" +"Issues" = "https://github.com/B612-Asteroid-Institute/adam-assist/issues" +"Source" = "https://github.com/B612-Asteroid-Institute/adam-assist" [project.optional-dependencies] diff --git a/src/adam_core/propagator/adam_assist.py b/src/adam_core/propagator/adam_assist.py index 70c55e5..e277891 100644 --- a/src/adam_core/propagator/adam_assist.py +++ b/src/adam_core/propagator/adam_assist.py @@ -1,16 +1,14 @@ import hashlib -import os -import pathlib from ctypes import c_uint32 from typing import Dict, List, Tuple, Union import assist import numpy as np +import numpy.typing as npt import pyarrow as pa import pyarrow.compute as pc import quivr as qv import rebound -import urllib3 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 @@ -19,6 +17,8 @@ from adam_core.orbits import Orbits from adam_core.orbits.variants import VariantOrbits from adam_core.time import Timestamp +from jpl_small_bodies_de441_n16 import de441_n16 +from naif_de440 import de440 from quivr.concat import concatenate from adam_core.propagator.propagator import OrbitType, Propagator, TimestampType @@ -30,38 +30,9 @@ except ImportError: __version__ = "0.0.0" -DATA_DIR = os.getenv("ASSIST_DATA_DIR", "~/.adam_assist_data") - # Use the Earth's equatorial radius as used in DE4XX ephemerides # adam_core defines it in au but we need it in km -EARTH_RADIUS_KM = c.R_EARTH * KM_P_AU - - -def download_jpl_ephemeris_files(data_dir: str = DATA_DIR) -> None: - ephemeris_urls = ( - "https://ssd.jpl.nasa.gov/ftp/eph/small_bodies/asteroids_de441/sb441-n16.bsp", - "https://ssd.jpl.nasa.gov/ftp/eph/planets/Linux/de440/linux_p1550p2650.440", - ) - data_dir_path = pathlib.Path(data_dir).expanduser() - data_dir_path.mkdir(parents=True, exist_ok=True) - for url in ephemeris_urls: - file_name = pathlib.Path(url).name - file_path = data_dir_path.joinpath(file_name) - if not file_path.exists(): - # use urllib3 - http = urllib3.PoolManager() - with ( - http.request("GET", url, preload_content=False) as r, - open(file_path, "wb") as out_file, - ): - if r.status != 200: - raise RuntimeError(f"Failed to download {url}") - while True: - data = r.read(1024) - if not data: - break - out_file.write(data) - r.release_conn() +EARTH_RADIUS_KM = c.R_EARTH_EQUATORIAL * KM_P_AU def uint32_hash(s: str) -> c_uint32: @@ -71,7 +42,8 @@ def uint32_hash(s: str) -> c_uint32: def hash_orbit_ids_to_uint32( - orbit_ids: np.ndarray[Tuple[np.dtype[np.int_]], np.dtype[np.str_]], + # orbit_ids: np.ndarray[Tuple[np.dtype[np.int_]], np.dtype[np.str_]], + orbit_ids: npt.NDArray[np.str_], ) -> Tuple[Dict[int, str], List[c_uint32]]: """ Derive uint32 hashes from orbit id strigns @@ -132,10 +104,9 @@ def _propagate_orbits_inner( """ Propagates one or more orbits with the same epoch to the specified times. """ - root_dir = pathlib.Path(DATA_DIR).expanduser() ephem = assist.Ephem( - planets_path=str(root_dir.joinpath("linux_p1550p2650.440")), - asteroids_path=str(root_dir.joinpath("sb441-n16.bsp")), + planets_path=de440, + asteroids_path=de441_n16, ) sim = None sim = rebound.Simulation() @@ -284,10 +255,9 @@ def _detect_impacts( coords = coords.set_column("time", input_orbit_times) orbits = orbits.set_column("coordinates", coords) - root_dir = pathlib.Path(DATA_DIR).expanduser() ephem = assist.Ephem( - planets_path=str(root_dir.joinpath("linux_p1550p2650.440")), - asteroids_path=str(root_dir.joinpath("sb441-n16.bsp")), + planets_path=de440, + asteroids_path=de441_n16, ) sim = None sim = rebound.Simulation() diff --git a/src/adam_core/propagator/adam_assist_version.py b/src/adam_core/propagator/adam_assist_version.py new file mode 100644 index 0000000..6b2c48a --- /dev/null +++ b/src/adam_core/propagator/adam_assist_version.py @@ -0,0 +1 @@ +__version__ = "0.1.3.dev4+g148ff70.d20241109" diff --git a/tests/test_ephemeris.py b/tests/test_ephemeris.py index 90a8d7e..e66d559 100644 --- a/tests/test_ephemeris.py +++ b/tests/test_ephemeris.py @@ -8,17 +8,13 @@ from astropy import units as u from numpy.testing import assert_allclose -from src.adam_core.propagator.adam_assist import ( - ASSISTPropagator, - download_jpl_ephemeris_files, -) +from src.adam_core.propagator.adam_assist import ASSISTPropagator def test_ephemeris(): """ Test the accurate of the ephemeris generator by comparing the propagated orbit to the JPL ephemeris """ - download_jpl_ephemeris_files() prop = ASSISTPropagator() OBJECT_IDS = [ "2020 AV2", diff --git a/tests/test_impacts.py b/tests/test_impacts.py index 31cd86f..ca57c57 100644 --- a/tests/test_impacts.py +++ b/tests/test_impacts.py @@ -1,13 +1,10 @@ import pytest from adam_core.dynamics.impacts import calculate_impacts from adam_core.orbits import Orbits -from adam_core.time import Timestamp from adam_core.orbits.query.horizons import query_horizons +from adam_core.time import Timestamp -from src.adam_core.propagator.adam_assist import ( - ASSISTPropagator, - download_jpl_ephemeris_files, -) +from src.adam_core.propagator.adam_assist import ASSISTPropagator # Contains a likely impactor with ~60% chance of impact in 30 days IMPACTOR_FILE_PATH_60 = "tests/data/I00007_orbit.parquet" @@ -20,7 +17,6 @@ @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_60)[0] propagator = ASSISTPropagator() variants, impacts = benchmark( @@ -39,7 +35,7 @@ def test_calculate_impacts_benchmark(benchmark, processes): @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( @@ -58,7 +54,6 @@ def test_calculate_impacts_benchmark(benchmark, processes): @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( @@ -77,7 +72,7 @@ def test_calculate_impacts_benchmark(benchmark, processes): def test_detect_impacts_time_direction(): start_time = Timestamp.from_mjd([60000], scale="utc") orbit = query_horizons(["1980 PA"], start_time) - download_jpl_ephemeris_files() + propagator = ASSISTPropagator() results, impacts = propagator._detect_impacts(orbit, 60) diff --git a/tests/test_propagate.py b/tests/test_propagate.py index d989b77..ef033b2 100644 --- a/tests/test_propagate.py +++ b/tests/test_propagate.py @@ -6,15 +6,10 @@ from adam_core.coordinates.residuals import Residuals from adam_core.orbits import Orbits from adam_core.orbits.query.horizons import query_horizons -from adam_core.orbits.query.sbdb import query_sbdb from adam_core.time import Timestamp from astropy import units as u -from numpy.testing import assert_allclose -from src.adam_core.propagator.adam_assist import ( - ASSISTPropagator, - download_jpl_ephemeris_files, -) +from src.adam_core.propagator.adam_assist import ASSISTPropagator DEFAULT_POSITION_TOLERANCE = (50 * u.m).to(u.au).value DEFAULT_VELOCITY_TOLERANCE = (1 * u.mm / u.s).to(u.au / u.day).value @@ -140,7 +135,7 @@ def test_propagate(): """ Test the accurate of the ephemeris generator by comparing the propagated orbit to the JPL ephemeris """ - download_jpl_ephemeris_files() + prop = ASSISTPropagator() millisecond_in_days = 1.1574074074074073e-8 @@ -189,7 +184,7 @@ def test_propagate_different_input_times(mocker): """ Ensure that we can pass in vectors with different epochs """ - download_jpl_ephemeris_files() + prop = ASSISTPropagator() watched_propagate_orbits_inner = mocker.spy(prop, "_propagate_orbits_inner") orbits = Orbits.from_kwargs(