Merge pull request #2873 from ChrisRackauckas-Claude/diffeqbase-initialization-algorithms

Update to use DiffEqBase 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.2"
 DocStringExtensions = "0.9.5"
 EnumX = "1.0.5"
 ExplicitImports = "1.13.1"

lib/OrdinaryDiffEqCore/Project.toml

@@ -85,7 +85,7 @@ FastPower = "1.1"
 Logging = "1.10"
 Mooncake = "0.4"
 AllocCheck = "0.2"
-DiffEqBase = "6.187"
+DiffEqBase = "6.190.2"
 FillArrays = "1.13"
 Adapt = "4.3"
 Reexport = "1.2"

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -27,6 +27,9 @@ import FastPower: fastpower
 import SciMLBase: solve!, step!, isadaptive
 import DiffEqBase: initialize!
 
+# DAE Initialization algorithms
+import DiffEqBase: DefaultInit, ShampineCollocationInit, BrownFullBasicInit
+
 # Internal utils
 import DiffEqBase: ODE_DEFAULT_NORM,
                    ODE_DEFAULT_ISOUTOFDOMAIN, ODE_DEFAULT_PROG_MESSAGE,

lib/OrdinaryDiffEqCore/src/initialize_dae.jl

@@ -1,25 +1,3 @@
-struct DefaultInit <: SciMLBase.DAEInitializationAlgorithm end
-
-struct ShampineCollocationInit{T, F} <: SciMLBase.DAEInitializationAlgorithm
-    initdt::T
-    nlsolve::F
-end
-function ShampineCollocationInit(; initdt = nothing, nlsolve = nothing)
-    ShampineCollocationInit(initdt, nlsolve)
-end
-function ShampineCollocationInit(initdt)
-    ShampineCollocationInit(; initdt = initdt, nlsolve = nothing)
-end
-
-struct BrownFullBasicInit{T, F} <: SciMLBase.DAEInitializationAlgorithm
-    abstol::T
-    nlsolve::F
-end
-function BrownFullBasicInit(; abstol = 1e-10, nlsolve = nothing)
-    BrownFullBasicInit(abstol, nlsolve)
-end
-BrownFullBasicInit(abstol) = BrownFullBasicInit(; abstol = abstol, nlsolve = nothing)
-
 ## Notes
 
 #=
@@ -44,7 +22,7 @@ end
 
 ## Default algorithms
 
