How can I implement a vector of StaticArrays as a variable in MTK?

Hi everyone. After trying for weeks I would like to ask the community if what I’m trying to do is possible or whether I have arrived to a dead end.

I’ve reduced it to a MWE. The system I want to solve in ModelingToolkit has two array variables, r and H. Each element of both arrays is a SMatrix from the StaticArrays package. First I define some vectors and initialize r to the value r0:

using ModelingToolkit, StaticArrays

@independent_variables x
D = Differential(x)

# Define some vectors and matrices
n = 5
ys = @SVector [0.26, 1.41, 3.6, 7.09, 12.64]
u = @SMatrix [0.93 0.31; -0.14 + 0.02im 0.81]

# Initialize variable
r0 = [SHermitianCompact{2, ComplexF64, 3}([exp(-ys[i]), 0, exp(-ys[i])]) for i in eachindex(ys)]

vars = @variables begin
    r(x)[1:n] = r0
    H(x)[1:n]
end

Next, I define a function called Hvar. This function takes several parameters, makes basic operations with matrices and returns the value of what the H variable should have. Notice that I want element-wise multiplication in the return value. I want each SMatrix of H be multiplied by each SMatrix of r.

# This computes the value of H as a function of other parameters
function Hvar(r, x, ys, U)
    H = [x^2 * U / ys[i] for i in eachindex(ys)]

    return H .* r
end

eqs = [
    H ~ Hvar(r, x, ys, u)
    D(r) ~ H
]

_sys = System(eqs, x, vars, []; name = :foo)
sys = mtkcompile(_sys)

When I try to compile the system with mtkcompile I get the following error:

ERROR: TypeError: in typeassert, expected Array{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymReal}}, got a value of type SizedVector{5, SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymReal}, Vector{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymReal}}}

I’ve tried converting the variables to floats with the .value field and then create a Num array back but it doesn’t work either.

Is there a way to define a function that takes a vector of SMatrix and other symbolic variables as input, do some work and return another vector of SMatrix? Note that my functions are general and outside of ModelingToolkit, they work without problems normal Julia arrays of with StaticArrays.

Did you try using @register_symbolic, see: Function Registration and Tracing · Symbolics.jl ?

Yes, I tried with either @register_symbolic or @register_array_symbolic and the error persisted.

For example:

@register_array_symbolic Hvar(r::AbstractVector, x, ys, U) begin
    size = size(r)
    eltype=eltype(r)
end

produces

ERROR: TypeError: in typeassert, expected Array{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymReal}}, got a value of type SizedVector{5, SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymReal}, Vector{SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymReal}}}