Merge pull request #503 from ChrisRackauckas-Claude/pr-464-rebased

feat: add support for OverrideInit and CheckInit (rebased)

Author: ChrisRackauckas

Project.toml

@@ -4,57 +4,69 @@ authors = ["Chris Rackauckas <[email protected]>"]
 version = "4.28.0"
 
 [deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Sundials_jll = "fb77eaff-e24c-56d4-86b1-d163f2edb164"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
 [compat]
+Accessors = "0.1.38"
+ADTypes = "1"
 AlgebraicMultigrid = "1"
 Aqua = "0.8"
+ArrayInterface = "7.17.1"
 CEnum = "0.5"
 DAEProblemLibrary = "0.1"
 DataStructures = "0.18, 0.19"
 DiffEqBase = "6.154"
 DiffEqCallbacks = "4"
+DifferentiationInterface = "0.6, 0.7"
 ExplicitImports = "1"
 ForwardDiff = "0.10"
 IncompleteLU = "0.2"
 Libdl = "1"
 LinearAlgebra = "1"
+LinearSolve = "3.40.0"
 Logging = "1"
 ModelingToolkit = "10"
 ODEProblemLibrary = "1"
 PrecompileTools = "1"
 Reexport = "1.0"
-SciMLBase = "2.9"
+SafeTestsets = "0.1"
+SciMLBase = "2.119.0"
 SparseArrays = "1"
-SparseConnectivityTracer = "0.6"
-SparseDiffTools = "2"
+SparseConnectivityTracer = "1"
 Sundials_jll = "7.4.1"
+SymbolicIndexingInterface = "0.3.35"
 Test = "1"
 julia = "1.10"
 
 [extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DAEProblemLibrary = "dfb8ca35-80a1-48ba-a605-84916a45b4f8"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "AlgebraicMultigrid", "Aqua", "DiffEqCallbacks", "ExplicitImports", "ODEProblemLibrary", "DAEProblemLibrary", "ForwardDiff", "SparseDiffTools", "SparseConnectivityTracer", "IncompleteLU", "ModelingToolkit"]
+test = ["Test", "ADTypes", "AlgebraicMultigrid", "Aqua", "DiffEqCallbacks", "ExplicitImports", "ODEProblemLibrary", "DAEProblemLibrary", "ForwardDiff", "DifferentiationInterface", "SparseConnectivityTracer", "IncompleteLU", "ModelingToolkit", "SafeTestsets"]

analyze_imports.jl

@@ -13,11 +13,16 @@ using DiffEqBase: DiffEqBase, NonlinearFunction, ODEFunction, add_saveat!,
                   update_coefficients!, warn_compat
 using SciMLBase: AbstractSciMLOperator, DAEProblem, ODEProblem, ReturnCode,
                  SciMLBase, SplitODEProblem, VectorContinuousCallback
+import Accessors: @reset
+import ArrayInterface
+import SymbolicIndexingInterface as SII
+import SymbolicIndexingInterface: ParameterIndexingProxy
 using DataStructures: DataStructures
 using Logging: Logging
 using SparseArrays: SparseArrays
 using LinearAlgebra: LinearAlgebra
 
+import LinearSolve # Required for initialization
 using Libdl: Libdl
 using CEnum: CEnum, @cenum
 
@@ -91,6 +96,7 @@ include("common_interface/verbosity.jl")
 include("common_interface/algorithms.jl")
 include("common_interface/integrator_types.jl")
 include("common_interface/integrator_utils.jl")
+include("common_interface/initialize_dae.jl")
 include("common_interface/solve.jl")
 
 import PrecompileTools

src/Sundials.jl

@@ -0,0 +1,87 @@
+struct SundialsDefaultInit <: DiffEqBase.DAEInitializationAlgorithm end
+
+function DiffEqBase.initialize_dae!(integrator::AbstractSundialsIntegrator, initializealg = integrator.initializealg)
+    _initialize_dae!(integrator, integrator.sol.prob, initializealg, Val(DiffEqBase.isinplace(integrator.sol.prob)))
+end
+
+struct IDADefaultInit <: DiffEqBase.DAEInitializationAlgorithm
+end
+
+function _initialize_dae!(integrator::IDAIntegrator, prob,
+    initializealg::IDADefaultInit, isinplace)
+    if integrator.u_modified
+        IDAReinit!(integrator)
+    end
+    integrator.f(integrator.tmp, integrator.du, integrator.u, integrator.p, integrator.t)
+    tstart, tend = integrator.sol.prob.tspan
+    if any(abs.(integrator.tmp) .>= integrator.opts.reltol)
+        if integrator.sol.prob.differential_vars === nothing && !integrator.alg.init_all
+            error("Must supply differential_vars argument to DAEProblem constructor to use IDA initial value solver.")
+        end
+        if integrator.alg.init_all
+            init_type = IDA_Y_INIT
+        else
+            init_type = IDA_YA_YDP_INIT
+            # Use preallocated NVector for differential_vars
+            if integrator.diff_vars_nvec !== nothing
+                integrator.flag = IDASetId(integrator.mem, integrator.diff_vars_nvec)
+            else
+                error("differential_vars NVector not preallocated but needed for IDASetId")
+            end
+        end
+        dt = integrator.dt == tstart ? tend : integrator.dt
+        integrator.flag = IDACalcIC(integrator.mem, init_type, dt)
+
+        # Reflect consistent initial conditions back into the integrator's
+        # shadow copy. N.B.: ({du, u}_nvec are aliased to {du, u}).
+        IDAGetConsistentIC(integrator.mem, integrator.u_nvec, integrator.du_nvec)
+    end
+    if integrator.t == tstart && integrator.flag < 0
+        integrator.sol = SciMLBase.solution_new_retcode(integrator.sol,
+            ReturnCode.InitialFailure)
+    end
+end
+
+function _initialize_dae!(integrator, prob, ::SundialsDefaultInit, isinplace)
+    if SciMLBase.has_initializeprob(prob.f)
+        _initialize_dae!(integrator, prob, SciMLBase.OverrideInit(), isinplace)
+    elseif integrator isa IDAIntegrator
+        _initialize_dae!(integrator, prob, IDADefaultInit(), isinplace)
+    end
+end
+
+function _initialize_dae!(integrator, prob, initalg::SciMLBase.NoInit, isinplace) end
+
+function _initialize_dae!(integrator, prob, initalg::SciMLBase.OverrideInit, isinplace::Union{Val{true}, Val{false}})
+    nlsolve_alg = KINSOL()
+    u0, p, success = SciMLBase.get_initial_values(prob, integrator, prob.f, initalg, isinplace; nlsolve_alg, abstol = integrator.opts.abstol, reltol = integrator.opts.reltol)
+
+    if isinplace === Val{true}()
+        integrator.u .= u0
+        if length(integrator.sol.u) == 1
+            integrator.sol.u[1] .= u0
+        end
+    else
+        integrator.u = u0
+        if length(integrator.sol.u) == 1
+            integrator.sol.u[1] = u0
+        end
+    end
+    integrator.p = p
+    sol = integrator.sol
+    @reset sol.prob.p = integrator.p
+    integrator.sol = sol
+
+    # For IDA, we need to reinitialize the solver after changing u, du, or p
+    if integrator isa IDAIntegrator && success
+        integrator.u_modified = true
+    end
+
+    if !success
+        integrator.sol = SciMLBase.solution_new_retcode(integrator.sol, ReturnCode.InitialFailure)
+    end
+end
+
+function _initialize_dae!(integrator, prob, initalg::SciMLBase.CheckInit, isinplace::Union{Val{true}, Val{false}})
+    SciMLBase.get_initial_values(prob, integrator, prob.f, initalg, isinplace; abstol = integrator.opts.abstol)
+end

src/common_interface/initialize_dae.jl

@@ -40,7 +40,8 @@ mutable struct CVODEIntegrator{N,
     oType,
     LStype,
     Atype,
-    CallbackCacheType} <: AbstractSundialsIntegrator{algType}
+    CallbackCacheType,
+    IA} <: AbstractSundialsIntegrator{algType}
     u::Array{Float64, N}
     u_nvec::NVector
     p::pType
