Merge pull request #2151 from SciML/overrideinit

Allow for using the initializeprob with OverrideInit

Author: ChrisRackauckas

Project.toml

@@ -65,14 +65,14 @@ Logging = "1.9"
 MacroTools = "0.5"
 MuladdMacro = "0.2.1"
 NLsolve = "4"
-NonlinearSolve = "3.3"
+NonlinearSolve = "3.7.3"
 Polyester = "0.7"
 PreallocationTools = "0.4.15"
 PrecompileTools = "1"
 Preferences = "1.3"
 RecursiveArrayTools = "2.36, 3"
 Reexport = "1.0"
-SciMLBase = "2.26"
+SciMLBase = "2.27.1"
 SciMLOperators = "0.3"
 SimpleNonlinearSolve = "1"
 SimpleUnPack = "1"

src/alg_utils.jl

@@ -291,6 +291,7 @@ end
 _alg_autodiff(::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
 _alg_autodiff(::DAEAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
 _alg_autodiff(::OrdinaryDiffEqImplicitAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
+_alg_autodiff(alg::CompositeAlgorithm) = _alg_autodiff(alg.algs[2])
 function _alg_autodiff(::Union{OrdinaryDiffEqExponentialAlgorithm{CS, AD},
         OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS, AD}
 }) where {

src/initialize_dae.jl

@@ -21,13 +21,34 @@ end
 BrownFullBasicInit(abstol) = BrownFullBasicInit(; abstol = abstol, nlsolve = nothing)
 
 default_nlsolve(alg, isinplace, u, autodiff = false) = alg
-function default_nlsolve(::Nothing, isinplace, u, autodiff = false)
-    TrustRegion(; autodiff = autodiff ? AutoForwardDiff() : AutoFiniteDiff())
+function default_nlsolve(::Nothing, isinplace, u, ::NonlinearProblem, autodiff = false)
+    FastShortcutNonlinearPolyalg(;
+        autodiff = autodiff ? AutoForwardDiff() : AutoFiniteDiff())
 end
-function default_nlsolve(::Nothing, isinplace::Val{false}, u::StaticArray, autodiff = false)
+function default_nlsolve(::Nothing, isinplace::Val{false}, u::StaticArray,
+        ::NonlinearProblem, autodiff = false)
     SimpleTrustRegion(autodiff = autodiff ? AutoForwardDiff() : AutoFiniteDiff())
 end
 
+function default_nlsolve(
+        ::Nothing, isinplace, u, ::NonlinearLeastSquaresProblem, autodiff = false)
+    FastShortcutNLLSPolyalg(; autodiff = autodiff ? AutoForwardDiff() : AutoFiniteDiff())
+end
+function default_nlsolve(::Nothing, isinplace::Val{false}, u::StaticArray,
+        ::NonlinearLeastSquaresProblem, autodiff = false)
+    SimpleGaussNewton(autodiff = autodiff ? AutoForwardDiff() : AutoFiniteDiff())
+end
+
+struct OverrideInit{T, F} <: DiffEqBase.DAEInitializationAlgorithm
+    abstol::T
+    nlsolve::F
+end
+
+function OverrideInit(; abstol = 1e-10, nlsolve = nothing)
+    OverrideInit(abstol, nlsolve)
+end
+OverrideInit(abstol) = OverrideInit(; abstol = abstol, nlsolve = nothing)
+
 ## Notes
 
 #=
@@ -54,19 +75,32 @@ end
 
 function _initialize_dae!(integrator, prob::ODEProblem,
         alg::DefaultInit, x::Val{true})
-    _initialize_dae!(integrator, prob,
-        BrownFullBasicInit(integrator.opts.abstol), x)
+    if SciMLBase.has_initializeprob(prob.f)
+        _initialize_dae!(integrator, prob,
+            OverrideInit(integrator.opts.abstol), x)
+    else
+        _initialize_dae!(integrator, prob,
+            BrownFullBasicInit(integrator.opts.abstol), x)
+    end
 end
 
 function _initialize_dae!(integrator, prob::ODEProblem,
         alg::DefaultInit, x::Val{false})
-    _initialize_dae!(integrator, prob,
-        BrownFullBasicInit(integrator.opts.abstol), x)
+    if SciMLBase.has_initializeprob(prob.f)
+        _initialize_dae!(integrator, prob,
+            OverrideInit(integrator.opts.abstol), x)
+    else
+        _initialize_dae!(integrator, prob,
+            BrownFullBasicInit(integrator.opts.abstol), x)
+    end
 end
 
 function _initialize_dae!(integrator, prob::DAEProblem,
         alg::DefaultInit, x::Val{false})
-    if prob.differential_vars === nothing
+    if SciMLBase.has_initializeprob(prob.f)
+        _initialize_dae!(integrator, prob,
+            OverrideInit(integrator.opts.abstol), x)
+    elseif prob.differential_vars === nothing
         _initialize_dae!(integrator, prob,
             ShampineCollocationInit(), x)
     else
@@ -77,7 +111,10 @@ end
 
 function _initialize_dae!(integrator, prob::DAEProblem,
         alg::DefaultInit, x::Val{true})
-    if prob.differential_vars === nothing
+    if SciMLBase.has_initializeprob(prob.f)
+        _initialize_dae!(integrator, prob,
+            OverrideInit(integrator.opts.abstol), x)
+    elseif prob.differential_vars === nothing
         _initialize_dae!(integrator, prob,
             ShampineCollocationInit(), x)
     else
@@ -92,6 +129,28 @@ function _initialize_dae!(integrator, prob::Union{ODEProblem, DAEProblem},
         alg::NoInit, x::Union{Val{true}, Val{false}})
 end
 
+## OverrideInit
+
+function _initialize_dae!(integrator, prob::Union{ODEProblem, DAEProblem},
+        alg::OverrideInit, isinplace::Union{Val{true}, Val{false}})
+    initializeprob = prob.f.initializeprob
+    isAD = alg_autodiff(integrator.alg) isa AutoForwardDiff
+    alg = default_nlsolve(alg.nlsolve, isinplace, initializeprob.u0, initializeprob, isAD)
+    nlsol = solve(initializeprob, alg)
+    if isinplace === Val{true}()
+        integrator.u .= prob.f.initializeprobmap(nlsol)
+    elseif isinplace === Val{false}()
+        integrator.u = prob.f.initializeprobmap(nlsol)
+    else
+        error("Unreachable reached. Report this error.")
+    end
+
+    if nlsol.retcode != ReturnCode.Success
+        integrator.sol = SciMLBase.solution_new_retcode(integrator.sol,
+            ReturnCode.InitialFailure)
+    end
+end
+
 ## ShampineCollocationInit
 
 #=
@@ -196,7 +255,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, isAD)
+        nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0, nlprob, isAD)
         nlsol = solve(nlprob, nlsolve; abstol = integrator.opts.abstol,
             reltol = integrator.opts.reltol)
         integrator.u .= nlsol.u
@@ -266,12 +325,12 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
             end
         end
 
-        nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0)
-
         nlfunc = NonlinearFunction(nlequation_oop;
             jac_prototype = f.jac_prototype,
             jac = jac)
         nlprob = NonlinearProblem(nlfunc, u0)
