Interpolate to tstops by step length and handle sub-dtmin stop spacing

src/ensemblegpukernel/integrators/integrator_utils.jl

@@ -14,30 +14,49 @@ function build_adaptive_controller_cache(::GPUTsit5IController, ::Type{T}) where
     return T(1 / 5), zero(T), T(5), T(1 / 5), T(9 / 10), T(1.0e-4), T(1.0e-4)
 end
 
+# Adaptive steps land on pending tstops and on `tf` by construction: `_bounded_step`
+# makes a trial step that would reach (or nearly reach) the next landing target exactly
+# the remaining distance, and the accept branch marks `t = tf` only for an accepted step
+# of length exactly `tf - t`. A step whose state is not finite is never accepted. Only a
+# landing on a stop before `tf` retries with the controller's covering step (landing on
+# the stop by interpolation), so no accepted step reaches past `tf`.
+# There is no tolerance-based endpoint snap. (At large |t| a step shortened by error
+# control can still round onto `tf` in `t += dt`; see
+# https://github.com/SciML/DiffEqGPU.jl/issues/567.)
+
 # `qold^beta2` of the PI controller. The I controller has `beta2 = 0`, so the factor is
 # exactly one and the device `pow` call is skipped.
 @inline _qold_factor(alg, qold, beta2) = qold^beta2
 @inline _qold_factor(::GPUTsit5IController, qold, beta2) = one(qold)
 
-# Whether a pending tstop before `tf` falls inside the step just accepted from `integ.t`.
+# Whether a pending tstop before `tf` lies inside the step just accepted from `integ.t`.
 @inline function _tstop_in_step(integ, tf, ::Type{T}) where {T}
     tstops = integ.tstops
     (tstops === nothing || integ.tstops_idx > length(tstops)) && return false
-    stop = @inbounds tstops[integ.tstops_idx]
-    return stop < tf && stop - integ.t - integ.dt - T(100) * eps(T) < T(0)
+    stop = convert(typeof(tf), @inbounds tstops[integ.tstops_idx])
+    return stop < tf && stop - integ.t <= integ.dt
 end
 
 # The next time the adaptive integrator has to land on exactly: the next pending
 # tstop after `integ.t` and before `tf`, otherwise `tf`.
-@inline function _next_stop(integ, tf)
+@inline function _next_stop(integ, tf, ::Type{T}) where {T}
     tstops = integ.tstops
     if tstops !== nothing && integ.tstops_idx <= length(tstops)
-        stop = @inbounds tstops[integ.tstops_idx]
+        stop = convert(typeof(tf), @inbounds tstops[integ.tstops_idx])
         integ.t < stop < tf && return stop
     end
     return tf
 end
 
