From 408a7a30fe0989a55b446240a18fb2a8c5c30993 Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Fri, 5 Jun 2026 10:41:42 -0400 Subject: [PATCH 1/9] disco optimizations --- lib/OrdinaryDiffEqBDF/src/controllers.jl | 21 ++- .../src/OrdinaryDiffEqCore.jl | 1 + lib/OrdinaryDiffEqCore/src/disco.jl | 12 +- .../src/integrators/controllers.jl | 89 +++++++--- .../src/integrators/type.jl | 2 + lib/OrdinaryDiffEqCore/src/solve.jl | 5 +- .../Integrators_I}/disco_tests.jl | 168 +++++++++--------- 7 files changed, 177 insertions(+), 121 deletions(-) rename {lib/OrdinaryDiffEqCore/test => test/Integrators_I}/disco_tests.jl (52%) diff --git a/lib/OrdinaryDiffEqBDF/src/controllers.jl b/lib/OrdinaryDiffEqBDF/src/controllers.jl index 171078d996b..f54acd8d141 100644 --- a/lib/OrdinaryDiffEqBDF/src/controllers.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl index e00cc98292c..190c6b0e576 100644 --- a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqCore/src/disco.jl b/lib/OrdinaryDiffEqCore/src/disco.jl index 5c18b9f8952..f4befea1371 100644 --- a/lib/OrdinaryDiffEqCore/src/disco.jl +++ b/lib/OrdinaryDiffEqCore/src/disco.jl @@ -38,13 +38,17 @@ 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 @@ -52,10 +56,14 @@ function find_discontinuity(u, uprev, integrator) 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 diff --git a/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl b/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl index 8b956cb9ebe..a4b0ecd74fd 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl +++ b/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 @@ -510,10 +517,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 @@ -567,18 +570,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 @@ -615,10 +626,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, @@ -708,14 +716,21 @@ 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 @@ -723,10 +738,13 @@ 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) @@ -781,11 +799,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` @@ -945,14 +958,23 @@ 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 @@ -960,7 +982,11 @@ 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 @@ -1025,10 +1051,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 @@ -1124,8 +1146,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 @@ -1135,7 +1164,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 diff --git a/lib/OrdinaryDiffEqCore/src/integrators/type.jl b/lib/OrdinaryDiffEqCore/src/integrators/type.jl index 6664f2e2d1c..66db13ffa96 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/type.jl +++ b/lib/OrdinaryDiffEqCore/src/integrators/type.jl @@ -176,6 +176,8 @@ mutable struct ODEIntegrator{ fsalfirst::FSALType fsallast::FSALType rng::RNGType + curr_discontinuity::Int + disco_checkpoint::tType W::WType P::PType sqdt::SqdtType diff --git a/lib/OrdinaryDiffEqCore/src/solve.jl b/lib/OrdinaryDiffEqCore/src/solve.jl index 10087e2c20f..361194531e4 100644 --- a/lib/OrdinaryDiffEqCore/src/solve.jl +++ b/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 @@ -689,6 +688,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) @@ -722,7 +723,7 @@ Base.@constprop :aggressive function _ode_init( isout, reeval_fsal, derivative_discontinuity, user_set_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 ) diff --git a/lib/OrdinaryDiffEqCore/test/disco_tests.jl b/test/Integrators_I/disco_tests.jl similarity index 52% rename from lib/OrdinaryDiffEqCore/test/disco_tests.jl rename to test/Integrators_I/disco_tests.jl index 02e9108e724..9d484d96a04 100644 --- a/lib/OrdinaryDiffEqCore/test/disco_tests.jl +++ b/test/Integrators_I/disco_tests.jl @@ -1,5 +1,8 @@ using OrdinaryDiffEqCore, DiffEqDevTools, Test, LinearAlgebra using OrdinaryDiffEqTsit5, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqFIRK +using Logging +global_logger(SimpleLogger(stderr, Logging.Error)) +using BenchmarkTools predictive_disco_controller(alg) = OrdinaryDiffEqCore.PredictiveController(alg; discontinuity_detection = true) PI_disco_controller(alg) = OrdinaryDiffEqCore.PIController(alg; discontinuity_detection = true) @@ -18,26 +21,25 @@ prob = ODEProblem(f, u0, tspan) #define callback condition(u, t, integrator) = u[1] - 1 cb = ContinuousCallback(condition, default_affect!; maybe_discontinuity = true) -cb2 = ContinuousCallback(condition, default_affect!; maybe_discontinuity = false) -sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1.0e-6, controller = predictive_disco_controller(RadauIIA5())) -# 253.208 μs (7657 allocations: 239.47 KiB) -sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb2, reltol = 1.0e-6) -# 347.917 μs (10308 allocations: 320.86 KiB) +sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) +# 272.333 μs (7961 allocations: 249.34 KiB) +sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) +# 369.000 μs (10315 allocations: 321.41 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject -sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1.0e-6, controller = PI_disco_controller(Rodas5P())) -# 456.292 μs (16771 allocations: 587.12 KiB) -sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb2, reltol = 1.0e-6) -# 589.666 μs (21804 allocations: 763.61 KiB) +sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) +# 419.750 μs (15586 allocations: 546.83 KiB) +sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) +# 587.042 μs (21827 allocations: 764.73 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject -sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1.0e-6, controller = PI_disco_controller(Tsit5())) -# 48.917 μs (7173 allocations: 227.61 KiB) -sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb2, reltol = 1.0e-6) -# 48.750 μs (7158 allocations: 227.05 KiB) +sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) +# 54.916 μs (7185 allocations: 228.16 KiB) +sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) +# 53.083 μs (7165 allocations: 227.38 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -60,25 +62,29 @@ prob = ODEProblem(f, u0, tspan) #define callbacks condition1(u, t, integrator) = u[1] - 1 cb1 = ContinuousCallback(condition1, default_affect!; maybe_discontinuity = true) -cb1f = ContinuousCallback(condition1, default_affect!; maybe_discontinuity = false) condition2(u, t, integrator) = u[1] - 2 cb2 = ContinuousCallback(condition2, default_affect!; maybe_discontinuity = true) -cb2f = ContinuousCallback(condition2, default_affect!; maybe_discontinuity = false) cb = CallbackSet(cb1, cb2) -cb2 = CallbackSet(cb1f, cb2f) -sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1.0e-6, controller = PI_disco_controller(Rodas5P())) -# 1.306 ms (45129 allocations: 1.54 MiB) -sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb2, reltol = 1.0e-6) -# 1.500 ms (52871 allocations: 1.80 MiB) +sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) +# 917.292 μs (25563 allocations: 795.66 KiB) +sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) +# 1.325 ms (36765 allocations: 1.11 MiB) +@test sol_disco_radau.retcode == ReturnCode.Success +@test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject + +sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) +# 1.371 ms (48673 allocations: 1.66 MiB) +sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) +# 1.465 ms (52921 allocations: 1.80 MiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject -sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1.0e-6, controller = PI_disco_controller(Tsit5())) -# 261.375 μs (34373 allocations: 1.06 MiB) -sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb2, reltol = 1.0e-6) -# 266.125 μs (38745 allocations: 1.19 MiB) +sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) +# 272.292 μs (34890 allocations: 1.07 MiB) +sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) +# 285.625 μs (38807 allocations: 1.19 MiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -101,71 +107,60 @@ prob_multi = ODEProblem(f_multi_exp!, u0_multi, tspan_multi) #define callbacks cond_multi_1(u, t, integrator) = u[1] - 0.3 cb_multi_1 = ContinuousCallback(cond_multi_1, default_affect!; maybe_discontinuity = true) -cb_multi_1f = ContinuousCallback(cond_multi_1, default_affect!; maybe_discontinuity = false) cond_multi_2(u, t, integrator) = u[1] - 0.8 cb_multi_2 = ContinuousCallback(cond_multi_2, default_affect!; maybe_discontinuity = true) -cb_multi_2f = ContinuousCallback(cond_multi_2, default_affect!; maybe_discontinuity = false) cb_multi = CallbackSet(cb_multi_1, cb_multi_2) -cb_multi2 = CallbackSet(cb_multi_1f, cb_multi_2f) -#disco solve -sol_disco_radau = solve(prob_multi, RadauIIA5(); callback = cb_multi, reltol = 1.0e-7, abstol = 1.0e-9, controller = predictive_disco_controller(RadauIIA5())) -# 145.542 μs (1641 allocations: 64.64 KiB) -sol_no_disco_radau = solve(prob_multi, RadauIIA5(); callback = cb_multi2, reltol = 1.0e-7, abstol = 1.0e-9) -# 125.667 μs (1044 allocations: 46.12 KiB) -@test sol_disco_radau.retcode == ReturnCode.Success -@test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject - -sol_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback = cb_multi, reltol = 1.0e-7, abstol = 1.0e-9, controller = PI_disco_controller(Rodas5P())) -# 253.291 μs (1758 allocations: 75.62 KiB) -sol_no_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback = cb_multi2, reltol = 1.0e-7, abstol = 1.0e-9) -# 239.292 μs (947 allocations: 43.89 KiB) +sol_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) +# 210.333 μs (1403 allocations: 58.92 KiB) +sol_no_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9) +# 238.250 μs (1024 allocations: 45.09 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject -sol_disco_tsit5 = solve(prob_multi, Tsit5(); callback = cb_multi, reltol = 1.0e-7, abstol = 1.0e-9, controller = PI_disco_controller(Tsit5())) -# 99.250 μs (1603 allocations: 59.85 KiB) -sol_no_disco_tsit5 = solve(prob_multi, Tsit5(); callback = cb_multi2, reltol = 1.0e-7, abstol = 1.0e-9) -# 88.000 μs (1105 allocations: 45.27 KiB) +sol_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) +# 98.084 μs (1488 allocations: 58.20 KiB) +sol_no_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol = 1e-7, abstol = 1e-9) +# 89.958 μs (1184 allocations: 47.15 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject -#TEST 4: STIFF DISCONTINUITY -# very stiff discontinuous system +#TEST 4: STIFF MULTI-COMPONENT DISCONTINUITY function f_stiff_disc!(du, u, p, t) - λ = p[1] # stiffness parameter - return if u[1] < 0.5 - du[1] = -λ * u[1] + λ * exp(-t) # stiff decay with forcing + λ = p[1] + if u[1] < 0.5 + du[1] = -λ * u[1] + λ * exp(-t) + du[2] = -λ * (u[2] - u[1]) + du[3] = -(u[3] - u[1]^2) else - du[1] = u[1] + du[1] = u[1] - u[2] + du[2] = -λ * (u[2] - u[1]) + du[3] = -(u[3] - u[1]) end end -u0_stiff = [0.1] -tspan_stiff = (0.0, 3.0) -prob_stiff = ODEProblem(f_stiff_disc!, u0_stiff, tspan_stiff, [100.0]) +u0_stiff = [0.1, 0.1, 0.01] +tspan_stiff = (0.0, 2.0) +prob_stiff = ODEProblem(f_stiff_disc!, u0_stiff, tspan_stiff, [500.0]) -#define callback cond_stiff(u, t, integrator) = u[1] - 0.5 cb_stiff = ContinuousCallback(cond_stiff, default_affect!; maybe_discontinuity = true) -cb_stiff_f = ContinuousCallback(cond_stiff, default_affect!; maybe_discontinuity = false) -#disco solve -sol_disco_radau = solve(prob_stiff, RadauIIA5(); callback = cb_stiff, reltol = 1.0e-9, abstol = 1.0e-11, controller = predictive_disco_controller(RadauIIA5())) -# 119.250 μs (1666 allocations: 66.52 KiB) -sol_no_disco_radau = solve(prob_stiff, RadauIIA5(); callback = cb_stiff_f, reltol = 1.0e-9, abstol = 1.0e-11) -# 124.167 μs (1476 allocations: 61.16 KiB) +sol_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = predictive_disco_controller(RadauIIA5())) +# 208.959 μs (2257 allocations: 101.69 KiB) +sol_no_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) +# 191.500 μs (2009 allocations: 93.05 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject -sol_disco_tsit5 = solve(prob_stiff, Tsit5(); callback = cb_stiff, reltol = 1.0e-9, abstol = 1.0e-11, controller = PI_disco_controller(Tsit5())) -# 80.375 μs (1870 allocations: 70.51 KiB) -sol_no_disco_tsit5 = solve(prob_stiff, Tsit5(); callback = cb_stiff_f, reltol = 1.0e-9, abstol = 1.0e-11) -# 79.583 μs (1783 allocations: 68.51 KiB) -@test sol_disco_tsit5.retcode == ReturnCode.Success -@test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject +sol_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = PI_disco_controller(Rodas5P())) +# 340.958 μs (2060 allocations: 105.09 KiB) +sol_no_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) +# 325.375 μs (1793 allocations: 79.44 KiB) +@test sol_disco_rosenbrock.retcode == ReturnCode.Success +@test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject #TEST 5: DISCONTINUOUS DAE # discontinuous DAE with mass matrix @@ -192,12 +187,11 @@ prob_dae = ODEProblem(f_dae_func, u0_dae, tspan_dae) cond_dae(u, t, integrator) = u[1] - 0.5 cb_dae = ContinuousCallback(cond_dae, default_affect!; maybe_discontinuity = true) -cb_daef = ContinuousCallback(cond_dae, default_affect!; maybe_discontinuity = false) -radau_disco = solve(prob_dae, RadauIIA5(); callback = cb_dae, reltol = 1.0e-8, abstol = 1.0e-10, controller = predictive_disco_controller(RadauIIA5())) -# 62.875 μs (754 allocations: 33.98 KiB) -radau_no_disco = solve(prob_dae, RadauIIA5(); callback = cb_daef, reltol = 1.0e-8, abstol = 1.0e-10) -# 58.417 μs (639 allocations: 30.61 KiB) +radau_disco = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10, controller = predictive_disco_controller(RadauIIA5())) +# 62.083 μs (743 allocations: 33.97 KiB) +radau_no_disco = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10) +# 58.458 μs (641 allocations: 30.88 KiB) @test radau_disco.retcode == ReturnCode.Success @test radau_disco.stats.nreject <= radau_no_disco.stats.nreject @@ -218,19 +212,18 @@ function condition_vec!(out, u, t, integrator) end default_affect_vec!(integrator, idx) = nothing cb_vec = VectorContinuousCallback(condition_vec!, default_affect_vec!, 2; maybe_discontinuity = true) -cb_vec2 = VectorContinuousCallback(condition_vec!, default_affect_vec!, 2; maybe_discontinuity = false) -sol_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol = 1.0e-7, abstol = 1.0e-9, controller = PI_disco_controller(Rodas5P())) -# 175.292 μs (1213 allocations: 54.08 KiB) -sol_no_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec2, reltol = 1.0e-7, abstol = 1.0e-9) -# 194.709 μs (888 allocations: 46.53 KiB) +sol_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) +# 197.417 μs (1289 allocations: 57.34 KiB) +sol_no_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9) +# 192.542 μs (936 allocations: 47.69 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject -sol_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol = 1.0e-7, abstol = 1.0e-9, controller = PI_disco_controller(Tsit5())) -# 73.333 μs (1349 allocations: 57.71 KiB) -sol_no_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec2, reltol = 1.0e-7, abstol = 1.0e-9) -# 74.708 μs (1256 allocations: 58.48 KiB) +sol_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) +# 87.583 μs (1537 allocations: 63.20 KiB) +sol_no_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9) +# 76.042 μs (1295 allocations: 59.73 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -251,18 +244,17 @@ prob = ODEProblem(f!, u, tspan) cond(u, t, integrator) = u[2] cb = ContinuousCallback(cond, default_affect!; maybe_discontinuity = true) -cb2 = ContinuousCallback(cond, default_affect!; maybe_discontinuity = false) -sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1.0e-8, abstol = 1.0e-10, controller = PI_disco_controller(Rodas5P())) -# 215.792 μs (1621 allocations: 61.61 KiB) -sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb2, reltol = 1.0e-8, abstol = 1.0e-10) -# 172.333 μs (845 allocations: 43.62 KiB) +sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Rodas5P())) +# 161.167 μs (1154 allocations: 51.39 KiB) +sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10) +# 177.250 μs (910 allocations: 45.09 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject -sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1.0e-8, abstol = 1.0e-10, controller = PI_disco_controller(Tsit5())) -# 66.292 μs (1450 allocations: 59.10 KiB) -sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb2, reltol = 1.0e-8, abstol = 1.0e-10) -# 52.500 μs (1117 allocations: 52.55 KiB) +sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Tsit5())) +# 64.958 μs (1357 allocations: 59.46 KiB) +sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10) +# 52.375 μs (1173 allocations: 54.07 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject From 268bbb0c468f5f68d62ce062477b5e8dcde28ccb Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Fri, 5 Jun 2026 11:23:06 -0400 Subject: [PATCH 2/9] addsteps optimization --- lib/OrdinaryDiffEqCore/src/disco.jl | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/lib/OrdinaryDiffEqCore/src/disco.jl b/lib/OrdinaryDiffEqCore/src/disco.jl index f4befea1371..755f7e82cbc 100644 --- a/lib/OrdinaryDiffEqCore/src/disco.jl +++ b/lib/OrdinaryDiffEqCore/src/disco.jl @@ -22,6 +22,7 @@ function find_discontinuity(u, uprev, integrator) 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 @@ -38,10 +39,13 @@ 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])) + 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 sol = solve(disco_prob, tspan = (0.0, breakpointθ == -one(dt) ? 1.0 : breakpointθ); save_everystep = false) tmp = sol[] @@ -56,8 +60,11 @@ function find_discontinuity(u, uprev, integrator) out_prev = i.condition(uprev, t, integrator) out_curr = i.condition(u, t + dt, integrator) if (out_prev * out_curr < zero(out_prev)) - _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) + 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 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θ)) From 690e356b6264ccfbb3595e3d482b086bfa822d74 Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Mon, 8 Jun 2026 14:11:41 -0400 Subject: [PATCH 3/9] Update disco_tests.jl --- test/Integrators_I/disco_tests.jl | 3 --- 1 file changed, 3 deletions(-) diff --git a/test/Integrators_I/disco_tests.jl b/test/Integrators_I/disco_tests.jl index 9d484d96a04..c3f5d968a7d 100644 --- a/test/Integrators_I/disco_tests.jl +++ b/test/Integrators_I/disco_tests.jl @@ -1,8 +1,5 @@ using OrdinaryDiffEqCore, DiffEqDevTools, Test, LinearAlgebra using OrdinaryDiffEqTsit5, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqFIRK -using Logging -global_logger(SimpleLogger(stderr, Logging.Error)) -using BenchmarkTools predictive_disco_controller(alg) = OrdinaryDiffEqCore.PredictiveController(alg; discontinuity_detection = true) PI_disco_controller(alg) = OrdinaryDiffEqCore.PIController(alg; discontinuity_detection = true) From a95e473561c939f01eaa4f22a609ef975d31ed4a Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Fri, 12 Jun 2026 09:23:09 -0400 Subject: [PATCH 4/9] refactor controller disco logic --- lib/OrdinaryDiffEqBDF/src/controllers.jl | 50 +++++------ lib/OrdinaryDiffEqCore/src/disco.jl | 16 ++-- .../src/integrators/controllers.jl | 88 ++++++++----------- .../src/integrators/integrator_utils.jl | 15 ++++ .../src/integrators/type.jl | 2 +- lib/OrdinaryDiffEqCore/src/solve.jl | 4 +- lib/OrdinaryDiffEqFIRK/src/controllers.jl | 10 ++- test/Integrators_I/disco_tests.jl | 84 +++++++++++------- 8 files changed, 143 insertions(+), 126 deletions(-) diff --git a/lib/OrdinaryDiffEqBDF/src/controllers.jl b/lib/OrdinaryDiffEqBDF/src/controllers.jl index f54acd8d141..37f1df5964e 100644 --- a/lib/OrdinaryDiffEqBDF/src/controllers.jl +++ b/lib/OrdinaryDiffEqBDF/src/controllers.jl @@ -83,14 +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 - if integrator.curr_discontinuity != -1 && - integrator.tdir * integrator.t >= integrator.tdir * integrator.disco_checkpoint - integrator.curr_discontinuity = -1 + 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 @@ -152,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 && q < 1.2 && q > 0.6 + 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) @@ -179,11 +179,7 @@ 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 no_prior_discontinuity && disco_dt > zero(disco_dt) - add_tstop!(integrator, integrator.disco_checkpoint) - end + disco_dt = set_discontinuity(integrator) if disco_dt > zero(disco_dt) integrator.dt = disco_dt return integrator.dt @@ -425,9 +421,9 @@ 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 + 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) @@ -441,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) @@ -593,9 +590,9 @@ 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 + 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) @@ -609,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 diff --git a/lib/OrdinaryDiffEqCore/src/disco.jl b/lib/OrdinaryDiffEqCore/src/disco.jl index 755f7e82cbc..37fac31aa64 100644 --- a/lib/OrdinaryDiffEqCore/src/disco.jl +++ b/lib/OrdinaryDiffEqCore/src/disco.jl @@ -1,5 +1,5 @@ -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 @@ -11,9 +11,11 @@ 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 + u = integrator.u + uprev = integrator.uprev cb === nothing && return -one(dt) isempty(cb.continuous_callbacks) && return -one(dt) p = integrator.p @@ -51,8 +53,8 @@ function find_discontinuity(u, uprev, integrator) tmp = sol[] if (!isnan(tmp) && (breakpointθ < zero(breakpointθ) || tmp < breakpointθ)) breakpointθ = tmp - integrator.curr_discontinuity = idx - integrator.disco_checkpoint = integrator.t + tmp * dt + 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 @@ -69,8 +71,8 @@ function find_discontinuity(u, uprev, integrator) tmp = sol[] if (!isnan(tmp) && (breakpointθ < zero(breakpointθ) || tmp < breakpointθ)) breakpointθ = tmp - integrator.curr_discontinuity = idx - integrator.disco_checkpoint = integrator.t + tmp * dt + integrator.is_disco_step = true + integrator.disco_checkpoint = integrator.t + dt #our prev rejected step, we shouldn't step past this end end end diff --git a/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl b/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl index a4b0ecd74fd..709892e864b 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl +++ b/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl @@ -572,24 +572,22 @@ end function step_accept_controller!(integrator, cache::IControllerCache, alg, q) (; qsteady_min, qsteady_max, discontinuity_detection) = cache.controller.basic + t = integrator.t + dt = integrator.dt 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 + if discontinuity_detection && integrator.is_disco_step + integrator.is_disco_step = false + return min((integrator.disco_checkpoint - t) / 4, dt / q) end - return integrator.dt / q # new dt + return 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 + disco_dt = set_discontinuity(integrator) if disco_dt > zero(disco_dt) integrator.dt = disco_dt return integrator.dt @@ -720,29 +718,25 @@ function step_accept_controller!(integrator, cache::PIControllerCache, alg, q) qoldinit = controller.qoldinit EEst = DiffEqBase.value(get_EEst(integrator)) + t = integrator.t + dt = integrator.dt 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 - cache.errold = qoldinit - else - cache.errold = max(EEst, qoldinit) + cache.errold = max(EEst, qoldinit) + if discontinuity_detection && integrator.is_disco_step + integrator.is_disco_step = false + return min((integrator.disco_checkpoint - t) / 4, dt / q) # new dt end - return integrator.dt / q # new dt + + return 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 + disco_dt = set_discontinuity(integrator) if disco_dt > zero(disco_dt) return integrator.dt = disco_dt end @@ -960,33 +954,26 @@ function step_accept_controller!(integrator, cache::PIDControllerCache, alg, dt_ (; controller) = cache (; qsteady_min, qsteady_max, discontinuity_detection) = controller.basic + t = integrator.t + dt = integrator.dt if qsteady_min <= inv(dt_factor) <= qsteady_max dt_factor = one(dt_factor) end - 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 + @inbounds begin + cache.err[3] = cache.err[2] + cache.err[2] = cache.err[1] + end + if discontinuity_detection && integrator.is_disco_step + integrator.is_disco_step = false + return min((integrator.disco_checkpoint - t) / 4, dt * dt_factor) end - return integrator.dt * dt_factor # new dt + return 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 + disco_dt = set_discontinuity(integrator) if disco_dt > zero(disco_dt) integrator.dt = disco_dt return integrator.dt @@ -1146,14 +1133,13 @@ function step_accept_controller!(integrator, cache::PredictiveControllerCache, a if qsteady_min <= qacc <= qsteady_max qacc = one(qacc) end - 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) + + cache.dtacc = DiffEqBase.value(integrator.dt) + cache.erracc = max(1.0e-2, EEst) + + if cache.controller.basic.discontinuity_detection && integrator.is_disco_step + integrator.is_disco_step = false + return min((integrator.disco_checkpoint - integrator.t) / 4, integrator.dt / qacc) end return integrator.dt / qacc @@ -1164,11 +1150,7 @@ 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 + disco_dt = set_discontinuity(integrator) if disco_dt > zero(disco_dt) integrator.dt = disco_dt return integrator.dt diff --git a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl b/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl index 527a6c3cc7f..87c46013990 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl +++ b/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl @@ -621,6 +621,18 @@ function fixed_t_for_tstop_error!(integrator, ttmp) end end +# Type-stable check: did the callback that fired have maybe_discontinuity = true? +@generated function _fired_cb_maybe_discontinuity( + cb_idx, + callbacks::NTuple{N, Union{ContinuousCallback, VectorContinuousCallback}} + ) where {N} + ex = :(false) + for i in N:-1:1 + ex = :(cb_idx == $i ? callbacks[$i].maybe_discontinuity : $ex) + end + return ex +end + # Use a generated function to call apply_callback! in a type-stable way @generated function apply_ith_callback!( integrator, @@ -682,6 +694,9 @@ function handle_callbacks!(integrator) idx, continuous_callbacks ) + if _fired_cb_maybe_discontinuity(idx, continuous_callbacks) + reinit_controller!(integrator, integrator.controller_cache) + end else integrator.event_last_time = 0 integrator.vector_event_last_time = 1 diff --git a/lib/OrdinaryDiffEqCore/src/integrators/type.jl b/lib/OrdinaryDiffEqCore/src/integrators/type.jl index 66db13ffa96..29a90571e73 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/type.jl +++ b/lib/OrdinaryDiffEqCore/src/integrators/type.jl @@ -176,7 +176,7 @@ mutable struct ODEIntegrator{ fsalfirst::FSALType fsallast::FSALType rng::RNGType - curr_discontinuity::Int + is_disco_step::Bool disco_checkpoint::tType W::WType P::PType diff --git a/lib/OrdinaryDiffEqCore/src/solve.jl b/lib/OrdinaryDiffEqCore/src/solve.jl index 361194531e4..18511cf014a 100644 --- a/lib/OrdinaryDiffEqCore/src/solve.jl +++ b/lib/OrdinaryDiffEqCore/src/solve.jl @@ -688,7 +688,7 @@ Base.@constprop :aggressive function _ode_init( controller_cache = setup_controller_cache(_alg, cache, controller, EEstT, disco_probs) - curr_discontinuity = -1 + is_disco_step = false disco_checkpoint = zero(tType) # Seed the initial EEst on the controller cache (was previously # `integrator.EEst = oneunit(EEstT)`). @@ -723,7 +723,7 @@ Base.@constprop :aggressive function _ode_init( isout, reeval_fsal, derivative_discontinuity, user_set_discontinuity, reinitialize, isdae, opts, stats, initializealg, differential_vars, - fsalfirst, fsallast, _rng, curr_discontinuity, disco_checkpoint, + fsalfirst, fsallast, _rng, is_disco_step, disco_checkpoint, W, P, sqdt, noise, c, rate_constants ) diff --git a/lib/OrdinaryDiffEqFIRK/src/controllers.jl b/lib/OrdinaryDiffEqFIRK/src/controllers.jl index 473c9be9dfe..46845d63739 100644 --- a/lib/OrdinaryDiffEqFIRK/src/controllers.jl +++ b/lib/OrdinaryDiffEqFIRK/src/controllers.jl @@ -21,8 +21,6 @@ function step_accept_controller!( if qsteady_min <= qacc <= qsteady_max qacc = one(qacc) end - ccache.dtacc = integrator.dt - ccache.erracc = max(1.0e-2, EEst) cache.step = step + 1 hist_iter = hist_iter * 0.8 + iter * 0.2 cache.hist_iter = hist_iter @@ -46,6 +44,12 @@ function step_accept_controller!( cache.hist_iter = iter end end + ccache.dtacc = integrator.dt + ccache.erracc = max(1.0e-2, EEst) + if controller.basic.discontinuity_detection && integrator.is_disco_step + integrator.is_disco_step = false + return min((integrator.disco_checkpoint - integrator.t) / 4, integrator.dt / qacc) + end return integrator.dt / qacc end @@ -63,7 +67,7 @@ function step_reject_controller!( min_stages = (alg.min_order - 1) ÷ 4 * 2 + 1 if (discontinuity_detection) - disco_dt = set_discontinuity(integrator.u, integrator.uprev, integrator) + disco_dt = set_discontinuity(integrator) if disco_dt > zero(dt) integrator.dt = disco_dt end diff --git a/test/Integrators_I/disco_tests.jl b/test/Integrators_I/disco_tests.jl index c3f5d968a7d..bea08e62b21 100644 --- a/test/Integrators_I/disco_tests.jl +++ b/test/Integrators_I/disco_tests.jl @@ -1,5 +1,7 @@ using OrdinaryDiffEqCore, DiffEqDevTools, Test, LinearAlgebra using OrdinaryDiffEqTsit5, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqFIRK +using Logging +global_logger(ConsoleLogger(stderr, Logging.Error)) predictive_disco_controller(alg) = OrdinaryDiffEqCore.PredictiveController(alg; discontinuity_detection = true) PI_disco_controller(alg) = OrdinaryDiffEqCore.PIController(alg; discontinuity_detection = true) @@ -20,23 +22,23 @@ condition(u, t, integrator) = u[1] - 1 cb = ContinuousCallback(condition, default_affect!; maybe_discontinuity = true) sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) -# 272.333 μs (7961 allocations: 249.34 KiB) +# 261.500 μs (7667 allocations: 240.53 KiB) sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) -# 369.000 μs (10315 allocations: 321.41 KiB) +# 357.834 μs (10329 allocations: 322.31 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) -# 419.750 μs (15586 allocations: 546.83 KiB) +# 440.167 μs (16190 allocations: 568.30 KiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) -# 587.042 μs (21827 allocations: 764.73 KiB) +# 583.750 μs (21828 allocations: 765.69 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) -# 54.916 μs (7185 allocations: 228.16 KiB) +# 54.791 μs (7183 allocations: 228.30 KiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) -# 53.083 μs (7165 allocations: 227.38 KiB) +# 55.875 μs (7177 allocations: 228.20 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -65,23 +67,23 @@ cb2 = ContinuousCallback(condition2, default_affect!; maybe_discontinuity = true cb = CallbackSet(cb1, cb2) sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) -# 917.292 μs (25563 allocations: 795.66 KiB) +# 947.334 μs (26352 allocations: 821.97 KiB) sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) -# 1.325 ms (36765 allocations: 1.11 MiB) +# 1.320 ms (36787 allocations: 1.11 MiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) -# 1.371 ms (48673 allocations: 1.66 MiB) +# 1.250 ms (45708 allocations: 1.56 MiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) -# 1.465 ms (52921 allocations: 1.80 MiB) +# 1.439 ms (52903 allocations: 1.80 MiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) -# 272.292 μs (34890 allocations: 1.07 MiB) +# 277.417 μs (35712 allocations: 1.10 MiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) -# 285.625 μs (38807 allocations: 1.19 MiB) +# 286.750 μs (38772 allocations: 1.19 MiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -110,17 +112,24 @@ cb_multi_2 = ContinuousCallback(cond_multi_2, default_affect!; maybe_discontinui cb_multi = CallbackSet(cb_multi_1, cb_multi_2) +sol_disco_radau = solve(prob_multi, RadauIIA5(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(RadauIIA5())) +# 142.917 μs (1576 allocations: 65.19 KiB) +sol_no_disco_radau = solve(prob_multi, RadauIIA5(); callback=cb_multi, reltol=1e-7, abstol=1e-9) +# 126.958 μs (1073 allocations: 48.44 KiB) +@test sol_disco_radau.retcode == ReturnCode.Success +@test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject + sol_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) -# 210.333 μs (1403 allocations: 58.92 KiB) +# 252.000 μs (1801 allocations: 71.69 KiB) sol_no_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9) -# 238.250 μs (1024 allocations: 45.09 KiB) +# 240.291 μs (969 allocations: 45.44 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) -# 98.084 μs (1488 allocations: 58.20 KiB) +# 93.334 μs (1430 allocations: 55.13 KiB) sol_no_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol = 1e-7, abstol = 1e-9) -# 89.958 μs (1184 allocations: 47.15 KiB) +# 85.542 μs (1133 allocations: 46.84 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -146,16 +155,16 @@ cond_stiff(u, t, integrator) = u[1] - 0.5 cb_stiff = ContinuousCallback(cond_stiff, default_affect!; maybe_discontinuity = true) sol_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = predictive_disco_controller(RadauIIA5())) -# 208.959 μs (2257 allocations: 101.69 KiB) +# 201.792 μs (2315 allocations: 103.89 KiB) sol_no_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) -# 191.500 μs (2009 allocations: 93.05 KiB) +# 191.791 μs (2022 allocations: 94.27 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = PI_disco_controller(Rodas5P())) -# 340.958 μs (2060 allocations: 105.09 KiB) +# 345.666 μs (1922 allocations: 81.44 KiB) sol_no_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) -# 325.375 μs (1793 allocations: 79.44 KiB) +# 328.875 μs (1765 allocations: 79.78 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject @@ -185,12 +194,19 @@ prob_dae = ODEProblem(f_dae_func, u0_dae, tspan_dae) cond_dae(u, t, integrator) = u[1] - 0.5 cb_dae = ContinuousCallback(cond_dae, default_affect!; maybe_discontinuity = true) -radau_disco = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10, controller = predictive_disco_controller(RadauIIA5())) -# 62.083 μs (743 allocations: 33.97 KiB) -radau_no_disco = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10) -# 58.458 μs (641 allocations: 30.88 KiB) -@test radau_disco.retcode == ReturnCode.Success -@test radau_disco.stats.nreject <= radau_no_disco.stats.nreject +sol_disco_radau = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10, controller = predictive_disco_controller(RadauIIA5())) +# 55.000 μs (739 allocations: 34.45 KiB) +sol_no_disco_radau = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10) +# 58.542 μs (655 allocations: 32.09 KiB) +@test sol_disco_radau.retcode == ReturnCode.Success +@test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject + +sol_disco_rosenbrock = solve(prob_dae, Rodas5P(); callback=cb_dae, reltol=1e-8, abstol=1e-10, controller = PI_disco_controller(Rodas5P())) +# 120.083 μs (819 allocations: 35.34 KiB) +sol_no_disco_rosenbrock = solve(prob_dae, Rodas5P(); callback=cb_dae, reltol=1e-8, abstol=1e-10) +# 119.708 μs (625 allocations: 31.69 KiB) +@test sol_disco_rosenbrock.retcode == ReturnCode.Success +@test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject #TEST 6: VECTOR CALLBACK function f_vec_disc!(du, u, p, t) @@ -211,16 +227,16 @@ default_affect_vec!(integrator, idx) = nothing cb_vec = VectorContinuousCallback(condition_vec!, default_affect_vec!, 2; maybe_discontinuity = true) sol_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) -# 197.417 μs (1289 allocations: 57.34 KiB) +# 191.416 μs (1197 allocations: 55.30 KiB) sol_no_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9) -# 192.542 μs (936 allocations: 47.69 KiB) +# 190.708 μs (915 allocations: 48.44 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) -# 87.583 μs (1537 allocations: 63.20 KiB) +# 83.083 μs (1498 allocations: 64.91 KiB) sol_no_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9) -# 76.042 μs (1295 allocations: 59.73 KiB) +# 74.916 μs (1282 allocations: 59.90 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -243,15 +259,15 @@ cond(u, t, integrator) = u[2] cb = ContinuousCallback(cond, default_affect!; maybe_discontinuity = true) sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Rodas5P())) -# 161.167 μs (1154 allocations: 51.39 KiB) +# 163.458 μs (1148 allocations: 50.66 KiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10) -# 177.250 μs (910 allocations: 45.09 KiB) +# 176.459 μs (866 allocations: 45.28 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Tsit5())) -# 64.958 μs (1357 allocations: 59.46 KiB) +# 58.375 μs (1227 allocations: 53.41 KiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10) -# 52.375 μs (1173 allocations: 54.07 KiB) +# 53.791 μs (1137 allocations: 53.68 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject From 96074665321510f12c7690f2795a7b8eb2e7ed0d Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Mon, 15 Jun 2026 14:26:50 -0400 Subject: [PATCH 5/9] add biasing --- lib/OrdinaryDiffEqCore/src/disco.jl | 18 +++--- lib/OrdinaryDiffEqCore/test/runtests.jl | 1 - test/Integrators_I/disco_tests.jl | 73 +++++++++++-------------- 3 files changed, 41 insertions(+), 51 deletions(-) diff --git a/lib/OrdinaryDiffEqCore/src/disco.jl b/lib/OrdinaryDiffEqCore/src/disco.jl index 37fac31aa64..2a7610e2847 100644 --- a/lib/OrdinaryDiffEqCore/src/disco.jl +++ b/lib/OrdinaryDiffEqCore/src/disco.jl @@ -1,8 +1,8 @@ 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 end return -one(dt) end @@ -16,12 +16,12 @@ function find_discontinuity(integrator) dt = integrator.dt u = integrator.u uprev = integrator.uprev - cb === nothing && return -one(dt) - isempty(cb.continuous_callbacks) && return -one(dt) + 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 @@ -49,9 +49,9 @@ function find_discontinuity(integrator) disco_zero.dt, disco_zero.f, disco_zero.p, disco_zero.cache, false, true, false) end disco_zero.ind = j - sol = solve(disco_prob, tspan = (0.0, breakpointθ == -one(dt) ? 1.0 : breakpointθ); save_everystep = false) + sol = solve(disco_prob, tspan = (0.0, breakpointθ); save_everystep = false) 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 @@ -67,9 +67,9 @@ function find_discontinuity(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) end - sol = solve(disco_prob, tspan = (0.0, breakpointθ == -one(dt) ? 1.0 : breakpointθ); save_everystep = false) + sol = solve(disco_prob, tspan = (0.0, breakpointθ); save_everystep = false) 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 diff --git a/lib/OrdinaryDiffEqCore/test/runtests.jl b/lib/OrdinaryDiffEqCore/test/runtests.jl index 46342deb64b..1f79736c7a8 100644 --- a/lib/OrdinaryDiffEqCore/test/runtests.jl +++ b/lib/OrdinaryDiffEqCore/test/runtests.jl @@ -33,5 +33,4 @@ end # Functional tests if TEST_GROUP == "Core" || TEST_GROUP == "ALL" @time @safetestset "Sparse isdiag Performance" include("sparse_isdiag_tests.jl") - @time @safetestset "Discontinuity Detection" include("disco_tests.jl") end diff --git a/test/Integrators_I/disco_tests.jl b/test/Integrators_I/disco_tests.jl index bea08e62b21..242454e05b8 100644 --- a/test/Integrators_I/disco_tests.jl +++ b/test/Integrators_I/disco_tests.jl @@ -1,7 +1,5 @@ using OrdinaryDiffEqCore, DiffEqDevTools, Test, LinearAlgebra using OrdinaryDiffEqTsit5, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqFIRK -using Logging -global_logger(ConsoleLogger(stderr, Logging.Error)) predictive_disco_controller(alg) = OrdinaryDiffEqCore.PredictiveController(alg; discontinuity_detection = true) PI_disco_controller(alg) = OrdinaryDiffEqCore.PIController(alg; discontinuity_detection = true) @@ -22,23 +20,23 @@ condition(u, t, integrator) = u[1] - 1 cb = ContinuousCallback(condition, default_affect!; maybe_discontinuity = true) sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) -# 261.500 μs (7667 allocations: 240.53 KiB) +# 310.958 μs (8429 allocations: 263.88 KiB) sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) -# 357.834 μs (10329 allocations: 322.31 KiB) +# 392.667 μs (10628 allocations: 331.61 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) -# 440.167 μs (16190 allocations: 568.30 KiB) +# 443.834 μs (15594 allocations: 547.55 KiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) -# 583.750 μs (21828 allocations: 765.69 KiB) +# 620.000 μs (21828 allocations: 765.69 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) -# 54.791 μs (7183 allocations: 228.30 KiB) +# 60.250 μs (7576 allocations: 240.67 KiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) -# 55.875 μs (7177 allocations: 228.20 KiB) +# 58.333 μs (7570 allocations: 240.58 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -67,23 +65,23 @@ cb2 = ContinuousCallback(condition2, default_affect!; maybe_discontinuity = true cb = CallbackSet(cb1, cb2) sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) -# 947.334 μs (26352 allocations: 821.97 KiB) +# 1.067 ms (28126 allocations: 875.64 KiB) sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) -# 1.320 ms (36787 allocations: 1.11 MiB) +# 1.311 ms (34747 allocations: 1.05 MiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) -# 1.250 ms (45708 allocations: 1.56 MiB) +# 1.128 ms (38426 allocations: 1.31 MiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) -# 1.439 ms (52903 allocations: 1.80 MiB) +# 1.464 ms (52903 allocations: 1.80 MiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) -# 277.417 μs (35712 allocations: 1.10 MiB) +# 271.125 μs (35982 allocations: 1.11 MiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) -# 286.750 μs (38772 allocations: 1.19 MiB) +# 279.166 μs (39645 allocations: 1.22 MiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -113,23 +111,23 @@ cb_multi_2 = ContinuousCallback(cond_multi_2, default_affect!; maybe_discontinui cb_multi = CallbackSet(cb_multi_1, cb_multi_2) sol_disco_radau = solve(prob_multi, RadauIIA5(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(RadauIIA5())) -# 142.917 μs (1576 allocations: 65.19 KiB) +# 137.542 μs (1470 allocations: 61.42 KiB) sol_no_disco_radau = solve(prob_multi, RadauIIA5(); callback=cb_multi, reltol=1e-7, abstol=1e-9) -# 126.958 μs (1073 allocations: 48.44 KiB) +# 127.209 μs (1073 allocations: 48.44 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) -# 252.000 μs (1801 allocations: 71.69 KiB) +# 229.167 μs (1593 allocations: 64.94 KiB) sol_no_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9) -# 240.291 μs (969 allocations: 45.44 KiB) +# 239.750 μs (969 allocations: 45.44 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) -# 93.334 μs (1430 allocations: 55.13 KiB) +# 96.917 μs (1469 allocations: 56.32 KiB) sol_no_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol = 1e-7, abstol = 1e-9) -# 85.542 μs (1133 allocations: 46.84 KiB) +# 90.250 μs (1133 allocations: 46.84 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -155,16 +153,16 @@ cond_stiff(u, t, integrator) = u[1] - 0.5 cb_stiff = ContinuousCallback(cond_stiff, default_affect!; maybe_discontinuity = true) sol_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = predictive_disco_controller(RadauIIA5())) -# 201.792 μs (2315 allocations: 103.89 KiB) +# 202.750 μs (2250 allocations: 101.73 KiB) sol_no_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) -# 191.791 μs (2022 allocations: 94.27 KiB) +# 198.958 μs (2054 allocations: 95.58 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = PI_disco_controller(Rodas5P())) -# 345.666 μs (1922 allocations: 81.44 KiB) +# 379.875 μs (2053 allocations: 85.81 KiB) sol_no_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) -# 328.875 μs (1765 allocations: 79.78 KiB) +# 335.709 μs (1765 allocations: 79.78 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject @@ -194,17 +192,10 @@ prob_dae = ODEProblem(f_dae_func, u0_dae, tspan_dae) cond_dae(u, t, integrator) = u[1] - 0.5 cb_dae = ContinuousCallback(cond_dae, default_affect!; maybe_discontinuity = true) -sol_disco_radau = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10, controller = predictive_disco_controller(RadauIIA5())) -# 55.000 μs (739 allocations: 34.45 KiB) -sol_no_disco_radau = solve(prob_dae, RadauIIA5(); callback=cb_dae, reltol=1e-8, abstol=1e-10) -# 58.542 μs (655 allocations: 32.09 KiB) -@test sol_disco_radau.retcode == ReturnCode.Success -@test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject - sol_disco_rosenbrock = solve(prob_dae, Rodas5P(); callback=cb_dae, reltol=1e-8, abstol=1e-10, controller = PI_disco_controller(Rodas5P())) -# 120.083 μs (819 allocations: 35.34 KiB) +# 117.500 μs (810 allocations: 35.44 KiB) sol_no_disco_rosenbrock = solve(prob_dae, Rodas5P(); callback=cb_dae, reltol=1e-8, abstol=1e-10) -# 119.708 μs (625 allocations: 31.69 KiB) +# 123.333 μs (625 allocations: 31.69 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject @@ -227,16 +218,16 @@ default_affect_vec!(integrator, idx) = nothing cb_vec = VectorContinuousCallback(condition_vec!, default_affect_vec!, 2; maybe_discontinuity = true) sol_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) -# 191.416 μs (1197 allocations: 55.30 KiB) +# 168.250 μs (1109 allocations: 51.56 KiB) sol_no_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9) -# 190.708 μs (915 allocations: 48.44 KiB) +# 204.292 μs (937 allocations: 49.41 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) -# 83.083 μs (1498 allocations: 64.91 KiB) +# 98.542 μs (1705 allocations: 69.88 KiB) sol_no_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9) -# 74.916 μs (1282 allocations: 59.90 KiB) +# 80.667 μs (1354 allocations: 62.98 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -259,15 +250,15 @@ cond(u, t, integrator) = u[2] cb = ContinuousCallback(cond, default_affect!; maybe_discontinuity = true) sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Rodas5P())) -# 163.458 μs (1148 allocations: 50.66 KiB) +# 152.583 μs (1074 allocations: 48.30 KiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10) -# 176.459 μs (866 allocations: 45.28 KiB) +# 181.875 μs (877 allocations: 45.72 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Tsit5())) -# 58.375 μs (1227 allocations: 53.41 KiB) +# 62.791 μs (1284 allocations: 55.26 KiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10) -# 53.791 μs (1137 allocations: 53.68 KiB) +# 54.708 μs (1155 allocations: 54.55 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject From 31da61e7b674bb857dd6cbf4ed51fc989b3b0768 Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Mon, 15 Jun 2026 14:28:00 -0400 Subject: [PATCH 6/9] Update disco.jl --- lib/OrdinaryDiffEqCore/src/disco.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/OrdinaryDiffEqCore/src/disco.jl b/lib/OrdinaryDiffEqCore/src/disco.jl index 2a7610e2847..ace16a72d04 100644 --- a/lib/OrdinaryDiffEqCore/src/disco.jl +++ b/lib/OrdinaryDiffEqCore/src/disco.jl @@ -49,7 +49,7 @@ function find_discontinuity(integrator) disco_zero.dt, disco_zero.f, disco_zero.p, disco_zero.cache, false, true, false) end disco_zero.ind = j - sol = solve(disco_prob, tspan = (0.0, breakpointθ); save_everystep = false) + sol = solve(disco_prob, tspan = (0.0, breakpointθ)) tmp = sol[] if (!isnan(tmp) && tmp < breakpointθ) breakpointθ = tmp @@ -67,7 +67,7 @@ function find_discontinuity(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) end - sol = solve(disco_prob, tspan = (0.0, breakpointθ); save_everystep = false) + sol = solve(disco_prob, tspan = (0.0, breakpointθ)) tmp = sol[] if (!isnan(tmp) && tmp < breakpointθ) breakpointθ = tmp From 80d7e74c82511e7670a67ee1cdf695231981882d Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Tue, 16 Jun 2026 13:06:09 -0400 Subject: [PATCH 7/9] code cleanup --- lib/OrdinaryDiffEqBDF/src/controllers.jl | 2 +- lib/OrdinaryDiffEqCore/src/disco.jl | 1 + .../src/integrators/controllers.jl | 84 +++++++------------ test/Integrators_I/disco_tests.jl | 64 +++++++------- 4 files changed, 62 insertions(+), 89 deletions(-) diff --git a/lib/OrdinaryDiffEqBDF/src/controllers.jl b/lib/OrdinaryDiffEqBDF/src/controllers.jl index 37f1df5964e..bfc3dfe9cbd 100644 --- a/lib/OrdinaryDiffEqBDF/src/controllers.jl +++ b/lib/OrdinaryDiffEqBDF/src/controllers.jl @@ -153,7 +153,7 @@ function step_accept_controller!(integrator, cache::Union{QNDFCache, QNDFConstan cache.order = kₙ q = integrator.dt / hₙ - if prefer_const_step && q < 1.2 && q > 0.6 + if prefer_const_step && 0.6 < q < 1.2 h elseif q <= get_qsteady_max(integrator) && q >= get_qsteady_min(integrator) h diff --git a/lib/OrdinaryDiffEqCore/src/disco.jl b/lib/OrdinaryDiffEqCore/src/disco.jl index ace16a72d04..4c9f7812520 100644 --- a/lib/OrdinaryDiffEqCore/src/disco.jl +++ b/lib/OrdinaryDiffEqCore/src/disco.jl @@ -3,6 +3,7 @@ function set_discontinuity(integrator) dt = integrator.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 diff --git a/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl b/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl index 709892e864b..b83666374a4 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl +++ b/lib/OrdinaryDiffEqCore/src/integrators/controllers.jl @@ -492,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 """ @@ -570,29 +586,18 @@ end # TODO change signature to remove the q input function step_accept_controller!(integrator, cache::IControllerCache, alg, q) - (; qsteady_min, qsteady_max, discontinuity_detection) = cache.controller.basic + (; qsteady_min, qsteady_max) = cache.controller.basic t = integrator.t dt = integrator.dt if qsteady_min <= q <= qsteady_max q = one(q) end - if discontinuity_detection && integrator.is_disco_step - integrator.is_disco_step = false - return min((integrator.disco_checkpoint - t) / 4, dt / q) - end - return 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) - 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 @@ -714,7 +719,7 @@ end function step_accept_controller!(integrator, cache::PIControllerCache, alg, q) (; controller) = cache - (; qsteady_min, qsteady_max, discontinuity_detection) = controller.basic + (; qsteady_min, qsteady_max) = controller.basic qoldinit = controller.qoldinit EEst = DiffEqBase.value(get_EEst(integrator)) @@ -724,23 +729,13 @@ function step_accept_controller!(integrator, cache::PIControllerCache, alg, q) q = one(q) end cache.errold = max(EEst, qoldinit) - if discontinuity_detection && integrator.is_disco_step - integrator.is_disco_step = false - return min((integrator.disco_checkpoint - t) / 4, dt / q) # new dt - end - - return 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) - if disco_dt > zero(disco_dt) - return integrator.dt = disco_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 @@ -952,7 +947,7 @@ end function step_accept_controller!(integrator, cache::PIDControllerCache, alg, dt_factor) (; controller) = cache - (; qsteady_min, qsteady_max, discontinuity_detection) = controller.basic + (; qsteady_min, qsteady_max) = controller.basic t = integrator.t dt = integrator.dt @@ -963,22 +958,11 @@ function step_accept_controller!(integrator, cache::PIDControllerCache, alg, dt_ cache.err[3] = cache.err[2] cache.err[2] = cache.err[1] end - if discontinuity_detection && integrator.is_disco_step - integrator.is_disco_step = false - return min((integrator.disco_checkpoint - t) / 4, dt * dt_factor) - end - return 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) - 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 @@ -1137,25 +1121,13 @@ function step_accept_controller!(integrator, cache::PredictiveControllerCache, a cache.dtacc = DiffEqBase.value(integrator.dt) cache.erracc = max(1.0e-2, EEst) - if cache.controller.basic.discontinuity_detection && integrator.is_disco_step - integrator.is_disco_step = false - return min((integrator.disco_checkpoint - integrator.t) / 4, integrator.dt / qacc) - end - - 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) - 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 diff --git a/test/Integrators_I/disco_tests.jl b/test/Integrators_I/disco_tests.jl index 242454e05b8..0310c1df5ed 100644 --- a/test/Integrators_I/disco_tests.jl +++ b/test/Integrators_I/disco_tests.jl @@ -20,23 +20,23 @@ condition(u, t, integrator) = u[1] - 1 cb = ContinuousCallback(condition, default_affect!; maybe_discontinuity = true) sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) -# 310.958 μs (8429 allocations: 263.88 KiB) +# 314.500 μs (8909 allocations: 278.73 KiB) sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) -# 392.667 μs (10628 allocations: 331.61 KiB) +# 376.125 μs (10628 allocations: 331.61 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) -# 443.834 μs (15594 allocations: 547.55 KiB) +# 444.500 μs (15596 allocations: 547.62 KiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) -# 620.000 μs (21828 allocations: 765.69 KiB) +# 620.709 μs (21830 allocations: 765.77 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) -# 60.250 μs (7576 allocations: 240.67 KiB) +# 55.625 μs (7576 allocations: 240.67 KiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) -# 58.333 μs (7570 allocations: 240.58 KiB) +# 55.083 μs (7570 allocations: 240.58 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -65,23 +65,23 @@ cb2 = ContinuousCallback(condition2, default_affect!; maybe_discontinuity = true cb = CallbackSet(cb1, cb2) sol_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6, controller = predictive_disco_controller(RadauIIA5())) -# 1.067 ms (28126 allocations: 875.64 KiB) +# 930.334 μs (25771 allocations: 802.09 KiB) sol_no_disco_radau = solve(prob, RadauIIA5(); callback = cb, reltol = 1e-6) -# 1.311 ms (34747 allocations: 1.05 MiB) +# 1.256 ms (34747 allocations: 1.05 MiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Rodas5P())) -# 1.128 ms (38426 allocations: 1.31 MiB) +# 1.224 ms (42698 allocations: 1.46 MiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-6) -# 1.464 ms (52903 allocations: 1.80 MiB) +# 1.551 ms (52905 allocations: 1.80 MiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6, controller = PI_disco_controller(Tsit5())) -# 271.125 μs (35982 allocations: 1.11 MiB) +# 267.167 μs (34737 allocations: 1.07 MiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-6) -# 279.166 μs (39645 allocations: 1.22 MiB) +# 279.416 μs (39645 allocations: 1.22 MiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -111,23 +111,23 @@ cb_multi_2 = ContinuousCallback(cond_multi_2, default_affect!; maybe_discontinui cb_multi = CallbackSet(cb_multi_1, cb_multi_2) sol_disco_radau = solve(prob_multi, RadauIIA5(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(RadauIIA5())) -# 137.542 μs (1470 allocations: 61.42 KiB) +# 148.875 μs (1564 allocations: 64.56 KiB) sol_no_disco_radau = solve(prob_multi, RadauIIA5(); callback=cb_multi, reltol=1e-7, abstol=1e-9) -# 127.209 μs (1073 allocations: 48.44 KiB) +# 129.875 μs (1073 allocations: 48.44 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) -# 229.167 μs (1593 allocations: 64.94 KiB) +# 209.000 μs (1474 allocations: 60.94 KiB) sol_no_disco_rosenbrock = solve(prob_multi, Rodas5P(); callback=cb_multi, reltol=1e-7, abstol=1e-9) -# 239.750 μs (969 allocations: 45.44 KiB) +# 237.250 μs (969 allocations: 45.44 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) -# 96.917 μs (1469 allocations: 56.32 KiB) +# 105.250 μs (1592 allocations: 60.34 KiB) sol_no_disco_tsit5 = solve(prob_multi, Tsit5(); callback=cb_multi, reltol = 1e-7, abstol = 1e-9) -# 90.250 μs (1133 allocations: 46.84 KiB) +# 88.000 μs (1133 allocations: 46.84 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -153,16 +153,16 @@ cond_stiff(u, t, integrator) = u[1] - 0.5 cb_stiff = ContinuousCallback(cond_stiff, default_affect!; maybe_discontinuity = true) sol_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = predictive_disco_controller(RadauIIA5())) -# 202.750 μs (2250 allocations: 101.73 KiB) +# 207.500 μs (2285 allocations: 101.72 KiB) sol_no_disco_radau = solve(prob_stiff, RadauIIA5(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) -# 198.958 μs (2054 allocations: 95.58 KiB) +# 196.458 μs (2054 allocations: 95.58 KiB) @test sol_disco_radau.retcode == ReturnCode.Success @test sol_disco_radau.stats.nreject <= sol_no_disco_radau.stats.nreject sol_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11, controller = PI_disco_controller(Rodas5P())) -# 379.875 μs (2053 allocations: 85.81 KiB) +# 364.833 μs (2030 allocations: 85.34 KiB) sol_no_disco_rosenbrock = solve(prob_stiff, Rodas5P(); callback=cb_stiff, reltol=1e-9, abstol=1e-11) -# 335.709 μs (1765 allocations: 79.78 KiB) +# 332.750 μs (1765 allocations: 79.78 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject @@ -193,9 +193,9 @@ cond_dae(u, t, integrator) = u[1] - 0.5 cb_dae = ContinuousCallback(cond_dae, default_affect!; maybe_discontinuity = true) sol_disco_rosenbrock = solve(prob_dae, Rodas5P(); callback=cb_dae, reltol=1e-8, abstol=1e-10, controller = PI_disco_controller(Rodas5P())) -# 117.500 μs (810 allocations: 35.44 KiB) +# 127.083 μs (831 allocations: 35.53 KiB) sol_no_disco_rosenbrock = solve(prob_dae, Rodas5P(); callback=cb_dae, reltol=1e-8, abstol=1e-10) -# 123.333 μs (625 allocations: 31.69 KiB) +# 121.625 μs (625 allocations: 31.69 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject @@ -218,16 +218,16 @@ default_affect_vec!(integrator, idx) = nothing cb_vec = VectorContinuousCallback(condition_vec!, default_affect_vec!, 2; maybe_discontinuity = true) sol_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Rodas5P())) -# 168.250 μs (1109 allocations: 51.56 KiB) +# 191.250 μs (1212 allocations: 53.62 KiB) sol_no_disco_rosenbrock = solve(prob_vec, Rodas5P(); callback = cb_vec, reltol=1e-7, abstol=1e-9) -# 204.292 μs (937 allocations: 49.41 KiB) +# 202.000 μs (939 allocations: 49.47 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9, controller = PI_disco_controller(Tsit5())) -# 98.542 μs (1705 allocations: 69.88 KiB) +# 94.459 μs (1626 allocations: 67.73 KiB) sol_no_disco_tsit5 = solve(prob_vec, Tsit5(); callback = cb_vec, reltol=1e-7, abstol=1e-9) -# 80.667 μs (1354 allocations: 62.98 KiB) +# 93.084 μs (1356 allocations: 63.04 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject @@ -250,15 +250,15 @@ cond(u, t, integrator) = u[2] cb = ContinuousCallback(cond, default_affect!; maybe_discontinuity = true) sol_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Rodas5P())) -# 152.583 μs (1074 allocations: 48.30 KiB) +# 169.208 μs (1190 allocations: 51.78 KiB) sol_no_disco_rosenbrock = solve(prob, Rodas5P(); callback = cb, reltol = 1e-8, abstol = 1e-10) -# 181.875 μs (877 allocations: 45.72 KiB) +# 184.792 μs (877 allocations: 45.72 KiB) @test sol_disco_rosenbrock.retcode == ReturnCode.Success @test sol_disco_rosenbrock.stats.nreject <= sol_no_disco_rosenbrock.stats.nreject sol_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10, controller = PI_disco_controller(Tsit5())) -# 62.791 μs (1284 allocations: 55.26 KiB) +# 59.625 μs (1248 allocations: 53.79 KiB) sol_no_disco_tsit5 = solve(prob, Tsit5(); callback = cb, reltol = 1e-8, abstol = 1e-10) -# 54.708 μs (1155 allocations: 54.55 KiB) +# 54.041 μs (1155 allocations: 54.55 KiB) @test sol_disco_tsit5.retcode == ReturnCode.Success @test sol_disco_tsit5.stats.nreject <= sol_no_disco_tsit5.stats.nreject From c2af3bc8eb794d1857f548ae5cb1f86eaa3d1965 Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Tue, 23 Jun 2026 10:14:05 -0400 Subject: [PATCH 8/9] cache reused vals --- lib/OrdinaryDiffEqCore/src/disco.jl | 12 +++++++----- lib/OrdinaryDiffEqCore/src/solve.jl | 10 ++++++---- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/lib/OrdinaryDiffEqCore/src/disco.jl b/lib/OrdinaryDiffEqCore/src/disco.jl index 4c9f7812520..593ee83be1d 100644 --- a/lib/OrdinaryDiffEqCore/src/disco.jl +++ b/lib/OrdinaryDiffEqCore/src/disco.jl @@ -50,7 +50,8 @@ function find_discontinuity(integrator) disco_zero.dt, disco_zero.f, disco_zero.p, disco_zero.cache, false, true, false) end disco_zero.ind = j - sol = solve(disco_prob, tspan = (0.0, breakpointθ)) + disco_prob.tspan[2] = breakpointθ + sol = solve(disco_prob) tmp = sol[] if (!isnan(tmp) && tmp < breakpointθ) breakpointθ = tmp @@ -60,15 +61,16 @@ function find_discontinuity(integrator) 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 - sol = solve(disco_prob, tspan = (0.0, breakpointθ)) + disco_prob.tspan[2] = breakpointθ + sol = solve(disco_prob) tmp = sol[] if (!isnan(tmp) && tmp < breakpointθ) breakpointθ = tmp diff --git a/lib/OrdinaryDiffEqCore/src/solve.jl b/lib/OrdinaryDiffEqCore/src/solve.jl index 18511cf014a..8d1de9fb0f1 100644 --- a/lib/OrdinaryDiffEqCore/src/solve.jl +++ b/lib/OrdinaryDiffEqCore/src/solve.jl @@ -16,7 +16,7 @@ determine_controller_datatype(u, internalnorm, ts::Tuple{<:Number, <:Number}) = determine_controller_datatype(u::AbstractVector{<:Number}, internalnorm, ts::Tuple{<:Integer, <:Integer}) = promote_type(typeof(DiffEqBase.value(internalnorm(u, ts[1]))), typeof(DiffEqBase.value(internalnorm(u, ts[2]))), eltype(float.(DiffEqBase.value(ts)))) determine_controller_datatype(u, internalnorm, ts::Tuple{<:Integer, <:Integer}) = promote_type(typeof(float(DiffEqBase.value(ts[1]))), typeof(float(DiffEqBase.value(ts[2])))) # This seems to be an assumption implicitly taken somewhere -mutable struct zero_func_struct{u1Type, uType, tType, kType, CacheType, idxsType, varsType, callbackType, outType, FunctionType, tType2, ParameterType} +mutable struct zero_func_struct{u1Type, uType, tType, kType, CacheType, idxsType, varsType, callbackType, outType, outCacheType, FunctionType, tType2, ParameterType} u₁::u1Type callback::callbackType dt::tType @@ -28,8 +28,8 @@ mutable struct zero_func_struct{u1Type, uType, tType, kType, CacheType, idxsType differential_vars::varsType ind::Int out::outType - out_low::outType - out_high::outType + out_low::outCacheType + out_high::outCacheType f::FunctionType tprev::tType2 p::ParameterType @@ -38,6 +38,8 @@ end parameter_values(z::zero_func_struct) = z.p function (z::zero_func_struct)(θ, p) + iszero(θ) && return z.out_low[z.ind] + isone(θ) && return z.out_high[z.ind] 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 @@ -672,7 +674,7 @@ Base.@constprop :aggressive function _ode_init( arr = (u isa AbstractArray) ? similar(u, i.len) : zeros(typeof(u), i.len) arr, similar(arr), similar(arr) else - nothing, nothing, nothing + nothing, zeros(Float64, 1), zeros(Float64, 1) end zero_func = zero_func_struct(u₁, i, _dt, uprev, u, k, cache, save_idxs, differential_vars, 1, out, out_low, out_high, f, tprev, p) zero_func_wrapped = FunctionWrapper{Float64, Tuple{Float64, typeof(p)}}(zero_func) From cba749578fae01a8fda5b0bbb59c8be74fe92254 Mon Sep 17 00:00:00 2001 From: Shreyas-Ekanathan Date: Mon, 6 Jul 2026 10:15:26 -0400 Subject: [PATCH 9/9] Update OrdinaryDiffEqCore.jl --- lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl index 3c48b486a72..8f818665c71 100644 --- a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl +++ b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl @@ -262,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,