feat: support symbolic arrays in @mtkmodels

src/systems/model_parsing.jl

@@ -83,7 +83,7 @@ function _model_macro(mod, name, expr, isconnector)
     :($name = $Model($f, $dict, $isconnector))
 end
 
-function parse_variable_def!(dict, mod, arg, varclass, kwargs, def = nothing)
+function parse_variable_def!(dict, mod, arg, varclass, kwargs; def = nothing, indices = 1:1)
     metatypes = [(:connection_type, VariableConnectType),
         (:description, VariableDescription),
         (:unit, VariableUnit),
@@ -104,20 +104,21 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, def = nothing)
     MLStyle.@match arg begin
         a::Symbol => begin
             push!(kwargs, Expr(:kw, a, nothing))
-            var = generate_var!(dict, a, varclass)
+            var = generate_var!(dict, a, varclass; indices)
             dict[:kwargs][getname(var)] = def
-            (var, nothing)
+            (var, def)
         end
         Expr(:call, a, b) => begin
             push!(kwargs, Expr(:kw, a, nothing))
-            var = generate_var!(dict, a, b, varclass)
+            var = generate_var!(dict, a, b, varclass; indices)
             dict[:kwargs][getname(var)] = def
-            (var, nothing)
+            (var, def)
         end
         Expr(:(=), a, b) => begin
             Base.remove_linenums!(b)
             def, meta = parse_default(mod, b)
-            var, _ = parse_variable_def!(dict, mod, a, varclass, kwargs, def)
+            var, def = parse_variable_def!(dict, mod, a, varclass, kwargs; def)
+            @info 128 var def meta
             dict[varclass][getname(var)][:default] = def
             if meta !== nothing
                 for (type, key) in metatypes
@@ -142,31 +143,33 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, def = nothing)
                 end
                 var = set_var_metadata(var, meta)
             end
-            (set_var_metadata(var, meta), def)
+            (var, def)
+        end
+        Expr(:ref, a, b) => begin
+            parse_variable_def!(dict, mod, a, varclass, kwargs; def, indices = eval(b))
         end
         _ => error("$arg cannot be parsed")
     end
 end
 
-function generate_var(a, varclass)
-    var = Symbolics.variable(a)
+
+function generate_var(a, varclass; indices = 1:1)
+    var = length(indices) > 1 ? (@variables $a[indices]) : Symbolics.variable(a)
     if varclass == :parameters
         var = toparam(var)
     end
-    var
+    var isa Vector ? var[1] : var
 end
 
-function generate_var!(dict, a, varclass)
-    #var = generate_var(Symbol("#", a), varclass)
-    var = generate_var(a, varclass)
+function generate_var!(dict, a, varclass; indices::UnitRange = 1:1)
     vd = get!(dict, varclass) do
         Dict{Symbol, Dict{Symbol, Any}}()
     end
     vd[a] = Dict{Symbol, Any}()
-    var
+    generate_var(a, varclass; indices)
 end
 
-function generate_var!(dict, a, b, varclass)
+function generate_var!(dict, a, b, varclass; indices::UnitRange = 1:1)
     iv = generate_var(b, :variables)
     prev_iv = get!(dict, :independent_variable) do
         iv
@@ -176,11 +179,15 @@ function generate_var!(dict, a, b, varclass)
         Dict{Symbol, Dict{Symbol, Any}}()
     end
     vd[a] = Dict{Symbol, Any}()
-    var = Symbolics.variable(a, T = SymbolicUtils.FnType{Tuple{Real}, Real})(iv)
+    var = if length(indices) > 1
+        @variables $a(iv)[indices]
+    else
+        Symbolics.variable(a, T = SymbolicUtils.FnType{Tuple{Real}, Real})(iv)
+    end
     if varclass == :parameters
         var = toparam(var)
     end
-    var
+    var isa Vector ? var[1] : var
 end
 
 function parse_default(mod, a)
@@ -215,7 +222,7 @@ end
 
 function set_var_metadata(a, ms)
     for (m, v) in ms
-        a = setmetadata(a, m, v)
+        a = a isa Symbolics.Arr ? collect(setmetadata.(a, m, v)) : setmetadata(a, m, v)
     end
     a
 end
@@ -406,11 +413,16 @@ end
 
 function parse_variable_arg!(expr, vs, dict, mod, arg, varclass, kwargs)
     vv, def = parse_variable_def!(dict, mod, arg, varclass, kwargs)
-    v = Num(vv)
-    name = getname(v)
-    push!(vs, name)
-    push!(expr.args,
-        :($name = $name === nothing ? $setdefault($vv, $def) : $setdefault($vv, $name)))
+    if !(vv isa Union{Symbolics.Arr, Vector})
+        vv = setdefault(vv, def)
+        name = getname(Num(vv))
+        push!(expr.args, :($name = $name === nothing ? $vv : $setdefault($vv, $name)))
+        push!(vs, name)
+    else
+        name = vv isa Vector ? getname(vv[1]) : getname(vv)
+        push!(expr.args, :($name = $name === nothing ? $setdefault.($vv, $def) : $setdefault.($vv, $name)))
+        push!(vs, :($name...))
+    end
 end
 
 function parse_variables!(exprs, vs, dict, mod, body, varclass, kwargs)

test/model_parsing.jl

@@ -173,15 +173,19 @@ end
     @mtkmodel MockModel begin
         @parameters begin
             a
+            a2[1:2]
             b(t)
+            b2(t)[1:2]
             cval
             jval
             kval
             c(t) = cval + jval
             d = 2
             e, [description = "e"]
+            e2[1:2], [description = "e2"]
             f = 3, [description = "f"]
             h(t), [description = "h(t)"]
+            h2(t)[1:2], [description = "h(t)"]
             i(t) = 4, [description = "i(t)"]
             j(t) = jval, [description = "j(t)"]
             k = kval, [description = "k"]
@@ -193,7 +197,12 @@ end
     end
 
     kval = 5
-    @named model = MockModel(; kval, cval = 1, func = identity)
+    @named model = MockModel(; b2 = 3, kval, cval = 1, func = identity)
+
+    @test lastindex(parameters(model)) == 21
+    @test lastindex(model.b2) == 2
+    @test getdefault.(model.b2) .== 3
+    @test getmetadata(model.e2[1], VariableDescription) == getmetadata(model.e2[2], VariableDescription) == "e2"
 
     @test hasmetadata(model.e, VariableDescription)
     @test hasmetadata(model.f, VariableDescription)
@@ -209,6 +218,8 @@ end
     @test_throws KeyError getdefault(model.e)
     @test getdefault(model.f) == 3
     @test getdefault(model.i) == 4
+    @test getdefault(model.b2[1]) == 3
+    @test getdefault(model.b2[2]) == 3
     @test isequal(getdefault(model.j), model.jval)
     @test isequal(getdefault(model.k), model.kval)
 end