diff --git a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl b/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl index a17ab531eb..3f243b408d 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl +++ b/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 diff --git a/lib/StochasticDiffEq/test/tstops_tests.jl b/lib/StochasticDiffEq/test/tstops_tests.jl index b897e3bc18..962cc1198d 100644 --- a/lib/StochasticDiffEq/test/tstops_tests.jl +++ b/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