Merge pull request #1012 from AayushSabharwal/as/oop-init-and-fixes

feat: allow `update_initializeprob!` to be oop, fix `SCCNonlinearProblem` dropping `u0::SVector` type

Author: ChrisRackauckas

src/SciMLBase.jl

@@ -25,6 +25,7 @@ import ADTypes: ADTypes, AbstractADType
 import Accessors: @set, @reset, @delete, @insert
 using Moshi.Data: @data
 using Moshi.Match: @match
+import StaticArraysCore
 
 using Reexport
 using SciMLOperators

src/initialization.jl

@@ -3,7 +3,8 @@
 
 A collection of all the data required for `OverrideInit`.
 """
-struct OverrideInitData{IProb, UIProb, IProbMap, IProbPmap, M}
+struct OverrideInitData{
+    IProb, UIProb, IProbMap, IProbPmap, M, OOP <: Union{Val{true}, Val{false}}}
     """
     The `AbstractNonlinearProblem` to solve for initialization.
     """
@@ -34,19 +35,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::OOP
 
     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::O) where {I, J, K, L, M, O}
         @assert initprob isa
                 Union{SCCNonlinearProblem, NonlinearProblem, NonlinearLeastSquaresProblem}
-        return new{I, J, K, L, M}(
-            initprob, update_initprob!, initprobmap, initprobpmap, metadata)
+        return new{I, J, K, L, M, O}(
+            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 +252,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)

src/problems/nonlinear_problems.jl

@@ -484,27 +484,37 @@ Note that this example aliases the parameters together for a memory-reduced repr
 * `probs`: the collection of problems to solve
 * `explictfuns!`: the explicit functions for mutating the parameter set
 """
-mutable struct SCCNonlinearProblem{uType, iip, P, E, F <: NonlinearFunction{iip}, Par} <:
+mutable struct SCCNonlinearProblem{uType, iip, P, E, F <: NonlinearFunction{iip},
+    Par, Palias <: Union{Val{true}, Val{false}}} <:
                AbstractNonlinearProblem{uType, iip}
     probs::P
     explicitfuns!::E
     # NonlinearFunction with `f = Returns(nothing)`
     f::F
     p::Par
-    parameters_alias::Bool
+    parameters_alias::Palias
 
     function SCCNonlinearProblem{P, E, F, Par}(probs::P, funs::E, f::F, pobj::Par,
-            alias::Bool) where {P, E, F <: NonlinearFunction, Par}
+            alias::Palias) where {P, E, F <: NonlinearFunction, Par, Palias}
+        init = state_values(first(probs))
+        if ArrayInterface.ismutable(init)
+            init = similar(init, 0)
+        else
+            init = StaticArraysCore.similar_type(init, StaticArraysCore.Size(0))()
+        end
         u0 = mapreduce(
-            state_values, vcat, probs; init = similar(state_values(first(probs)), 0))
+            state_values, vcat, probs; init = init)
         uType = typeof(u0)
-        new{uType, false, P, E, F, Par}(probs, funs, f, pobj, alias)
+        new{uType, false, P, E, F, Par, Palias}(probs, funs, f, pobj, alias)
     end
 end
 
 function SCCNonlinearProblem(probs, explicitfuns!, parameter_object = nothing,
-        parameters_alias = false; kwargs...)
+        parameters_alias::Union{Bool, Val{true}, Val{false}} = Val(false); kwargs...)
     f = NonlinearFunction{false}(Returns(nothing); kwargs...)
+    if parameters_alias isa Bool
+        parameters_alias = Val(parameters_alias)
+    end
     return SCCNonlinearProblem{typeof(probs), typeof(explicitfuns!),
         typeof(f), typeof(parameter_object)}(
         probs, explicitfuns!, f, parameter_object, parameters_alias)
@@ -526,8 +536,13 @@ function SymbolicIndexingInterface.parameter_values(prob::SCCNonlinearProblem)
     prob.p
 end
 function SymbolicIndexingInterface.state_values(prob::SCCNonlinearProblem)