+        nlsolve = default_nlsolve(alg.nlsolve, isinplace, nlprob, u0)
+
         nlsol = solve(nlprob, nlsolve; abstol = integrator.opts.abstol,
             reltol = integrator.opts.reltol)
         integrator.u = nlsol.u
@@ -351,7 +410,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, isAD)
+    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0, nlprob, isAD)
     nlsol = solve(nlprob, nlsolve; abstol = integrator.opts.abstol,
         reltol = integrator.opts.reltol)
 
@@ -395,7 +454,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, u0)
+    nlsolve = default_nlsolve(alg.nlsolve, isinplace, nlprob, u0)
 
     nlfunc = NonlinearFunction(nlequation; jac_prototype = f.jac_prototype)
     nlprob = NonlinearProblem(nlfunc, u0)
@@ -493,11 +552,10 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     end
 
     J = algebraic_jacobian(f.jac_prototype, algebraic_eqs, algebraic_vars)
-
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u, isAD)
-
     nlfunc = NonlinearFunction(nlequation!; jac_prototype = J)
     nlprob = NonlinearProblem(nlfunc, alg_u, p)
+    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u, nlprob, isAD)
+
     nlsol = solve(nlprob, nlsolve; abstol = alg.abstol, reltol = integrator.opts.reltol)
     alg_u .= nlsol
 
@@ -554,11 +612,10 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     end
 
     J = algebraic_jacobian(f.jac_prototype, algebraic_eqs, algebraic_vars)
-
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0, isAD)
-
     nlfunc = NonlinearFunction(nlequation; jac_prototype = J)
     nlprob = NonlinearProblem(nlfunc, u0[algebraic_vars])
+    nlsolve = default_nlsolve(alg.nlsolve, isinplace, u0, nlprob, isAD)
+
     nlsol = solve(nlprob, nlsolve)
 
     u[algebraic_vars] .= nlsol.u
@@ -689,11 +746,11 @@ function _initialize_dae!(integrator, prob::DAEProblem,
         f(du, u, p, t)
     end
 
-    nlsolve = default_nlsolve(alg.nlsolve, isinplace, integrator.u)
-
     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)
+
     nlsol = solve(nlprob, nlsolve)
 
     du = ifelse.(differential_vars, nlsol.u, du)

test/interface/dae_initialize_integration.jl

@@ -13,6 +13,7 @@ n2 = ODESystem(single_neuron_eqs, t, [v, w, F], [g, e, b], name = :n2)
 connections = [0 ~ n1.F - D * Dk * max(n1.v - n2.v, 0)
                0 ~ n2.F - D * max(n2.v - n1.v, 0)]
 connected = ODESystem(connections, t, [], [D, Dk], systems = [n1, n2], name = :connected)
+connected = complete(connected)
 
 u0 = [
     n1.v => -2,

test/interface/jacobian_tests.jl

@@ -55,8 +55,8 @@ function lotka(du, u, p, t)
 end
 
 prob = ODEProblem(lotka, [1.0, 1.0], (0.0, 1.0), [1.5, 1.0, 3.0, 1.0])
-de = ModelingToolkit.modelingtoolkitize(prob)
-prob2 = remake(prob, f = ODEFunction(de; jac = true))
+de = ModelingToolkit.modelingtoolkitize(prob) |> complete
+prob2 = ODEProblem(de; jac = true)
 
 sol = solve(prob, TRBDF2())
 

test/regression/psos_and_energy_conservation.jl

@@ -90,7 +90,7 @@ end
 
 # energy conserving callback:
 # important to use save = false, I dont want rescaling points
-cb = ManifoldProjection(ghh, nlopts = Dict(:ftol => 1e-13), save = false)
+cb = ManifoldProjection(ghh, abstol = 1e-13, save = false)
 
 # Callback for Poincare surface of section
 function psos_callback(j, direction = +1, offset::Real = 0,