Skip to content

Big slowdown of ODE with matrix equations #3708

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hersle opened this issue Jun 7, 2025 · 7 comments
Open

Big slowdown of ODE with matrix equations #3708

hersle opened this issue Jun 7, 2025 · 7 comments
Labels
bug Something isn't working

Comments

@hersle
Copy link
Contributor

hersle commented Jun 7, 2025

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 hersle added the bug Something isn't working label Jun 7, 2025
@hersle
Copy link
Contributor Author

hersle commented Jun 7, 2025

Profiling (after running the first code block above):

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

Image

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
Copy link
Member

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
Copy link
Contributor Author

hersle commented Jun 8, 2025

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
Copy link
Member

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

@RobbesU
Copy link

RobbesU commented Jun 9, 2025

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

@hersle
Copy link
Contributor Author

hersle commented Jun 9, 2025

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
Copy link
Member

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

collect is scalarizing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

No branches or pull requests

3 participants