Implement Discontinuity Detection Algorithm

lib/OrdinaryDiffEqBDF/src/OrdinaryDiffEqBDF.jl

@@ -22,7 +22,7 @@ import OrdinaryDiffEqCore: alg_order, calculate_residuals!,
     DAEAlgorithm, _unwrap_val, DummyController,
     get_fsalfirstlast, generic_solver_docstring, _ad_chunksize_int, _ad_fdtype, _fixup_ad,
     _ode_interpolant, _ode_interpolant!, has_stiff_interpolation,
-    _ode_addsteps!, DerivativeOrderNotPossibleError
+    _ode_addsteps!, DerivativeOrderNotPossibleError, set_discontinuity
 using OrdinaryDiffEqSDIRK: ImplicitEulerConstantCache, ImplicitEulerCache
 
 using TruncatedStacktraces: @truncate_stacktrace

lib/OrdinaryDiffEqBDF/src/controllers.jl

@@ -97,6 +97,16 @@ function bdf_step_reject_controller!(integrator, cache, EEst1)
     h = integrator.dt
     cache.consfailcnt += 1
     cache.nconsteps = 0
+
+    discontinuity_detection = cache.controller.discontinuity_detection
+    if discontinuity_detection
+        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator, integrator.cache)
+        if disco_dt != -1
+            integrator.dt = disco_dt
+            return integrator.dt
+        end
+    end
+
     if cache.consfailcnt > 1
         h = h / 2
     end
@@ -495,4 +505,4 @@ function step_accept_controller!(
         cache.qwait -= 1 # countdown
     end
     return integrator.dt / q
-end
+end 

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -152,6 +152,7 @@ include("cache_utils.jl")
 include("initialize_dae.jl")
 
 include("perform_step/composite_perform_step.jl")
+include("disco.jl")
 
 include("dense/generic_dense.jl")
 

lib/OrdinaryDiffEqCore/src/disco.jl

@@ -0,0 +1,74 @@
+function set_discontinuity(u, uprev, integrator, cache) 
+    breakpointθ = find_discontinuity(u, uprev, integrator, cache) 
+    dt = integrator.dt
+    t = integrator.t
+    if !isnan(breakpointθ) && 1e-6 < breakpointθ < 1.0
+        #println("Discontinuity detected at t = ", t + breakpointθ * dt)
+        return breakpointθ * dt
+    end
+    return -1
+end
+
+function find_discontinuity(u, uprev, integrator, cache)
+    cb = integrator.opts.callback
+    cb === nothing && return -1
+    isempty(cb.continuous_callbacks) && return -1
+    p = integrator.p
+    t = integrator.t
+    dt = integrator.dt
+    save_idxs = integrator.opts.save_idxs
+    k = integrator.k
+    cache = integrator.cache
+    differential_vars = integrator.differential_vars
+    θlo = zero(dt)
+    θhi = one(dt)
+    bracket = [θlo, θhi]
+    breakpointθ = -one(dt)
+    idx = 1
+    for i in cb.continuous_callbacks
+        if (!(i.maybe_discontinuity)) 
+            continue 
+        end
+        disco_prob = integrator.disco_probs[idx]
+        disco_zero = disco_prob.f.f
+        disco_zero.dt = dt
+        disco_zero.uprev = uprev
+        disco_zero.u = u
+        disco_zero.k = k
+        disco_zero.cache = cache
+        disco_zero.differential_vars = differential_vars
+        disco_zero.idxs = save_idxs
+        disco_zero.tprev = t
+        disco_zero.f = integrator.f
+        disco_zero.p = p
+        if (i isa VectorContinuousCallback)
+            len_cb = i.len
+            out_prev = similar(u)
+            out_curr = similar(u)
+            i.condition(out_prev, uprev, t, integrator)
+            i.condition(out_curr, u, t + dt, integrator)
+            for j in 1:len_cb
+                if (out_prev[j] * out_curr[j] < zero(out_prev[j]))
+                    disco_zero.ind = j
+                    sol = solve(disco_prob; bracket = bracket)
+                    tmp = sol[]
+                    if (!isnan(tmp) && (breakpointθ == -1 || tmp < breakpointθ)) 
+                        breakpointθ = tmp 
+                    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; bracket = bracket)
+                tmp = sol[]
+                if (!isnan(tmp) && (breakpointθ == -1 || tmp < breakpointθ)) 
+                    breakpointθ = tmp 
+                end 
+            end
+        end
+        idx += 1
+    end
+    breakpointθ
+end

lib/OrdinaryDiffEqCore/src/integrators/controllers.jl

@@ -185,14 +185,20 @@ Controller cache used by algorithms that manage step-size selection themselves
 (BDF, Nordsieck, Leaping, …). Holds the scalar error estimate exposed through
 `get_EEst(integrator)` and a reference to the algorithm cache so existing dispatch
 on the algorithm cache continues to work.
