feat: add fields to OverrideInit, better nlsolve_alg handling

Author: AayushSabharwal

src/SciMLBase.jl

@@ -21,7 +21,7 @@ import CommonSolve: solve, init, step!, solve!
 import FunctionWrappersWrappers
 import RuntimeGeneratedFunctions
 import EnumX
-import ADTypes: AbstractADType
+import ADTypes: ADTypes, AbstractADType
 import Accessors: @set, @reset
 using Expronicon.ADT: @match
 
@@ -351,7 +351,15 @@ struct CheckInit <: DAEInitializationAlgorithm end
 """
 $(TYPEDEF)
 """
-struct OverrideInit <: DAEInitializationAlgorithm end
+struct OverrideInit{T, F} <: DAEInitializationAlgorithm
+    abstol::T
+    nlsolve::F
+end
+
+function OverrideInit(; abstol = 1e-10, nlsolve = nothing)
+    OverrideInit(abstol, nlsolve)
+end
+OverrideInit(abstol) = OverrideInit(; abstol = abstol, nlsolve = nothing)
 
 # PDE Discretizations
 

src/initialization.jl

@@ -160,7 +160,7 @@ argument, failing which this function will throw an error. The success value ret
 depends on the success of the nonlinear solve.
 """
 function get_initial_values(prob, valp, f, alg::OverrideInit,
-        isinplace::Union{Val{true}, Val{false}}; nlsolve_alg = nothing, kwargs...)
+        iip::Union{Val{true}, Val{false}}; nlsolve_alg = nothing, kwargs...)
     u0 = state_values(valp)
     p = parameter_values(valp)
 
@@ -171,15 +171,16 @@ function get_initial_values(prob, valp, f, alg::OverrideInit,
     initdata::OverrideInitData = f.initialization_data
     initprob = initdata.initializeprob
 
-    if nlsolve_alg === nothing
+    nlsolve_alg = something(nlsolve_alg, alg.nlsolve, Some(nothing))
+    if nlsolve_alg === nothing && state_values(initprob) !== nothing
         throw(OverrideInitMissingAlgorithm())
     end
 
     if initdata.update_initializeprob! !== nothing
         initdata.update_initializeprob!(initprob, valp)
     end
 
-    nlsol = solve(initprob, nlsolve_alg)
+    nlsol = solve(initprob, nlsolve_alg; abstol = alg.abstol)
 
     u0 = initdata.initializeprobmap(nlsol)
     if initdata.initializeprobpmap !== nothing

test/initialization.jl

@@ -101,15 +101,43 @@ end
     end
 
     @testset "Solves" begin
-        u0, p, success = SciMLBase.get_initial_values(
-            prob, integ, fn, SciMLBase.OverrideInit(),
-            Val(false); nlsolve_alg = NewtonRaphson())
+        @testset "with explicit alg" begin
+            u0, p, success = SciMLBase.get_initial_values(
+                prob, integ, fn, SciMLBase.OverrideInit(),
+                Val(false); nlsolve_alg = NewtonRaphson())
 
-        @test u0 ≈ [2.0, 2.0]
-        @test p ≈ 1.0
-        @test success
+            @test u0 ≈ [2.0, 2.0]
+            @test p ≈ 1.0
+            @test success
 
-        initprob.p[1] = 1.0
+            initprob.p[1] = 1.0
+        end
+        @testset "with alg in `OverrideInit`" begin
+            u0, p, success = SciMLBase.get_initial_values(
+                prob, integ, fn, SciMLBase.OverrideInit(nlsolve = NewtonRaphson()),
+                Val(false))
+
+            @test u0 ≈ [2.0, 2.0]
+            @test p ≈ 1.0
+            @test success
+
+            initprob.p[1] = 1.0
+        end
+        @testset "with trivial problem and no alg" begin
+            iprob = NonlinearProblem((u, p) -> 0.0, nothing, 1.0)
+            iprobmap = (_) -> [1.0, 1.0]
+            initdata = SciMLBase.OverrideInitData(iprob, nothing, iprobmap, nothing)
+            _fn = ODEFunction(rhs2; initialization_data = initdata)
+            _prob = ODEProblem(_fn, [2.0, 0.0], (0.0, 1.0), 1.0)
+            _integ = init(_prob; initializealg = NoInit())
+
+            u0, p, success = SciMLBase.get_initial_values(
+                _prob, _integ, _fn, SciMLBase.OverrideInit(), Val(false))
+
+            @test u0 ≈ [1.0, 1.0]
+            @test p ≈ 1.0
+            @test success
+        end
     end
 
     @testset "Solves with non-integrator value provider" begin