Merge pull request #2869 from ChrisRackauckas-Claude/fix-tstop-overshoot-with-flag

Fix tstop overshoot error and super dense time event triggers

Author: ChrisRackauckas

lib/OrdinaryDiffEqCore/ext/OrdinaryDiffEqCoreMooncakeExt.jl

@@ -11,7 +11,7 @@ Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{
     typeof(OrdinaryDiffEqCore.increment_nf!), Vararg,
 }
 Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{
-    typeof(OrdinaryDiffEqCore.fixed_t_for_floatingpoint_error!), Vararg,
+    typeof(OrdinaryDiffEqCore.fixed_t_for_tstop_error!), Vararg,
 }
 Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{
     typeof(OrdinaryDiffEqCore.increment_accept!), Vararg,

lib/OrdinaryDiffEqCore/src/enzyme_rules.jl

@@ -5,7 +5,7 @@ function EnzymeCore.EnzymeRules.inactive_noinl(
 end
 
 function EnzymeCore.EnzymeRules.inactive_noinl(
-        ::typeof(fixed_t_for_floatingpoint_error!), args...
+        ::typeof(fixed_t_for_tstop_error!), args...
     )
     return true
 end

lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl

@@ -94,21 +94,67 @@ function last_step_failed(integrator::ODEIntegrator)
     return integrator.last_stepfail && !integrator.opts.adaptive
 end
 
+# Accessor functions for tstop flag fields with fallbacks for non-ODE integrators
+# (e.g. DDEIntegrator in DelayDiffEq.jl which doesn't have these fields)
+_get_next_step_tstop(integrator::ODEIntegrator) = integrator.next_step_tstop
+_get_next_step_tstop(integrator) = false
+
+function _set_tstop_flag!(integrator::ODEIntegrator, is_tstop::Bool, target = nothing)
+    integrator.next_step_tstop = is_tstop
+    if is_tstop && target !== nothing
+        integrator.tstop_target = target
+    end
+    return nothing
+end
+_set_tstop_flag!(integrator, is_tstop::Bool, target = nothing) = nothing
+
+_get_tstop_target(integrator::ODEIntegrator) = integrator.tstop_target
+
 function modify_dt_for_tstops!(integrator)
     return if has_tstop(integrator)
         tdir_t = integrator.tdir * integrator.t
         tdir_tstop = first_tstop(integrator)
+        distance_to_tstop = abs(tdir_tstop - tdir_t)
+
         if integrator.opts.adaptive
-            integrator.dt = integrator.tdir *
-                min(abs(integrator.dt), abs(tdir_tstop - tdir_t)) # step! to the end
+            original_dt = abs(integrator.dt)
+            integrator.dtpropose = original_dt
+            if original_dt < distance_to_tstop
+                _set_tstop_flag!(integrator, false)
+            else
+                _set_tstop_flag!(
+                    integrator, true, integrator.tdir * tdir_tstop)
+            end
+            integrator.dt = integrator.tdir * min(original_dt, distance_to_tstop)
         elseif iszero(integrator.dtcache) && integrator.dtchangeable
-            integrator.dt = integrator.tdir * abs(tdir_tstop - tdir_t)
+            integrator.dt = integrator.tdir * distance_to_tstop
+            _set_tstop_flag!(
+                integrator, true, integrator.tdir * tdir_tstop)
         elseif integrator.dtchangeable && !integrator.force_stepfail
             # always try to step! with dtcache, but lower if a tstop
             # however, if force_stepfail then don't set to dtcache, and no tstop worry
+            if abs(integrator.dtcache) < distance_to_tstop
+                _set_tstop_flag!(integrator, false)
+            else
+                _set_tstop_flag!(
+                    integrator, true, integrator.tdir * tdir_tstop)
+            end
             integrator.dt = integrator.tdir *
-                min(abs(integrator.dtcache), abs(tdir_tstop - tdir_t)) # step! to the end
+                min(abs(integrator.dtcache), distance_to_tstop)
+        else
+            _set_tstop_flag!(integrator, false)
         end
+    else
+        _set_tstop_flag!(integrator, false)
+    end
+end
+
+function handle_tstop_step!(integrator)
+    return if integrator.t isa AbstractFloat && abs(integrator.dt) < eps(abs(integrator.t))
+        # Skip perform_step! entirely for tiny dt
+        integrator.accept_step = true
+    else
+        perform_step!(integrator, integrator.cache)
     end
 end
 
@@ -183,7 +229,7 @@ function _savevalues!(integrator, force_save, reduce_size)::Tuple{Bool, Bool}
             integrator.opts.save_everystep &&
                 (
                 isempty(integrator.sol.t) ||
-                    (integrator.t !== integrator.sol.t[end]) &&
+                    (integrator.t !== integrator.sol.t[end] || iszero(integrator.dt)) &&
                     (integrator.opts.save_end || integrator.t !== integrator.sol.prob.tspan[2])
             )
         )
@@ -344,6 +390,15 @@ function _loopfooter!(integrator)
         if integrator.accept_step # Accept
             increment_accept!(integrator.stats)
             integrator.last_stepfail = false
+            integrator.tprev = integrator.t
+
+            if _get_next_step_tstop(integrator)
+                # Step controller dt is overly pessimistic, since dt = time to tstop.
+                # Restore the original dt so the controller proposes a reasonable next step.
+                integrator.dt = integrator.dtpropose
+            end
+            integrator.t = fixed_t_for_tstop_error!(integrator, ttmp)
+
             dtnew = DiffEqBase.value(
                 step_accept_controller!(
                     integrator,
@@ -352,8 +407,6 @@ function _loopfooter!(integrator)
                 )
             ) *
                 oneunit(integrator.dt)
-            integrator.tprev = integrator.t
-            integrator.t = fixed_t_for_floatingpoint_error!(integrator, ttmp)
             calc_dt_propose!(integrator, dtnew)
             handle_callbacks!(integrator)
         else # Reject
@@ -362,7 +415,7 @@ function _loopfooter!(integrator)
     elseif !integrator.opts.adaptive #Not adaptive
         increment_accept!(integrator.stats)
         integrator.tprev = integrator.t
-        integrator.t = fixed_t_for_floatingpoint_error!(integrator, ttmp)
+        integrator.t = fixed_t_for_tstop_error!(integrator, ttmp)
         integrator.last_stepfail = false
         integrator.accept_step = true
         integrator.dtpropose = integrator.dt
@@ -406,18 +459,12 @@ function log_step!(progress_name, progress_id, progress_message, dt, u, p, t, ts
     )
 end
 
-function fixed_t_for_floatingpoint_error!(integrator, ttmp)
-    return if has_tstop(integrator)
-        tstop = integrator.tdir * first_tstop(integrator)
-        if abs(ttmp - tstop) <
-                100eps(float(max(integrator.t, tstop) / oneunit(integrator.t))) *
-                oneunit(integrator.t)
-            tstop
-        else
-            ttmp
-        end
+function fixed_t_for_tstop_error!(integrator, ttmp)
+    if _get_next_step_tstop(integrator)
+        _set_tstop_flag!(integrator, false)
+        return _get_tstop_target(integrator)
     else
-        ttmp
+        return ttmp
     end
 end
 

lib/OrdinaryDiffEqCore/src/integrators/type.jl

@@ -128,6 +128,8 @@ mutable struct ODEIntegrator{
     force_stepfail::Bool
     last_stepfail::Bool
     just_hit_tstop::Bool
+    next_step_tstop::Bool
+    tstop_target::tType
     do_error_check::Bool
     event_last_time::Int
     vector_event_last_time::Int

lib/OrdinaryDiffEqCore/src/solve.jl

@@ -629,6 +629,8 @@ function SciMLBase.__init(
     u_modified = false
     EEst = oneunit(EEstT) # https://github.com/JuliaPhysics/Measurements.jl/pull/135
     just_hit_tstop = false
+    next_step_tstop = false
+    tstop_target = zero(t)
     isout = false
     accept_step = false
     force_stepfail = false
@@ -678,7 +680,7 @@ function SciMLBase.__init(
         callback_cache,
         kshortsize, force_stepfail,
         last_stepfail,
-        just_hit_tstop, do_error_check,
+        just_hit_tstop, next_step_tstop, tstop_target, do_error_check,
         event_last_time,
         vector_event_last_time,
         last_event_error, accept_step,
@@ -741,14 +743,24 @@ end
 
 function SciMLBase.solve!(integrator::ODEIntegrator)
     @inbounds while !isempty(integrator.opts.tstops)
-        while integrator.tdir * integrator.t < first(integrator.opts.tstops)
+        first_tstop = first(integrator.opts.tstops)
+        while integrator.tdir * integrator.t < first_tstop
             loopheader!(integrator)
             if integrator.do_error_check && check_error!(integrator) != ReturnCode.Success
                 return integrator.sol
             end
-            perform_step!(integrator, integrator.cache)
+
+            # Use special tstop handling if flag is set, otherwise normal stepping
+            if integrator.next_step_tstop
+                handle_tstop_step!(integrator)
+            else
+                perform_step!(integrator, integrator.cache)
+            end
+
+            should_exit = integrator.next_step_tstop
+
             loopfooter!(integrator)
-            if isempty(integrator.opts.tstops)
+            if isempty(integrator.opts.tstops) || should_exit
                 break
             end
         end
@@ -821,11 +833,11 @@ end
 
     for t in tstops
         tdir_t = tdir * t
-        tdir_t0 < tdir_t ≤ tdir_tf && push!(tstops_internal, tdir_t)
+        tdir_t0 < tdir_t < tdir_tf && push!(tstops_internal, tdir_t)
     end
     for t in d_discontinuities
         tdir_t = tdir * t
-        tdir_t0 < tdir_t ≤ tdir_tf && push!(tstops_internal, tdir_t)
+        tdir_t0 < tdir_t < tdir_tf && push!(tstops_internal, tdir_t)
     end
     push!(tstops_internal, tdir_tf)
 
@@ -842,11 +854,11 @@ function reinit_tstops!(::Type{T}, tstops_internal, tstops, d_discontinuities, t
 
     for t in tstops
         tdir_t = tdir * t
-        tdir_t0 < tdir_t ≤ tdir_tf && push!(tstops_internal, tdir_t)
+        tdir_t0 < tdir_t < tdir_tf && push!(tstops_internal, tdir_t)
     end
     for t in d_discontinuities
         tdir_t = tdir * t
-        tdir_t0 < tdir_t ≤ tdir_tf && push!(tstops_internal, tdir_t)
+        tdir_t0 < tdir_t < tdir_tf && push!(tstops_internal, tdir_t)
     end
     return push!(tstops_internal, tdir_tf)
 end

test/interface/ode_tstops_tests.jl

@@ -1,4 +1,4 @@
-using OrdinaryDiffEq, Test, Random
+using OrdinaryDiffEq, Test, Random, StaticArrays, DiffEqCallbacks
 import ODEProblemLibrary: prob_ode_linear
 Random.seed!(100)
 
@@ -100,3 +100,120 @@ end
     sol2 = solve(prob2, Tsit5())
     @test 0.0:0.07:1.0 ⊆ sol2.t
 end
+
+# Tests for issue #2752: tstop overshoot errors with StaticArrays
+
+@testset "StaticArrays vs Arrays with extreme precision" begin
+    function precise_dynamics(u, p, t)
+        x = @view u[1:2]
+        v = @view u[3:4]
+        dv = -0.01 * x + 1.0e-6 * sin(100 * t) * SVector{2}(1, 1)
+        return SVector{4}(v[1], v[2], dv[1], dv[2])
+    end
+
+    function precise_dynamics_array!(du, u, p, t)
+        x = @view u[1:2]
+        v = @view u[3:4]
+        dv = -0.01 * x + 1.0e-6 * sin(100 * t) * [1, 1]
+        du[1] = v[1]
+        du[2] = v[2]
+        du[3] = dv[1]
+        du[4] = dv[2]
+    end
+
+    u0_static = SVector{4}(1.0, -0.5, 0.01, 0.01)
+    u0_array = [1.0, -0.5, 0.01, 0.01]
+    tspan = (0.0, 2.0)
+    tstops = [0.5, 1.0, 1.5]
+
+    prob_static = ODEProblem(precise_dynamics, u0_static, tspan)
+    sol_static = solve(
+        prob_static, Vern9(); reltol = 1.0e-12, abstol = 1.0e-15,
+        tstops = tstops
+    )
+    @test SciMLBase.successful_retcode(sol_static)
+    for tstop in tstops
+        @test tstop ∈ sol_static.t
+    end
+
+    prob_array = ODEProblem(precise_dynamics_array!, u0_array, tspan)
+    sol_array = solve(
+        prob_array, Vern9(); reltol = 1.0e-12, abstol = 1.0e-15,
+        tstops = tstops
+    )
+    @test SciMLBase.successful_retcode(sol_array)
+    for tstop in tstops
+        @test tstop ∈ sol_array.t
+    end
+
+    @test isapprox(sol_static(2.0), sol_array(2.0), rtol = 1.0e-10)
+end
+
+@testset "Backward integration with tstop flags" begin
+    function decay_ode(u, p, t)
+        SA[-0.1 * u[1]]
+    end
+
+    u0 = SVector{1}(1.0)
+    tspan = (2.0, 0.0)
+    tstops = [1.5, 1.0, 0.5]
+
+    prob = ODEProblem(decay_ode, u0, tspan)
+    sol = solve(prob, Vern9(); tstops = tstops, reltol = 1.0e-12, abstol = 1.0e-15)
+    @test SciMLBase.successful_retcode(sol)
+    for tstop in tstops
+        @test tstop ∈ sol.t
+    end
+end
+
+@testset "PresetTimeCallback with tstop flags" begin
+    callback_times = Float64[]
+
+    function affect_preset!(integrator)
+        push!(callback_times, integrator.t)
+        integrator.u += 0.1 * integrator.u
+    end
+
+    function simple_growth(u, p, t)
+        SA[0.1 * u[1]]
+    end
+
+    u0 = SA[1.0]
+    tspan = (0.0, 3.0)
+    critical_times = [0.5, 1.0, 1.5, 2.0, 2.5]
+
+    preset_cb = PresetTimeCallback(critical_times, affect_preset!)
+
+    prob = ODEProblem(simple_growth, u0, tspan)
+    sol = solve(
+        prob, Vern9(); tstops = critical_times, callback = preset_cb,
+        reltol = 1.0e-10, abstol = 1.0e-12
+    )
+
+    @test SciMLBase.successful_retcode(sol)
+    for time in critical_times
+        @test any(abs.(sol.t .- time) .< 1.0e-10)
+    end
+    @test length(callback_times) == length(critical_times)
+end
+
+@testset "Multiple close tstops with StaticArrays" begin
+    function oscillator(u, p, t)
+        SVector{2}(u[2], -u[1])
+    end
+
+    u0 = SVector{2}(1.0, 0.0)
+    tspan = (0.0, 4.0)
+    close_tstops = [
+        1.0, 1.0 + 1.0e-14, 1.0 + 2.0e-14, 1.0 + 5.0e-14,
+        2.0, 2.0 + 1.0e-15, 2.0 + 1.0e-14,
+        3.0, 3.0 + 1.0e-13,
+    ]
+
+    prob = ODEProblem(oscillator, u0, tspan)
+    sol = solve(prob, Vern9(); tstops = close_tstops, reltol = 1.0e-12, abstol = 1.0e-15)
+    @test SciMLBase.successful_retcode(sol)
+    for time in [1.0, 2.0, 3.0]
+        @test any(abs.(sol.t .- time) .< 1.0e-10)
+    end
+end