Point to fork and use data packages (#12)

* Point to fork and use data packages

* linting

* typing

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

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]

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()

src/adam_core/propagator/adam_assist_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.3.dev4+g148ff70.d20241109"

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",

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)

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(