-function _initialize_dae!(integrator, prob::ODEProblem,
+function _initialize_dae!(integrator::ODEIntegrator, prob::ODEProblem,
         alg::DefaultInit, x::Union{Val{true}, Val{false}})
     if SciMLBase.has_initializeprob(prob.f)
         _initialize_dae!(integrator, prob,
@@ -58,7 +36,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     end
 end
 
-function _initialize_dae!(integrator, prob::DAEProblem,
+function _initialize_dae!(integrator::ODEIntegrator, prob::DAEProblem,
         alg::DefaultInit, x::Union{Val{true}, Val{false}})
     if SciMLBase.has_initializeprob(prob.f)
         _initialize_dae!(integrator, prob,
@@ -77,7 +55,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     end
 end
 
-function _initialize_dae!(integrator, prob::DiscreteProblem,
+function _initialize_dae!(integrator::ODEIntegrator, prob::DiscreteProblem,
         alg::DefaultInit, x::Union{Val{true}, Val{false}})
     if SciMLBase.has_initializeprob(prob.f)
         # integrator.opts.abstol is `false` for `DiscreteProblem`.
@@ -124,13 +102,13 @@ end
 
 ## NoInit
 
-function _initialize_dae!(integrator, prob::AbstractDEProblem,
+function _initialize_dae!(integrator::ODEIntegrator, prob::AbstractDEProblem,
         alg::NoInit, x::Union{Val{true}, Val{false}})
 end
 
 ## OverrideInit
 
-function _initialize_dae!(integrator, prob::AbstractDEProblem,
+function _initialize_dae!(integrator::ODEIntegrator, prob::AbstractDEProblem,
         alg::OverrideInit, isinplace::Union{Val{true}, Val{false}})
     initializeprob = prob.f.initialization_data.initializeprob
 
@@ -172,7 +150,7 @@ function _initialize_dae!(integrator, prob::AbstractDEProblem,
 end
 
 ## CheckInit
-function _initialize_dae!(integrator, prob::AbstractDEProblem, alg::CheckInit,
+function _initialize_dae!(integrator::ODEIntegrator, prob::AbstractDEProblem, alg::CheckInit,
         isinplace::Union{Val{true}, Val{false}})
     SciMLBase.get_initial_values(
         prob, integrator, prob.f, alg, isinplace; abstol = integrator.opts.abstol)

lib/OrdinaryDiffEqNonlinearSolve/Project.toml

@@ -63,7 +63,7 @@ RecursiveArrayTools = "3.36"
 ODEProblemLibrary = "0.1.8"
 PreallocationTools = "0.4"
 AllocCheck = "0.2"
-DiffEqBase = "6.176"
+DiffEqBase = "6.190.2"
 SafeTestsets = "0.1.0"
 SciMLOperators = "1.4"
 SciMLStructures = "1.7"

lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl

@@ -53,19 +53,20 @@ Solve for `u`
 
 =#
 
-function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocationInit,
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::ODEProblem, alg::DiffEqBase.ShampineCollocationInit,
         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 = alg.initdt
+    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)]
@@ -168,18 +169,19 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
     return
 end
 
-function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocationInit,
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::ODEProblem, alg::DiffEqBase.ShampineCollocationInit,
         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 = alg.initdt
+    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)]
@@ -246,7 +248,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
     return
 end
 
-function _initialize_dae!(integrator, prob::DAEProblem,
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::DAEProblem,
         alg::ShampineCollocationInit, isinplace::Val{true})
     @unpack p, t, f = integrator
     u0 = integrator.u
@@ -323,7 +325,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     return
 end
 
-function _initialize_dae!(integrator, prob::DAEProblem,
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::DAEProblem,
         alg::ShampineCollocationInit, isinplace::Val{false})
     @unpack p, t, f = integrator
     u0 = integrator.u
@@ -387,8 +389,8 @@ function algebraic_jacobian(jac_prototype::T, algebraic_eqs,
     jac_prototype[algebraic_eqs, algebraic_vars]
 end
 
-function _initialize_dae!(integrator, prob::ODEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{true})
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::ODEProblem,
+        alg::DiffEqBase.BrownFullBasicInit, isinplace::Val{true})
     @unpack p, t, f = integrator
     u = integrator.u
     M = integrator.f.mass_matrix
@@ -468,8 +470,8 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     return
 end
 
-function _initialize_dae!(integrator, prob::ODEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{false})
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::ODEProblem,
+        alg::DiffEqBase.BrownFullBasicInit, isinplace::Val{false})
     @unpack p, t, f = integrator
 
     u0 = integrator.u
@@ -536,8 +538,8 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     return
 end
 
-function _initialize_dae!(integrator, prob::DAEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{true})
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::DAEProblem,
+        alg::DiffEqBase.BrownFullBasicInit, isinplace::Val{true})
     @unpack p, t, f = integrator
     differential_vars = prob.differential_vars
     u = integrator.u
@@ -592,8 +594,9 @@ function _initialize_dae!(integrator, prob::DAEProblem,
         f(out, du_tmp, uu, p, t)
     end
 
-    if alg.nlsolve !== nothing
-        nlsolve = alg.nlsolve
+    nlsolve_alg = alg.nlsolve
+    if nlsolve_alg !== nothing
+        nlsolve = nlsolve_alg
     else
         nlsolve = NewtonRaphson(autodiff = alg_autodiff(integrator.alg))
     end
@@ -617,8 +620,8 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     return
 end
 
-function _initialize_dae!(integrator, prob::DAEProblem,
-        alg::BrownFullBasicInit, isinplace::Val{false})
+function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::DAEProblem,
+        alg::DiffEqBase.BrownFullBasicInit, isinplace::Val{false})
     @unpack p, t, f = integrator
     differential_vars = prob.differential_vars