B612-Asteroid-Institute / precovery

Fast precovery of small body observations at scale
BSD 3-Clause "New" or "Revised" License
6 stars 2 forks source link

Add parallel precovery_db implementation that uses adam_core orbits/propagation #99

Closed ntellis closed 1 month ago

ntellis commented 1 month ago

This adds a new class, PrecoveryDatabaseQv which uses adam_core propagation and Orbits / ephemerides.

Currently on first run in a new environment, precovery is pretty slow. This appears to be due to jax, where compile time is expensive for the @jit-ed propagate_2body in adam core. My understanding of this is limited, but there is some indication that this is also related to the input size. Just based on the runtimes profiled, it seems that recompilation needs to occur to propagate to 500 frames, even if the function has been compiled to propagate to say, 300. For this reason, precovery over the whole of NSC can take 5 minutes for orbit A, then 1 minute if run a second time, but it may still take 5 minutes to run for orbit B, all in the same environment. This is the main bottleneck right now, and hamstrings precovery pretty significantly if it's run in a fresh environment (like a k8s pod, for example).

The other current expense is in creating the observers class repeatedly via Observers.from_code, which involves hitting the spice files. This could be sped up.

Also included in this PR are the transition to PDM, removal of pyoorb as a dependency.

ntellis commented 1 month ago

The rest of precovery has now been migrated to use the adam_core Orbits class and the Precovery/FrameCandidates quivr tables. Was able to remove a lot of now unused code, including the entire old Orbit class, the precovery candidate and frame candidate dataclasses, the conversions from the same to the quivr types, etc.

The test dataset (in /tests/data/index/) has been regenerated with Assist ephemeris generation, get_orbits.py now generates /tests/data/sample_orbits.parquet, which is what the tests read from.

There is now a new test decorator that checks for the existence of ASSIST_DATA_DIR in the env - this is the location of the jpl ephemeris files. We skip index, precovery search tests without it. test_index.py only runs a precovery search on one orbit for now, due to how long it takes (~5 min, still mostly dominated by 2-body propagation time).

Core precovery is propagator-agnostic, but all of the tests initialize an AssistPropagator before doing any propagation or ephemeris generation.

Consumers will need to be updated:

ntellis commented 1 month ago

Incorporated those suggestions @akoumjian, think this is ready to merge?