use std::sync::Arc;
use super::{Euler, OdeError, OdeSolverConfig, OdeSystem, Rkf45, Rosenbrock23};
pub type SharedOdeSystem = Arc<dyn OdeSystem + Send + Sync>;
#[derive(Debug, Clone)]
pub enum OdeSolver {
Euler(Euler),
Rkf45(Rkf45),
Rosenbrock23(Rosenbrock23),
}
impl OdeSolver {
#[must_use]
pub fn euler(n: usize, abs_tol: f64, rel_tol: f64) -> Self {
Self::Euler(Euler::new(n, abs_tol, rel_tol))
}
#[must_use]
pub fn rkf45(n: usize, abs_tol: f64, rel_tol: f64) -> Self {
Self::Rkf45(Rkf45::new(n, abs_tol, rel_tol))
}
#[must_use]
pub fn rosenbrock23(n: usize, abs_tol: f64, rel_tol: f64) -> Self {
Self::Rosenbrock23(Rosenbrock23::new(n, abs_tol, rel_tol))
}
#[must_use]
pub const fn name(&self) -> &'static str {
match self {
Self::Euler(_) => "Euler",
Self::Rkf45(_) => "RKF45",
Self::Rosenbrock23(_) => "Rosenbrock23",
}
}
#[must_use]
pub const fn requires_jacobian(&self) -> bool {
matches!(self, Self::Rosenbrock23(_))
}
#[must_use]
pub fn config(&self) -> &OdeSolverConfig {
match self {
Self::Euler(s) => &s.config,
Self::Rkf45(s) => &s.config,
Self::Rosenbrock23(s) => &s.config,
}
}
pub fn config_mut(&mut self) -> &mut OdeSolverConfig {
match self {
Self::Euler(s) => &mut s.config,
Self::Rkf45(s) => &mut s.config,
Self::Rosenbrock23(s) => &mut s.config,
}
}
pub fn solve_step<Sys: OdeSystem + ?Sized>(
&mut self,
ode: &Sys,
x: &mut f64,
y: &mut Vec<f64>,
dx_try: &mut f64,
) -> Result<(), OdeError> {
match self {
Self::Euler(s) => s.solve_step(ode, x, y, dx_try),
Self::Rkf45(s) => s.solve_step(ode, x, y, dx_try),
Self::Rosenbrock23(s) => s.solve_step(ode, x, y, dx_try),
}
}
pub fn integrate<Sys: OdeSystem + ?Sized>(
&mut self,
ode: &Sys,
x_start: f64,
x_end: f64,
y: &mut Vec<f64>,
dx_est: &mut f64,
) -> Result<(), OdeError> {
match self {
Self::Euler(s) => s.integrate(ode, x_start, x_end, y, dx_est),
Self::Rkf45(s) => s.integrate(ode, x_start, x_end, y, dx_est),
Self::Rosenbrock23(s) => s.integrate(ode, x_start, x_end, y, dx_est),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct NoTypedSystem;
impl OdeSystem for NoTypedSystem {
fn n_eqns(&self) -> usize {
0
}
fn derivatives(&self, _x: f64, _y: &[f64], _dydx: &mut Vec<f64>) {}
}
#[derive(Debug)]
pub struct TypedStateIntegrator<S: OdeSystem> {
pub solver: OdeSolver,
pub system: S,
}
impl<S: OdeSystem> TypedStateIntegrator<S> {
pub fn new(system: S, solver: OdeSolver) -> Self {
Self { solver, system }
}
pub fn solve_step(
&mut self,
x: &mut f64,
y: &mut Vec<f64>,
dx_try: &mut f64,
) -> Result<(), OdeError> {
let Self { solver, system } = self;
solver.solve_step(&*system, x, y, dx_try)
}
pub fn integrate(
&mut self,
x_start: f64,
x_end: f64,
y: &mut Vec<f64>,
dx_est: &mut f64,
) -> Result<(), OdeError> {
let Self { solver, system } = self;
solver.integrate(&*system, x_start, x_end, y, dx_est)
}
pub fn into_system(self) -> S {
self.system
}
}
#[derive(Clone)]
pub struct DynSystemIntegrator {
pub solver: OdeSolver,
pub system: SharedOdeSystem,
}
impl std::fmt::Debug for DynSystemIntegrator {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DynSystemIntegrator")
.field("solver", &self.solver)
.field(
"system",
&format_args!("<dyn OdeSystem, {} eqns>", self.system.n_eqns()),
)
.finish()
}
}
impl DynSystemIntegrator {
pub fn new(system: SharedOdeSystem, solver: OdeSolver) -> Self {
Self { solver, system }
}
pub fn solve_step(
&mut self,
x: &mut f64,
y: &mut Vec<f64>,
dx_try: &mut f64,
) -> Result<(), OdeError> {
let Self { solver, system } = self;
solver.solve_step(system.as_ref(), x, y, dx_try)
}
pub fn integrate(
&mut self,
x_start: f64,
x_end: f64,
y: &mut Vec<f64>,
dx_est: &mut f64,
) -> Result<(), OdeError> {
let Self { solver, system } = self;
solver.integrate(system.as_ref(), x_start, x_end, y, dx_est)
}
}
#[derive(Debug)]
pub enum OdeIntegrator<S: OdeSystem = NoTypedSystem> {
TypedState(TypedStateIntegrator<S>),
DynSystem(DynSystemIntegrator),
}
impl<S: OdeSystem> OdeIntegrator<S> {
pub fn typed(system: S, solver: OdeSolver) -> Self {
Self::TypedState(TypedStateIntegrator::new(system, solver))
}
#[must_use]
pub fn n_eqns(&self) -> usize {
match self {
Self::TypedState(i) => i.system.n_eqns(),
Self::DynSystem(i) => i.system.n_eqns(),
}
}
#[must_use]
pub fn solver(&self) -> &OdeSolver {
match self {
Self::TypedState(i) => &i.solver,
Self::DynSystem(i) => &i.solver,
}
}
pub fn solver_mut(&mut self) -> &mut OdeSolver {
match self {
Self::TypedState(i) => &mut i.solver,
Self::DynSystem(i) => &mut i.solver,
}
}
#[must_use]
pub const fn is_typed_state(&self) -> bool {
matches!(self, Self::TypedState(_))
}
#[must_use]
pub fn typed_system(&self) -> Option<&S> {
match self {
Self::TypedState(i) => Some(&i.system),
Self::DynSystem(_) => None,
}
}
#[must_use]
pub fn shared_system(&self) -> Option<&SharedOdeSystem> {
match self {
Self::TypedState(_) => None,
Self::DynSystem(i) => Some(&i.system),
}
}
pub fn solve_step(
&mut self,
x: &mut f64,
y: &mut Vec<f64>,
dx_try: &mut f64,
) -> Result<(), OdeError> {
match self {
Self::TypedState(i) => i.solve_step(x, y, dx_try),
Self::DynSystem(i) => i.solve_step(x, y, dx_try),
}
}
pub fn integrate(
&mut self,
x_start: f64,
x_end: f64,
y: &mut Vec<f64>,
dx_est: &mut f64,
) -> Result<(), OdeError> {
match self {
Self::TypedState(i) => i.integrate(x_start, x_end, y, dx_est),
Self::DynSystem(i) => i.integrate(x_start, x_end, y, dx_est),
}
}
}
impl OdeIntegrator<NoTypedSystem> {
pub fn shared(system: SharedOdeSystem, solver: OdeSolver) -> Self {
Self::DynSystem(DynSystemIntegrator::new(system, solver))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::matrix::SquareMatrix;
struct Decay;
impl OdeSystem for Decay {
fn n_eqns(&self) -> usize {
1
}
fn derivatives(&self, _x: f64, y: &[f64], dydx: &mut Vec<f64>) {
dydx[0] = -y[0];
}
fn jacobian(&self, _x: f64, _y: &[f64], dfdx: &mut Vec<f64>, dfdy: &mut SquareMatrix) {
dfdx[0] = 0.0;
dfdy.set(0, 0, -1.0);
}
}
#[test]
fn typed_variant_reaches_every_stepper() {
let exact = (-1.0_f64).exp();
for (solver, tol) in [
(OdeSolver::euler(1, 1e-3, 1e-2), 1e-2),
(OdeSolver::rkf45(1, 1e-10, 1e-8), 1e-6),
(OdeSolver::rosenbrock23(1, 1e-10, 1e-8), 1e-6),
] {
let name = solver.name();
let mut integrator = OdeIntegrator::typed(Decay, solver);
let mut y = vec![1.0_f64];
let mut dx = 0.1;
integrator.integrate(0.0, 1.0, &mut y, &mut dx).unwrap();
let err = (y[0] - exact).abs();
println!("{name}: y = {:.16}, error = {err:.4e}", y[0]);
assert!(err < tol, "{name}: y = {:.16}, error = {err:.4e}", y[0]);
}
}
#[test]
fn dyn_variant_matches_typed_bit_for_bit() {
let mut typed = OdeIntegrator::typed(Decay, OdeSolver::rkf45(1, 1e-10, 1e-8));
let mut y_typed = vec![1.0_f64];
let mut dx = 0.1;
typed.integrate(0.0, 1.0, &mut y_typed, &mut dx).unwrap();
let shared: SharedOdeSystem = Arc::new(Decay);
let mut dynamic = OdeIntegrator::shared(shared, OdeSolver::rkf45(1, 1e-10, 1e-8));
let mut y_dyn = vec![1.0_f64];
let mut dx2 = 0.1;
dynamic.integrate(0.0, 1.0, &mut y_dyn, &mut dx2).unwrap();
println!("typed = {:.17}, shared = {:.17}", y_typed[0], y_dyn[0]);
assert_eq!(
y_typed[0], y_dyn[0],
"typed = {:.17}, shared = {:.17}",
y_typed[0], y_dyn[0]
);
}
#[test]
fn integrator_is_storable_without_a_lifetime() {
struct Holder {
integrator: OdeIntegrator<Decay>,
state: Vec<f64>,
}
let mut holder = Holder {
integrator: OdeIntegrator::typed(Decay, OdeSolver::rkf45(1, 1e-10, 1e-8)),
state: vec![1.0],
};
let mut dx = 0.1;
holder
.integrator
.integrate(0.0, 1.0, &mut holder.state, &mut dx)
.unwrap();
println!(
"stored integrator: n_eqns = {}, y = {:.16}",
holder.integrator.n_eqns(),
holder.state[0]
);
assert_eq!(holder.integrator.n_eqns(), 1);
assert!(
(holder.state[0] - (-1.0_f64).exp()).abs() < 1e-6,
"y = {:.16}",
holder.state[0]
);
}
#[test]
fn accessors_report_the_variant_truthfully() {
assert!(!OdeSolver::euler(1, 1e-6, 1e-4).requires_jacobian());
assert!(!OdeSolver::rkf45(1, 1e-6, 1e-4).requires_jacobian());
assert!(OdeSolver::rosenbrock23(1, 1e-6, 1e-4).requires_jacobian());
assert_eq!(OdeSolver::euler(1, 1e-6, 1e-4).name(), "Euler");
assert_eq!(OdeSolver::rkf45(1, 1e-6, 1e-4).name(), "RKF45");
assert_eq!(
OdeSolver::rosenbrock23(1, 1e-6, 1e-4).name(),
"Rosenbrock23"
);
let typed = OdeIntegrator::typed(Decay, OdeSolver::rkf45(1, 1e-6, 1e-4));
assert!(typed.typed_system().is_some());
assert!(typed.shared_system().is_none());
let shared: SharedOdeSystem = Arc::new(Decay);
let dynamic = OdeIntegrator::shared(shared, OdeSolver::rkf45(1, 1e-6, 1e-4));
assert!(dynamic.typed_system().is_none());
assert!(dynamic.shared_system().is_some());
}
#[test]
fn placeholder_system_integrates_as_a_no_op() {
let mut integrator = OdeIntegrator::typed(NoTypedSystem, OdeSolver::euler(0, 1e-6, 1e-4));
let mut y: Vec<f64> = Vec::new();
let mut dx = 0.5;
integrator.integrate(0.0, 1.0, &mut y, &mut dx).unwrap();
assert_eq!(integrator.n_eqns(), 0);
assert_eq!(y.len(), 0);
}
}