@@ -66,6 +67,7 @@ mutable struct CVODEIntegrator{N,
     vector_event_last_time::Int
     callback_cache::CallbackCacheType
     last_event_error::Float64
+    initializealg::IA
     ctx_handle::ContextHandle
 end
 
@@ -102,7 +104,8 @@ mutable struct ARKODEIntegrator{N,
     MLStype,
     Mtype,
     CallbackCacheType,
-    MemType} <: AbstractSundialsIntegrator{ARKODE}
+    MemType,
+    IA} <: AbstractSundialsIntegrator{ARKODE}
     u::Array{Float64, N}
     u_nvec::NVector
     p::pType
@@ -130,15 +133,16 @@ mutable struct ARKODEIntegrator{N,
     vector_event_last_time::Int
     callback_cache::CallbackCacheType
     last_event_error::Float64
+    initializealg::IA
     ctx_handle::ContextHandle
 end
 
 function (integrator::ARKODEIntegrator{
         N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, ARKStepMem})(t::Number,
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, ARKStepMem, IA})(t::Number,
         deriv::Type{Val{T}} = Val{0};
         idxs = nothing) where {N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, T}
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, IA, T}
     out = similar(integrator.u)
     out_nvec = NVector(vec(out), integrator.ctx_handle.ctx)
     integrator.flag = @checkflag ARKStepGetDky(integrator.mem, t, Cint(T), out_nvec)
@@ -148,10 +152,10 @@ end
 
 function (integrator::ARKODEIntegrator{
         N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, ERKStepMem})(t::Number,
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, ERKStepMem, IA})(t::Number,
         deriv::Type{Val{T}} = Val{0};
         idxs = nothing) where {N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, T}
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, IA, T}
     out = similar(integrator.u)
     out_nvec = NVector(vec(out), integrator.ctx_handle.ctx)
     integrator.flag = @checkflag ERKStepGetDky(integrator.mem, t, Cint(T), out_nvec)
