fix: allow specifying type of buffers inside MTKParameters

Author: AayushSabharwal

src/systems/parameter_buffer.jl

@@ -28,7 +28,11 @@ the default behavior).
 """
 function MTKParameters(
         sys::AbstractSystem, p, u0 = Dict(); tofloat = false,
-        t0 = nothing, substitution_limit = 1000, floatT = nothing)
+        t0 = nothing, substitution_limit = 1000, floatT = nothing,
+        container_type = Vector)
+    if !(container_type <: AbstractArray)
+        container_type = Array
+    end
     ic = if has_index_cache(sys) && get_index_cache(sys) !== nothing
         get_index_cache(sys)
     else
@@ -133,18 +137,23 @@ function MTKParameters(
             end
         end
     end
-    tunable_buffer = narrow_buffer_type(tunable_buffer)
+    tunable_buffer = narrow_buffer_type(tunable_buffer; container_type)
     if isempty(tunable_buffer)
         tunable_buffer = SizedVector{0, Float64}()
     end
-    initials_buffer = narrow_buffer_type(initials_buffer)
+    initials_buffer = narrow_buffer_type(initials_buffer; container_type)
     if isempty(initials_buffer)
         initials_buffer = SizedVector{0, Float64}()
     end
-    disc_buffer = narrow_buffer_type.(disc_buffer)
-    const_buffer = narrow_buffer_type.(const_buffer)
+    disc_buffer = narrow_buffer_type.(disc_buffer; container_type)
+    const_buffer = narrow_buffer_type.(const_buffer; container_type)
     # Don't narrow nonnumeric types
-    nonnumeric_buffer = nonnumeric_buffer
+    if !isempty(nonnumeric_buffer)
+        nonnumeric_buffer = map(nonnumeric_buffer) do buf
+            SymbolicUtils.Code.create_array(
+                container_type, nothing, Val(1), Val(length(buf)), buf...)
+        end
+    end
 
     mtkps = MTKParameters{
         typeof(tunable_buffer), typeof(initials_buffer), typeof(disc_buffer),
@@ -160,21 +169,44 @@ function rebuild_with_caches(p::MTKParameters, cache_templates::BufferTemplate..
     @set p.caches = buffers
 end
 
-function narrow_buffer_type(buffer::AbstractArray)
+function narrow_buffer_type(buffer::AbstractArray; container_type = typeof(buffer))
     type = Union{}
     for x in buffer
         type = promote_type(type, typeof(x))
     end
-    return convert.(type, buffer)
+    return SymbolicUtils.Code.create_array(
+        container_type, type, Val(ndims(buffer)), Val(length(buffer)), buffer...)
 end
 
-function narrow_buffer_type(buffer::AbstractArray{<:AbstractArray})
-    buffer = narrow_buffer_type.(buffer)
+function narrow_buffer_type(
+        buffer::AbstractArray{<:AbstractArray}; container_type = typeof(buffer))
+    type = Union{}
+    for arr in buffer
+        for x in arr
+            type = promote_type(type, typeof(x))
+        end
+    end
+    buffer = map(buffer) do buf
+        SymbolicUtils.Code.create_array(
+            container_type, type, Val(ndims(buf)), Val(size(buf)), buf...)
+    end
+    return SymbolicUtils.Code.create_array(
+        container_type, nothing, Val(ndims(buffer)), Val(size(buffer)), buffer...)
+end
+
+function narrow_buffer_type(buffer::BlockedArray; container_type = typeof(parent(buffer)))
     type = Union{}
     for x in buffer
-        type = promote_type(type, eltype(x))
+        type = promote_type(type, typeof(x))
+    end
+    tmp = SymbolicUtils.Code.create_array(
+        container_type, type, Val(ndims(buffer)), Val(size(buffer)), buffer...)
+    blocks = ntuple(Val(ndims(buffer))) do i
+        bsizes = blocksizes(buffer, i)
+        SymbolicUtils.Code.create_array(
+            container_type, Int, Val(1), Val(length(bsizes)), bsizes...)
     end
-    return broadcast.(convert, type, buffer)
+    return BlockedArray(tmp, blocks...)
 end
 
 function buffer_to_arraypartition(buf)

src/systems/problem_utils.jl

@@ -1150,7 +1150,7 @@ function process_SciMLProblem(
     end
     evaluate_varmap!(op, ps; limit = substitution_limit)
     if is_split(sys)
-        p = MTKParameters(sys, op; floatT = floatT)
+        p = MTKParameters(sys, op; floatT = floatT, container_type = pType)
     else
         p = better_varmap_to_vars(op, ps; tofloat, container_type = pType)
     end

test/initial_values.jl

@@ -2,6 +2,7 @@ using ModelingToolkit
 using ModelingToolkit: t_nounits as t, D_nounits as D, get_u0
 using OrdinaryDiffEq
 using DataInterpolations
+using StaticArrays
 using SymbolicIndexingInterface: getu
 
 @variables x(t)[1:3]=[1.0, 2.0, 3.0] y(t) z(t)[1:2]
@@ -281,3 +282,28 @@ end
     @test prob.p isa Vector{Float64}
     @test length(prob.p) == 5
 end
+
+@testset "MTKParameters uses given `pType` for inner buffers" begin
+    @parameters σ ρ β
+    @variables x(t) y(t) z(t)
+
+    eqs = [D(D(x)) ~ σ * (y - x),
+        D(y) ~ x * (ρ - z) - y,
+        D(z) ~ x * y - β * z]
+
+    @mtkbuild sys = ODESystem(eqs, t)
+
+    u0 = SA[D(x) => 2.0f0,
+        x => 1.0f0,
+        y => 0.0f0,
+        z => 0.0f0]
+
+    p = SA[σ => 28.0f0,
+        ρ => 10.0f0,
+        β => 8.0f0 / 3]
+
+    tspan = (0.0f0, 100.0f0)
+    prob = ODEProblem(sys, u0, tspan, p)
+    @test prob.p.tunable isa SVector
+    @test prob.p.initials isa SVector
+end

test/mtkparameters.jl

@@ -1,8 +1,8 @@
 using ModelingToolkit
 using ModelingToolkit: t_nounits as t, D_nounits as D, MTKParameters
-using SymbolicIndexingInterface
+using SymbolicIndexingInterface, StaticArrays
 using SciMLStructures: SciMLStructures, canonicalize, Tunable, Discrete, Constants
-using BlockArrays: BlockedArray, Block
+using BlockArrays: BlockedArray, BlockedVector, Block
 using OrdinaryDiffEq
 using ForwardDiff
 using JET
@@ -27,6 +27,15 @@ end
 @test getp(sys, a)(ps) == getp(sys, b)(ps) == getp(sys, c)(ps) == 0.0
 @test getp(sys, d)(ps) isa Int
 
+@testset "`container_type`" begin
+    ps2 = MTKParameters(sys, ivs; container_type = SVector)
+    @test ps2.tunable isa SVector
+    @test ps2.initials isa SVector
+    @test ps2.discrete isa Tuple{<:BlockedVector{Float64, <:SVector}}
+    @test ps2.constant isa Tuple{<:SVector, <:SVector, <:SVector{1, <:SMatrix}}
+    @test ps2.nonnumeric isa Tuple{<:SVector}
+end
+
 ivs[a] = 1.0
 ps = MTKParameters(sys, ivs)
 for (p, val) in ivs