Update to use DiffEqBase initialization algorithms

- Import DefaultInit, ShampineCollocationInit, BrownBasicInit from DiffEqBase
- Add extended versions for OrdinaryDiffEq-specific options
- Update OrdinaryDiffEqNonlinearSolve to accept both base and extended types
- Add helper functions to handle algorithm parameter differences
- Bump DiffEqBase to 6.190 for new initialization algorithms

Author: ChrisRackauckas

Project.toml

@@ -112,7 +112,7 @@ Adapt = "4.3"
 ArrayInterface = "7.19"
 CommonSolve = "0.2.4"
 DataStructures = "0.18.22, 0.19"
-DiffEqBase = "6.186"
+DiffEqBase = "6.190"
 DocStringExtensions = "0.9.5"
 EnumX = "1.0.5"
 ExplicitImports = "1.13.1"

lib/OrdinaryDiffEqCore/src/initialize_dae.jl

@@ -1,16 +1,31 @@
-struct DefaultInit <: SciMLBase.DAEInitializationAlgorithm end
+# Import standard initialization algorithms from DiffEqBase
+import DiffEqBase: DefaultInit, ShampineCollocationInit, BrownBasicInit
 
-struct ShampineCollocationInit{T, F} <: SciMLBase.DAEInitializationAlgorithm
+# Re-export for backward compatibility
+export DefaultInit, ShampineCollocationInit, BrownBasicInit
+
+# Extended versions with OrdinaryDiffEq-specific options
+struct ShampineCollocationInitExt{T, F} <: SciMLBase.DAEInitializationAlgorithm
     initdt::T
     nlsolve::F
 end
-function ShampineCollocationInit(; initdt = nothing, nlsolve = nothing)
-    ShampineCollocationInit(initdt, nlsolve)
+function ShampineCollocationInitExt(; initdt = nothing, nlsolve = nothing)
+    ShampineCollocationInitExt(initdt, nlsolve)
 end
-function ShampineCollocationInit(initdt)
-    ShampineCollocationInit(; initdt = initdt, nlsolve = nothing)
+function ShampineCollocationInitExt(initdt)
+    ShampineCollocationInitExt(; initdt = initdt, nlsolve = nothing)
 end
 
+# Constructor that delegates to extended version for backward compatibility
+function ShampineCollocationInit(; initdt = nothing, nlsolve = nothing)
+    if initdt !== nothing || nlsolve !== nothing
+        ShampineCollocationInitExt(initdt, nlsolve)
+    else
+        DiffEqBase.ShampineCollocationInit()
+    end
+end
+ShampineCollocationInit(initdt) = ShampineCollocationInitExt(; initdt = initdt, nlsolve = nothing)
+
 struct BrownFullBasicInit{T, F} <: SciMLBase.DAEInitializationAlgorithm
     abstol::T
     nlsolve::F
@@ -20,6 +35,18 @@ function BrownFullBasicInit(; abstol = 1e-10, nlsolve = nothing)
 end
 BrownFullBasicInit(abstol) = BrownFullBasicInit(; abstol = abstol, nlsolve = nothing)
 
+# Alias for consistency with DiffEqBase naming
+const BrownBasicInitExt = BrownFullBasicInit
+
+# Constructor that delegates for backward compatibility
+function BrownBasicInit(; abstol = nothing, nlsolve = nothing)
+    if abstol !== nothing || nlsolve !== nothing
+        BrownFullBasicInit(something(abstol, 1e-10), nlsolve)
+    else
+        DiffEqBase.BrownBasicInit()
+    end
+end
+
 ## Notes
 
 #=
