Merge pull request #2466 from SciML/checkinit2

Implement CheckInit

Author: ChrisRackauckas

Project.toml

@@ -133,7 +133,7 @@ PrecompileTools = "1"
 Preferences = "1.3"
 RecursiveArrayTools = "2.36, 3"
 Reexport = "1.0"
-SciMLBase = "2.50.4"
+SciMLBase = "2.53"
 SciMLOperators = "0.3"
 SciMLStructures = "1"
 SimpleNonlinearSolve = "1"

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -58,7 +58,7 @@ using DiffEqBase: check_error!, @def, _vec, _reshape
 
 using FastBroadcast: @.., True, False
 
-using SciMLBase: NoInit, _unwrap_val
+using SciMLBase: NoInit, CheckInit, _unwrap_val
 
 import SciMLBase: alg_order
 

lib/OrdinaryDiffEqCore/src/initialize_dae.jl

@@ -176,3 +176,83 @@ function _initialize_dae!(integrator, prob::Union{ODEProblem, DAEProblem},
             ReturnCode.InitialFailure)
     end
 end
+
+## CheckInit
+struct CheckInitFailureError <: Exception
+    normresid
+    abstol
+end
+
+function Base.showerror(io::IO, e::CheckInitFailureError)
+    print(io, "CheckInit specified but initialization not satisifed. normresid = $(e.normresid) > abstol = $(e.abstol)")
+end
+
+function _initialize_dae!(integrator, prob::ODEProblem, alg::CheckInit,
+        isinplace::Val{true})
+    @unpack p, t, f = integrator
+    M = integrator.f.mass_matrix
+    tmp = first(get_tmp_cache(integrator))
+    u0 = integrator.u
+
+    algebraic_vars = [all(iszero, x) for x in eachcol(M)]
+    algebraic_eqs = [all(iszero, x) for x in eachrow(M)]
+    (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return
+    update_coefficients!(M, u0, p, t)
+    f(tmp, u0, p, t)
+    tmp .= ArrayInterface.restructure(tmp, algebraic_eqs .* _vec(tmp))
+
+    normresid = integrator.opts.internalnorm(tmp, t)
+    if normresid > integrator.opts.abstol 
+        throw(CheckInitFailureError(normresid, integrator.opts.abstol))
+    end
+end
+
+function _initialize_dae!(integrator, prob::ODEProblem, alg::CheckInit,
+        isinplace::Val{false})
+    @unpack p, t, f = integrator
+    u0 = integrator.u
+    M = integrator.f.mass_matrix
+
+    algebraic_vars = [all(iszero, x) for x in eachcol(M)]
+    algebraic_eqs = [all(iszero, x) for x in eachrow(M)]
+    (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return
+    update_coefficients!(M, u0, p, t)
+    du = f(u0, p, t)
+    resid = _vec(du)[algebraic_eqs]
+
+    normresid = integrator.opts.internalnorm(resid, t)
+    if normresid > integrator.opts.abstol 
+        throw(CheckInitFailureError(normresid, integrator.opts.abstol))
+    end
+end
+
+function _initialize_dae!(integrator, prob::DAEProblem,
+        alg::CheckInit, isinplace::Val{true})
+    @unpack p, t, f = integrator
+    u0 = integrator.u
+    resid = get_tmp_cache(integrator)[2]
+
+    f(resid, integrator.du, u0, p, t)
+    normresid = integrator.opts.internalnorm(resid, t)
+    if normresid > integrator.opts.abstol 
+        throw(CheckInitFailureError(normresid, integrator.opts.abstol))
+    end
+end
+
+function _initialize_dae!(integrator, prob::DAEProblem,
+        alg::CheckInit, isinplace::Val{false})
+    @unpack p, t, f = integrator
+    u0 = integrator.u
+
+    nlequation_oop = u -> begin
+        f((u - u0) / dt, u, p, t)
+    end
+
+    nlequation = (u, _) -> nlequation_oop(u)
+
+    resid = f(integrator.du, u0, p, t)
+    normresid = integrator.opts.internalnorm(resid, t)
+    if normresid > integrator.opts.abstol 
+        throw(CheckInitFailureError(normresid, integrator.opts.abstol))
+    end
+end

lib/OrdinaryDiffEqNonlinearSolve/src/OrdinaryDiffEqNonlinearSolve.jl

@@ -34,8 +34,8 @@ end
 using OrdinaryDiffEqCore: resize_nlsolver!, _initialize_dae!,
                           AbstractNLSolverAlgorithm, AbstractNLSolverCache,
                           AbstractNLSolver, NewtonAlgorithm, @unpack,
-                          OverrideInit, ShampineCollocationInit, BrownFullBasicInit, _vec,
-                          _unwrap_val, DAEAlgorithm,
+                          OverrideInit, ShampineCollocationInit, BrownFullBasicInit,
+                          _vec, _unwrap_val, DAEAlgorithm,
                           _reshape, calculate_residuals, calculate_residuals!,
                           has_special_newton_error, isadaptive,
                           TryAgain, DIRK, COEFFICIENT_MULTISTEP, NORDSIECK_MULTISTEP, GLM,

test/interface/checkinit_tests.jl

@@ -0,0 +1,50 @@
+using OrdinaryDiffEqBDF, OrdinaryDiffEqRosenbrock, LinearAlgebra, ForwardDiff, Test
+using OrdinaryDiffEqCore
+
+function rober(du, u, p, t)
+    y₁, y₂, y₃ = u
+    k₁, k₂, k₃ = p
+    du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
+    du[2] = k₁ * y₁ - k₃ * y₂ * y₃ - k₂ * y₂^2
+    du[3] = y₁ + y₂ + y₃ - 1
+    nothing
+end
+function rober(u, p, t)
+    y₁, y₂, y₃ = u
+    k₁, k₂, k₃ = p
+    [-k₁ * y₁ + k₃ * y₂ * y₃,
+        k₁ * y₁ - k₃ * y₂ * y₃ - k₂ * y₂^2,
+        y₁ + y₂ + y₃ - 1]
+end
+M = [1.0 0 0
+     0 1.0 0
+     0 0 0]
+roberf = ODEFunction(rober, mass_matrix = M)
+roberf_oop = ODEFunction{false}(rober, mass_matrix = M)
+prob_mm = ODEProblem(roberf, [1.0, 0.0, 0.2], (0.0, 1e5), (0.04, 3e7, 1e4))
+prob_mm_oop = ODEProblem(roberf_oop, [1.0, 0.0, 0.2], (0.0, 1e5), (0.04, 3e7, 1e4))
+
+@test_throws OrdinaryDiffEqCore.CheckInitFailureError solve(prob_mm, Rodas5P(), reltol = 1e-8, abstol = 1e-8, initializealg = SciMLBase.CheckInit())
+@test_throws OrdinaryDiffEqCore.CheckInitFailureError solve(prob_mm_oop, Rodas5P(), reltol = 1e-8, abstol = 1e-8, initializealg = SciMLBase.CheckInit())
+
+f_oop = function (du, u, p, t)
+    out1 = -0.04u[1] + 1e4 * u[2] * u[3] - du[1]
+    out2 = +0.04u[1] - 3e7 * u[2]^2 - 1e4 * u[2] * u[3] - du[2]
+    out3 = u[1] + u[2] + u[3] - 1.0
+    [out1, out2, out3]
+end
+
+f = function (resid, du, u, p, t)
+    resid[1] = -0.04u[1] + 1e4 * u[2] * u[3] - du[1]
+    resid[2] = +0.04u[1] - 3e7 * u[2]^2 - 1e4 * u[2] * u[3] - du[2]
+    resid[3] = u[1] + u[2] + u[3] - 1.0
+end
+
+u₀ = [1.0, 0, 0.2]
+du₀ = [0.0, 0.0, 0.0]
+tspan = (0.0, 100000.0)
+differential_vars = [true, true, false]
+prob = DAEProblem(f, du₀, u₀, tspan, differential_vars = differential_vars)
+prob_oop = DAEProblem(f_oop, du₀, u₀, tspan, differential_vars = differential_vars)
+@test_throws OrdinaryDiffEqCore.CheckInitFailureError solve(prob, DFBDF(), reltol = 1e-8, abstol = 1e-8, initializealg = SciMLBase.CheckInit())
+@test_throws OrdinaryDiffEqCore.CheckInitFailureError solve(prob_oop, DFBDF(), reltol = 1e-8, abstol = 1e-8, initializealg = SciMLBase.CheckInit())

test/runtests.jl

@@ -65,7 +65,8 @@ end
         @time @safetestset "Linear Solver Split ODE Tests" include("interface/linear_solver_split_ode_test.jl")
         @time @safetestset "Sparse Diff Tests" include("interface/sparsediff_tests.jl")
         @time @safetestset "Enum Tests" include("interface/enums.jl")
-        @time @safetestset "Enum Tests" include("interface/get_du.jl")
+        @time @safetestset "CheckInit Tests" include("interface/checkinit_tests.jl")
+        @time @safetestset "Get du Tests" include("interface/get_du.jl")
         @time @safetestset "Mass Matrix Tests" include("interface/mass_matrix_tests.jl")
         @time @safetestset "W-Operator prototype tests" include("interface/wprototype_tests.jl")
     end