fix: retain type of u0 in SCCNonlinearProblem impl for state_values and remake

Author: AayushSabharwal

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/problems/nonlinear_problems.jl

@@ -495,8 +495,14 @@ mutable struct SCCNonlinearProblem{uType, iip, P, E, F <: NonlinearFunction{iip}
 
     function SCCNonlinearProblem{P, E, F, Par}(probs::P, funs::E, f::F, pobj::Par,
             alias::Bool) where {P, E, F <: NonlinearFunction, Par}
+        init = state_values(first(probs))
+        if ArrayInterface.ismutable(init)
+            init = similar(init)
+        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)
     end
@@ -526,8 +532,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)
+    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)

src/remake.jl

@@ -834,13 +834,16 @@ function remake(prob::SCCNonlinearProblem; u0 = missing, p = missing, probs = mi
     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 = newu0[(offset + 1):(offset + length(state_values(subprob)))],
-                    p = newp)
+                remake(subprob; u0 = _u0, p = newp)
             else
                 remake(subprob;
-                    u0 = newu0[(offset + 1):(offset + length(state_values(subprob)))])
+                    u0 = _u0)
             end
             offset += length(state_values(subprob))
             return subprob

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,47 @@ 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 = 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