From 96fad713f40083980a8a02e12b767795f91e61c6 Mon Sep 17 00:00:00 2001 From: Markus Date: Wed, 29 Jul 2026 21:04:34 +0200 Subject: [PATCH] feat: Add static or dynamic stepsize --- src/dsl/native.rs | 34 ++++++- src/lib.rs | 3 +- src/simulator/equation/sde/em.rs | 143 ++++++++++++++++++++++++------ src/simulator/equation/sde/mod.rs | 48 +++++++++- 4 files changed, 197 insertions(+), 31 deletions(-) diff --git a/src/dsl/native.rs b/src/dsl/native.rs index 64fe8d60..1362ad58 100644 --- a/src/dsl/native.rs +++ b/src/dsl/native.rs @@ -34,7 +34,7 @@ use crate::{ equation::{ metadata::ValidatedRoute, ode::{closure_helpers::PMProblem, ExplicitRkTableau, OdeSolver, SdirkTableau}, - sde::simulate_sde_event_with, + sde::{simulate_sde_event_with, SdeStepSize}, EqnKind, Equation, EquationPriv, EquationTypes, Predictions, }, likelihood::{Prediction, SubjectPredictions}, @@ -1057,6 +1057,7 @@ pub struct NativeSdeModel { shared: Arc, nparticles: usize, cache: Option, + step_size: SdeStepSize, } #[derive(Clone, Debug)] @@ -2013,6 +2014,7 @@ impl NativeSdeModel { shared: Arc::new(SharedNativeModel::new(info, artifact)?), nparticles, cache: Some(SdeLikelihoodCache::new(DEFAULT_CACHE_SIZE)), + step_size: SdeStepSize::default(), }) } @@ -2036,6 +2038,35 @@ impl NativeSdeModel { self } + /// Configure the solver's step size strategy: either a fixed step size or + /// adaptive step size control (see [`SdeStepSize`]). + pub fn with_step_size(mut self, step_size: SdeStepSize) -> Self { + self.step_size = step_size; + self + } + + /// Use a fixed step size `dt` for every integration step, disabling + /// adaptive error control. + pub fn with_fixed_step_size(mut self, dt: f64) -> Self { + self.step_size = SdeStepSize::Fixed(dt); + self + } + + /// Divide every interval between consecutive bolus/observation events + /// into exactly `n` equal steps (infusions are applied continuously + /// within the drift function and don't create extra intervals). + pub fn with_event_steps(mut self, n: usize) -> Self { + self.step_size = SdeStepSize::EventSteps(n); + self + } + + /// Use adaptive step size control, targeting the given relative (`rtol`) + /// and absolute (`atol`) tolerances. + pub fn with_adaptive_step_size(mut self, rtol: f64, atol: f64) -> Self { + self.step_size = SdeStepSize::adaptive(rtol, atol); + self + } + pub fn info(&self) -> &NativeModelInfo { self.shared.info.as_ref() } @@ -2280,6 +2311,7 @@ impl NativeSdeModel { drift_state, start_time, end_time, + self.step_size, ); if let Some(error) = function_error.into_inner() { return Err(error); diff --git a/src/lib.rs b/src/lib.rs index 815d6993..da821810 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -151,7 +151,8 @@ pub use crate::simulator::equation::{ self, ode::{ExplicitRkTableau, OdeSolver, SdirkTableau}, Analytical, AnalyticalKernel, Cache, Equation, ModelKind, ModelMetadata, ModelMetadataError, - NameDomain, Predictions, RouteInputPolicy, RouteKind, State, ValidatedModelMetadata, ODE, SDE, + NameDomain, Predictions, RouteInputPolicy, RouteKind, SdeStepSize, State, + ValidatedModelMetadata, ODE, SDE, }; pub use error::PharmsolError; #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] diff --git a/src/simulator/equation/sde/em.rs b/src/simulator/equation/sde/em.rs index b80cfc07..64bfb46d 100644 --- a/src/simulator/equation/sde/em.rs +++ b/src/simulator/equation/sde/em.rs @@ -2,11 +2,50 @@ use nalgebra::DVector; use rand::rng; use rand_distr::{Distribution, Normal}; +/// Step size strategy for the Euler-Maruyama SDE solver. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum SdeStepSize { + /// Advance the solver in fixed-size increments of `dt`. + Fixed(f64), + /// Divide each solved interval (i.e. the span between two consecutive + /// bolus/observation events; infusions are applied continuously within + /// the drift function and don't create extra intervals) into exactly + /// `n` equal steps, regardless of the interval's length. + EventSteps(usize), + /// Adapt the step size between `min_step` and `max_step` so that the + /// estimated local error stays within `rtol`/`atol`. + Adaptive { + rtol: f64, + atol: f64, + min_step: f64, + max_step: f64, + }, +} + +impl SdeStepSize { + /// Convenience constructor for the adaptive strategy using the crate's + /// historical default `min_step`/`max_step` bounds. + pub fn adaptive(rtol: f64, atol: f64) -> Self { + SdeStepSize::Adaptive { + rtol, + atol, + min_step: 1e-6, + max_step: 0.1, + } + } +} + +impl Default for SdeStepSize { + fn default() -> Self { + SdeStepSize::adaptive(1e-2, 1e-2) + } +} + /// Implementation of the Euler-Maruyama method for solving stochastic differential equations. /// /// This structure holds the SDE system parameters and state, providing a numerical method -/// for approximating solutions to stochastic differential equations with adaptive step size -/// control for improved accuracy. +/// for approximating solutions to stochastic differential equations, using either a fixed +/// or an adaptive step size (see [`SdeStepSize`]). pub struct EM where D: Fn(f64, &DVector, &mut DVector), @@ -15,10 +54,7 @@ where drift: D, diffusion: G, state: DVector, - rtol: f64, - atol: f64, - max_step: f64, - min_step: f64, + step_size: SdeStepSize, } impl EM @@ -32,25 +68,23 @@ where /// /// * `drift` - Function defining the deterministic component of the SDE /// * `diffusion` - Function defining the stochastic component of the SDE - /// * `params` - Vector of model parameters /// * `initial_state` - Initial state vector of the system - /// * `cov` - Covariates that may influence the system dynamics - /// * `infusions` - Vector of infusion events to be applied during simulation - /// * `rtol` - Relative tolerance for adaptive step size control - /// * `atol` - Absolute tolerance for adaptive step size control + /// * `step_size` - Fixed or adaptive step size configuration /// /// # Returns /// /// A new instance of the Euler-Maruyama solver configured with the given parameters. - pub fn new(drift: D, diffusion: G, initial_state: DVector, rtol: f64, atol: f64) -> Self { + pub fn new( + drift: D, + diffusion: G, + initial_state: DVector, + step_size: SdeStepSize, + ) -> Self { Self { drift, diffusion, state: initial_state, - rtol, - atol, - max_step: 0.1, - min_step: 1e-6, + step_size, } } @@ -64,12 +98,12 @@ where /// # Returns /// /// The maximum normalized error between the two approximations. - fn calculate_error(&self, y1: &DVector, y2: &DVector) -> f64 { + fn calculate_error(&self, y1: &DVector, y2: &DVector, rtol: f64, atol: f64) -> f64 { let n = y1.len(); let mut err = 0.0f64; for i in 0..n { - let tol = self.atol + self.rtol * self.state[i].abs(); + let tol = atol + rtol * self.state[i].abs(); let e = (y1[i] - y2[i]).abs() / tol; err = err.max(e); } @@ -87,9 +121,17 @@ where /// # Returns /// /// The adjusted step size for the next iteration. - fn compute_new_step(&self, dt: f64, error: f64, safety: f64) -> f64 { + #[allow(clippy::too_many_arguments)] + fn compute_new_step( + &self, + dt: f64, + error: f64, + safety: f64, + min_step: f64, + max_step: f64, + ) -> f64 { let mut new_dt = dt * safety * (1.0 / error).powf(0.5); - new_dt = new_dt.clamp(self.min_step, self.max_step); + new_dt = new_dt.clamp(min_step, max_step); new_dt } @@ -117,9 +159,9 @@ where } } - /// Solves the SDE system over the specified time interval. - /// - /// Uses adaptive step size control to balance accuracy and performance. + /// Solves the SDE system over the specified time interval, using either a + /// fixed step size or adaptive step size control, depending on how the + /// solver was configured (see [`SdeStepSize`]). /// /// # Arguments /// @@ -132,8 +174,55 @@ where /// * Vector of time points where solutions were computed /// * Vector of state vectors corresponding to each time point pub fn solve(&mut self, t0: f64, tf: f64) -> (Vec, Vec>) { + match self.step_size { + SdeStepSize::Fixed(dt) => self.solve_fixed(t0, tf, dt), + SdeStepSize::EventSteps(n) => { + let dt = (tf - t0) / n.max(1) as f64; + self.solve_fixed(t0, tf, dt) + } + SdeStepSize::Adaptive { + rtol, + atol, + min_step, + max_step, + } => self.solve_adaptive(t0, tf, rtol, atol, min_step, max_step), + } + } + + /// Advances the system from `t0` to `tf` in fixed increments of `dt`, + /// with no local error control (the final step is truncated to land exactly on `tf`). + fn solve_fixed(&mut self, t0: f64, tf: f64, dt: f64) -> (Vec, Vec>) { + let mut t = t0; + let mut times = vec![t0]; + let mut solution = vec![self.state.clone()]; + + while t < tf { + let step = dt.min(tf - t); + let mut next = self.state.clone(); + self.euler_maruyama_step(t, step, &mut next); + self.state = next; + t += step; + times.push(t); + solution.push(self.state.clone()); + } + + (times, solution) + } + + /// Advances the system from `t0` to `tf`, adapting the step size so that the + /// estimated local error (via step-doubling) stays within `rtol`/`atol`. + #[allow(clippy::too_many_arguments)] + fn solve_adaptive( + &mut self, + t0: f64, + tf: f64, + rtol: f64, + atol: f64, + min_step: f64, + max_step: f64, + ) -> (Vec, Vec>) { let mut t = t0; - let mut dt = self.max_step; + let mut dt = max_step; let safety = 0.9; let mut times = vec![t0]; let mut solution = vec![self.state.clone()]; @@ -149,17 +238,17 @@ where self.euler_maruyama_step(t, dt / 2.0, &mut y2); self.euler_maruyama_step(t + dt / 2.0, dt / 2.0, &mut y2); - let error = self.calculate_error(&y1, &y2); + let error = self.calculate_error(&y1, &y2, rtol, atol); if error <= 1.0 { t += dt; self.state = y2; // Use more accurate solution times.push(t); solution.push(self.state.clone()); - dt = self.compute_new_step(dt, error, safety); + dt = self.compute_new_step(dt, error, safety, min_step, max_step); dt = dt.min(tf - t); // Don't step beyond tf } else { - dt = self.compute_new_step(dt, error, safety); + dt = self.compute_new_step(dt, error, safety, min_step, max_step); } } diff --git a/src/simulator/equation/sde/mod.rs b/src/simulator/equation/sde/mod.rs index 726daad6..ba3bcac6 100644 --- a/src/simulator/equation/sde/mod.rs +++ b/src/simulator/equation/sde/mod.rs @@ -1,5 +1,7 @@ mod em; +pub use em::SdeStepSize; + use diffsol::{NalgebraContext, Vector}; use nalgebra::DVector; use ndarray::{concatenate, Array2, Axis}; @@ -109,6 +111,7 @@ fn simulate_sde_event( ndrugs: usize, ti: f64, tf: f64, + step_size: SdeStepSize, ) -> V { if ti == tf { return x; @@ -151,7 +154,15 @@ fn simulate_sde_event( out.copy_from(out_v.inner()); }; - simulate_sde_event_with(drift_closure, diffusion_closure, x.inner().clone(), ti, tf).into() + simulate_sde_event_with( + drift_closure, + diffusion_closure, + x.inner().clone(), + ti, + tf, + step_size, + ) + .into() } pub(crate) fn simulate_sde_event_with( @@ -160,6 +171,7 @@ pub(crate) fn simulate_sde_event_with( initial_state: DVector, ti: f64, tf: f64, + step_size: SdeStepSize, ) -> DVector where D: Fn(f64, &DVector, &mut DVector), @@ -169,7 +181,7 @@ where return initial_state; } - let mut sde = em::EM::new(drift, diffusion, initial_state, 1e-2, 1e-2); + let mut sde = em::EM::new(drift, diffusion, initial_state, step_size); let (_time, solution) = sde.solve(ti, tf); solution.last().unwrap().clone() } @@ -195,6 +207,7 @@ pub struct SDE { injected_bolus_mappings: InjectedBolusMappings, cache: Option, error_model_cache: Option, + step_size: SdeStepSize, } impl SDE { @@ -231,6 +244,7 @@ impl SDE { error_model_cache: Some(BoundErrorModelCache::new( DEFAULT_BOUND_ERROR_MODEL_CACHE_SIZE, )), + step_size: SdeStepSize::default(), } } @@ -255,6 +269,35 @@ impl SDE { self } + /// Configure the solver's step size strategy: either a fixed step size or + /// adaptive step size control (see [`SdeStepSize`]). + pub fn with_step_size(mut self, step_size: SdeStepSize) -> Self { + self.step_size = step_size; + self + } + + /// Use a fixed step size `dt` for every integration step, disabling + /// adaptive error control. + pub fn with_fixed_step_size(mut self, dt: f64) -> Self { + self.step_size = SdeStepSize::Fixed(dt); + self + } + + /// Divide every interval between consecutive bolus/observation events + /// into exactly `n` equal steps (infusions are applied continuously + /// within the drift function and don't create extra intervals). + pub fn with_event_steps(mut self, n: usize) -> Self { + self.step_size = SdeStepSize::EventSteps(n); + self + } + + /// Use adaptive step size control, targeting the given relative (`rtol`) + /// and absolute (`atol`) tolerances. + pub fn with_adaptive_step_size(mut self, rtol: f64, atol: f64) -> Self { + self.step_size = SdeStepSize::adaptive(rtol, atol); + self + } + /// Attach validated handwritten-model metadata to this SDE model. pub fn with_metadata(mut self, metadata: ModelMetadata) -> Result { let metadata = metadata.validate_for_with_particles(ModelKind::Sde, self.nparticles)?; @@ -509,6 +552,7 @@ impl EquationPriv for SDE { ndrugs, ti, tf, + self.step_size, ) .inner() .clone();