feat: allow update_initializeprob! to be out-of-place in OverrideInitData

Author: AayushSabharwal

src/initialization.jl

@@ -34,19 +34,26 @@ struct OverrideInitData{IProb, UIProb, IProbMap, IProbPmap, M}
     Additional metadata required by the creator of the initialization.
     """
     metadata::M
+    """
+    If this flag is `Val{true}`, `update_initializeprob!` is treated as an out-of-place
+    function which returns the updated `initializeprob`.
+    """
+    is_update_oop::Union{Type{Val{true}}, Type{Val{false}}}
 
     function OverrideInitData(initprob::I, update_initprob!::J, initprobmap::K,
-            initprobpmap::L, metadata::M) where {I, J, K, L, M}
+            initprobpmap::L, metadata::M, is_update_oop) where {I, J, K, L, M}
         @assert initprob isa
                 Union{SCCNonlinearProblem, NonlinearProblem, NonlinearLeastSquaresProblem}
         return new{I, J, K, L, M}(
-            initprob, update_initprob!, initprobmap, initprobpmap, metadata)
+            initprob, update_initprob!, initprobmap, initprobpmap, metadata, is_update_oop)
     end
 end
 
 function OverrideInitData(
-        initprob, update_initprob!, initprobmap, initprobpmap; metadata = nothing)
-    OverrideInitData(initprob, update_initprob!, initprobmap, initprobpmap, metadata)
+        initprob, update_initprob!, initprobmap, initprobpmap;
+        metadata = nothing, is_update_oop = Val{false})
+    OverrideInitData(
+        initprob, update_initprob!, initprobmap, initprobpmap, metadata, is_update_oop)
 end
 
 """
@@ -244,7 +251,11 @@ function get_initial_values(prob, valp, f, alg::OverrideInit,
     initprob = initdata.initializeprob
 
     if initdata.update_initializeprob! !== nothing
-        initdata.update_initializeprob!(initprob, valp)
+        if initdata.is_update_oop == Val{true}
+            initprob = initdata.update_initializeprob!(initprob, valp)
+        else
+            initdata.update_initializeprob!(initprob, valp)
+        end
     end
 
     if is_trivial_initialization(initdata)

test/initialization.jl

@@ -240,6 +240,24 @@ end
         @test success
     end
 
+    @testset "`is_update_oop` flag" begin
+        initprob = remake(initprob; u0 = ones(2), p = ones(1))
+        update_initializeprob = function (initprob, valp)
+            return remake(initprob; p = [valp.u[1]])
+        end
+        initdata = SciMLBase.OverrideInitData(initprob, update_initializeprob, initprobmap,
+            initprobpmap; is_update_oop = Val{true})
+        fn = ODEFunction(rhs2; initialization_data = initdata)
+        prob = ODEProblem(fn, [2.0, 0.0], (0.0, 1.0), 0.0)
+        integ = init(prob; initializealg = NoInit())
+        u0, p, success = SciMLBase.get_initial_values(
+            prob, integ, fn, SciMLBase.OverrideInit(), Val(false);
+            nlsolve_alg = NewtonRaphson(), abstol, reltol)
+        @test u0 ≈ [2.0, 2.0]
+        @test p ≈ 1.0
+        @test success
+    end
+
     @testset "Solves without `initializeprobmap`" begin
         initdata = SciMLBase.@set initialization_data.initializeprobmap = nothing
         fn = ODEFunction(rhs2; initialization_data = initdata)