Parameter splitting breaks runtime generated jacobians

[david-hofmann]

🐞: Setting split=true breaks the call of a runtime generated Jacobian:

Minimal Reproducible Example 👇

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

sys = let
           @parameters σ ρ β a::Bool
           @variables x(t) y(t) z(t)

           eqs = [D(x) ~ a * σ * (y - x),
                  D(y) ~ x * (ρ - z) - y,
                  D(z) ~ x * y - β * z]

           @named sys = ODESystem(eqs, t)
           structural_simplify(sys; split=true)
end;

jac = generate_jacobian(sys, expression=Val{false})[1]
jac(unknowns(sys), parameters(sys), t)

Error & Stacktrace ⚠️

ERROR: BoundsError: attempt to access Tuple{Vector{SymbolicUtils.BasicSymbolic{Real}}, Vector{Any}, Num} at index [4]
Stacktrace:
 [1] getindex
   @ ./tuple.jl:31 [inlined]
 [2] macro expansion
   @ ~/.julia/packages/RuntimeGeneratedFunctions/M9ZX8/src/RuntimeGeneratedFunctions.jl:162 [inlined]
 [3] macro expansion
   @ ./none:0 [inlined]
 [4] generated_callfunc
   @ ./none:0 [inlined]
 [5] (::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{…})(::Vector{…}, ::Vector{…}, ::Num)
   @ RuntimeGeneratedFunctions ~/.julia/packages/RuntimeGeneratedFunctions/M9ZX8/src/RuntimeGeneratedFunctions.jl:150
 [6] top-level scope
   @ REPL[42]:1
Some type information was truncated. Use `show(err)` to see complete types.

Maybe relevant to note that with split=false the generated Jacobian has 3 parameters:

RuntimeGeneratedFunction(#=in Symbolics=#, #=using Symbolics=#, :((ˍ₋arg1, ˍ₋arg2, t)->#= /home/david/.julia/packages/Symbolics/PxO3a/src/build_function.jl:340 =# @inbounds(begin
              #= /home/david/.julia/packages/Symbolics/PxO3a/src/build_function.jl:340 =#
              begin
                  #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:385 =#
                  #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:386 =#
                  #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:387 =#
                  begin
                      begin
                          begin
                              begin
                                  begin
                                      #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:480 =#
                                      (SymbolicUtils.Code.create_array)(typeof(ˍ₋arg1), nothing, Val{2}(), Val{(3, 3)}(), (*)((*)(-1, ˍ₋arg2[1]), ˍ₋arg2[2]), (+)((*)(-1, ˍ₋arg1[3]), ˍ₋arg2[3]), ˍ₋arg1[2], (*)(ˍ₋arg2[1], ˍ₋arg2[2]), -1, ˍ₋arg1[1], 0, (*)(-1, ˍ₋arg1[1]), (*)(-1, ˍ₋arg2[4]))
                                  end
                              end
                          end
                      end
                  end
              end
          end)))

whereas with split=true it has 4:

RuntimeGeneratedFunction(#=in Symbolics=#, #=using Symbolics=#, :((ˍ₋arg1, ˍ₋arg2, ˍ₋arg3, t)->#= /home/david/.julia/packages/Symbolics/PxO3a/src/build_function.jl:340 =# @inbounds(begin
              #= /home/david/.julia/packages/Symbolics/PxO3a/src/build_function.jl:340 =#
              begin
                  #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:385 =#
                  #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:386 =#
                  #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:387 =#
                  begin
                      begin
                          begin
                              begin
                                  begin
                                      #= /home/david/.julia/packages/SymbolicUtils/jf8aQ/src/code.jl:480 =#
                                      (SymbolicUtils.Code.create_array)(typeof(ˍ₋arg1), nothing, Val{2}(), Val{(3, 3)}(), (*)((*)(-1, ˍ₋arg3[1]), ˍ₋arg2[3]), (+)((*)(-1, ˍ₋arg1[3]), ˍ₋arg2[1]), ˍ₋arg1[2], (*)(ˍ₋arg3[1], ˍ₋arg2[3]), -1, ˍ₋arg1[1], 0, (*)(-1, ˍ₋arg1[1]), (*)(-1, ˍ₋arg2[2]))
                                  end
                              end
                          end
                      end
                  end
              end
          end)))

It would be nice if the Jacobian treated the parameters the same way as parameters(sys) does. If that is not possible then having a function that returns the parameters in their split structure would help.

Thank you!!

[vyudu]

Creating an ODEProblem from the system and using prob.p gets the parameters in their split form, jac(unknowns(sys), prob.p, t), does this work for your workflow?

EDIT: constructing a MTKParameters(sys, p) also works and is probably preferable to the ODEProblem

[AayushSabharwal]

We also have calculate_jacobian(sys) to get a symbolic jacobian

[david-hofmann]

Thank you very much! MTKParameters solved it also in my workflow!

Just a note: calculate_jacobian(sys) decreases performance since we need to update parameter values on runtime. This is why we need to use generate_jacobian.