Vectorize Propagation working again (#89)

* Add support for long term orbit propagation

* J2 calculation fixed for Jupiter and Sun

* typo in GR and J2 calculation, GMS double applied

* documentation and changelog

CHANGELOG.md

@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Added support for long term integrations, on the scale of less than a mega-year.
+
 ### Changed
 
 - Optimizations to SPICE kernel queries leading to a 20% speedup in orbit propagation,
@@ -16,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Moved python wrappers over propagation into the rust backend entirely.
 - Simplified polymorphic support in the rust wrappers, removing some intermediate
   objects which were over-complicating the interface.
+- Refactored the computation of gravitational forces, improving maintainability.
 
 ### Fixed
 

docs/conf.py

@@ -63,7 +63,7 @@
 
 # -- Sphinx gallery settings --------------------------------------------------
 sphinx_gallery_conf = {
-    "run_stale_examples": False,
+    "run_stale_examples": True,
     "filename_pattern": "",
     "examples_dirs": ["../src/examples"],
 }

src/examples/README.rst

@@ -1,4 +1,4 @@
 Examples
 ========
 
-Here are a collection of examples which demonstrate different parts of the neospy.
+A collection of examples which demonstrate different parts of the neospy.

src/neospy/propagation.py

@@ -10,11 +10,17 @@
 
 from .vector import State
 from . import spice
-from ._core import NonGravModel, propagate_n_body, propagate_two_body
+from ._core import (
+    NonGravModel,
+    propagate_n_body,
+    propagate_n_body_long,
+    propagate_two_body,
+)
 
 
 __all__ = [
     "propagate_n_body",
+    "propagate_n_body_long",
     "propagate_two_body",
     "NonGravModel",
     "moid",

src/neospy/rust/lib.rs

@@ -62,6 +62,7 @@ fn _core(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(kepler::propagation_kepler_py, m)?)?;
 
     m.add_function(wrap_pyfunction!(propagation::propagation_n_body_spk_py, m)?)?;
+    m.add_function(wrap_pyfunction!(propagation::propagation_n_body_py, m)?)?;
 
     m.add_function(wrap_pyfunction!(fovs::fov_checks_py, m)?)?;
     m.add_function(wrap_pyfunction!(fovs::fov_spk_checks_py, m)?)?;

src/neospy/rust/propagation.rs

@@ -1,7 +1,7 @@
 use itertools::Itertools;
 use neospy_core::{
     errors::Error,
-    propagation,
+    propagation::{self, NonGravModel},
     spice::{self, get_spk_singleton},
     state::State,
     time::{scales::TDB, Time},
@@ -136,3 +136,97 @@ pub fn propagation_n_body_spk_py(
 
     Ok(res)
 }
+
+/// It is *STRONGLY* recommended to use `propagate_n_body` instead of this function
+/// wherever possible. This function is specifically meant for kilo-year or longer
+/// simulations, it is slower and less accurate than `propagate_n_body`, but that
+/// function only works for as long as there are SPICE kernels for the planets
+/// available.
+///
+/// This is designed to not require SPICE kernels to function, and is meant for long
+/// term simulations on the other of less than a mega-year. It is not recommended
+/// for orbits longer than this.
+///
+/// Propagation using this will treat the Earth and Moon as a single object for
+/// performance reasons.
+///
+/// Propagate the provided :class:`~neospy.State` using N body mechanics to the
+/// specified times, very few approximations are made, this can be very CPU intensive.
+///
+/// This does not compute light delay, however it does include corrections for general
+/// relativity due to the Sun.
+///
+/// This returns two lists of states:
+/// - First one contains the states of the objects at the end of the integration
+/// - Second contains the states of the planets at the end of the integration, which may
+///   be used as input for continuing the integration.
+///
+/// Parameters
+/// ----------
+/// states:
+///     The initial states, this is a list of multiple State objects.
+/// jd:
+///     A JD to propagate the initial states to.
+/// planet_states:
+///     Optional list of the planet's states at the same time as the provided states.
+///     If this is not provided, the planets positions are loaded from the Spice kernels
+///     if that information is available.
+/// non_gravs:
+///     A list of non-gravitational terms for each object. If provided, then every
+///     object must have an associated :class:`~NonGravModel`.
+/// batch_size:
+///     Number of objects to propagate at once with the planets. This is used to break
+///     up the simulation for multi-core support. It additionally has effects on the
+///     integrator stepsize which is difficult to predict before running. This can be
+///     manually tuned for increased performance, it should have no other effects than
+///     performance.
+///
+/// Returns
+/// -------
+/// Iterable
+///     A :class:`~neospy.State` at the new time.
+#[pyfunction]
+#[pyo3(name = "propagate_n_body_long", signature = (states, jd_final, planet_states=None, non_gravs=None, batch_size=10))]
+pub fn propagation_n_body_py(
+    states: Vec<PyState>,
+    jd_final: PyTime,
+    planet_states: Option<Vec<PyState>>,
+    non_gravs: Option<Vec<Option<PyNonGravModel>>>,
+    batch_size: usize,
+) -> PyResult<(Vec<PyState>, Vec<PyState>)> {
+    let states: Vec<State> = states.into_iter().map(|x| x.0).collect();
+    let planet_states: Option<Vec<State>> =
+        planet_states.map(|s| s.into_iter().map(|x| x.0).collect());
+
+    let non_gravs = non_gravs.unwrap_or(vec![None; states.len()]);
+    let non_gravs: Vec<Option<NonGravModel>> =
+        non_gravs.into_iter().map(|y| y.map(|z| z.0)).collect();
+
+    let jd = jd_final.jd();
+    let res = states
+        .into_iter()
+        .zip(non_gravs.into_iter())
+        .collect_vec()
+        .par_chunks(batch_size)
+        .map(|chunk| {
+            let (chunk_state, chunk_nongrav): (Vec<State>, Vec<Option<NonGravModel>>) =
+                chunk.iter().cloned().collect();
+
+            propagation::propagate_n_body_vec(chunk_state, jd, planet_states.clone(), chunk_nongrav)
+                .map(|(states, planets)| {
+                    (
+                        states.into_iter().map(PyState::from).collect::<Vec<_>>(),
+                        planets.into_iter().map(PyState::from).collect::<Vec<_>>(),
+                    )
+                })
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+
+    let mut final_states = Vec::new();
+    let mut final_planets = Vec::new();
+    for (mut state, planet) in res.into_iter() {
+        final_states.append(&mut state);
+        final_planets = planet;
+    }
+    Ok((final_states, final_planets))
+}

src/neospy/rust/spice/spk.rs

@@ -1,5 +1,5 @@
 use neospy_core::spice::{get_spk_singleton, try_name_from_id};
-use pyo3::{pyfunction, PyResult};
+use pyo3::{pyfunction, PyResult, Python};
 
 use crate::frame::PyFrames;
 use crate::state::PyState;
@@ -8,12 +8,13 @@ use crate::time::PyTime;
 /// Load all specified files into the SPK shared memory singleton.
 #[pyfunction]
 #[pyo3(name = "spk_load")]
-pub fn spk_load_py(filenames: Vec<String>) -> PyResult<()> {
+pub fn spk_load_py(py: Python<'_>, filenames: Vec<String>) -> PyResult<()> {
     let mut singleton = get_spk_singleton().write().unwrap();
     if filenames.len() > 100 {
         eprintln!("Loading {} spk files...", filenames.len());
     }
     for filename in filenames.iter() {
+        py.check_signals()?;
         let load = (*singleton).load_file(filename);
         if let Err(err) = load {
             eprintln!("{} failed to load. {}", filename, err);

src/neospy_core/benches/propagation.rs

@@ -1,4 +1,6 @@
 extern crate criterion;
+use std::time::Duration;
+
 use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
 use lazy_static::lazy_static;
 use neospy_core::prelude::*;
@@ -54,6 +56,16 @@ fn prop_n_body_radau(state: State, dt: f64) {
     propagation::propagate_n_body_spk(state, jd, false, None).unwrap();
 }
 
+fn prop_n_body_vec_radau(mut state: State, dt: f64) {
+    let spk = get_spk_singleton().read().unwrap();
+    spk.try_change_center(&mut state, 10).unwrap();
+    state.try_change_frame_mut(Frame::Ecliptic).unwrap();
+    let states = vec![state.clone(); 100];
+    let non_gravs = vec![None; 100];
+    let jd = state.jd + dt;
+    propagation::propagate_n_body_vec(states, jd, None, non_gravs).unwrap();
+}
+
 fn prop_n_body_radau_par(state: State, dt: f64) {
     let states: Vec<State> = (0..100).map(|_| state.clone()).collect();
     let _tmp: Vec<State> = states
@@ -87,7 +99,7 @@ pub fn two_body_numeric(c: &mut Criterion) {
             _ => panic!(),
         };
         twobody_num_group.bench_with_input(BenchmarkId::new("Single", name), &state, |b, s| {
-            b.iter(|| prop_2_body_radau(s.clone(), black_box(1000.0)))
+            b.iter(|| prop_2_body_radau(black_box(s.clone()), black_box(1000.0)))
         });
     }
 }
@@ -106,14 +118,32 @@ pub fn n_body_prop(c: &mut Criterion) {
             _ => panic!(),
         };
         nbody_group.bench_with_input(BenchmarkId::new("Single", name), &state, |b, s| {
-            b.iter(|| prop_n_body_radau(s.clone(), black_box(1000.0)))
+            b.iter(|| prop_n_body_radau(black_box(s.clone()), black_box(1000.0)))
         });
 
         nbody_group.bench_with_input(BenchmarkId::new("Parallel", name), &state, |b, s| {
             b.iter(|| prop_n_body_radau_par(black_box(s.clone()), black_box(1000.0)))
         });
     }
 }
+pub fn n_body_prop_vec(c: &mut Criterion) {
+    let mut nbody_group = c.benchmark_group("N-Body-Vec");
+
+    for state in [
+        CIRCULAR.clone(),
+        ELLIPTICAL.clone(),
+        PARABOLIC.clone(),
+        HYPERBOLIC.clone(),
+    ] {
+        let name = match &state.desig {
+            Desig::Name(n) => n,
+            _ => panic!(),
+        };
+        nbody_group.bench_with_input(BenchmarkId::new("Single", name), &state, |b, s| {
+            b.iter(|| prop_n_body_vec_radau(black_box(s.clone()), black_box(1000.0)))
+        });
+    }
+}
 
 pub fn two_body_analytic(c: &mut Criterion) {
     let mut twobody_group = c.benchmark_group("2-Body-Analytic");
@@ -135,6 +165,6 @@ pub fn two_body_analytic(c: &mut Criterion) {
 }
 
 criterion_group!(name=benches;
-                 config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
-                 targets=two_body_analytic, n_body_prop, two_body_numeric);
+                 config = Criterion::default().sample_size(30).measurement_time(Duration::from_secs(15)).with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
+                 targets=n_body_prop_vec, two_body_analytic, n_body_prop, two_body_numeric);
 criterion_main!(benches);