@@ -177,3 +204,20 @@ function _initialize_dae!(integrator, prob::AbstractDEProblem, alg::CheckInit,
     SciMLBase.get_initial_values(
         prob, integrator, prob.f, alg, isinplace; abstol = integrator.opts.abstol)
 end
+
+# Delegate base DiffEqBase types to extended versions with default options
+function _initialize_dae!(integrator, prob::AbstractDEProblem,
+        alg::DiffEqBase.ShampineCollocationInit, isinplace::Union{Val{true}, Val{false}})
+    _initialize_dae!(integrator, prob, ShampineCollocationInitExt(nothing, nothing), isinplace)
+end
+
+function _initialize_dae!(integrator, prob::AbstractDEProblem,
+        alg::DiffEqBase.BrownBasicInit, isinplace::Union{Val{true}, Val{false}})
+    _initialize_dae!(integrator, prob, BrownFullBasicInit(integrator.opts.abstol, nothing), isinplace)
+end
+
+# Handle DiffEqBase.DefaultInit same as our DefaultInit
+function _initialize_dae!(integrator, prob::AbstractDEProblem,
+        alg::DiffEqBase.DefaultInit, isinplace::Union{Val{true}, Val{false}})
+    _initialize_dae!(integrator, prob, DefaultInit(), isinplace)
+end

lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl

@@ -16,6 +16,33 @@ end
     end
 end
 
+# Helper to get abstol from either the algorithm or integrator opts
+@inline function get_abstol(alg::Union{BrownFullBasicInit, OrdinaryDiffEqCore.BrownFullBasicInit}, integrator)
+    return alg.abstol
+end
+
+@inline function get_abstol(alg::DiffEqBase.BrownBasicInit, integrator)
+    return integrator.opts.abstol
+end
+
+# Helper to get nlsolve from either the algorithm or nothing
+@inline function get_nlsolve(alg::Union{BrownFullBasicInit, OrdinaryDiffEqCore.BrownFullBasicInit, OrdinaryDiffEqCore.ShampineCollocationInitExt})
+    return alg.nlsolve
+end
+
+@inline function get_nlsolve(alg::Union{DiffEqBase.BrownBasicInit, DiffEqBase.ShampineCollocationInit})
+    return nothing
+end
+
+# Helper to get initdt from either the algorithm or nothing
+@inline function get_initdt(alg::OrdinaryDiffEqCore.ShampineCollocationInitExt)
+    return alg.initdt
+end
+
+@inline function get_initdt(alg::Union{DiffEqBase.ShampineCollocationInit, OrdinaryDiffEqCore.ShampineCollocationInit})
+    return nothing
+end
+
 function default_nlsolve(
         ::Nothing, isinplace::Val{true}, u, ::AbstractNonlinearProblem, autodiff = false)
     FastShortcutNonlinearPolyalg(;
@@ -53,19 +80,20 @@ Solve for `u`
 
 =#
 
-function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocationInit,
+function _initialize_dae!(integrator, prob::ODEProblem, alg::Union{ShampineCollocationInit, OrdinaryDiffEqCore.ShampineCollocationInitExt},
         isinplace::Val{true})
     @unpack p, t, f = integrator
     M = integrator.f.mass_matrix
     dtmax = integrator.opts.dtmax
     tmp = first(get_tmp_cache(integrator))
     u0 = integrator.u
 
-    dt = if alg.initdt === nothing
+    initdt = get_initdt(alg)
+    dt = if initdt === nothing
         integrator.dt != 0 ? min(integrator.dt / 5, dtmax) :
         (prob.tspan[end] - prob.tspan[begin]) / 1000 # Haven't implemented norm reduction
     else
-        alg.initdt
+        initdt
     end
 
     algebraic_vars = [all(iszero, x) for x in eachcol(M)]
@@ -77,7 +105,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
 
     check_dae_tolerance(integrator, tmp, integrator.opts.abstol, t, isinplace) && return
 
-    if isdefined(integrator.cache, :nlsolver) && !isnothing(alg.nlsolve)
+    if isdefined(integrator.cache, :nlsolver) && !isnothing(get_nlsolve(alg))
         # backward Euler
         nlsolver = integrator.cache.nlsolver
         oldγ, oldc, oldmethod,
@@ -149,7 +177,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
             jac_prototype = f.jac_prototype,
             jac = jac)
         nlprob = NonlinearProblem(nlfunc, integrator.u, p)
-        nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0, nlprob, isAD)
+        nlsolve = default_nlsolve(get_nlsolve(alg), isinplace, u0, nlprob, isAD)
         nlsol = solve(nlprob, nlsolve; abstol = integrator.opts.abstol,
             reltol = integrator.opts.reltol)
         integrator.u .= nlsol.u
@@ -168,18 +196,19 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
     return
 end
 
