Resync SDE noise when late add_tstop! shortens dt

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

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

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
 

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)

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

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!,
 )