Big slowdown of ODE with matrix equations

[hersle]

This "matrix ODE" is much slower and allocates much more than it used to:

using ModelingToolkit, OrdinaryDiffEq
using ModelingToolkit: t_nounits as t, D_nounits as D

@variables x(t)[1:2,1:1]
eqs = [D(x[1,1]) ~ x[1,1], x[2,1] ~ 0.0]
@named sys = System(eqs, t)
ssys = mtkcompile(sys)
prob = ODEProblem(ssys, [x[1,1] => 1.0], (0.0, 1.0))
@time sol = solve(prob, Tsit5(); dt = 1e-5, adaptive = false, save_everystep = false)

0.031453 seconds (1.80 M allocations: 64.097 MiB, 45.73% gc time)

For comparison, here is the equivalent "scalar ODE":

@variables x11(t) x21(t)
eqs = [D(x11) ~ x11, x21 ~ 0.0]
@named sys = System(eqs, t)
ssys = mtkcompile(sys)
prob = ODEProblem(ssys, [x11 => 1.0], (0.0, 1.0))
@time sol = solve(prob, Tsit5(); dt = 1e-5, adaptive = false, save_everystep = false)

0.004658 seconds (144 allocations: 8.594 KiB)

This was for a 2x1 matrix. The problem is much worse for larger sizes.

[hersle]

Profiling (after running the first code block above):

using ProfileCanvas
@profview sol = solve(prob, Tsit5(); dt = 1e-5, adaptive = false, save_everystep = false)

The allocations come from create_array in the generated ODE function:

prob.f.f.f_iip # see var"##cse#5" = ...

RuntimeGeneratedFunction(#=in ModelingToolkit=#, #=using ModelingToolkit=#, :((ˍ₋out, __mtk_arg_1, ___mtkparameters___, t)->#= /home/hermasl/.julia/packages/Symbolics/JCopU/src/build_function.jl:368 =# @inbounds(begin
              #= /home/hermasl/.julia/packages/Symbolics/JCopU/src/build_function.jl:368 =#
              begin
                  #= /home/hermasl/.julia/packages/SymbolicUtils/aooYZ/src/code.jl:409 =#
                  #= /home/hermasl/.julia/packages/SymbolicUtils/aooYZ/src/code.jl:410 =#
                  #= /home/hermasl/.julia/packages/SymbolicUtils/aooYZ/src/code.jl:411 =#
                  begin
                      __mtk_arg_2 = ___mtkparameters___[1]
                      begin
                          var"##cse#1" = (view)(__mtk_arg_2, 1:2)
                          var"##cse#2" = (reshape)(var"##cse#1", (2, 1))
                          var"Initial(x(t))" = var"##cse#2"
                          var"##cse#3" = (view)(__mtk_arg_2, 3:4)
                          var"##cse#4" = (reshape)(var"##cse#3", (2, 1))
                          var"Initial(xˍt(t))" = var"##cse#4"
                          var"(x(t))[2, 1]" = 0.0
                          var"##cse#5" = begin
                                  #= /home/hermasl/.julia/packages/SymbolicUtils/aooYZ/src/code.jl:510 =#
                                  (SymbolicUtils.Code.create_array)(Matrix{SymbolicUtils.BasicSymbolic{Real}}, nothing, Val{2}(), Val{(2, 1)}(), __mtk_arg_1[1], var"(x(t))[2, 1]")
                              end
                          var"##cse#6" = (ModelingToolkit.StructuralTransformations.change_origin)(CartesianIndex(1, 1), var"##cse#5")
                          var"x(t)" = var"##cse#6"
                          #= /home/hermasl/.julia/packages/SymbolicUtils/aooYZ/src/code.jl:464 =# @inbounds begin
                                  #= /home/hermasl/.julia/packages/SymbolicUtils/aooYZ/src/code.jl:460 =#
                                  ˍ₋out[1] = __mtk_arg_1[1]
                                  #= /home/hermasl/.julia/packages/SymbolicUtils/aooYZ/src/code.jl:462 =#
                                  ˍ₋out
                              end
                      end
                  end
              end
          end)))

[ChrisRackauckas]

So this is known, but a bit necessary as an intermediate step. Think of just a matrix multiplication. If you scalarize the equations, then:

du[1] = A[1,1]*x[1] + A[2,1]*x[2]
du[2] = A[1,2]*x[1] + A[2,2]*x[2]

That's not allocating

tmp = A*x
du .= tmp

That is allocating. But, it will also be O(1) in compile time. So we got array functions working and we're keeping in that direction. Though we probably need to start setting up the functions with a bump allocator or something.

That said, the "workaround" for now is just to manually scalarize. The "regression" is now that isn't done automatically, so you can just force it if that's the behavior you need. But down the line, we'll be improving Symbolics.jl's build_function for this kind of behavior, using a bump allocator and Reactant.jl

[hersle]

Thank you, I see and think that is an excellent direction to go.

By scalarizing manually, do you mean to only define scalar variables and write out all the equations for every index? Or is there a way to "scalarize-transform" the equations/system after creating it with array variables? My equations are created by looping over the indices with a user-configurable maximum index, so I really wish for some automatic handling.

I do not have success with e.g.

sys = ModelingToolkit.scalarize(sys) # this would be very convenient

or

@named sys = System(ModelingToolkit.scalarize(eqs), t, vcat(ModelingToolkit.scalarize(x)...), [])

[ChrisRackauckas]

We can probably make an overload that does that. It wouldn't be so hard to do.

[RobbesU]

Rather than manually scalarizing, doesn't the collect( D(u) .~ A*x)... still work for this? Seems to for me.

[hersle]

I think that would work if my equations were of that form, with u being fully unknown. But I have equations for x where x[1,1] is unknown and x[1,2] is observed, which I think may be what triggers the allocations.

[ChrisRackauckas]

Rather than manually scalarizing, doesn't the collect( D(u) .~ A*x)... still work for this? Seems to for me.

collect is scalarizing.