disco optimizations

Author: Shreyas-Ekanathan

lib/OrdinaryDiffEqBDF/src/controllers.jl

@@ -83,6 +83,11 @@ 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
+    if integrator.curr_discontinuity != -1 &&
+            integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint
+        integrator.curr_discontinuity = -1
+        cache.nconsteps = 0
+    end
     cache.nconsteps += 1
     if iszero(OrdinaryDiffEqCore.get_EEst(integrator))
         return integrator.dt * get_current_qmax(integrator, get_qmax(integrator))
@@ -174,8 +179,12 @@ function bdf_step_reject_controller!(integrator, cache, EEst1)
     end
 
     if discontinuity_detection
+        no_prior_discontinuity = (integrator.curr_discontinuity == -1)
         disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
-        if disco_dt != -1
+        if no_prior_discontinuity && disco_dt > zero(disco_dt)
+            add_tstop!(integrator, integrator.disco_checkpoint)
+        end
+        if disco_dt > zero(disco_dt)
             integrator.dt = disco_dt
             return integrator.dt
         end
@@ -416,6 +425,11 @@ function step_accept_controller!(
         q
     ) where {max_order}
     cache.consfailcnt = 0
+    if integrator.curr_discontinuity != -1 &&
+            integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint
+        integrator.curr_discontinuity = -1
+        cache.nconsteps = 0
+    end
     if q <= get_qsteady_max(integrator) && q >= get_qsteady_min(integrator)
         q = one(q)
     end
@@ -579,6 +593,11 @@ function step_accept_controller!(
         q
     ) where {max_order}
     cache.consfailcnt = 0
+    if integrator.curr_discontinuity != -1 &&
+            integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint
+        integrator.curr_discontinuity = -1
+        cache.nconsteps = 0
+    end
     if q <= get_qsteady_max(integrator) && q >= get_qsteady_min(integrator)
         q = one(q)
     end

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -112,6 +112,7 @@ import EnzymeCore
     Tableau        # tableau-derived predictor (α / const_stage_guess)
 end
 export Predictor
+export set_discontinuity
 
 const CompiledFloats = Union{Float32, Float64}
 import Preferences

lib/OrdinaryDiffEqCore/src/disco.jl

@@ -38,24 +38,32 @@ function find_discontinuity(u, uprev, integrator)
             len_cb = i.len
             i.condition(disco_zero.out_low, uprev, t, integrator)
             i.condition(disco_zero.out_high, u, t + dt, integrator)
+            _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)
             for j in 1:len_cb
                 if (disco_zero.out_low[j] * disco_zero.out_high[j] < zero(disco_zero.out_low[j]))
                     disco_zero.ind = j
-                    sol = solve(disco_prob)
+                    sol = solve(disco_prob, tspan = (0.0, breakpointθ == -one(dt) ? 1.0 : breakpointθ); save_everystep = false)
                     tmp = sol[]
                     if (!isnan(tmp) && (breakpointθ < zero(breakpointθ) || tmp < breakpointθ))
                         breakpointθ = tmp
+                        integrator.curr_discontinuity = idx
+                        integrator.disco_checkpoint = integrator.t + tmp * dt
                     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))
-                sol = solve(disco_prob)
+                _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)
+                sol = solve(disco_prob, tspan = (0.0, breakpointθ == -one(dt) ? 1.0 : breakpointθ); save_everystep = false)
                 tmp = sol[]
                 if (!isnan(tmp) && (breakpointθ < zero(breakpointθ) || tmp < breakpointθ))
                     breakpointθ = tmp
+                    integrator.curr_discontinuity = idx
+                    integrator.disco_checkpoint = integrator.t + tmp * dt
                 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
@@ -506,10 +513,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
 
@@ -563,18 +566,26 @@ end
 
 # TODO change signature to remove the q input
 function step_accept_controller!(integrator, cache::IControllerCache, alg, q)
-    (; qsteady_min, qsteady_max) = cache.controller.basic
+    (; qsteady_min, qsteady_max, discontinuity_detection) = cache.controller.basic
 
     if qsteady_min <= q <= qsteady_max
         q = one(q)
     end
+    if discontinuity_detection && integrator.curr_discontinuity != -1 &&
+            integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint
+        integrator.curr_discontinuity = -1
+    end
     return integrator.dt / q # new dt
 end
 
 function step_reject_controller!(integrator, cache::IControllerCache, alg)
     discontinuity_detection = cache.controller.basic.discontinuity_detection
     if discontinuity_detection
+        no_prior_discontinuity = (integrator.curr_discontinuity == -1)
         disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
+        if no_prior_discontinuity && disco_dt > zero(disco_dt)
+            add_tstop!(integrator, integrator.disco_checkpoint)
+        end
         if disco_dt > zero(disco_dt)
             integrator.dt = disco_dt
             return integrator.dt
@@ -611,10 +622,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,
@@ -703,25 +711,35 @@ end
 
 function step_accept_controller!(integrator, cache::PIControllerCache, alg, q)
     (; controller) = cache
-    (; qsteady_min, qsteady_max) = controller.basic
+    (; qsteady_min, qsteady_max, discontinuity_detection) = controller.basic
     qoldinit = controller.qoldinit
     EEst = DiffEqBase.value(get_EEst(integrator))
 
     if qsteady_min <= q <= qsteady_max
         q = one(q)
     end