-    mapreduce(
-        state_values, vcat, prob.probs; init = similar(state_values(first(prob.probs)), 0))
+    init = state_values(first(prob.probs))
+    if ArrayInterface.ismutable(init)
+        init = similar(init, 0)
+    else
+        init = StaticArraysCore.similar_type(init, StaticArraysCore.Size(0))()
+    end
+    mapreduce(state_values, vcat, prob.probs; init)
 end
 
 function SymbolicIndexingInterface.set_state!(prob::SCCNonlinearProblem, val, idx)
@@ -542,7 +557,7 @@ end
 function SymbolicIndexingInterface.set_parameter!(prob::SCCNonlinearProblem, val, idx)
     if prob.p !== nothing
         set_parameter!(prob.p, val, idx)
-        prob.parameters_alias && return
+        prob.parameters_alias === Val(true) && return
     end
     for scc in prob.probs
         is_parameter(scc, idx) || continue

src/remake.jl

@@ -803,6 +803,61 @@ function remake(prob::NonlinearLeastSquaresProblem; f = missing, u0 = missing, p
     return prob
 end
 
+function scc_update_subproblems(probs::Vector, newu0, newp, parameters_alias)
+    offset = Ref(0)
+    return map(probs) do subprob
+        # N should be inferred if `prob` is type-stable and `subprob.u0 isa StaticArray`
+        N = length(state_values(subprob))
+        if ArrayInterface.ismutable(newu0)
+            _u0 = newu0[(offset[] + 1):(offset[] + N)]
+        else
+            _u0 = StaticArraysCore.similar_type(
+                newu0, StaticArraysCore.Size(N))(newu0[(offset[] + 1):(offset[] + N)])
+        end
+        subprob = if parameters_alias === Val(true)
+            remake(subprob; u0 = _u0, p = newp)
+        else
+            remake(subprob; u0 = _u0)
+        end
+        offset[] += length(state_values(subprob))
+        return subprob
+    end
+end
+
+@generated function scc_update_subproblems(probs::Tuple, newu0, newp, parameters_alias)
+    function get_expr(i::Int)
+        subprob_name = Symbol(:subprob, i)
+        quote
+            $subprob_name = probs[$i]
+            # N should be inferred if `prob` is type-stable and `subprob.u0 isa StaticArray`
+            N = length(state_values($subprob_name))
+            if ArrayInterface.ismutable(newu0)
+                _u0 = newu0[(offset + 1):(offset + N)]
+            else
+                _u0 = StaticArraysCore.similar_type(
+                    newu0, StaticArraysCore.Size(N))(newu0[(offset + 1):(offset + N)])
+            end
+            $subprob_name = if parameters_alias === Val(true)
+                remake($subprob_name; u0 = _u0, p = newp)
+            else
+                remake($subprob_name; u0 = _u0)
+            end
+            offset += N
+        end, subprob_name
+    end
+    expr = quote
+        offset = 0
+    end
+    subprob_names = []
+    for i in 1:fieldcount(probs)
+        subexpr, spname = get_expr(i)
+        push!(expr.args, subexpr)
+        push!(subprob_names, spname)
+    end
+    push!(expr.args, Expr(:tuple, subprob_names...))
+    return expr
+end
+
 """
     remake(prob::SCCNonlinearProblem; u0 = missing, p = missing, probs = missing,
         parameters_alias = prob.parameters_alias, sys = missing, explicitfuns! = missing)
@@ -818,7 +873,10 @@ override the values in `probs`. `sys` is the index provider for the full system.
 function remake(prob::SCCNonlinearProblem; u0 = missing, p = missing, probs = missing,
         parameters_alias = prob.parameters_alias, f = missing, sys = missing,
         interpret_symbolicmap = true, use_defaults = false, explicitfuns! = missing)
-    if p !== missing && !parameters_alias && probs === missing
+    if parameters_alias isa Bool
+        parameters_alias = Val(parameters_alias)
+    end
+    if p !== missing && parameters_alias === Val(false) && probs === missing
         throw(ArgumentError("`parameters_alias` is `false` for the given `SCCNonlinearProblem`. Please provide the subproblems using the keyword `probs` with the parameters updated appropriately in each."))
     end
     newu0, newp = updated_u0_p(prob, u0, p; interpret_symbolicmap, use_defaults)
@@ -831,28 +889,14 @@ function remake(prob::SCCNonlinearProblem; u0 = missing, p = missing, probs = mi
     if sys === missing
         sys = prob.f.sys
     end
-    offset = 0
-    if u0 !== missing || p !== missing && parameters_alias
-        probs = map(probs) do subprob
-            subprob = if parameters_alias
-                remake(subprob;
-                    u0 = newu0[(offset + 1):(offset + length(state_values(subprob)))],
-                    p = newp)
-            else
-                remake(subprob;
-                    u0 = newu0[(offset + 1):(offset + length(state_values(subprob)))])
-            end
-            offset += length(state_values(subprob))
-            return subprob
-        end
+    if u0 !== missing || p !== missing && parameters_alias === Val(true)
+        probs = scc_update_subproblems(probs, newu0, newp, parameters_alias)
     end
     f = coalesce(f, prob.f)
     f = remake(f; sys)
-    props = getproperties(f)
-    props = @delete props.f
 
-    return SCCNonlinearProblem(
-        probs, explicitfuns!, newp, parameters_alias; props...)
+    return SCCNonlinearProblem{typeof(probs), typeof(explicitfuns!), typeof(f), typeof(newp)}(
+        probs, explicitfuns!, f, newp, parameters_alias)
 end
 
 function varmap_has_var(varmap, var)

test/downstream/modelingtoolkit_remake.jl

@@ -330,7 +330,7 @@ end
     newp = remake_buffer(sccprob.f.sys, sccprob.p, [σ], [3.0])
     sccprob4 = remake(sccprob; parameters_alias = false, p = newp,
         probs = [remake(sccprob.probs[1]; p = deepcopy(newp)), sccprob.probs[2]])
-    @test !sccprob4.parameters_alias
+    @test sccprob4.parameters_alias === Val(false)
     @test sccprob4.p !== sccprob4.probs[1].p
     @test sccprob4.p !== sccprob4.probs[2].p
 end

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)

test/problem_building_test.jl

@@ -1,4 +1,4 @@
-using Test, SciMLBase, SymbolicIndexingInterface, Accessors
+using Test, SciMLBase, SymbolicIndexingInterface, Accessors, StaticArrays
 
 function simplependulum!(du, u, p, t)
     θ = u[1]
@@ -98,3 +98,48 @@ let
     @test dprob8.u0 === u₀
     @test dprob8.tspan === tspan
 end
+
+@testset "SCCNonlinearProblem with static arrays" begin
+    function f1(u, p)
+        y = u[1]
+        x = p[1]
+        return SA[1 - y^2 - x^2]
+    end
+
+    function f2(u, p)
+        yt = u[1]
+        x, xt, y = p
+        return SA[-2y * yt - 2x * xt]
+    end
+
+    function f3(u, p)
+        lam = u[1]
+        x, xt, y, yt = p
+        return SA[-2xt^2 - 2yt^2 - 2y * (-1 + y * lam) - 2x^2 * lam]
+    end
+
+    explicit1 = Returns(nothing)
+    function explicit2(p, sols)
+        p[3] = sols[1].u[1]
+    end
+    function explicit3(p, sols)
+        p[4] = sols[2].u[1]
+    end
+
+    p = [1.0, 0.0, NaN, NaN]
+    prob1 = NonlinearProblem(f1, SA[1.0], p)
+    prob2 = NonlinearProblem(f2, SA[0.0], p)
+    prob3 = NonlinearProblem(f3, SA[1.0], p)
+    sccprob = SCCNonlinearProblem(
+        (prob1, prob2, prob3), (explicit1, explicit2, explicit3), p, true)
+
+    @test !SciMLBase.isinplace(sccprob)
+    @test sccprob isa SCCNonlinearProblem{SVector{3, Float64}}
+    @test state_values(sccprob) isa SVector{3, Float64}
+    @test sccprob.p === prob1.p === prob2.p === prob3.p
+
+    sccprob2 = @inferred remake(sccprob; u0 = SA[2.0, 1.0, 2.0])
+    @test !SciMLBase.isinplace(sccprob2)
+    @test sccprob2 isa SCCNonlinearProblem{SVector{3, Float64}}
+    @test state_values(sccprob2) isa SVector{3, Float64}
+end