+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. 
 """
 mutable struct DummyControllerCache{T, C} <: AbstractControllerCache
     EEst::T
     cache::C
+    discontinuity_detection::Bool
 end
 
 function setup_controller_cache(alg, cache, controller::DummyController, ::Type{E}) where {E}
-    return DummyControllerCache{E, typeof(cache)}(oneunit(E), cache)
+    discontinuity_detection = false
+    return DummyControllerCache{E, typeof(cache)}(oneunit(E), cache, discontinuity_detection)
 end
 
 # Algorithms with integrated controllers (BDF, Nordsieck, …) only define their
@@ -236,6 +242,10 @@ 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
 
@@ -250,31 +260,35 @@ struct IController{T} <: AbstractController
     gamma::T
     qsteady_min::T
     qsteady_max::T
+    discontinuity_detection::Bool
 end
 
 function IController(; qmin = 1 // 5, qmax = 10 // 1, qmax_first_step = 10000 // 1, gamma = 9 // 10, qsteady_min = 1 // 1, qsteady_max = 6 // 5)
+    discontinuity_detection = false
     return IController{typeof(qmin)}( # FIXME combined promoted type
         qmin,
         qmax,
         qmax_first_step,
         gamma,
         qsteady_min,
         qsteady_max,
+        discontinuity_detection
     )
 end
 
 function IController(alg; kwargs...)
     return IController(Float64, alg; kwargs...)
 end
 
-function IController(QT, alg; qmin = nothing, qmax = nothing, qmax_first_step = nothing, gamma = nothing, qsteady_min = nothing, qsteady_max = nothing)
+function IController(QT, alg; qmin = nothing, qmax = nothing, qmax_first_step = nothing, gamma = nothing, qsteady_min = nothing, qsteady_max = nothing, discontinuity_detection = nothing)
     return IController{QT}(
         qmin === nothing ? qmin_default(alg) : qmin,
         qmax === nothing ? qmax_default(alg) : qmax,
         qmax_first_step === nothing ? QT(10000) : QT(qmax_first_step),
         gamma === nothing ? gamma_default(alg) : gamma,
         qsteady_min === nothing ? qsteady_min_default(alg) : qsteady_min,
         qsteady_max === nothing ? qsteady_max_default(alg) : qsteady_max,
+        discontinuity_detection === nothing ? false : discontinuity_detection
     )
 end
 
@@ -320,7 +334,15 @@ function step_accept_controller!(integrator, cache::IControllerCache, alg, q)
 end
 
 function step_reject_controller!(integrator, cache::IControllerCache, alg)
-    return integrator.dt = cache.dtreject
+    discontinuity_detection = cache.controller.discontinuity_detection
+    if discontinuity_detection
+        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator, integrator.cache)
+        if disco_dt != -1
+            integrator.dt = disco_dt
+            return integrator.dt
+        end
+    end
+    return integrator.dt = cache.dtreject # TODO this does not look right.
 end
 
 reinit_controller!(integrator::SciMLBase.DEIntegrator, cache::IControllerCache) = nothing
@@ -351,7 +373,10 @@ 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,
@@ -377,9 +402,11 @@ mutable struct PIController{T} <: AbstractController # TODO remove the mutable o
     qsteady_min::T
     qsteady_max::T
     qoldinit::T
+    discontinuity_detection::Bool
 end
 
 function PIController(beta1::Real, beta2::Real; qmin = 1 // 5, qmax = 10 // 0, qmax_first_step = 10000 // 1, gamma = 9 // 10, qsteady_min = 1 // 1, qsteady_max = 6 // 5, qoldinit = 1 // 10^4)
+    discontinuity_detection = false
     return PIController{typeof(beta1)}(
         beta1,
         beta2,
@@ -390,14 +417,15 @@ function PIController(beta1::Real, beta2::Real; qmin = 1 // 5, qmax = 10 // 0, q
         qsteady_min,
         qsteady_max,
         qoldinit,
+        discontinuity_detection
     )
 end
 
 function PIController(alg; kwargs...)
     return PIController(Float64, alg; kwargs...)
 end
 
-function PIController(QT, alg; beta1 = nothing, beta2 = nothing, qmin = nothing, qmax = nothing, qmax_first_step = nothing, gamma = nothing, qsteady_min = nothing, qsteady_max = nothing, qoldinit = nothing)
+function PIController(QT, alg; beta1 = nothing, beta2 = nothing, qmin = nothing, qmax = nothing, qmax_first_step = nothing, gamma = nothing, qsteady_min = nothing, qsteady_max = nothing, qoldinit = nothing, discontinuity_detection = nothing)
     beta2 = beta2 === nothing ? beta2_default(alg) : beta2
     beta1 = beta1 === nothing ? beta1_default(alg, beta2) : beta1
     qoldinit = qoldinit === nothing ? 1 // 10^4 : qoldinit
@@ -410,7 +438,8 @@ function PIController(QT, alg; beta1 = nothing, beta2 = nothing, qmin = nothing,
         gamma === nothing ? gamma_default(alg) : gamma,
         qsteady_min === nothing ? qsteady_min_default(alg) : qsteady_min,
         qsteady_max === nothing ? qsteady_max_default(alg) : qsteady_max,
-        qoldinit,
+        qoldinit, 
+        discontinuity_detection === nothing ? false : discontinuity_detection
     )
 end
 
@@ -465,6 +494,14 @@ end
 function step_reject_controller!(integrator, cache::PIControllerCache, alg)
     (; controller, q11) = cache
     (; qmin, gamma) = controller
+    discontinuity_detection = cache.controller.discontinuity_detection
+    if discontinuity_detection
+        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator, integrator.cache)
+        if disco_dt != -1
+            integrator.dt = disco_dt
+            return integrator.dt
+        end
+    end
     return integrator.dt /= min(inv(qmin), q11 / gamma)
 end
 
@@ -517,6 +554,11 @@ 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`
@@ -552,18 +594,21 @@ struct PIDController{QT, Limiter} <: AbstractController
     limiter::Limiter    # limiter of the dt factor (before clipping)
     qsteady_min::QT
     qsteady_max::QT
+    discontinuity_detection::Bool
 end
 
 @inline default_dt_factor_limiter(x) = one(x) + atan(x - one(x))
 
 function PIDController(beta1::Real, beta2::Real, beta3::Real = zero(beta1); accept_safety = 0.81, limiter = default_dt_factor_limiter, qsteady_min = 1 // 1, qsteady_max = 6 // 5)
     beta = map(float, promote(beta1, beta2, beta3))
+    discontinuity_detection = false
     return PIDController{typeof(beta1), typeof(limiter)}(
         beta,
         accept_safety,
         limiter,
         qsteady_min,
         qsteady_max,
+        discontinuity_detection
     )
 end
 
@@ -586,6 +631,7 @@ function PIDController(QT, alg; beta = nothing, accept_safety = 0.81, limiter =
         limiter,
         QT(qsteady_min === nothing ? qsteady_min_default(alg) : qsteady_min),
         QT(qsteady_max === nothing ? qsteady_max_default(alg) : qsteady_max),
+        discontinuity_detection === nothing ? false : discontinuity_detection
     )
 end
 
@@ -683,6 +729,14 @@ function step_accept_controller!(integrator, cache::PIDControllerCache, alg, dt_
 end
 
 function step_reject_controller!(integrator, cache::PIDControllerCache, alg)
+    discontinuity_detection = cache.controller.discontinuity_detection
+    if discontinuity_detection
+        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator, integrator.cache)
+        if disco_dt != -1
+            integrator.dt = disco_dt
+            return integrator.dt
+        end
+    end
     return integrator.dt *= cache.dt_factor
 end
 
@@ -738,7 +792,10 @@ 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
@@ -754,31 +811,35 @@ struct PredictiveController{T} <: AbstractController
     gamma::T
     qsteady_min::T
     qsteady_max::T
+    discontinuity_detection::Bool
 end
 
 function PredictiveController(; qmin = float(1 // 5), qmax = 10 // 1, qmax_first_step = 10000 // 1, gamma = 9 // 10, qsteady_min = 1 // 1, qsteady_max = 6 // 5)
+    discontinuity_detection = false
     return PredictiveController{typeof(qmin)}( # FIXME combined promoted type
         qmin,
         qmax,
         qmax_first_step,
         gamma,
         qsteady_min,
         qsteady_max,
+        discontinuity_detection
     )
 end
 
 function PredictiveController(alg; kwargs...)
     return PredictiveController(Float64, alg; kwargs...)
 end
 
-function PredictiveController(QT, alg; qmin = nothing, qmax = nothing, qmax_first_step = nothing, gamma = nothing, qsteady_min = nothing, qsteady_max = nothing)
+function PredictiveController(QT, alg; qmin = nothing, qmax = nothing, qmax_first_step = nothing, gamma = nothing, qsteady_min = nothing, qsteady_max = nothing, discontinuity_detection = nothing)
     return PredictiveController{QT}(
         qmin === nothing ? qmin_default(alg) : qmin,
         qmax === nothing ? qmax_default(alg) : qmax,
         qmax_first_step === nothing ? QT(10000) : QT(qmax_first_step),
         gamma === nothing ? gamma_default(alg) : gamma,
         qsteady_min === nothing ? qsteady_min_default(alg) : qsteady_min,
         qsteady_max === nothing ? qsteady_max_default(alg) : qsteady_max,
+        discontinuity_detection === nothing ? false : discontinuity_detection
     )
 end
 
@@ -868,6 +929,14 @@ end
 function step_reject_controller!(integrator, cache::PredictiveControllerCache, alg)
     (; dt, success_iter) = integrator
     (; qold) = cache
+    discontinuity_detection = cache.controller.discontinuity_detection
+    if discontinuity_detection
+        disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator, integrator.cache)
+        if disco_dt != -1
+            integrator.dt = disco_dt
+            return integrator.dt
+        end
+    end
     return integrator.dt = success_iter == 0 ? 0.1 * dt : dt / qold
 end
 

lib/OrdinaryDiffEqCore/src/integrators/type.jl

@@ -175,6 +175,8 @@ mutable struct ODEIntegrator{
     fsalfirst::FSALType
     fsallast::FSALType
     rng::RNGType
+    #disco_prob::IntervalNonlinearProblem
+    disco_probs::Vector{IntervalNonlinearProblem} #should we change this?
     W::WType
     P::PType
     sqdt::SqdtType