@@ -161,11 +165,11 @@ end
 
 function (integrator::ARKODEIntegrator{
         N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, ARKStepMem})(out,
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, ARKStepMem, IA})(out,
         t::Number,
         deriv::Type{Val{T}} = Val{0};
         idxs = nothing) where {N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, T}
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, IA, T}
     out_nvec = NVector(vec(out), integrator.ctx_handle.ctx)
     integrator.flag = @checkflag ARKStepGetDky(integrator.mem, t, Cint(T), out_nvec)
     copyto!(out, out_nvec.v)
@@ -174,11 +178,11 @@ end
 
 function (integrator::ARKODEIntegrator{
         N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, ERKStepMem})(out,
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, ERKStepMem, IA})(out,
         t::Number,
         deriv::Type{Val{T}} = Val{0};
         idxs = nothing) where {N, pType, solType, algType, fType, UFType, JType, oType,
-        LStype, Atype, MLStype, Mtype, CallbackCacheType, T}
+        LStype, Atype, MLStype, Mtype, CallbackCacheType, IA, T}
     out_nvec = NVector(vec(out), integrator.ctx_handle.ctx)
     integrator.flag = @checkflag ERKStepGetDky(integrator.mem, t, Cint(T), out_nvec)
     copyto!(out, out_nvec.v)

src/common_interface/integrator_types.jl

@@ -157,7 +157,7 @@ end
 
 function handle_callback_modifiers!(integrator::IDAIntegrator)
     # Implicitly does IDAReinit!
-    DiffEqBase.initialize_dae!(integrator, IDADefaultInit())
+    DiffEqBase.initialize_dae!(integrator)
 end
 
 function DiffEqBase.add_tstop!(integrator::AbstractSundialsIntegrator, t)
@@ -211,6 +211,8 @@ end
 @inline function Base.getproperty(integrator::AbstractSundialsIntegrator, sym::Symbol)
     if sym == :dt
         return integrator.t - integrator.tprev
+    elseif sym == :ps
+        return ParameterIndexingProxy(integrator)
     else
         return getfield(integrator, sym)
     end
@@ -228,46 +230,6 @@ end
 # Required for callbacks
 DiffEqBase.set_proposed_dt!(i::AbstractSundialsIntegrator, dt) = nothing
 
-DiffEqBase.initialize_dae!(integrator::AbstractSundialsIntegrator) = nothing
-
-struct IDADefaultInit <: DiffEqBase.DAEInitializationAlgorithm
-end
-
-function DiffEqBase.initialize_dae!(integrator::IDAIntegrator,
-        initializealg::IDADefaultInit)
-    if integrator.u_modified
-        IDAReinit!(integrator)
-    end
-    integrator.f(integrator.tmp, integrator.du, integrator.u, integrator.p, integrator.t)
-    tstart, tend = integrator.sol.prob.tspan
-    if any(abs.(integrator.tmp) .>= integrator.opts.reltol)
-        if integrator.sol.prob.differential_vars === nothing && !integrator.alg.init_all
-            error("Must supply differential_vars argument to DAEProblem constructor to use IDA initial value solver.")
-        end
-        if integrator.alg.init_all
-            init_type = IDA_Y_INIT
-        else
-            init_type = IDA_YA_YDP_INIT
-            # Use preallocated NVector for differential_vars
-            if integrator.diff_vars_nvec !== nothing
-                integrator.flag = IDASetId(integrator.mem, integrator.diff_vars_nvec)
-            else
-                error("differential_vars NVector not preallocated but needed for IDASetId")
-            end
-        end
-        dt = integrator.dt == tstart ? tend : integrator.dt
-        integrator.flag = IDACalcIC(integrator.mem, init_type, dt)
-
-        # Reflect consistent initial conditions back into the integrator's
-        # shadow copy. N.B.: ({du, u}_nvec are aliased to {du, u}).
-        IDAGetConsistentIC(integrator.mem, integrator.u_nvec, integrator.du_nvec)
-    end
-    if integrator.t == tstart && integrator.flag < 0
-        integrator.sol = SciMLBase.solution_new_retcode(integrator.sol,
-            ReturnCode.InitialFailure)
-    end
-end
-
 DiffEqBase.has_reinit(integrator::AbstractSundialsIntegrator) = true
 function DiffEqBase.reinit!(integrator::AbstractSundialsIntegrator,
         u0 = integrator.sol.prob.u0;
@@ -343,3 +305,6 @@ DiffEqBase.get_tstops_array(integ::AbstractSundialsIntegrator) = get_tstops(inte
 function DiffEqBase.get_tstops_max(integ::AbstractSundialsIntegrator)
     maximum(get_tstops_array(integ))
 end
+
+# SII
+SII.symbolic_container(integ::AbstractSundialsIntegrator) = integ.sol