-    cache.errold = max(EEst, qoldinit)
+
+    if discontinuity_detection && integrator.curr_discontinuity != -1 &&
+            integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint
+        integrator.curr_discontinuity = -1
+        cache.errold = qoldinit
+    else
+        cache.errold = max(EEst, qoldinit)
+    end
     return integrator.dt / q # new dt
 end
 
 function step_reject_controller!(integrator, cache::PIControllerCache, alg)
     (; controller, q11) = cache
     (; qmin, gamma, discontinuity_detection) = controller.basic
     if discontinuity_detection
+        no_prior_discontinuity = (integrator.curr_discontinuity == -1)
         disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
+        if no_prior_discontinuity && disco_dt > zero(disco_dt) #found a discontinuity
+            add_tstop!(integrator, integrator.disco_checkpoint)
+        end
         if disco_dt > zero(disco_dt)
-            integrator.dt = disco_dt
-            return integrator.dt
+            return integrator.dt = disco_dt
         end
     end
     return integrator.dt /= min(inv(qmin), q11 / gamma)
@@ -776,11 +794,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`
@@ -940,22 +953,35 @@ end
 
 function step_accept_controller!(integrator, cache::PIDControllerCache, alg, dt_factor)
     (; controller) = cache
-    (; qsteady_min, qsteady_max) = controller.basic
+    (; qsteady_min, qsteady_max, discontinuity_detection) = controller.basic
 
     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]
+    if discontinuity_detection && integrator.curr_discontinuity != -1 &&
+            integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint
+        integrator.curr_discontinuity = -1
+        @inbounds begin
+            cache.err[3] = one(eltype(cache.err))
+            cache.err[2] = cache.err[1]
+        end
+    else
+        @inbounds begin
+            cache.err[3] = cache.err[2]
+            cache.err[2] = cache.err[1]
+        end
     end
     return integrator.dt * dt_factor # new dt
 end
 
 function step_reject_controller!(integrator, cache::PIDControllerCache, alg)
     discontinuity_detection = cache.controller.basic.discontinuity_detection
     if discontinuity_detection
+        no_prior_discontinuity = (integrator.curr_discontinuity == -1)
         disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
+        if no_prior_discontinuity && disco_dt > zero(disco_dt)
+            add_tstop!(integrator, integrator.disco_checkpoint)
+        end
         if disco_dt > zero(disco_dt)
             integrator.dt = disco_dt
             return integrator.dt
@@ -1020,10 +1046,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
@@ -1119,8 +1141,15 @@ function step_accept_controller!(integrator, cache::PredictiveControllerCache, a
     if qsteady_min <= qacc <= qsteady_max
         qacc = one(qacc)
     end
-    cache.dtacc = DiffEqBase.value(integrator.dt)
-    cache.erracc = max(1.0e-2, EEst)
+    if cache.controller.basic.discontinuity_detection && integrator.curr_discontinuity != -1 &&
+            integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint
+        integrator.curr_discontinuity = -1
+        cache.dtacc = DiffEqBase.value(integrator.dt)
+        cache.erracc = one(EEst)
+    else
+        cache.dtacc = DiffEqBase.value(integrator.dt)
+        cache.erracc = max(1.0e-2, EEst)
+    end
 
     return integrator.dt / qacc
 end
@@ -1130,7 +1159,11 @@ function step_reject_controller!(integrator, cache::PredictiveControllerCache, a
     (; qold) = cache
     discontinuity_detection = cache.controller.basic.discontinuity_detection
     if discontinuity_detection
+        no_prior_discontinuity = (integrator.curr_discontinuity == -1)
         disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator)
+        if no_prior_discontinuity && disco_dt > zero(disco_dt)
+            add_tstop!(integrator, integrator.disco_checkpoint)
+        end
         if disco_dt > zero(disco_dt)
             integrator.dt = disco_dt
             return integrator.dt

lib/OrdinaryDiffEqCore/src/integrators/type.jl

@@ -175,6 +175,8 @@ mutable struct ODEIntegrator{
     fsalfirst::FSALType
     fsallast::FSALType
     rng::RNGType
+    curr_discontinuity::Int
+    disco_checkpoint::tType
     W::WType
     P::PType
     sqdt::SqdtType

lib/OrdinaryDiffEqCore/src/solve.jl

@@ -38,7 +38,6 @@ end
 parameter_values(z::zero_func_struct) = z.p
 
 function (z::zero_func_struct)(θ, p)
-    _ode_addsteps!(z.k, z.tprev, z.uprev, z.u, z.dt, z.f, z.p, z.cache, false, true, false)
     ode_interpolant!(z.u₁, θ, z.dt, z.uprev, z.u, z.k, z.cache, z.idxs, Val{0}, z.differential_vars)
     return zero_condition(z.callback, z.out, z.u₁, z.tprev + θ * z.dt, z, z.ind)
 end
@@ -688,6 +687,8 @@ Base.@constprop :aggressive function _ode_init(
 
     controller_cache = setup_controller_cache(_alg, cache, controller, EEstT, disco_probs)
 
+    curr_discontinuity = -1
+    disco_checkpoint = zero(tType)
     # Seed the initial EEst on the controller cache (was previously
     # `integrator.EEst = oneunit(EEstT)`).
     set_EEst!(controller_cache, EEst)
@@ -721,7 +722,7 @@ Base.@constprop :aggressive function _ode_init(
         isout, reeval_fsal,
         derivative_discontinuity, reinitialize, isdae,
         opts, stats, initializealg, differential_vars,
-        fsalfirst, fsallast, _rng,
+        fsalfirst, fsallast, _rng, curr_discontinuity, disco_checkpoint,
         W, P, sqdt,
         noise, c, rate_constants
     )