fix: improve type-stability of SCCNonlinearProblem

Author: AayushSabharwal

src/problems/nonlinear_problems.jl

@@ -484,33 +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)
+            init = similar(init, 0)
         else
             init = StaticArraysCore.similar_type(init, StaticArraysCore.Size(0))()
         end
         u0 = mapreduce(
             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)
@@ -534,7 +538,7 @@ end
 function SymbolicIndexingInterface.state_values(prob::SCCNonlinearProblem)
     init = state_values(first(prob.probs))
     if ArrayInterface.ismutable(init)
-        init = similar(init)
+        init = similar(init, 0)
     else
         init = StaticArraysCore.similar_type(init, StaticArraysCore.Size(0))()
     end

src/remake.jl

@@ -803,6 +803,49 @@ 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
+
+function scc_update_subproblems(probs::Tuple, newu0, newp, parameters_alias)
+    offset = Ref(0)
+    return ntuple(Val(length(probs))) do i
+        subprob = probs[i]
+        # 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[] += N
+        return subprob
+    end
+end
+
 """
     remake(prob::SCCNonlinearProblem; u0 = missing, p = missing, probs = missing,
         parameters_alias = prob.parameters_alias, sys = missing, explicitfuns! = missing)
@@ -831,23 +874,8 @@ 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
-            _u0 = newu0[(offset + 1):(offset + length(state_values(subprob)))]
-            if !ArrayInterface.ismutable(newu0)
-                _u0 = StaticArraysCore.similar_type(
-                    newu0, StaticArraysCore.Size(length(_u0)))(_u0)
-            end
-            subprob = if parameters_alias
-                remake(subprob; u0 = _u0, p = newp)
-            else
-                remake(subprob;
-                    u0 = _u0)
-            end
-            offset += length(state_values(subprob))
-            return subprob
-        end
+        probs = scc_update_subproblems(probs, newu0, newp, parameters_alias)
     end
     f = coalesce(f, prob.f)
     f = remake(f; sys)

test/problem_building_test.jl

@@ -130,14 +130,15 @@ end
     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)
+    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 = remake(sccprob; u0 = SA[2.0, 1.0, 2.0])
+    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}