diff --git a/CHANGELOG.md b/CHANGELOG.md index 34d9977..af79f88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Saving`SimultaneousStates` to parquet files can now optionally include a column containing the TDB JD of when the state information was last updated. This allows users to selectively update state vectors only when necessary. +- Added multi-core propagation support to rust backend. ### Changed diff --git a/src/kete/rust/propagation.rs b/src/kete/rust/propagation.rs index 3c4cc17..c9d85f3 100644 --- a/src/kete/rust/propagation.rs +++ b/src/kete/rust/propagation.rs @@ -125,7 +125,7 @@ pub fn propagation_n_body_spk_py( .into_iter() .zip(non_gravs.into_iter()) .collect_vec() - .chunks(1000) + .chunks(500) { py.check_signals()?; diff --git a/src/kete_core/src/propagation/mod.rs b/src/kete_core/src/propagation/mod.rs index 084fec3..d24aef4 100644 --- a/src/kete_core/src/propagation/mod.rs +++ b/src/kete_core/src/propagation/mod.rs @@ -35,11 +35,13 @@ use crate::constants::GravParams; use crate::errors::Error; use crate::frames::Equatorial; -use crate::prelude::{Desig, KeteResult}; +use crate::prelude::{Desig, KeteResult, SimultaneousStates}; use crate::spice::LOADED_SPK; use crate::state::State; use crate::time::{TDB, Time}; +use itertools::Itertools; use nalgebra::{DVector, SMatrix, SVector, Vector3}; +use rayon::prelude::*; mod acceleration; mod kepler; @@ -110,6 +112,84 @@ pub fn propagate_n_body_spk( Ok(new_state) } +/// Propagate the provided [`Vec>`] using N body mechanics to the +/// specified times, no approximations are made, this can be very CPU intensive. +/// +/// This uses rayon to use as many cores as are available. +/// +/// This does not compute light delay, however it does include corrections for general +/// relativity due to the Sun. +/// +/// # Parameters +/// +/// states: +/// The initial states, this is a list of multiple State objects. +/// jd: +/// A JD to propagate the initial states to. +/// ``include_asteroids``: +/// If this is true, the computation will include the largest 5 asteroids. +/// The asteroids are: Ceres, Pallas, Interamnia, Hygiea, and Vesta. +/// ``non_gravs``: +/// A list of non-gravitational terms for each object. If provided, then every +/// object must have an associated :class:`~NonGravModel` or `None`. +/// ``suppress_errors``: +/// If True, errors during propagation will return NaN for the relevant state +/// vectors, but propagation will continue. +/// +/// # Errors +/// +/// Integration may fail for a large number of reasons. This function accepts a +/// ``suppress_errors`` parameter which may be used to continue propagation as long +/// as possible for the most objects as possible. +pub fn propagation_n_body_spk_par( + states: Vec>, + jd: Time, + include_asteroids: bool, + non_gravs: Option>>, + suppress_errors: bool, +) -> KeteResult { + let non_gravs = non_gravs.unwrap_or(vec![None; states.len()]); + + if states.len() != non_gravs.len() { + Err(Error::ValueError( + "non_gravs must be the same length as states.".into(), + ))?; + } + + let res: KeteResult> = states + .into_iter() + .zip(non_gravs) + .collect_vec() + .into_par_iter() + .with_min_len(10) + .map(|(state, model)| { + let center = state.center_id; + let desig = state.desig.clone(); + + // if the input has a NAN in it, skip the propagation entirely and return + // the nans. + if !state.is_finite() { + if !suppress_errors { + Err(Error::ValueError("Input state contains NaNs.".into()))?; + } + return Ok(State::::new_nan(desig, jd, center)); + } + match propagate_n_body_spk(state, jd, include_asteroids, model) { + Ok(state) => Ok(state), + Err(er) => { + if suppress_errors { + Ok(State::::new_nan(desig, jd, center)) + } else { + Err(er)? + } + } + } + }) + .collect(); + + SimultaneousStates::new_exact(res?, None) +} + /// Initialization function for the picard integrator which initializes the state /// using two body mechanics. fn picard_two_body_init( diff --git a/src/kete_core/src/simult_states.rs b/src/kete_core/src/simult_states.rs index 557e28e..cffbf74 100644 --- a/src/kete_core/src/simult_states.rs +++ b/src/kete_core/src/simult_states.rs @@ -153,4 +153,14 @@ impl SimultaneousStates { }) .collect()) } + + /// Number of states + #[must_use] + #[allow( + clippy::len_without_is_empty, + reason = "Cannot be constructed as empty" + )] + pub fn len(&self) -> usize { + self.states.len() + } }