-function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocationInit,
+function _initialize_dae!(integrator, prob::ODEProblem, alg::Union{ShampineCollocationInit, OrdinaryDiffEqCore.ShampineCollocationInitExt},
         isinplace::Val{false})
     @unpack p, t, f = integrator
     u0 = integrator.u
     M = integrator.f.mass_matrix
     dtmax = integrator.opts.dtmax
 
-    dt = if alg.initdt === nothing
+    initdt = get_initdt(alg)
+    dt = if initdt === nothing
         integrator.dt != 0 ? min(integrator.dt / 5, dtmax) :
         (prob.tspan[end] - prob.tspan[begin]) / 1000 # Haven't implemented norm reduction
     else
-        alg.initdt
+        initdt
     end
 
     algebraic_vars = [all(iszero, x) for x in eachcol(M)]
@@ -191,7 +220,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
 
     check_dae_tolerance(integrator, resid, integrator.opts.abstol, t, isinplace) && return
 
-    if isdefined(integrator.cache, :nlsolver) && !isnothing(alg.nlsolve)
+    if isdefined(integrator.cache, :nlsolver) && !isnothing(get_nlsolve(alg))
         # backward Euler
         nlsolver = integrator.cache.nlsolver
         oldγ, oldc, oldmethod,
@@ -225,7 +254,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
             jac_prototype = f.jac_prototype,
             jac = jac)
         nlprob = NonlinearProblem(nlfunc, u0)
-        nlsolve = default_nlsolve(alg.nlsolve, isinplace, nlprob, u0)
+        nlsolve = default_nlsolve(get_nlsolve(alg), isinplace, nlprob, u0)
 
         nlsol = solve(nlprob, nlsolve; abstol = integrator.opts.abstol,
             reltol = integrator.opts.reltol)
@@ -306,7 +335,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
         jac_prototype = f.jac_prototype,
         jac = jac)
     nlprob = NonlinearProblem(nlfunc, u0, p)
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0, nlprob, isAD)
+    nlsolve = default_nlsolve(get_nlsolve(alg), isinplace, u0, nlprob, isAD)
     nlsol = solve(nlprob, nlsolve; abstol = integrator.opts.abstol,
         reltol = integrator.opts.reltol)
 
@@ -350,7 +379,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     nlfunc = NonlinearFunction(nlequation; jac_prototype = f.jac_prototype,
         jac = jac)
     nlprob = NonlinearProblem(nlfunc, u0)
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, nlprob, u0)
+    nlsolve = default_nlsolve(get_nlsolve(alg), isinplace, nlprob, u0)
 
     nlfunc = NonlinearFunction(nlequation; jac_prototype = f.jac_prototype)
     nlprob = NonlinearProblem(nlfunc, u0)
@@ -388,7 +417,7 @@ function algebraic_jacobian(jac_prototype::T, algebraic_eqs,
 end
 
 function _initialize_dae!(integrator, prob::ODEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{true})
+        alg::Union{BrownFullBasicInit, DiffEqBase.BrownBasicInit}, isinplace::Val{true})
     @unpack p, t, f = integrator
     u = integrator.u
     M = integrator.f.mass_matrix
@@ -403,7 +432,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
 
     tmp .= ArrayInterface.restructure(tmp, algebraic_eqs .* _vec(tmp))
 
-    check_dae_tolerance(integrator, tmp, alg.abstol, t, isinplace) && return
+    check_dae_tolerance(integrator, tmp, get_abstol(alg, integrator), t, isinplace) && return
     alg_u = @view u[algebraic_vars]
 
     # These non-dual values are thus used to make the caches
@@ -451,9 +480,9 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     J = algebraic_jacobian(f.jac_prototype, algebraic_eqs, algebraic_vars)
     nlfunc = NonlinearFunction(nlequation!; jac_prototype = J)
     nlprob = NonlinearProblem(nlfunc, alg_u, p)
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u, nlprob, isAD)
+    nlsolve = default_nlsolve(get_nlsolve(alg), isinplace, u, nlprob, isAD)
 
-    nlsol = solve(nlprob, nlsolve; abstol = alg.abstol, reltol = integrator.opts.reltol)
+    nlsol = solve(nlprob, nlsolve; abstol = get_abstol(alg, integrator), reltol = integrator.opts.reltol)
     alg_u .= nlsol
 
     recursivecopy!(integrator.uprev, integrator.u)
