Indexing and iterating in ODEProblem and modelingtoolkitize / error messages

[DanielVandH]

I am working on some complicated 2D finite volume code and am trying to refactor it to allow for modelingtoolkitize. In doing so I am having some difficulties dealing with iteration and indexing, and the error messages don't really seem so clear to me in how I can fix it.

The following code covers the type of iteration I am dealing with:

using DifferentialEquations, ModelingToolkit, Symbolics
@inline do_something(k) = sin(k^2)
@register_symbolic do_something(k) 
@inline lookup_val(dict, idx) = dict[idx] # getindex
@register_symbolic lookup_val(dict, idx)
@inline set_val!(dict, val, idx) = dict[idx] = val # setindex!
@register_symbolic set_val!(dict, val, idx)
function ode_f_sym(du, u, p, t)
    vals, u_idx, itr_idx = p 
    for (i, j) in itr_idx
        set_val!(du, lookup_val(vals, i) * lookup_val(u, j), i)
    end
    for k in u_idx 
        set_val!(du, lookup_val(du, k) + do_something(lookup_val(u, k)), k)
    end
    nothing 
end 
a, b, c, d, e = 1.7, 2.3, -2.0, 10.0, 7.19
u_idx = [2, 3, 5, 4, 1] 
itr_idx = Dict(1 => 3, 2 => 5, 3 => 1, 4 => 4, 5 => 2)
prob = ODEProblem(ode_f_sym, [0.0, 1.0, 2.3, -2.0, 1.0], (0.0, 1.0), ((a, b, c, d, e), u_idx, itr_idx))
modelingtoolkitize(prob)

I define additional functions for getindex and setindex! so that they are not traced (as someone helped me with here https://discourse.julialang.org/t/modelingtoolkit-indexing-a-parameter-by-another-parameter-in-an-odeproblem/83569/2?u=legola18), along with some additional do_something function. When I run modelingtoolkitize as shown, I get:

ERROR: BoundsError: attempt to access Num at index [2]
Stacktrace:
 [1] indexed_iterate(I::Num, i::Int64, state::Nothing)
   @ Base .\tuple.jl:98
 [2] ode_f_sym(du::Vector{Num}, u::Vector{Num}, p::Tuple{Num, Num, Num}, t::Num)
   @ Main c:\Users\licer\.julia\dev\FVM\dev\sparsity.jl:104
 [3] (::ODEFunction{true, typeof(ode_f_sym), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing})(::Vector{Num}, ::Vararg{Any})
   @ SciMLBase C:\Users\licer\.julia\packages\SciMLBase\IJbT7\src\scimlfunctions.jl:1624
 [4] modelingtoolkitize(prob::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Tuple{NTuple{5, Float64}, Vector{Int64}, Dict{Int64, Int64}}, ODEFunction{true, typeof(ode_f_sym), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ ModelingToolkit C:\Users\licer\.julia\packages\ModelingToolkit\tMgaW\src\systems\diffeqs\modelingtoolkitize.jl:48
 [5] modelingtoolkitize(prob::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Tuple{NTuple{5, Float64}, Vector{Int64}, Dict{Int64, Int64}}, ODEFunction{true, typeof(ode_f_sym), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem})
   @ ModelingToolkit C:\Users\licer\.julia\packages\ModelingToolkit\tMgaW\src\systems\diffeqs\modelingtoolkitize.jl:7
 [6] top-level scope
   @ c:\Users\licer\.julia\dev\FVM\dev\sparsity.jl:116

How can I change the code further to allow for this type of work? It would be nice if there were e.g. a suggestion in this error message for what could be done for this type of loop, too - I know improving these error messages is a current work in progress.

Just for reference, this is what the code above looks like with hard-coded numbers (to make it easier to compare to for what 'works') and the Jacobian that I actually want to compute:

@inline do_something(k) = sin(k^2)
function ode_f_basic(du, u, p, t)
    vals= p 
    du[1] = vals[1] * u[3] 
    du[2] = vals[2] * u[5] 
    du[3] = vals[3] * u[1] 
    du[4] = vals[4] * u[4] 
    du[5] = vals[5] * u[2]
    du[2] += do_something(u[2])
    du[3] += do_something(u[3])
    du[5] += do_something(u[5])
    nothing 
end 
a, b, c, d, e = 1.7, 2.3, -2.0, 10.0, 7.19
u_idx = [2, 3, 5] 
itr_idx = Dict(1 => 3, 2 => 5, 3 => 1, 4 => 4, 5 => 2)
prob = ODEProblem(ode_f_basic, [0.0, 1.0, 2.3, -2.0, 1.0], (0.0, 1.0), ((a, b, c, d, e)))
de = modelingtoolkitize(prob)
calculate_jacobian(de)

5×5 Matrix{Num}:
  0   0                                        α₁                                         0   0
  0  Differential(x₂(t))(do_something(x₂(t)))   0                                         0  α₂
 α₃   0                                        Differential(x₃(t))(do_something(x₃(t)))   0   0
  0   0                                         0                                        α₄   0
  0  α₅                                         0                                         0  Differential(x₅(t))(do_something(x₅(t)))

[ChrisRackauckas]

https://github.com/SciML/ModelingToolkit.jl/pull/1682 at least improves the error message.

[DanielVandH]

Yeah. So is it just not expected for MTK to work with these kinds of variables, then? i.e. with a p that is a combination of tuples, Dicts, arrays, and so on. (My actual example has dicts, functions, tuples, ntuples, arrays, and dualcaches defined in p actually.. even more difficult.)

If there are some ways to work around these errors that you know of, maybe that's also useful to add/link to in the new error message.

[ChrisRackauckas]

https://github.com/SciML/DifferentialEquations.jl/issues/881 is turning to a more general solution.