+# Trial step from `integ.t`: never past the next landing target, and exactly the
+# remaining distance when the proposal reaches it or would leave less than 1% of a
+# step before it.
+@inline function _bounded_step(integ, dt, tf, ::Type{T}) where {T}
+    target = _next_stop(integ, tf, T)
+    remaining = target - integ.t
+    return remaining <= dt + dt / 100 || integ.t + dt >= target ? remaining : dt
+end
+
 @inline function savevalues!(
         integrator::SciMLBase.AbstractODEIntegrator{
             AlgType,
@@ -62,21 +81,35 @@ end
     elseif saveat !== nothing
         saved = true
         savedexactly = true
-        while integrator.cur_t <= length(saveat) && saveat[integrator.cur_t] <= integrator.t
-            savet = saveat[integrator.cur_t]
-            Θ = (savet - integrator.tprev) / integrator.dt
-            @inbounds us[integrator.cur_t] = _ode_interpolant(
-                Θ, integrator.dt,
-                integrator.uprev, integrator
-            )
-            @inbounds ts[integrator.cur_t] = savet
-            integrator.cur_t += 1
-        end
+        _save_saveat!(integrator, ts, us, integrator.t)
     end
 
     return saved, savedexactly
 end
 
+# Save the pending `saveat` points up to `tsave` from the current step. The dense output
+# of several methods uses `integrator.u` as the step's end state, so this must run
+# before `integrator.u` is moved back to an event time.
+@inline function _save_saveat!(integrator, ts, us, tsave)
+    saveat = integrator.saveat
+    while integrator.cur_t <= length(saveat) && saveat[integrator.cur_t] <= tsave
+        savet = saveat[integrator.cur_t]
+        # The step's endpoints are known exactly; only interior points need the
+        # dense output, which is not exact at Θ = 0 or 1 for every method.
+        @inbounds us[integrator.cur_t] = if savet == integrator.t
+            integrator.u
+        elseif savet == integrator.tprev
+            integrator.uprev
+        else
+            Θ = (savet - integrator.tprev) / integrator.dt
+            _ode_interpolant(Θ, integrator.dt, integrator.uprev, integrator)
+        end
+        @inbounds ts[integrator.cur_t] = savet
+        integrator.cur_t += 1
+    end
+    return nothing
+end
+
 @inline function DiffEqBase.terminate!(
         integrator::SciMLBase.AbstractODEIntegrator{
             AlgType,
@@ -263,6 +296,7 @@ end
         cb_time, prev_sign, event_idx, ts,
         us
     ) where {AlgType <: GPUODEAlgorithm, IIP, S, T}
+    integrator.saveat === nothing || _save_saveat!(integrator, ts, us, cb_time)
     DiffEqBase.change_t_via_interpolation!(integrator, cb_time)
 
     # The new absolute-time callback handling can leave dtnew ≈ 0 when

src/ensemblegpukernel/linalg/linsolve.jl

@@ -14,9 +14,8 @@ end
     @inbounds return similar_type(b, typeof(a[1] \ b[1]))(a[1] \ b[1])
 end
 
-@inline function _linear_solve(
+@inline function _cramer_solve(
         ::Size{(2, 2)},
-        ::Size{(2,)},
         a::StaticMatrix{<:Any, <:Any, Ta},
         b::StaticVector{<:Any, Tb}
     ) where {Ta, Tb}
@@ -28,9 +27,8 @@ end
     )
 end
 
-@inline function _linear_solve(
+@inline function _cramer_solve(
         ::Size{(3, 3)},
-        ::Size{(3,)},
         a::StaticMatrix{<:Any, <:Any, Ta},
         b::StaticVector{<:Any, Tb}
     ) where {Ta, Tb}
@@ -56,6 +54,26 @@ end
     )
 end
 
+# Cramer's rule forms products of up to n matrix entries, which overflow (or underflow)
+# for badly scaled systems such as a Rosenbrock W = I/(γ dt) - J with a tiny dt. When
+# that happens, solve the system divided by its largest entry instead: the solution is
+# unchanged, and well-scaled systems keep their exact unscaled arithmetic.
+for Sa in [(2, 2), (3, 3)]
+    @eval @inline function _linear_solve(
+            ::Size{$Sa},
+            ::Size{($Sa[1],)},
+            a::StaticMatrix{<:Any, <:Any, Ta},
+            b::StaticVector{<:Any, Tb}
+        ) where {Ta, Tb}
+        d = det(a)
+        x = _cramer_solve(Size($Sa), a, b)
+        isfinite(d) && !iszero(d) && all(isfinite, x) && return x
+        s = maximum(abs, a)
+        (isfinite(s) && !iszero(s)) || return x
+        return _cramer_solve(Size($Sa), a / s, b / s)
+    end
+end
+
 for Sa in [(2, 2), (3, 3)]  # not needed for Sa = (1, 1);
     @eval begin
         @inline function _linear_solve(

src/ensemblegpukernel/perform_step/gpu_kvaerno3_perform_step.jl

@@ -99,6 +99,9 @@ end
     t = integ.t
     p = integ.p
     tf = integ.tf
+    dtprop = dt
+    dt = _bounded_step(integ, dt, tf, T)
+    shortened_to_land = dt < dtprop && _next_stop(integ, tf, T) < tf
 
     tmp = integ.tmp
     f = integ.f
@@ -122,7 +125,9 @@ end
     EEst = convert(T, Inf)
 
     while EEst > convert(T, 1.0)
-        dt < convert(T, 1.0f-14) && error("dt<dtmin")
+        (dt < convert(T, 1.0f-14) || t + dt == t) &&
+            dt != _next_stop(integ, tf, T) - t &&
+            error("dt<dtmin")
 
         ## Steps
 
@@ -181,6 +186,9 @@ end
         tmp = (err) ./
             (abstol .+ max.(abs.(uprev), abs.(u)) * reltol)
         EEst = DiffEqBase.ODE_DEFAULT_NORM(tmp, t)
+        # `ODE_DEFAULT_NORM` is computed with fast-math, so NaN is checked on its inputs.
+        finite = mapreduce(isfinite, &, u) & mapreduce(isfinite, &, tmp)
+        finite || (EEst = T(Inf))
 
         q11 = EEst^beta1
         if iszero(EEst)
@@ -190,11 +198,15 @@ end
         end
 
         if EEst > 1
-            dt = dt / min(inv(qmin), q11 / gamma)
+            # A landing step on a stop that overflowed retries with the controller's covering
+            # step, which lands on the stop by interpolation; it never reaches past `tf`.
+            dt = !finite && shortened_to_land ? dtprop : dt / min(inv(qmin), q11 / gamma)
+            shortened_to_land = false
         else # EEst <= 1
             q = max(inv(qmax), min(inv(qmin), q / gamma))
             qold = max(EEst, qoldinit)
             dtnew = dt / q #dtnew
+            shortened_to_land && (dtnew = max(dtnew, dtprop))
 
             @inbounds begin # Necessary for interpolation
                 integ.k1 = k1
@@ -208,18 +220,19 @@ end
 
             if _tstop_in_step(integ, tf, T)
                 integ.t = integ.tstops[integ.tstops_idx]
-                if abs(integ.t - (t + dt)) > eps(integ.t)
+                if integ.t - t != dt
                     integ.u = integ(integ.t)
+                    mapreduce(isfinite, &, integ.u) || error("non-finite state at tstop")
                 end
                 dt = integ.t - integ.tprev
                 integ.tstops_idx += 1
-            elseif (tf - t - dt) < convert(T, 1.0f-14)
+            elseif dt == tf - t
                 integ.t = tf
             else
                 ##Advance the integrator
                 integ.t += dt
             end
-            integ.dtnew = min(abs(dtnew), abs(_next_stop(integ, tf) - integ.t))
+            integ.dtnew = dtnew
         end
     end
 

src/ensemblegpukernel/perform_step/gpu_kvaerno5_perform_step.jl

@@ -130,6 +130,9 @@ end
     t = integ.t
     p = integ.p
     tf = integ.tf
+    dtprop = dt
+    dt = _bounded_step(integ, dt, tf, T)
+    shortened_to_land = dt < dtprop && _next_stop(integ, tf, T) < tf
 
     tmp = integ.tmp
     f = integ.f
@@ -155,7 +158,9 @@ end
     EEst = convert(T, Inf)
 
     while EEst > convert(T, 1.0)
-        dt < convert(T, 1.0f-14) && error("dt<dtmin")
+        (dt < convert(T, 1.0f-14) || t + dt == t) &&
+            dt != _next_stop(integ, tf, T) - t &&
+            error("dt<dtmin")
 
         ## Steps
 
@@ -249,6 +254,9 @@ end
         tmp = (err) ./
             (abstol .+ max.(abs.(uprev), abs.(u)) * reltol)
         EEst = DiffEqBase.ODE_DEFAULT_NORM(tmp, t)
+        # `ODE_DEFAULT_NORM` is computed with fast-math, so NaN is checked on its inputs.
+        finite = mapreduce(isfinite, &, u) & mapreduce(isfinite, &, tmp)
+        finite || (EEst = T(Inf))
 
         q11 = EEst^beta1
         if iszero(EEst)
@@ -258,11 +266,15 @@ end
         end
 
         if EEst > 1
-            dt = dt / min(inv(qmin), q11 / gamma)
+            # A landing step on a stop that overflowed retries with the controller's covering
+            # step, which lands on the stop by interpolation; it never reaches past `tf`.
+            dt = !finite && shortened_to_land ? dtprop : dt / min(inv(qmin), q11 / gamma)
+            shortened_to_land = false
         else # EEst <= 1
             q = max(inv(qmax), min(inv(qmin), q / gamma))
             qold = max(EEst, qoldinit)
             dtnew = dt / q #dtnew
+            shortened_to_land && (dtnew = max(dtnew, dtprop))
 
             @inbounds begin # Necessary for interpolation
                 integ.k1 = k1
@@ -276,18 +288,19 @@ end
 
             if _tstop_in_step(integ, tf, T)
                 integ.t = integ.tstops[integ.tstops_idx]
-                if abs(integ.t - (t + dt)) > eps(integ.t)
+                if integ.t - t != dt
                     integ.u = integ(integ.t)
+                    mapreduce(isfinite, &, integ.u) || error("non-finite state at tstop")
                 end
                 dt = integ.t - integ.tprev
                 integ.tstops_idx += 1
-            elseif (tf - t - dt) < convert(T, 1.0f-14)
+            elseif dt == tf - t
                 integ.t = tf
             else
                 ##Advance the integrator
                 integ.t += dt
             end
-            integ.dtnew = min(abs(dtnew), abs(_next_stop(integ, tf) - integ.t))
+            integ.dtnew = dtnew
         end
     end
 

src/ensemblegpukernel/perform_step/gpu_rodas4_perform_step.jl

@@ -126,6 +126,9 @@ end
     t = integ.t
     p = integ.p
     tf = integ.tf
+    dtprop = dt
+    dt = _bounded_step(integ, dt, tf, T)
+    shortened_to_land = dt < dtprop && _next_stop(integ, tf, T) < tf
 
     tmp = integ.tmp
     f = integ.f
@@ -158,7 +161,9 @@ end
     EEst = convert(T, Inf)
 
     while EEst > convert(T, 1.0)
-        dt < convert(T, 1.0f-14) && error("dt<dtmin")
+        (dt < convert(T, 1.0f-14) || t + dt == t) &&
+            dt != _next_stop(integ, tf, T) - t &&
+            error("dt<dtmin")
 
         # Precalculations
         dtC21 = C21 / dt
@@ -225,6 +230,9 @@ end
 
         tmp = k6 ./ (abstol .+ max.(abs.(uprev), abs.(u)) * reltol)
         EEst = DiffEqBase.ODE_DEFAULT_NORM(tmp, t)
+        # `ODE_DEFAULT_NORM` is computed with fast-math, so NaN is checked on its inputs.
+        finite = mapreduce(isfinite, &, u) & mapreduce(isfinite, &, tmp)
+        finite || (EEst = T(Inf))
 
         q11 = EEst^beta1
         if iszero(EEst)
@@ -234,11 +242,15 @@ end
         end
 
         if EEst > 1
-            dt = dt / min(inv(qmin), q11 / gamma)
+            # A landing step on a stop that overflowed retries with the controller's covering
+            # step, which lands on the stop by interpolation; it never reaches past `tf`.
+            dt = !finite && shortened_to_land ? dtprop : dt / min(inv(qmin), q11 / gamma)
+            shortened_to_land = false
         else # EEst <= 1
             q = max(inv(qmax), min(inv(qmin), q / gamma))
             qold = max(EEst, qoldinit)
             dtnew = dt / q #dtnew
+            shortened_to_land && (dtnew = max(dtnew, dtprop))
 
             @inbounds begin # Necessary for interpolation
                 @unpack h21, h22, h23, h24, h25, h31, h32, h33, h34, h35 = integ.tab
@@ -253,18 +265,19 @@ end
 
             if _tstop_in_step(integ, tf, T)
                 integ.t = integ.tstops[integ.tstops_idx]
-                if abs(integ.t - (t + dt)) > eps(integ.t)
+                if integ.t - t != dt
                     integ.u = integ(integ.t)
+                    mapreduce(isfinite, &, integ.u) || error("non-finite state at tstop")
                 end
                 dt = integ.t - integ.tprev
                 integ.tstops_idx += 1
-            elseif (tf - t - dt) < convert(T, 1.0f-14)
+            elseif dt == tf - t
                 integ.t = tf
             else
                 ##Advance the integrator
                 integ.t += dt
             end
-            integ.dtnew = min(abs(dtnew), abs(_next_stop(integ, tf) - integ.t))
+            integ.dtnew = dtnew
         end
     end