Bridge SDE noise when a late tstop shortens dt

lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl

@@ -50,6 +50,25 @@ reinit_noise!(::Nothing, dt) = nothing
 @inline _get_W(integrator) = hasfield(typeof(integrator), :W) ? getfield(integrator, :W) : nothing
 @inline _get_P(integrator) = hasfield(typeof(integrator), :P) ? getfield(integrator, :P) : nothing
 
+# `fix_dt_at_bounds!`/`modify_dt_for_tstops!` can shorten `dt` after the pending
+# noise increment was already drawn for the longer step (e.g. `add_tstop!` called
+# between `init` and the first `step!`). Shrinking the step from the same start
+# time is exactly the situation a step rejection describes, so reuse the rejection
+# path: it bridges the drawn increment down to `dt` and keeps the remainder of the
+# path on the RSWM stack. Growing the step cannot be bridged, so the increment is
+# left alone; `perform_step!` then sees the noise the process was already
+# committed to.
+function shrink_noise_to_integrator_dt!(integrator)
+    W = _get_W(integrator)
+    isnothing(W) && return nothing
+    if abs(integrator.dt) < abs(W.dt)
+        reject_noise!(W, integrator.dt, integrator.u, integrator.p)
+        reject_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 +114,7 @@ function loopheader!(integrator)
     choose_algorithm!(integrator, integrator.cache)
     fix_dt_at_bounds!(integrator)
     modify_dt_for_tstops!(integrator)
+    shrink_noise_to_integrator_dt!(integrator)
     integrator.force_stepfail = false
     return nothing
 end

lib/StochasticDiffEq/test/tstops_tests.jl

@@ -1,4 +1,5 @@
 using StochasticDiffEq, Test, Random
+using DiffEqNoiseProcess: WienerProcess, RSWM
 using SDEProblemLibrary: prob_sde_linear
 Random.seed!(100)
 prob = prob_sde_linear
@@ -36,3 +37,70 @@ 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)
+
+# The noise has to follow the solver onto the shortened step. Without bridging,
+# W keeps stepping on the original dt grid and drifts away from sol.t for the
+# rest of the integration.
+for alg in (EM(), LambaEM())
+    i = init(prob3175, alg; dt = 0.02, save_noise = true)
+    add_tstop!(i, 0.01)
+    solve!(i)
+    @test i.W.t == i.sol.t
+    @test i.W.curt == i.t
+end
+
+# ... and bridging is what it must do: committing the pending increment instead
+# would make the noise take a step the solver never took.
+for alg in (EM(), LambaEM())
+    i = init(prob3175, alg; dt = 0.02)
+    add_tstop!(i, 0.01)
+    solve!(i)
+    @test i.W.iter == i.iter
+end
+
+# For du = dW the Euler step is exact, so the endpoints must agree with the
+# Brownian path the solution was built from.
+fexact(u, p, t) = 0.0
+gexact(u, p, t) = 1.0
+probexact = SDEProblem(fexact, gexact, [0.0], (0.0, 1.0))
+
+for alg in (EM(), LambaEM())
+    i = init(probexact, alg; dt = 0.02, save_noise = true)
+    add_tstop!(i, 0.01)
+    solve!(i)
+    @test i.sol.u[end][1] ≈ only(i.W.curW) atol = 1.0e-12
+end
+
+# RSwM1 hands a stack chunk back through W.dt, so the tstop shortening has to
+# bridge that chunk rather than relabel it.
+for adaptivealg in (:RSwM1, :RSwM2, :RSwM3)
+    W1 = WienerProcess(0.0, 0.0; rswm = RSWM(adaptivealg = adaptivealg))
+    prob1 = SDEProblem(fexact, gexact, 0.0, (0.0, 1.0), noise = W1)
+    i = init(prob1, LambaEM(); dt = 0.02, tstops = [0.331], save_noise = true)
+    add_tstop!(i, 0.01)
+    solve!(i)
+    @test 0.01 ∈ i.sol.t
+    @test 0.331 ∈ i.sol.t
+    @test i.sol.u[end] ≈ i.W.curW atol = 1.0e-12
+end