How to specify vector MTK states or parameters to ODESystem call? (ODESystem(eq,t,var,par))

Hey folks,

I’m working to implement a simple vectorized step function using callbacks as a learning exercise. I was able to create a simple scalar step function:

function step(time;name)    
    x = @variables u(t)=0 [output = true,irreducible = true]    
    p = @parameters time=time height=0        
    eq = [0 ~ u - height]
    ODESystem(eq, t, x, p; name=name, continuous_events = [(t - time ~ 0) => (height ~ 1)])
end

@named s = step(1)
_s = structural_simplify(s)
prob = ODEProblem(_s,[],(0.,10.))
sol = solve(prob)

This seemed to work okay:

julia> sol = solve(prob)
retcode: Success
Interpolation: specialized 3rd order "free" stiffness-aware interpolation
t: 11-element Vector{Float64}:
  0.0
  1.0e-6
  1.1e-5
  0.00011099999999999999
  0.0011109999999999998
  0.011110999999999996
  0.11111099999999996
  0.9999999999999999
  0.9999999999999999
  1.9999999999999996
 10.0
u: 11-element Vector{Vector{Float64}}:
 [0.0]
 [0.0]
 [0.0]
 [0.0]
 [0.0]
 [0.0]
 [0.0]
 [0.0]
 [1.0000000000039475]
 [1.0]
 [1.0]

However, if I don’t specify states and parameters in my ODESystem call, parameter “time” does not show up as a parameter. This seems to be because time is not part of eq (though it is part of the callbacks).

function step(time;name)    
    x = @variables u(t)=0 [output = true,irreducible = true]    
    p = @parameters time=time height=0        
    eq = [0 ~ u - height]
    ODESystem(eq, t; name=name, continuous_events = [(t - time ~ 0) => (height ~ 1)])
end

julia> @named s = step(1)
Model s with 1 equations
States (1):
  u(t) [defaults to 0]
Parameters (1):
  height [defaults to 0]

I tested this by arbitrarily adding "time " into the eq.

function step(time;name)    
    x = @variables u(t)=0 [output = true,irreducible = true]    
    p = @parameters time=time height=0        
    eq = [0 ~ u - height + time]
    ODESystem(eq, t; name=name, continuous_events = [(t - time ~ 0) => (height ~ 1)])
end

julia> @named s = step(1)
Model s with 1 equations
States (1):
  u(t) [defaults to 0]
Parameters (2):
  height [defaults to 0]
  time [defaults to 1]

Solving the system without time as a parameter expectedly produces an error:

julia> sol = solve(prob)
ERROR: MethodError: no method matching *(::Int64, ::typeof(time))
Closest candidates are:
  *(::Any, ::Any, ::Any, ::Any...) at operators.jl:560
  *(::T, ::T) where T<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8} at int.jl:88
  *(::Union{AbstractIrrational, AbstractFloat, Integer, Rational}, ::Static.StaticInt{0})

Now, if I move to the slightly more complicated vectorized step function, I run into the following problem:

function steps(times;name)        
    n = length(times)
    x = @variables u(t)[1:n]=zeros(n) [output = true,irreducible = true]
    p = @parameters times[1:n]=times heights[1:n]=zeros(n)        

    eq = [0 ~ u[i] - heights[i] for i in 1:n]
        
    cond = [t - times[i] ~ 0 for i in 1:n]
    affect = [heights[i] ~ 1 for i in 1:n]
    ODESystem(eq,t,x,p; name=name, continuous_events = cond => affect)
end
julia> @named s2 = steps([1,3])
ERROR: MethodError: no method matching hasmetadata(::Vector{Num}, ::Type{Symbolics.VariableDefaultValue})
Closest candidates are:
  hasmetadata(::SymbolicUtils.Symbolic, ::Any) at ..\.julia\packages\SymbolicUtils\qulQp\src\types.jl:15
  hasmetadata(::Complex{Num}, ::Any) at ..\.julia\packages\Symbolics\FGTCH\src\Symbolics.jl:163
  hasmetadata(::Num, ::Any) at ..\.julia\packages\Symbolics\FGTCH\src\Symbolics.jl:163
Stacktrace:

Seems to fail because hasmetadata has not been extended to Vector{Symbolics.Arr}

julia> x = @variables u(t)[1:2]=zeros(2) [output = true,irreducible = true]
1-element Vector{Symbolics.Arr{Num, 1}}:
 (u(t))[1:2]

Removing x and p from the ODESystem call expectedly (from the scalar case) removes “times” from the parameters of the system, and the error more blatantly tells me that “times” is not defined.

function steps(times;name)        
    n = length(times)
    x = @variables u(t)[1:n]=zeros(n) [output = true,irreducible = true]
    p = @parameters times[1:n]=times heights[1:n]=zeros(n)        

    eq = [0 ~ u[i] - heights[i] for i in 1:n]
        
    cond = [t - times[i] ~ 0 for i in 1:n]
    affect = [heights[i] ~ 1 for i in 1:n]
    ODESystem(eq,t; name=name, continuous_events = cond => affect)
end
julia> @named s2 = steps([1,3])
Model s2 with 2 equations
States (2):
  (u(t))[1] [defaults to 0.0]
  (u(t))[2] [defaults to 0.0]
Parameters (2):
  heights[1] [defaults to 0.0]
  heights[2] [defaults to 0.0]

julia> _s2 = structural_simplify(s2)
Model s2 with 2 equations
States (2):
  (u(t))[1] [defaults to 0.0]
  (u(t))[2] [defaults to 0.0]
Parameters (2):
  heights[1] [defaults to 0.0]
  heights[2] [defaults to 0.0]
Incidence matrix:2×2 SparseArrays.SparseMatrixCSC{Num, Int64} with 2 stored entries:
 ×  ⋅
 ⋅  ×

julia> prob = ODEProblem(_s2,[],(0.,10.))
ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 10.0)
u0: 2-element Vector{Float64}:
 0.0
 0.0

julia> sol = solve(prob)
ERROR: UndefVarError: times not defined
Stacktrace:

TL;DR, I can successfully create a continuous_event based scalar step function, but only if I supply states and variables in the ODESystem call. Attempting the same for a vector step function with vector states and parameters fails due to Symbolics.hasmetadata.

From the tutorials, I understand that vectorization in MTK is still immature. I’m wondering if there’s a workaround (or of course if my simple brain is just missing something obvious) for supplying vector states and parameters to the ODESystem to get around hasmetadata?

It should be ODESystem(eq,t,[x...;],[p...;]; name=name, continuous_events = cond => affect) because x is a vector of vectors and p is a vector of vectors. Maybe MTK should do this automatically.

Works perfectly. Agree that it’s probably more intuitive to just supply x,p and let MTK handle it, but this form gets the job done. Thanks for the help!