Merge pull request #3720 from Shreyas-Ekanathan/disco-optimizations

Optimizations to Discontinuity Detection Scheme

Author: oscardssmith

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,36 @@ 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
+    if is_disco
+        return min((integrator.disco_checkpoint - integrator.t) / 4, new_dt)
     end
-    return integrator.dt / q
+    return new_dt
+end
+
+"""
+    bdf_restart_estimates!(cache)
+
+Restart the order and constant-step estimates for a step that straddles a derivative
+discontinuity, where the history behind it and the solution ahead belong to different
+regimes. We do this because the BDF step-size and order logic is based on the history of the solution, 
+and we have effectively entered a new regime where old estimates no longer apply.
+"""
+function bdf_restart_estimates!(cache)
+    cache.order = 1
+    cache.nconsteps = 0
+    if hasfield(typeof(cache), :qwait)
+        cache.qwait = 3
+    end
+    return nothing
 end
 
 function bdf_step_reject_controller!(integrator, cache, EEst1)
@@ -174,8 +199,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 +441,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 +457,12 @@ function step_accept_controller!(
     elseif cache.qwait > 0
         cache.qwait -= 1 # countdown
     end
-    return integrator.dt / q
+    new_dt = integrator.dt / q
+    if is_disco
+        bdf_restart_estimates!(cache)
+        return min((integrator.disco_checkpoint - integrator.t) / 4, new_dt)
+    end
+    return new_dt
 end
 
 function step_reject_controller!(integrator, alg::DFBDF)
@@ -579,6 +614,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 +630,10 @@ function step_accept_controller!(
     elseif cache.qwait > 0
         cache.qwait -= 1 # countdown
     end
-    return integrator.dt / q
+    new_dt = integrator.dt / q
+    if is_disco
+        bdf_restart_estimates!(cache)
+        return min((integrator.disco_checkpoint - integrator.t) / 4, new_dt)
+    end
+    return new_dt
 end

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -149,6 +149,7 @@ methods.
     Tableau        # tableau-derived predictor (α / const_stage_guess)
 end
 export Predictor
+export set_discontinuity
 
 """
     CompiledFloats
@@ -476,7 +477,7 @@ include("precompilation_setup.jl")
             # Integrator step / cache / initialization hooks.
             :_ode_init, :_determine_initdt, :ode_determine_initdt, :_initialize_dae!,
             :find_algebraic_vars_eqs, :postamble!, :apply_step!, :last_step_failed, :reset_alg_dependent_opts!,
-            :set_discontinuity, :resolve_basic,
+            :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,15 +1,16 @@
 """
-    set_discontinuity(u, uprev, integrator) -> dt
+    set_discontinuity(integrator) -> dt
 
 Return the sub-step `Δt` at which a continuous callback's root lies within the
 current step (a discontinuity to stop at), or a negative value if there is none in
 `(0, 1)·dt`.
 """
-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θ)
+        # below formula is equivalent to (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
@@ -18,17 +19,21 @@ 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
+    integrator.is_disco_step = false
+    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
@@ -47,22 +52,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