@@ -469,7 +498,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
 end
 
 function _initialize_dae!(integrator, prob::ODEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{false})
+        alg::Union{BrownFullBasicInit, DiffEqBase.BrownBasicInit}, isinplace::Val{false})
     @unpack p, t, f = integrator
 
     u0 = integrator.u
@@ -482,7 +511,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     du = f(u0, p, t)
     resid = _vec(du)[algebraic_eqs]
 
-    check_dae_tolerance(integrator, resid, alg.abstol, t, isinplace) && return
+    check_dae_tolerance(integrator, resid, get_abstol(alg, integrator), t, isinplace) && return
 
     isAD = alg_autodiff(integrator.alg) isa AutoForwardDiff
     if isAD
@@ -511,7 +540,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     J = algebraic_jacobian(f.jac_prototype, algebraic_eqs, algebraic_vars)
     nlfunc = NonlinearFunction(nlequation; jac_prototype = J)
     nlprob = NonlinearProblem(nlfunc, u0[algebraic_vars])
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0, nlprob, isAD)
+    nlsolve = default_nlsolve(get_nlsolve(alg), isinplace, u0, nlprob, isAD)
 
     nlsol = solve(nlprob, nlsolve)
 
@@ -537,7 +566,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
 end
 
 function _initialize_dae!(integrator, prob::DAEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{true})
+        alg::Union{BrownFullBasicInit, DiffEqBase.BrownBasicInit}, isinplace::Val{true})
     @unpack p, t, f = integrator
     differential_vars = prob.differential_vars
     u = integrator.u
@@ -561,7 +590,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     normtmp = get_tmp_cache(integrator)[1]
     f(normtmp, du, u, p, t)
 
-    if check_dae_tolerance(integrator, normtmp, alg.abstol, t, isinplace)
+    if check_dae_tolerance(integrator, normtmp, get_abstol(alg, integrator), t, isinplace)
         return
     elseif differential_vars === nothing
         error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.")
@@ -592,15 +621,16 @@ function _initialize_dae!(integrator, prob::DAEProblem,
         f(out, du_tmp, uu, p, t)
     end
 
-    if alg.nlsolve !== nothing
-        nlsolve = alg.nlsolve
+    nlsolve_alg = get_nlsolve(alg)
+    if nlsolve_alg !== nothing
+        nlsolve = nlsolve_alg
     else
         nlsolve = NewtonRaphson(autodiff = alg_autodiff(integrator.alg))
     end
 
     nlfunc = NonlinearFunction(nlequation!; jac_prototype = f.jac_prototype)
     nlprob = NonlinearProblem(nlfunc, ifelse.(differential_vars, du, u), p)
-    nlsol = solve(nlprob, nlsolve; abstol = alg.abstol, reltol = integrator.opts.reltol)
+    nlsol = solve(nlprob, nlsolve; abstol = get_abstol(alg, integrator), reltol = integrator.opts.reltol)
 
     @. du = ifelse(differential_vars, nlsol.u, du)
     @. u = ifelse(differential_vars, u, nlsol.u)
@@ -618,12 +648,12 @@ function _initialize_dae!(integrator, prob::DAEProblem,
 end
 
 function _initialize_dae!(integrator, prob::DAEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{false})
+        alg::Union{BrownFullBasicInit, DiffEqBase.BrownBasicInit}, isinplace::Val{false})
     @unpack p, t, f = integrator
     differential_vars = prob.differential_vars
 
     if check_dae_tolerance(
-        integrator, f(integrator.du, integrator.u, p, t), alg.abstol, t, isinplace)
+        integrator, f(integrator.du, integrator.u, p, t), get_abstol(alg, integrator), t, isinplace)
         return
     elseif differential_vars === nothing
         error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.")
@@ -647,7 +677,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     nlfunc = NonlinearFunction(nlequation; jac_prototype = f.jac_prototype)
     nlprob = NonlinearProblem(nlfunc, ifelse.(differential_vars, du, u))
 
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, nlprob, integrator.u)
+    nlsolve = default_nlsolve(get_nlsolve(alg), isinplace, nlprob, integrator.u)
 
     @show nlsolve