pub mod euler;
pub mod integrator;
pub mod parallel;
pub mod rkf45;
pub mod rosenbrock23;
use petir::mathf::RealMath;
pub use euler::Euler;
pub use integrator::{
DynSystemIntegrator, NoTypedSystem, OdeIntegrator, OdeSolver, SharedOdeSystem,
TypedStateIntegrator,
};
pub use parallel::{
adaptive_quadrature_batch, ensemble_backend_for, integrate_ensemble, integrate_ensemble_mixed,
quadrature_backend_for, quadrature_batch, AdaptiveSettings, GaussOrder, OdeEnsemble,
OdeEnsembleFailure, OdeLane, OdeLaneSolution, OdeLaneStatus, QuadratureBatch,
QuadratureBatchFailure, QuadratureInterval, QuadratureRule, QuadratureSolution,
QuadratureStatus,
};
pub use rkf45::Rkf45;
pub use rosenbrock23::Rosenbrock23;
use crate::matrix::SquareMatrix;
pub trait OdeSystem {
fn n_eqns(&self) -> usize;
fn derivatives(&self, x: f64, y: &[f64], dydx: &mut Vec<f64>);
fn jacobian(&self, _x: f64, _y: &[f64], _dfdx: &mut Vec<f64>, _dfdy: &mut SquareMatrix) {
unimplemented!("jacobian not implemented for this ODE system");
}
}
#[derive(Debug, Clone)]
pub struct OdeSolverConfig {
pub abs_tol: f64,
pub rel_tol: f64,
pub safe_scale: f64,
pub alpha_inc: f64,
pub alpha_dec: f64,
pub min_scale: f64,
pub max_scale: f64,
pub max_steps: usize,
}
impl Default for OdeSolverConfig {
fn default() -> Self {
Self {
abs_tol: 1e-6,
rel_tol: 1e-4,
safe_scale: 0.9,
alpha_inc: 0.2,
alpha_dec: 0.25,
min_scale: 0.2,
max_scale: 10.0,
max_steps: 10_000,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum OdeError {
StepSizeUnderflow,
MaxStepsExceeded(usize),
NonFiniteState,
}
impl std::fmt::Display for OdeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::StepSizeUnderflow => write!(f, "ODE step size underflow"),
Self::MaxStepsExceeded(n) => write!(f, "ODE exceeded {n} steps"),
Self::NonFiniteState => write!(
f,
"ODE produced a non-finite error estimate (NaN or infinite state)"
),
}
}
}
impl std::error::Error for OdeError {}
pub(crate) fn normalize_error(
y0: &[f64],
y: &[f64],
err: &[f64],
abs_tol: f64,
rel_tol: f64,
) -> f64 {
err.iter()
.zip(y0)
.zip(y)
.map(|((e, &a), &b)| {
let tol = abs_tol + rel_tol * a.abs().max(b.abs());
e.abs() / tol
})
.fold(0.0_f64, |acc, e| {
if e.is_nan() || acc.is_nan() {
f64::INFINITY
} else {
acc.max(e)
}
})
}
pub(crate) fn adaptive_step<Sys: OdeSystem + ?Sized>(
cfg: &OdeSolverConfig,
mut inner_step: impl FnMut(f64, &[f64], &[f64], f64, &mut Vec<f64>) -> f64,
ode: &Sys,
x: &mut f64,
y: &mut Vec<f64>,
dydx0: &mut Vec<f64>,
y_temp: &mut Vec<f64>,
dx_try: &mut f64,
) -> Result<(), OdeError> {
let mut dx = *dx_try;
ode.derivatives(*x, y, dydx0);
let err = loop {
let err = inner_step(*x, y, dydx0, dx, y_temp);
if !err.is_finite() {
return Err(OdeError::NonFiniteState);
}
if err <= 1.0 {
break err;
}
let scale = (cfg.safe_scale * err.r_powf(-cfg.alpha_dec)).max(cfg.min_scale);
dx *= scale;
if dx.abs() < f64::EPSILON {
return Err(OdeError::StepSizeUnderflow);
}
};
*x += dx;
std::mem::swap(y, y_temp);
let threshold = (cfg.max_scale / cfg.safe_scale).r_powf(-1.0 / cfg.alpha_inc);
*dx_try = if err > threshold {
let scale = (cfg.safe_scale * err.r_powf(-cfg.alpha_inc)).clamp(cfg.min_scale, cfg.max_scale);
dx * scale
} else {
dx * cfg.safe_scale * cfg.max_scale
};
Ok(())
}
pub(crate) fn integrate_interval(
cfg: &OdeSolverConfig,
step_fn: &mut dyn FnMut(&mut f64, &mut Vec<f64>, &mut f64) -> Result<(), OdeError>,
x_start: f64,
x_end: f64,
y: &mut Vec<f64>,
dx_est: &mut f64,
) -> Result<(), OdeError> {
let mut x = x_start;
let mut dx = *dx_est;
let mut n_steps = 0usize;
while x < x_end {
let dx_try = (x + dx).min(x_end) - x;
let mut dx_limited = dx_try;
step_fn(&mut x, y, &mut dx_limited)?;
dx = dx_limited;
n_steps += 1;
if n_steps > cfg.max_steps {
return Err(OdeError::MaxStepsExceeded(n_steps));
}
}
*dx_est = dx;
Ok(())
}