From 97f52ad04cb2bfe5b1e3ca3212f43b3a9da195d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juli=C3=A1n=20D=2E=20Ot=C3=A1lvaro?= Date: Thu, 20 Aug 2026 20:14:38 +0100 Subject: [PATCH] fix: handle close ODE stop times safely --- src/dsl/native.rs | 80 +++--- src/simulator/equation/ode/closure.rs | 18 ++ src/simulator/equation/ode/mod.rs | 96 +++---- tests/bolus_reinit_stop_time.rs | 383 ++++++++++++++++++++++++++ 4 files changed, 499 insertions(+), 78 deletions(-) create mode 100644 tests/bolus_reinit_stop_time.rs diff --git a/src/dsl/native.rs b/src/dsl/native.rs index f2fc1955..c215f4dc 100644 --- a/src/dsl/native.rs +++ b/src/dsl/native.rs @@ -234,7 +234,7 @@ impl FunctionSession for NativeFunctionSession<'_> { )) })?; - function(time, states, params, covariates, routes, derived, out); + unsafe { function(time, states, params, covariates, routes, derived, out) }; Ok(()) } } @@ -1350,11 +1350,8 @@ impl NativeOdeModel { S: OdeSolverMethod<'a, PMProblem<'a, F>>, { // Mirror the closure-based ODE event loop: stop at every infusion - // start and end boundary in addition to subject events, using the - // left-continuous rate while integrating toward a boundary and the - // right-continuous rate after reaching it. This keeps the JIT - // implementation numerically consistent with the reference [`ODE`] - // path (see `ode::run_events`). + // boundary in addition to subject events. This keeps compiled models + // numerically consistent with the reference [`ODE`] path. let infusion_boundary_times = solver.problem().eqn.infusion_boundary_times(); let mut infusion_boundary_cursor = 0usize; let mut index = 0usize; @@ -1448,6 +1445,10 @@ impl NativeOdeModel { Ok(OdeSolverStopReason::InternalTimestep) => continue, Ok(OdeSolverStopReason::TstopReached) => { solver.problem().eqn.set_left_continuity_time(None); + if solver.state().t != stop_time { + *solver.state_mut().t = stop_time; + pending_reinit = true; + } if is_infusion_boundary { pending_reinit = true; } @@ -1472,32 +1473,49 @@ impl NativeOdeModel { OdeSolverError::StopTimeAtCurrentTime, )) => { solver.problem().eqn.set_left_continuity_time(None); - let state_t = solver.state().t; - let stop_reached = crate::simulator::equation::ode::stop_time_reached( - stop_time, state_t, - ); - - if stop_reached { - if is_infusion_boundary { - pending_reinit = true; - } - // The requested stop is the current time within - // a small relative tolerance. If it is an - // infusion boundary before the next subject - // event, keep integrating toward the event; - // break only when the reached stop is the - // event time itself. - if stop_time < next_event_time { - continue; - } - break; + // Close forward stops are coalesced by contract after checking + // that doing so cannot omit material infusion input. An actually + // earlier stop uses a distinct error. Snap the logical time before the next + // event or RHS evaluation. Otherwise a state a few ULPs before + // an infusion boundary can be restarted with the old rate. + let state_time = solver.state().t; + if state_time > stop_time { + return Err(PharmsolError::from_solver_error( + diffsol::error::DiffsolError::OdeSolverError( + OdeSolverError::StopTimeBeforeCurrentTime { + stop_time, + state_time, + }, + ), + stop_time, + )); + } + let skipped_infusion = solver + .problem() + .eqn + .infusion_amount_between(state_time, stop_time); + let state_scale = solver + .state() + .y + .as_slice() + .iter() + .fold(0.0_f64, |scale, value| scale.max(value.abs())); + let material_tolerance = + self.atol.abs() + self.rtol.abs() * state_scale; + if skipped_infusion > material_tolerance { + return Err(PharmsolError::OtherError(format!( + "coalescing stop times from t = {state_time:.16e} to \ + t = {stop_time:.16e} would skip infusion amount \ + {skipped_infusion:.6e}, above tolerance \ + {material_tolerance:.6e}" + ))); + } + *solver.state_mut().t = stop_time; + pending_reinit = true; + if stop_time < next_event_time { + continue; } - return Err(PharmsolError::from_solver_error( - diffsol::error::DiffsolError::OdeSolverError( - OdeSolverError::StopTimeAtCurrentTime, - ), - stop_time, - )); + break; } Err(err) => { solver.problem().eqn.set_left_continuity_time(None); diff --git a/src/simulator/equation/ode/closure.rs b/src/simulator/equation/ode/closure.rs index 14ae22e5..bbf4f888 100644 --- a/src/simulator/equation/ode/closure.rs +++ b/src/simulator/equation/ode/closure.rs @@ -182,6 +182,20 @@ impl InfusionSchedule { &self.boundary_times } + /// Absolute infusion input omitted when two adjacent stops are coalesced. + /// The event loop calls this only before the next infusion boundary, so the + /// active rate is constant over the interval. + fn infusion_amount_between(&self, from: f64, to: f64) -> f64 { + if to <= from { + return 0.0; + } + let duration = to - from; + self.tracks + .iter() + .map(|track| track.rate_at_left(to).abs() * duration) + .sum() + } + fn fill_rate_vector(&self, time: f64, rateiv: &mut V) { let left_continuity_time = self.left_continuity_time.get(); rateiv.fill(0.0); @@ -414,6 +428,10 @@ where self.infusion_schedule.infusion_boundary_times() } + pub(crate) fn infusion_amount_between(&self, from: f64, to: f64) -> f64 { + self.infusion_schedule.infusion_amount_between(from, to) + } + /// Evaluate the full RHS (including the currently scheduled infusion /// rates) at time `t` into `dx`. /// diff --git a/src/simulator/equation/ode/mod.rs b/src/simulator/equation/ode/mod.rs index c90128ab..ffdc634d 100644 --- a/src/simulator/equation/ode/mod.rs +++ b/src/simulator/equation/ode/mod.rs @@ -313,7 +313,7 @@ fn _simulate_subject_dense( Some(error_models) => Some(ode.bind_error_models(error_models)?), None => None, }; - let bound_error_models = bound_error_models.as_ref().map(|models| &**models); + let bound_error_models = bound_error_models.as_deref(); let mut output = SubjectPredictions::new(ode.nparticles()); @@ -585,24 +585,6 @@ where state.dy.copy_from(dy_scratch); } -/// Whether a requested solver stop is effectively at the current state time. -/// -/// diffsol reports `StopTimeAtCurrentTime` not only for a stop exactly at the -/// current time, but also when its internal state time has landed a few ULPs -/// past the requested stop (adaptive steps may end slightly beyond a stop). -/// Dense output grids built with floating-point arithmetic routinely place -/// requested times a few ULPs away from event times (e.g. a `t += dt` -/// accumulation puts a point ~16 ULPs after a bolus at `t = 12`), so accept a -/// stop within a small relative tolerance of the current time instead of -/// erroring. The tolerance stays far below any meaningful time difference: -/// ~64-128 ULPs of the current time, i.e. at most ~1e-13 at `t = 12`. -/// -/// Shared with the DSL/JIT ODE path ([`crate::dsl::native::NativeOdeModel`]). -pub(crate) fn stop_time_reached(stop_time: f64, state_t: f64) -> bool { - let tolerance = f64::EPSILON * state_t.abs().max(1.0) * 64.0; - (stop_time - state_t).abs() <= tolerance -} - impl ODE { /// Generic event-loop runner, parameterized over the concrete solver type. #[allow(clippy::too_many_arguments)] @@ -758,6 +740,10 @@ impl ODE { Ok(OdeSolverStopReason::InternalTimestep) => continue, Ok(OdeSolverStopReason::TstopReached) => { solver.problem().eqn.set_left_continuity_time(None); + if solver.state().t != stop_time { + *solver.state_mut().t = stop_time; + pending_reinit = true; + } if is_infusion_boundary { pending_reinit = true; } @@ -782,33 +768,49 @@ impl ODE { OdeSolverError::StopTimeAtCurrentTime, )) => { solver.problem().eqn.set_left_continuity_time(None); - let state_t = solver.state().t; - let stop_reached = stop_time_reached(stop_time, state_t); - - if stop_reached { - if is_infusion_boundary { - pending_reinit = true; - } - // The requested stop is the current time within - // a small relative tolerance. If it is an - // infusion boundary before the next subject - // event, keep integrating toward the event; - // break only when the reached stop is the - // event time itself. Breaking early would skip - // the remaining interval and leave the solver - // state at the boundary when the observation is - // evaluated. - if stop_time < next_event_time { - continue; - } - break; + // Close forward stops are coalesced by contract after checking + // that doing so cannot omit material infusion input. An actually + // earlier stop uses a distinct error. Snap the logical time before the next + // event or RHS evaluation. Otherwise a state a few ULPs before + // an infusion boundary can be restarted with the old rate. + let state_time = solver.state().t; + if state_time > stop_time { + return Err(PharmsolError::from_solver_error( + diffsol::error::DiffsolError::OdeSolverError( + OdeSolverError::StopTimeBeforeCurrentTime { + stop_time, + state_time, + }, + ), + stop_time, + )); + } + let skipped_infusion = solver + .problem() + .eqn + .infusion_amount_between(state_time, stop_time); + let state_scale = solver + .state() + .y + .as_slice() + .iter() + .fold(0.0_f64, |scale, value| scale.max(value.abs())); + let material_tolerance = + self.atol.abs() + self.rtol.abs() * state_scale; + if skipped_infusion > material_tolerance { + return Err(PharmsolError::OtherError(format!( + "coalescing stop times from t = {state_time:.16e} to \ + t = {stop_time:.16e} would skip infusion amount \ + {skipped_infusion:.6e}, above tolerance \ + {material_tolerance:.6e}" + ))); + } + *solver.state_mut().t = stop_time; + pending_reinit = true; + if stop_time < next_event_time { + continue; } - return Err(PharmsolError::from_solver_error( - diffsol::error::DiffsolError::OdeSolverError( - OdeSolverError::StopTimeAtCurrentTime, - ), - stop_time, - )); + break; } Err(err) => { solver.problem().eqn.set_left_continuity_time(None); @@ -1349,8 +1351,8 @@ mod tests { // The infusion ends exactly one ULP after the observation at t = 10. // After landing on the observation stop the solver is already within // diffsol's round-off of the end boundary, so `set_stop_time` reports - // `StopTimeAtCurrentTime` and the loop accepts it through the - // same-ULP check. The reached boundary is *before* the observation at + // `StopTimeAtCurrentTime`, which confirms the boundary is reached. + // The reached boundary is *before* the observation at // t = 20, so the event loop must keep integrating toward it; breaking // early would evaluate the observation with the state frozen at the // boundary and miss the exponential decay. diff --git a/tests/bolus_reinit_stop_time.rs b/tests/bolus_reinit_stop_time.rs new file mode 100644 index 00000000..95f4c292 --- /dev/null +++ b/tests/bolus_reinit_stop_time.rs @@ -0,0 +1,383 @@ +//! Regression coverage for accepted stop times around state and RHS discontinuities. +//! +//! A solver restart can land a few ULPs from a requested event or infusion +//! boundary while diffsol still correctly reports that stop as reached. The +//! event loop must accept diffsol's `StopTimeAtCurrentTime`, align the logical +//! state time with the accepted stop, and restart with the post-boundary RHS. + +use pharmsol::prelude::*; + +#[cfg(feature = "dsl-jit")] +use pharmsol::dsl::{ + compile_module_source_to_runtime, CompiledRuntimeModel, RuntimeCompilationTarget, +}; + +const OBSERVATION_TIMES: [f64; 15] = [ + 0.0, + 0.35, + 0.516666666666667, + 0.983333333333333, + 1.48333333333333, + 2.0, + 2.5, + 3.0, + 4.0, + 4.98333333333333, + 6.98333333333333, + 7.98333333333333, + 10.0, + 11.0, + 12.0, +]; + +fn solver_cases() -> [(&'static str, OdeSolver); 4] { + [ + ("BDF", OdeSolver::Bdf), + ("TSIT45", OdeSolver::ExplicitRk(ExplicitRkTableau::Tsit45)), + ("TRBDF2", OdeSolver::Sdirk(SdirkTableau::TrBdf2)), + ("ESDIRK34", OdeSolver::Sdirk(SdirkTableau::Esdirk34)), + ] +} + +const PARAMETERS: [(&str, f64); 5] = [ + ("ka", 3.6156922578811646), + ("cl0", 1.0289061069488525), + ("vc0", 187.13204860687256), + ("q0", 2.4602913856506348), + ("vp0", 58.32162380218506), +]; + +fn subject_with_bolus_history() -> Subject { + let mut builder = Subject::builder("g34"); + for dose_index in -16..=0 { + builder = builder.bolus(f64::from(dose_index) * 12.0, 750.0, "input_1"); + } + for time in OBSERVATION_TIMES { + builder = builder.missing_observation(time, "outeq_1"); + } + builder.build() +} + +fn closure_model(solver: OdeSolver) -> equation::ODE { + equation::ODE::new( + |x, p, _t, dx, b, rateiv, _cov| { + fetch_params!(p, ka, cl0, vc0, q0, vp0); + let ke = cl0 / vc0; + let k23 = q0 / vc0; + let k32 = q0 / vp0; + + dx[0] = b[0] - x[0] * ka; + dx[1] = rateiv[0] + x[0] * ka + x[2] * k32 - x[1] * (ke + k23); + dx[2] = x[1] * k23 - x[2] * k32; + }, + |_p, _t, _cov| lag! {}, + |_p, _t, _cov| fa! {}, + |_p, _t, _cov, _x| {}, + |x, p, _t, _cov, y| { + fetch_params!(p, _ka, _cl0, vc0, _q0, _vp0); + y[0] = x[1] / vc0; + }, + ) + .with_nstates(3) + .with_ndrugs(1) + .with_nout(1) + .with_solver(solver) + .with_metadata( + equation::metadata::new("bolus_reinit_stop_time") + .parameters(["ka", "cl0", "vc0", "q0", "vp0"]) + .states(["x1", "x2", "x3"]) + .outputs(["outeq_1"]) + .routes([equation::Route::bolus("input_1") + .to_state("x1") + .expect_explicit_input()]), + ) + .expect("regression model metadata should validate") +} + +#[test] +fn closure_solvers_accept_reached_stop_after_bolus_restarts( +) -> Result<(), Box> { + for (label, solver) in solver_cases() { + let model = closure_model(solver); + let parameters = Parameters::with_model(&model, PARAMETERS)?; + let predictions = model + .estimate_predictions_dense(&subject_with_bolus_history(), parameters.as_slice()) + .unwrap_or_else(|error| panic!("{label}: bolus-history simulation failed: {error}")); + + assert_eq!(predictions.predictions().len(), OBSERVATION_TIMES.len()); + assert!(predictions + .predictions() + .iter() + .all(|prediction| prediction.prediction().is_finite())); + } + Ok(()) +} + +#[cfg(feature = "dsl-jit")] +const DSL_MODEL: &str = r#" +name = bolus_reinit_stop_time +kind = ode +params = ka, cl0, vc0, q0, vp0 +states = x1, x2, x3 +outputs = outeq_1 + +bolus(input_1) -> x1 +infusion(input_1) -> x2 + +cl = cl0 +vc = vc0 +q = q0 +vp = vp0 +ke = cl / vc +k23 = q / vc +k32 = q / vp + +dx(x1) = -(x1 * ka) +dx(x2) = x1 * ka + (x3 * k32) - (x2 * (ke + k23)) +dx(x3) = x2 * k23 - (x3 * k32) + +out(outeq_1) = x2 / vc +"#; + +fn infusion_boundary_subject() -> Subject { + Subject::builder("accepted_infusion_boundary") + .infusion(5.0, 100.0, "input_1", 10.0_f64.next_up() - 5.0) + .missing_observation(10.0, "cp") + .missing_observation(20.0, "cp") + .build() +} + +fn stepped_infusion_boundary_subject() -> Subject { + let mut builder = + Subject::builder("stepped_infusion_boundary").infusion(0.0, 100.0, "input_1", 12.0); + for time in OBSERVATION_TIMES.into_iter().filter(|time| *time < 12.0) { + builder = builder.missing_observation(time, "cp"); + } + builder.missing_observation(20.0, "cp").build() +} + +fn closure_infusion_model(solver: OdeSolver) -> equation::ODE { + equation::ODE::new( + |x, _p, _t, dx, _b, rateiv, _cov| { + dx[0] = rateiv[0] - 0.5 * x[0]; + }, + |_p, _t, _cov| lag! {}, + |_p, _t, _cov| fa! {}, + |_p, _t, _cov, _x| {}, + |x, _p, _t, _cov, y| y[0] = x[0], + ) + .with_nstates(1) + .with_ndrugs(1) + .with_nout(1) + .with_solver(solver) + .with_metadata( + equation::metadata::new("accepted_infusion_boundary") + .states(["central"]) + .outputs(["cp"]) + .routes([equation::Route::infusion("input_1").to_state("central")]), + ) + .expect("infusion metadata should validate") +} + +fn accepted_boundary_expected() -> f64 { + let delivered = 100.0 * (1.0 - (-2.5_f64).exp()) / 2.5; + delivered * (-5.0_f64).exp() +} + +fn stepped_boundary_expected() -> f64 { + (100.0 / 12.0) / 0.5 * (1.0 - (-6.0_f64).exp()) * (-4.0_f64).exp() +} + +fn assert_post_infusion_decay( + label: &str, + predictions: &SubjectPredictions, + expected: f64, + maximum_relative_error: f64, +) { + assert!(!predictions.predictions().is_empty(), "{label}"); + let actual = predictions.predictions().last().unwrap().prediction(); + let relative_error = (actual - expected).abs() / expected; + assert!( + relative_error < maximum_relative_error, + "{label}: post-infusion prediction {actual:.16e}, expected {expected:.16e}, relative error {relative_error:.3e}" + ); +} + +#[test] +fn closure_solvers_restart_with_post_infusion_rhs() { + for (label, solver) in solver_cases() { + let model = closure_infusion_model(solver); + for (scenario, subject, expected) in [ + ( + "accepted", + infusion_boundary_subject(), + accepted_boundary_expected(), + ), + ( + "stepped", + stepped_infusion_boundary_subject(), + stepped_boundary_expected(), + ), + ] { + let predictions = model + .estimate_predictions_dense(&subject, &[]) + .unwrap_or_else(|error| { + panic!("closure {label} {scenario} infusion failed: {error}") + }); + let maximum_relative_error = if label == "TSIT45" { 1.0e-3 } else { 1.0e-2 }; + assert_post_infusion_decay( + &format!("closure {label} {scenario}"), + &predictions, + expected, + maximum_relative_error, + ); + } + } +} + +#[cfg(feature = "dsl-jit")] +const INFUSION_DSL_MODEL: &str = r#" +name = accepted_infusion_boundary +kind = ode +states = central +outputs = cp +infusion(input_1) -> central +dx(central) = -(0.5 * central) +out(cp) = central +"#; + +#[test] +#[cfg(feature = "dsl-jit")] +fn jit_solvers_restart_with_post_infusion_rhs() -> Result<(), Box> { + for (label, solver) in solver_cases() { + let compiled = compile_module_source_to_runtime( + INFUSION_DSL_MODEL, + Some("accepted_infusion_boundary"), + RuntimeCompilationTarget::Jit, + |_, _| {}, + )?; + let model = match compiled { + CompiledRuntimeModel::Ode(model) => { + CompiledRuntimeModel::Ode(model.with_solver(solver)) + } + _ => return Err("expected an ODE model".into()), + }; + for (scenario, subject, expected) in [ + ( + "accepted", + infusion_boundary_subject(), + accepted_boundary_expected(), + ), + ( + "stepped", + stepped_infusion_boundary_subject(), + stepped_boundary_expected(), + ), + ] { + let predictions = match &model { + CompiledRuntimeModel::Ode(model) => model + .estimate_predictions_dense(&subject, &[]) + .unwrap_or_else(|error| { + panic!("JIT {label} {scenario} infusion failed: {error}") + }), + _ => unreachable!(), + }; + let maximum_relative_error = if label == "TSIT45" { 1.0e-3 } else { 1.0e-2 }; + assert_post_infusion_decay( + &format!("JIT {label} {scenario}"), + &predictions, + expected, + maximum_relative_error, + ); + } + } + Ok(()) +} + +fn material_short_infusion_subject() -> Subject { + Subject::builder("material_short_infusion") + .infusion(1.0, 100.0, "input_1", 1.0_f64.next_up() - 1.0) + .missing_observation(2.0, "cp") + .build() +} + +fn assert_material_infusion_error(label: &str, error: PharmsolError) { + let message = error.to_string(); + assert!( + message.contains("would skip infusion amount"), + "{label}: unexpected error: {message}" + ); +} + +#[test] +fn closure_solvers_reject_material_coalesced_infusions() { + for (label, solver) in solver_cases() { + let model = closure_infusion_model(solver); + let error = model + .estimate_predictions_dense(&material_short_infusion_subject(), &[]) + .unwrap_err(); + assert_material_infusion_error(&format!("closure {label}"), error); + } +} + +#[test] +#[cfg(feature = "dsl-jit")] +fn jit_solvers_reject_material_coalesced_infusions() -> Result<(), Box> { + for (label, solver) in solver_cases() { + let compiled = compile_module_source_to_runtime( + INFUSION_DSL_MODEL, + Some("accepted_infusion_boundary"), + RuntimeCompilationTarget::Jit, + |_, _| {}, + )?; + let model = match compiled { + CompiledRuntimeModel::Ode(model) => { + CompiledRuntimeModel::Ode(model.with_solver(solver)) + } + _ => return Err("expected an ODE model".into()), + }; + let error = match &model { + CompiledRuntimeModel::Ode(model) => model + .estimate_predictions_dense(&material_short_infusion_subject(), &[]) + .unwrap_err(), + _ => unreachable!(), + }; + assert_material_infusion_error(&format!("JIT {label}"), error); + } + Ok(()) +} + +#[test] +#[cfg(feature = "dsl-jit")] +fn jit_solvers_accept_reached_stop_after_bolus_restarts() -> Result<(), Box> +{ + for (label, solver) in solver_cases() { + let compiled = compile_module_source_to_runtime( + DSL_MODEL, + Some("bolus_reinit_stop_time"), + RuntimeCompilationTarget::Jit, + |_, _| {}, + )?; + let model = match compiled { + CompiledRuntimeModel::Ode(model) => { + CompiledRuntimeModel::Ode(model.with_solver(solver)) + } + _ => return Err("expected an ODE model".into()), + }; + let parameters = Parameters::with_model(&model, PARAMETERS)?; + + let predictions = match &model { + CompiledRuntimeModel::Ode(model) => model + .estimate_predictions_dense(&subject_with_bolus_history(), parameters.as_slice()) + .unwrap_or_else(|error| panic!("{label}: JIT bolus-history failed: {error}")), + _ => unreachable!(), + }; + + assert_eq!(predictions.predictions().len(), OBSERVATION_TIMES.len()); + assert!(predictions + .predictions() + .iter() + .all(|prediction| prediction.prediction().is_finite())); + } + Ok(()) +}