From 40b662695b9d04b506fa1fca694e2d08cde0f2c5 Mon Sep 17 00:00:00 2001 From: Niyati Jain Date: Mon, 27 Jul 2026 22:44:50 +0530 Subject: [PATCH] Resync SDE noise when late add_tstop! shortens dt (#3175) Fixed-step EM could step past a tstop added after init because the Wiener increment still matched the original dt. After modify_dt_for_tstops!, resample noise to the shortened dt, and keep RSWM from expanding dt while tstops remain. --- docs/src/devtools/internals/public_api.md | 1 + .../src/OrdinaryDiffEqCore.jl | 2 +- .../src/integrators/integrator_utils.jl | 34 ++++++++++++++++++- lib/StochasticDiffEq/test/tstops_tests.jl | 21 ++++++++++++ .../src/integrators/integrator_utils.jl | 7 ++++ lib/StochasticDiffEqCore/test/qa/qa.jl | 3 +- 6 files changed, 65 insertions(+), 3 deletions(-) diff --git a/docs/src/devtools/internals/public_api.md b/docs/src/devtools/internals/public_api.md index ccbd7139878..740501c8ce1 100644 --- a/docs/src/devtools/internals/public_api.md +++ b/docs/src/devtools/internals/public_api.md @@ -295,6 +295,7 @@ Used by SDE/RODE solver sublibraries; no-ops for pure ODEs. ```@docs OrdinaryDiffEqCore.accept_noise! OrdinaryDiffEqCore.reject_noise! +OrdinaryDiffEqCore.resync_noise! OrdinaryDiffEqCore.save_noise! OrdinaryDiffEqCore.reinit_noise! OrdinaryDiffEqCore.noise_curt diff --git a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl index 3faf0ecb3b2..82f1cc475e3 100644 --- a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl +++ b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl @@ -479,7 +479,7 @@ include("precompilation_setup.jl") :set_discontinuity, :resolve_basic, # Noise hooks used by the SDE/RODE solver sublibs. :accept_noise!, :reinit_noise!, :reject_noise!, :save_noise!, :noise_curt, - :is_noise_saveable, + :is_noise_saveable, :resync_noise!, # Docstring builder used by solver sublibs. :differentiation_rk_docstring, # SDE/RODE abstract cache hierarchy. The SDE analogue of the already-public diff --git a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl b/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl index 0c1b95a5467..eb44eb91bbd 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl +++ b/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl @@ -44,12 +44,35 @@ Reset the noise process `W` to its initial state for a fresh integration with st `dt` (used by `reinit!`). No-op when `W` is `nothing`. """ reinit_noise!(::Nothing, dt) = nothing +""" + resync_noise!(W, dt, u, p) + +Re-align the noise process `W` to a newly shortened integrator `dt` (typically +after `modify_dt_for_tstops!`) and resample the pending increment without +advancing the noise clock. No-op when `W` is `nothing`. Extended by +StochasticDiffEq for `NoiseProcess` types. +""" +resync_noise!(::Nothing, args...) = nothing # Noise field accessors — safe for any integrator type. # ODEIntegrator has W/P/sqdt; other integrators (DDEIntegrator) don't. @inline _get_W(integrator) = hasfield(typeof(integrator), :W) ? getfield(integrator, :W) : nothing @inline _get_P(integrator) = hasfield(typeof(integrator), :P) ? getfield(integrator, :P) : nothing +# After `modify_dt_for_tstops!` shortens `dt`, the pending `W.dW` may still match +# the previous step size. Resample so the forthcoming `perform_step!` uses a +# consistent increment (needed for late `add_tstop!` on fixed-step SDE methods). +function resync_noise_to_integrator_dt!(integrator) + W = _get_W(integrator) + isnothing(W) && return nothing + if W.dt != integrator.dt + resync_noise!(W, integrator.dt, integrator.u, integrator.p) + resync_noise!(_get_P(integrator), integrator.dt, integrator.u, integrator.p) + integrator.sqdt = integrator.tdir * sqrt(abs(integrator.dt)) + end + return nothing +end + # Trait: does the integrator+solution support dense output k-array storage? # True for ODEIntegrator (has integrator.k and sol.k), false for SDEIntegrator # (no integrator.k) and RODESolution/DAESolution (no sol.k). @@ -95,6 +118,7 @@ function loopheader!(integrator) choose_algorithm!(integrator, integrator.cache) fix_dt_at_bounds!(integrator) modify_dt_for_tstops!(integrator) + resync_noise_to_integrator_dt!(integrator) integrator.force_stepfail = false return nothing end @@ -166,7 +190,15 @@ function apply_step!(integrator) accept_noise!(W, integrator.dt, integrator.u, integrator.p, true) accept_noise!(_get_P(integrator), integrator.dt, integrator.u, integrator.p, true) if !isnothing(W) - integrator.dt = W.dt # RSWM readback + # RSWM may propose a new dt via W.dt, but must not expand past a pending + # tstop that already shortened integrator.dt (SciML/OrdinaryDiffEq.jl#3175). + if has_tstop(integrator) + W.dt = integrator.dt + P = _get_P(integrator) + !isnothing(P) && (P.dt = integrator.dt) + else + integrator.dt = W.dt # RSWM readback + end integrator.sqdt = @fastmath integrator.tdir * sqrt(abs(integrator.dt)) end diff --git a/lib/StochasticDiffEq/test/tstops_tests.jl b/lib/StochasticDiffEq/test/tstops_tests.jl index b897e3bc18c..f87af0e540d 100644 --- a/lib/StochasticDiffEq/test/tstops_tests.jl +++ b/lib/StochasticDiffEq/test/tstops_tests.jl @@ -36,3 +36,24 @@ for (i, tdir) in enumerate([-1.0; 1.0]) @test tstop ∈ integrator.sol.t end end + +# SciML/OrdinaryDiffEq.jl#3175 / StochasticDiffEq.jl#413: +# late add_tstop! after init must shorten the first EM step to the tstop. +f3175(u, p, t) = 1.0 +g3175(u, p, t) = 0.1 +prob3175 = SDEProblem(f3175, g3175, [1.0], (0.0, 1.0)) + +i_em = init(prob3175, EM(); dt = 0.02) +add_tstop!(i_em, 0.01) +step!(i_em) +@test i_em.t == 0.01 + +i_lamba = init(prob3175, LambaEM(); dt = 0.02) +add_tstop!(i_lamba, 0.01) +step!(i_lamba) +@test i_lamba.t == 0.01 + +# Past-time add_tstop! is rejected for both fixed-step and adaptive EM. +i_past = init(prob3175, EM(); dt = 0.02) +step!(i_past) +@test_throws ErrorException add_tstop!(i_past, 0.01) diff --git a/lib/StochasticDiffEqCore/src/integrators/integrator_utils.jl b/lib/StochasticDiffEqCore/src/integrators/integrator_utils.jl index 3634e4ccc89..552ae1aa4ec 100644 --- a/lib/StochasticDiffEqCore/src/integrators/integrator_utils.jl +++ b/lib/StochasticDiffEqCore/src/integrators/integrator_utils.jl @@ -80,6 +80,13 @@ function OrdinaryDiffEqCore.accept_noise!(W::SciMLBase.AbstractNoiseProcess, dt, return DiffEqNoiseProcess.accept_step!(W, dt, u, p, setup) end +function OrdinaryDiffEqCore.resync_noise!(W::SciMLBase.AbstractNoiseProcess, dt, u, p) + W.dt = dt + DiffEqNoiseProcess.accept_step!(W, dt, u, p, false) + DiffEqNoiseProcess.setup_next_step!(W, u, p) + return nothing +end + function OrdinaryDiffEqCore.reject_noise!(W::SciMLBase.AbstractNoiseProcess, dt, u, p) return DiffEqNoiseProcess.reject_step!(W, dt, u, p) end diff --git a/lib/StochasticDiffEqCore/test/qa/qa.jl b/lib/StochasticDiffEqCore/test/qa/qa.jl index 1b22d5c4d58..8576519aea6 100644 --- a/lib/StochasticDiffEqCore/test/qa/qa.jl +++ b/lib/StochasticDiffEqCore/test/qa/qa.jl @@ -18,7 +18,7 @@ const ODEC_INTERNAL = ( :_initialize_dae!, :_ode_init, :accept_noise!, :concrete_jac, :get_chunksize, :handle_callback_modifiers!, :has_autodiff, :is_noise_saveable, :noise_curt, :ode_determine_initdt, :qsteady_max_default, :qsteady_min_default, - :reinit_noise!, :reject_noise!, :save_noise!, :standardtag, + :reinit_noise!, :reject_noise!, :resync_noise!, :save_noise!, :standardtag, ) # Still-internal SciMLBase interface names (pending SciMLBase#1412 round-5 @@ -70,6 +70,7 @@ const ODEC_STOCHASTIC_SURFACE = ( OrdinaryDiffEqCore.noise_curt, OrdinaryDiffEqCore.reinit_noise!, OrdinaryDiffEqCore.reject_noise!, + OrdinaryDiffEqCore.resync_noise!, OrdinaryDiffEqCore.save_noise!, )