Help porting to @mtkmodel DSL and a structural_simplify error

Specific Domains Modelling & Simulations
1

I have an MTK model (ported from the epidemics R package where I build the model using a loop. How are symbolic arrays dealt with in the @mtkmodel DSL? I’ve tried a few things, but I can’t find what I need to port this model over.

using Symbolics
using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using LinearAlgebra
using OrdinaryDiffEq
using Plots

K = 3 # Number of age groups
@parameters σ γ N[1:K] B[1:K, 1:K] # exposure, recovery, population sizes, and contact matrix
@variables S(t)[1:K] E(t)[1:K] I(t)[1:K] R(t)[1:K] λ(t)[1:K] # SEIR and force of infection

eqs = Vector{ModelingToolkit.Equation}()
for k in 1:K
    push!(eqs, λ[k] ~ sum([B[i,k]*I[i]/N[i] for i in 1:K]))
    push!(eqs, D(S[k]) ~ -λ[k] * S[k])
    push!(eqs, D(E[k]) ~ λ[k] * S[k] - σ*E[k])
    push!(eqs, D(I[k]) ~ σ*E[k] - γ*I[k])
    push!(eqs, D(R[k]) ~ γ*I[k])
end

@mtkbuild sys = ODESystem(eqs, t) 
#sys = structural_simplify(sys) # This fails

uktotal = [14799290, 16526302, 28961159]
i₀ = 1e-6
inits = [1.0-i₀, 0.0, i₀, 0.0]
ukpop = hcat(repeat([uktotal], 4)...)
u0_vector = vec(hcat(repeat([inits], K)...)' .* ukpop)
cm_orig = hcat([[7.883663, 2.794154, 1.565665],
           [3.120220, 4.854839, 2.624868],
           [3.063895, 4.599893, 5.005571]]...)'
cm_norm = cm_orig/maximum(eigvals(cm_orig))

p = [sys.B => 1.3/7 * cm_norm,
     sys.σ => 1.0/2,
     sys.γ => 1.0/7,
     sys.N => [14799290.0, 16526302.0, 28961159.0]]

u0 = [sys.S => u0_vector[1:K],
      sys.E => u0_vector[K+1:2K],
      sys.I => u0_vector[2K+1:3K],
      sys.R => u0_vector[3K+1:4K]]

prob = ODEProblem(sys, u0, (0.0, 600.0), p)
sol = solve(prob, Tsit5(), saveat=1.0)

u = sol.u
S_out = permutedims(hcat([[u[i][j] for j in [1, 5, 9]] for i in 1:length(u)]...))
E_out = permutedims(hcat([[u[i][j] for j in [2, 6, 10]] for i in 1:length(u)]...))
I_out = permutedims(hcat([[u[i][j] for j in [3, 7, 11]] for i in 1:length(u)]...))
R_out = permutedims(hcat([[u[i][j] for j in [4, 8, 12]] for i in 1:length(u)]...))
plot(sol.t, I_out, label="Infectious", ylim=(0,500000))
plot!(sol.t, E_out, label="Exposed")

In addition, structural_simplify throws an error on the system, which I found unexpected.

julia> structural_simplify(sys) 
ERROR: ExtraVariablesSystemException: The system is unbalanced. There are 15 highest order derivative variables and 12 equations.
More variables than equations, here are the potential extra variable(s):
 (S(t))[1]
 (E(t))[1]
 (I(t))[1]
 (R(t))[1]
 (S(t))[2]
 (E(t))[2]
 (I(t))[2]
 (R(t))[2]
 (S(t))[3]
 (E(t))[3]
 (I(t))[3]
 (R(t))[3]
 (λ(t))[1]
 (λ(t))[2]
 (λ(t))[3]
Stacktrace:
  [1] error_reporting(state::TearingState{…}, bad_idxs::Vector{…}, n_highest_vars::Int64, iseqs::Bool, orig_inputs::Set{…})
    @ ModelingToolkit.StructuralTransformations ~/.julia/packages/ModelingToolkit/cAhZr/src/structural_transformation/utils.jl:50
  [2] check_consistency(state::TearingState{ODESystem}, orig_inputs::Set{Any})
    @ ModelingToolkit.StructuralTransformations ~/.julia/packages/ModelingToolkit/cAhZr/src/structural_transformation/utils.jl:85
  [3] _structural_simplify!(state::TearingState{…}, io::Nothing; simplify::Bool, check_consistency::Bool, fully_determined::Bool, warn_initialize_determined::Bool, dummy_derivative::Bool, kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/systemstructure.jl:689
  [4] structural_simplify!(state::TearingState{…}, io::Nothing; simplify::Bool, check_consistency::Bool, fully_determined::Bool, warn_initialize_determined::Bool, kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/systemstructure.jl:635
  [5] __structural_simplify(sys::ODESystem, io::Nothing; simplify::Bool, kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/systems.jl:0
  [6] __structural_simplify
    @ ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/systems.jl:55 [inlined]
  [7] structural_simplify(sys::ODESystem, io::Nothing; simplify::Bool, split::Bool, kwargs::@Kwargs{})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/systems.jl:22
  [8] structural_simplify
    @ ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/systems.jl:19 [inlined]
  [9] structural_simplify(sys::ODESystem)
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/systems.jl:19
 [10] top-level scope
    @ REPL[1]:1
2

If you already simplified via @mtkbuild, you should not simplify a second time.

3

Thanks @ChrisRackauckas! I didn’t know @mtkbuild also simplified. It was a bit confusing to have an error though (rather than a graceful ‘this is already simplified’).

Any tips on the porting? I’ve tried for loops in the @equations block, but that doesn’t seem to work.

4

I can add a better error to that.

I’m not sure I understand this question.

5

I think this is almost the same model, B is transposed to make the matrix multiplication simpler (adding an explicit permutedims doesn’t seem to work with symbolic matrices: ERROR: Don't know how to index by CartesianIndex()).

Looks like you can’t use structural_parameters to control the size of arrays so I’m not sure how to port that idea (LoadError: promotion of types Int64 and Symbol failed to change any argument)

The loops aren’t necessary, just create the array equations and let MTK scalarize.

@mtkmodel Example begin
    @structural_parameters begin
        K=3
    end
    @parameters begin
        σ=1.0/2, [description="exposure"]
        γ=1.0/7, [description="recovery"]
        N[1:3], [description="population sizes"]
        B[1:3, 1:3], [description="contact matrix"]
    end
    @variables begin
        S(t)[1:3]
        E(t)[1:3]
        I(t)[1:3]
        R(t)[1:3]
        λ(t)[1:3]
    end
    @equations begin
        λ ~ (B * I) ./ N
        D(S) ~ -λ .* S
        D(E) ~ λ .* S .- σ .* E
        D(I) ~ σ .* E - λ .* I
        D(R) ~ γ .* I
    end
end
6

Thanks for the tips @contradict! It would be useful to have permutedims working, as well as being able to specify the size of arrays (one for @ChrisRackauckas). While the loops aren’t necessary, I would like to know how to specify them explicitly - is this possible?

7

I think the only things that can be in an @equations block are a list of equations or a conditional that evaluates to a list of equations.

Do you have an example of something you would like to do with loops that isn’t expressible with array variables?