Optimizations to Discontinuity Detection Scheme

lib/OrdinaryDiffEqBDF/src/controllers.jl

@@ -83,9 +83,14 @@ end
 function step_accept_controller!(integrator, cache::Union{QNDFCache, QNDFConstantCache}, alg::QNDF{max_order}, q) where {max_order}
     #step is accepted, reset count of consecutive failed steps
     cache.consfailcnt = 0
+    is_disco = integrator.is_disco_step
+    if is_disco
+        integrator.is_disco_step = false
+        cache.nconsteps = 0
+    end
     cache.nconsteps += 1
-    if iszero(OrdinaryDiffEqCore.get_EEst(integrator))
-        return integrator.dt * get_current_qmax(integrator, get_qmax(integrator))
+    new_dt = if iszero(OrdinaryDiffEqCore.get_EEst(integrator))
+        integrator.dt * get_current_qmax(integrator, get_qmax(integrator))
     else
         est = OrdinaryDiffEqCore.get_EEst(integrator)
         estₖ₋₁ = cache.EEst1
@@ -147,16 +152,16 @@ function step_accept_controller!(integrator, cache::Union{QNDFCache, QNDFConstan
         end
         cache.order = kₙ
         q = integrator.dt / hₙ
-    end
-    if prefer_const_step
-        if q < 1.2 && q > 0.6
-            return integrator.dt
+
+        if prefer_const_step && 0.6 < q < 1.2
+            h
+        elseif q <= get_qsteady_max(integrator) && q >= get_qsteady_min(integrator)
+            h
+        else
+            h / q
         end
     end
-    if q <= get_qsteady_max(integrator) && q >= get_qsteady_min(integrator)
-        return integrator.dt
-    end
-    return integrator.dt / q
+    return is_disco ? min((integrator.disco_checkpoint - integrator.t) / 4, new_dt) : new_dt
 end
 
 function bdf_step_reject_controller!(integrator, cache, EEst1)
@@ -174,8 +179,8 @@ function bdf_step_reject_controller!(integrator, cache, EEst1)
     end
 
     if discontinuity_detection
-        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
-        if disco_dt != -1
+        disco_dt = set_discontinuity(integrator)
+        if disco_dt > zero(disco_dt)
             integrator.dt = disco_dt
             return integrator.dt
         end
@@ -416,6 +421,11 @@ function step_accept_controller!(
         q
     ) where {max_order}
     cache.consfailcnt = 0
+    is_disco = integrator.is_disco_step
+    if is_disco
+        integrator.is_disco_step = false
+        cache.nconsteps = 0
+    end
     if q <= get_qsteady_max(integrator) && q >= get_qsteady_min(integrator)
         q = one(q)
     end
@@ -427,7 +437,8 @@ function step_accept_controller!(
     elseif cache.qwait > 0
         cache.qwait -= 1 # countdown
     end
-    return integrator.dt / q
+    new_dt = integrator.dt / q
+    return is_disco ? min((integrator.disco_checkpoint - integrator.t) / 4, new_dt) : new_dt
 end
 
 function step_reject_controller!(integrator, alg::DFBDF)
@@ -579,6 +590,11 @@ function step_accept_controller!(
         q
     ) where {max_order}
     cache.consfailcnt = 0
+    is_disco = integrator.is_disco_step
+    if is_disco
+        integrator.is_disco_step = false
+        cache.nconsteps = 0
+    end
     if q <= get_qsteady_max(integrator) && q >= get_qsteady_min(integrator)
         q = one(q)
     end
@@ -590,5 +606,6 @@ function step_accept_controller!(
     elseif cache.qwait > 0
         cache.qwait -= 1 # countdown
     end
-    return integrator.dt / q
+    new_dt = integrator.dt / q
+    return is_disco ? min((integrator.disco_checkpoint - integrator.t) / 4, new_dt) : new_dt
 end

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -122,6 +122,7 @@ import EnzymeCore
     Tableau        # tableau-derived predictor (α / const_stage_guess)
 end
 export Predictor
+export set_discontinuity
 
 const CompiledFloats = Union{Float32, Float64}
 import Preferences
@@ -261,7 +262,7 @@ include("precompilation_setup.jl")
         # Integrator step / cache / initialization hooks.
         :_ode_init, :_determine_initdt, :ode_determine_initdt, :_initialize_dae!,
         :postamble!, :apply_step!, :last_step_failed, :reset_alg_dependent_opts!,
-        :handle_callback_modifiers!, :set_discontinuity, :resolve_basic,
+        :handle_callback_modifiers!, :resolve_basic,
         # Noise hooks used by the SDE/RODE solver sublibs.
         :accept_noise!, :reinit_noise!, :reject_noise!, :save_noise!, :noise_curt,
         :is_noise_saveable,

lib/OrdinaryDiffEqCore/src/disco.jl

@@ -1,8 +1,9 @@
-function set_discontinuity(u, uprev, integrator)
-    breakpointθ = find_discontinuity(u, uprev, integrator)
+function set_discontinuity(integrator)
+    breakpointθ = find_discontinuity(integrator)
     dt = integrator.dt
-    if 1.0e-10 < breakpointθ < 1.0
-        return breakpointθ * dt
+    if breakpointθ < one(breakpointθ)
+        return (0.5 + 0.4 * sin(π * (breakpointθ - 0.5))) * dt
+        return (0.5 + 1.2372 * (breakpointθ - 0.5) - 1.7487 * (breakpointθ - 0.5)^3) * dt
     end
     return -one(dt)
 end
@@ -11,17 +12,20 @@ get_disco_probs(cache::AbstractControllerCache) = cache.controller.basic.disco_p
 get_disco_probs(cache::DummyControllerCache) = cache.disco_probs
 get_disco_probs(cache::CompositeControllerCache) = get_disco_probs(first(cache.caches))
 
-function find_discontinuity(u, uprev, integrator)
+function find_discontinuity(integrator)
     cb = integrator.opts.callback
     dt = integrator.dt
-    cb === nothing && return -one(dt)
-    isempty(cb.continuous_callbacks) && return -one(dt)
+    u = integrator.u
+    uprev = integrator.uprev
+    cb === nothing && return one(dt)
+    isempty(cb.continuous_callbacks) && return one(dt)
     p = integrator.p
     t = integrator.t
     k = integrator.k
-    breakpointθ = -one(dt)
+    breakpointθ = one(dt)
     disco_probs = get_disco_probs(integrator.controller_cache)
     idx = 1
+    addsteps_called = false
     for i in cb.continuous_callbacks
         if (!(i.maybe_discontinuity))
             continue
@@ -40,22 +44,38 @@ function find_discontinuity(u, uprev, integrator)
             i.condition(disco_zero.out_high, u, t + dt, integrator)
             for j in 1:len_cb
                 if (disco_zero.out_low[j] * disco_zero.out_high[j] < zero(disco_zero.out_low[j]))
+                    if (!addsteps_called)
+                        addsteps_called = true
+                        _ode_addsteps!(disco_zero.k, disco_zero.tprev, disco_zero.uprev, disco_zero.u,
+                                    disco_zero.dt, disco_zero.f, disco_zero.p, disco_zero.cache, false, true, false)
+                    end
                     disco_zero.ind = j
+                    disco_prob.tspan[2] = breakpointθ
                     sol = solve(disco_prob)
                     tmp = sol[]
-                    if (!isnan(tmp) && (breakpointθ < zero(breakpointθ) || tmp < breakpointθ))
+                    if (!isnan(tmp) && tmp < breakpointθ)
                         breakpointθ = tmp
+                        integrator.is_disco_step = true
+                        integrator.disco_checkpoint = integrator.t + dt #our prev rejected step, we shouldn't step too far past this
                     end
                 end
             end
         else
-            out_prev = i.condition(uprev, t, integrator)
-            out_curr = i.condition(u, t + dt, integrator)
-            if (out_prev * out_curr < zero(out_prev))
+            disco_zero.out_low[1] = i.condition(uprev, t, integrator)
+            disco_zero.out_high[1] = i.condition(u, t + dt, integrator)
+            if (disco_zero.out_low[1] * disco_zero.out_high[1] < zero(disco_zero.out_low[1]))
+                if (!addsteps_called)
+                    addsteps_called = true
+                    _ode_addsteps!(disco_zero.k, disco_zero.tprev, disco_zero.uprev, disco_zero.u,
+                                disco_zero.dt, disco_zero.f, disco_zero.p, disco_zero.cache, false, true, false)
+                end
+                disco_prob.tspan[2] = breakpointθ
                 sol = solve(disco_prob)
                 tmp = sol[]
-                if (!isnan(tmp) && (breakpointθ < zero(breakpointθ) || tmp < breakpointθ))
+                if (!isnan(tmp) && tmp < breakpointθ)
                     breakpointθ = tmp
+                    integrator.is_disco_step = true
+                    integrator.disco_checkpoint = integrator.t + dt #our prev rejected step, we shouldn't step past this
                 end
             end
         end

lib/OrdinaryDiffEqCore/src/integrators/controllers.jl

@@ -43,6 +43,13 @@ The fields are:
   inside this interval, dt is held constant.
 - `failfactor`: post-Newton-failure shrink factor used by
   [`post_newton_controller!`](@ref).
+- 'discontinuity_detection': If `discontinuity_detection` is set to true, the algorithm will run the autonomous
+    discontinuity detection to predict the best next timestep after step rejection. 
+    Otherwise, it follows the default step rejection algorithm. This feature is currently
+    defaulted off. 
+- `disco_probs`: If `discontinuity_detection` is set to true, this field holds the vector of 
+`IntervalNonlinearProblem`s used for discontinuity detection. Otherwise, it can be left empty.
+
 
 User-supplied overrides flow through controllers as a `NamedTuple` of
 keyword arguments (whatever subset the user passed). At
@@ -485,6 +492,22 @@ for (accessor, default) in (
     end
 end
 
+# Discontinuity-detection helpers 
+# Shared logic in every step_accept/reject_controller! below.
+
+function handle_disco_accept!(integrator, controller_basic, t, nominal_new_dt)
+    controller_basic.discontinuity_detection && integrator.is_disco_step || return nominal_new_dt
+    integrator.is_disco_step = false
+    return min((integrator.disco_checkpoint - t) / 4, nominal_new_dt)
+end
+
+function handle_disco_reject!(integrator, controller_basic)
+    controller_basic.discontinuity_detection || return false
+    disco_dt = set_discontinuity(integrator)
+    disco_dt > zero(disco_dt) || return false
+    integrator.dt = disco_dt
+    return true
+end
 
 # Standard integral (I) step size controller
 """
@@ -510,10 +533,6 @@ the interval `[qmin, qmax]`.
 A step will be accepted whenever the estimated error `get_EEst(integrator)` is
 less than or equal to unity. Otherwise, the step is rejected and re-tried with
 the predicted step size.
-If `discontinuity_detection` is set to true, the algorithm will run the autonomous
-discontinuity detection to predict the best next timestep after step rejection. 
-Otherwise, it follows the default step rejection algorithm. This feature is currently
-defaulted off. 
 
 ## References
 
@@ -569,21 +588,16 @@ end
 function step_accept_controller!(integrator, cache::IControllerCache, alg, q)
     (; qsteady_min, qsteady_max) = cache.controller.basic
 
+    t = integrator.t
+    dt = integrator.dt
     if qsteady_min <= q <= qsteady_max
         q = one(q)
     end
-    return integrator.dt / q # new dt
+    return handle_disco_accept!(integrator, cache.controller.basic, t, dt / q)
 end
 
 function step_reject_controller!(integrator, cache::IControllerCache, alg)
-    discontinuity_detection = cache.controller.basic.discontinuity_detection
-    if discontinuity_detection
-        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
-        if disco_dt > zero(disco_dt)
-            integrator.dt = disco_dt
-            return integrator.dt
-        end
-    end
+    handle_disco_reject!(integrator, cache.controller.basic) && return integrator.dt
     return integrator.dt = cache.dtreject # TODO this does not look right.
 end
 
@@ -615,10 +629,7 @@ the interval `[qmin, qmax]`.
 A step will be accepted whenever the estimated error `get_EEst(integrator)` is
 less than or equal to unity. Otherwise, the step is rejected and re-tried with
 the predicted step size.
-If `discontinuity_detection` is set to true, the algorithm will run the autonomous
-discontinuity detection to predict the best next timestep after step rejection. 
-Otherwise, it follows the default step rejection algorithm. This feature is currently
-defaulted off. 
+
 !!! note
 
     The coefficients `beta1, beta2` are not scaled by the order of the method,
@@ -712,23 +723,19 @@ function step_accept_controller!(integrator, cache::PIControllerCache, alg, q)
     qoldinit = controller.qoldinit
     EEst = SciMLBase.value(get_EEst(integrator))
 
+    t = integrator.t
+    dt = integrator.dt
     if qsteady_min <= q <= qsteady_max
         q = one(q)
     end
     cache.errold = max(EEst, qoldinit)
-    return integrator.dt / q # new dt
+    return handle_disco_accept!(integrator, controller.basic, t, dt / q)
 end
 
 function step_reject_controller!(integrator, cache::PIControllerCache, alg)
     (; controller, q11) = cache
-    (; qmin, gamma, discontinuity_detection) = controller.basic
-    if discontinuity_detection
-        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
-        if disco_dt > zero(disco_dt)
-            integrator.dt = disco_dt
-            return integrator.dt
-        end
-    end
+    (; qmin, gamma) = controller.basic
+    handle_disco_reject!(integrator, controller.basic) && return integrator.dt
     return integrator.dt /= min(inv(qmin), q11 / gamma)
 end
 
@@ -781,11 +788,6 @@ Some standard controller parameters suggested in the literature are
 | H211PI     | `1//6`  | `1//6`  | `0`     |
 | H312PID    | `1//18` | `1//9`  | `1//18` |
 
-If `discontinuity_detection` is set to true, the algorithm will run the autonomous
-discontinuity detection to predict the best next timestep after step rejection. 
-Otherwise, it follows the default step rejection algorithm. This feature is currently
-defaulted off. 
-
 !!! note
 
     In contrast to the [`PIController`](@ref), the coefficients `beta1, beta2, beta3`
@@ -947,25 +949,20 @@ function step_accept_controller!(integrator, cache::PIDControllerCache, alg, dt_
     (; controller) = cache
     (; qsteady_min, qsteady_max) = controller.basic
 
+    t = integrator.t
+    dt = integrator.dt
     if qsteady_min <= inv(dt_factor) <= qsteady_max
         dt_factor = one(dt_factor)
     end
     @inbounds begin
         cache.err[3] = cache.err[2]
         cache.err[2] = cache.err[1]
     end
-    return integrator.dt * dt_factor # new dt
+    return handle_disco_accept!(integrator, controller.basic, t, dt * dt_factor)
 end
 
 function step_reject_controller!(integrator, cache::PIDControllerCache, alg)
-    discontinuity_detection = cache.controller.basic.discontinuity_detection
-    if discontinuity_detection
-        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
-        if disco_dt > zero(disco_dt)
-            integrator.dt = disco_dt
-            return integrator.dt
-        end
-    end
+    handle_disco_reject!(integrator, cache.controller.basic) && return integrator.dt
     return integrator.dt *= cache.dt_factor
 end
 
@@ -1025,10 +1022,6 @@ integrator.dt / qacc
 ```
 
 When it rejects, it's the same as the [`IController`](@ref):
-If `discontinuity_detection` is set to true, the algorithm will run the autonomous
-discontinuity detection to predict the best next timestep after step rejection. 
-Otherwise, it follows the default step rejection algorithm. This feature is currently
-defaulted off. 
 ```julia
 if integrator.success_iter == 0
     integrator.dt *= 0.1
@@ -1127,20 +1120,13 @@ function step_accept_controller!(integrator, cache::PredictiveControllerCache, a
     cache.dtacc = SciMLBase.value(integrator.dt)
     cache.erracc = max(1.0e-2, EEst)
 
-    return integrator.dt / qacc
+    return handle_disco_accept!(integrator, cache.controller.basic, integrator.t, integrator.dt / qacc)
 end
 
 function step_reject_controller!(integrator, cache::PredictiveControllerCache, alg)
     (; dt, success_iter) = integrator
     (; qold) = cache
-    discontinuity_detection = cache.controller.basic.discontinuity_detection
-    if discontinuity_detection
-        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
-        if disco_dt > zero(disco_dt)
-            integrator.dt = disco_dt
-            return integrator.dt
-        end
-    end
+    handle_disco_reject!(integrator, cache.controller.basic) && return integrator.dt
     return integrator.dt = success_iter == 0 ? 0.1 * dt : dt / qold
 end