Lambdifying a symbolic matrix in Julia

If I have a symbolic matrix defined as T below, is there any way to lambdify this as function of variables, say σ..., and return a matrix in a way that does not allocate (besides the allocation for the matrix that is returned)?

using SymPy, BenchmarkTools

σ = [symbols("σ_$$i$$j") for i in 1:4, j in 1:4];
T = Array{Sym}(undef,4,4);
for i in 1:4
for j in 1:4
T[i,j] = 1 + σ[i,j];
end
end


Directly using SymPy's lambdify on T results in allocation, whereas lambdify on the individual elements of T does not:

f_mat = lambdify(T, invoke_latest=false);
@benchmark \$f_mat(1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1)
BenchmarkTools.Trial:
memory estimate:  5.20 KiB
allocs estimate:  76
--------------
minimum time:     17.900 μs (0.00% GC)
median time:      19.699 μs (0.00% GC)
mean time:        25.533 μs (22.15% GC)
maximum time:     44.889 ms (99.92% GC)
--------------
samples:          10000
evals/sample:     1

f = lambdify(T[1,1], invoke_latest=false)
@benchmark f(1)
BenchmarkTools.Trial:
memory estimate:  0 bytes
allocs estimate:  0
--------------
minimum time:     13.726 ns (0.00% GC)
median time:      14.729 ns (0.00% GC)
mean time:        14.931 ns (0.00% GC)
maximum time:     75.751 ns (0.00% GC)
--------------
samples:          10000
evals/sample:     998


For context, I am trying to implement a set of ODEs with DifferentialEquations.jl and am looking for a fast way to update a matrix representing the ODEs (specifically, I wish to update the ODEs as they change in time). I could of course lambdify all the individual functions of T and simply loop over the ODEs, but my understanding is that this approach in general will scale linearly with the number of ODEs (or, equivalently, the number of elements in T). One caveat is that the individual functions in T in general could be very different, so I am not sure any vectorization can be utilized here, and so the non-vectorized loop might be the best I can do regardless. I am quite new to both Julia and metaprogramming of this sort, so it is quite possible I am thinking of this the wrong way, and it would be great if someone could shed some light on my issue here. Thanks!

It's hard to get around the allocations implicit to SymPy in this case. It wants to allocate the matrix, so the easiest thing to do would be, as you show, build individual scalar functions. But then composing those together can be a bit of a hassle, since you don't want to put them into an array since they are all different types and that would then ruin the potential optimizations of looping over the calls. This means you'd have to create big tuples of functions and metaprogram over that as well in order to fully get rid of the allocations by directly building the indexed tuples... it's not fun so I don't plan on showing the code there but if you really need to do that explanation should be sufficient.

An easier way to do this kind of metaprogramming but end up at a non-allocating ODE function is via ModelingToolkit. The DifferentialEquations.jl DSLs have recently switched from SymPy/SymEngine to ModelingToolkit as the backend, and the DifferentialEquations.jl documentation will soon be suggesting this approach for handling this kind of problem, so I'll explain that suggested workflow. Similar to SymPy, you'd just create the symbolic expressions programmatically:

using ModelingToolkit, BenchmarkTools

@variables t σ[1:4,1:4](t)
@derivatives D'~t
eqs = Array{Equation}(undef,4,4)
for i in 1:4
for j in 1:4
eqs[i,j] =  D(σ[i,j]) ~ 1 + σ[i,j]
end
end


From there you can tell it to build an ODESystem out of the equations and tell it to build primitives for DifferentialEquations.jl:

sys = ODESystem(vec(eqs))
f = ODEFunction(sys)


And these primitives are made to be non-allocating fast functions:

u = rand(4,4)
du = similar(u)
@benchmark f(du,u,nothing,0.0)

BenchmarkTools.Trial:
memory estimate:  0 bytes
allocs estimate:  0
--------------
minimum time:     78.041 ns (0.00% GC)
median time:      91.340 ns (0.00% GC)
mean time:        108.952 ns (0.00% GC)
maximum time:     829.795 ns (0.00% GC)
--------------
samples:          10000
evals/sample:     970


And there you go, there's the You can directly investigate the generated code as well:

generate_function(sys)[2]

## Generated:
:((var"##MTIIPVar#353", var"##MTKArg#349", var"##MTKArg#350", var"##MTKArg#351")->begin
@inbounds begin
let (σ₁ˏ₁, σ₂ˏ₁, σ₃ˏ₁, σ₄ˏ₁, σ₁ˏ₂, σ₂ˏ₂, σ₃ˏ₂, σ₄ˏ₂, σ₁ˏ₃, σ₂ˏ₃, σ₃ˏ₃, σ₄ˏ₃, σ₁ˏ₄, σ₂ˏ₄, σ₃ˏ₄, σ₄ˏ₄, t) = (var"##MTKArg#349"[1], var"##MTKArg#349"[2], var"##MTKArg#349"[3], var"##MTKArg#349"[4], var"##MTKArg#349"[5], var"##MTKArg#349"[6], var"##MTKArg#349"[7], var"##MTKArg#349"[8], var"##MTKArg#349"[9], var"##MTKArg#349"[10], var"##MTKArg#349"[11], var"##MTKArg#349"[12], var"##MTKArg#349"[13], var"##MTKArg#349"[14], var"##MTKArg#349"[15], var"##MTKArg#349"[16], var"##MTKArg#351")
var"##MTIIPVar#353"[1] = 1 + σ₁ˏ₁
var"##MTIIPVar#353"[2] = 1 + σ₂ˏ₁
var"##MTIIPVar#353"[3] = 1 + σ₃ˏ₁
var"##MTIIPVar#353"[4] = 1 + σ₄ˏ₁
var"##MTIIPVar#353"[5] = 1 + σ₁ˏ₂
var"##MTIIPVar#353"[6] = 1 + σ₂ˏ₂
var"##MTIIPVar#353"[7] = 1 + σ₃ˏ₂
var"##MTIIPVar#353"[8] = 1 + σ₄ˏ₂
var"##MTIIPVar#353"[9] = 1 + σ₁ˏ₃
var"##MTIIPVar#353"[10] = 1 + σ₂ˏ₃
var"##MTIIPVar#353"[11] = 1 + σ₃ˏ₃
var"##MTIIPVar#353"[12] = 1 + σ₄ˏ₃
var"##MTIIPVar#353"[13] = 1 + σ₁ˏ₄
var"##MTIIPVar#353"[14] = 1 + σ₂ˏ₄
var"##MTIIPVar#353"[15] = 1 + σ₃ˏ₄
var"##MTIIPVar#353"[16] = 1 + σ₄ˏ₄
end
end
nothing
end)


And just for demonstration, you can tell it to calculate the sparse Jacobian and multithread the code:

generate_jacobian(sys,sparse=true,multithread=true)[2]

## Generated:
:((var"##MTIIPVar#363", var"##MTKArg#359", var"##MTKArg#360", var"##MTKArg#361")->begin
@inbounds begin
let (σ₁ˏ₁, σ₂ˏ₁, σ₃ˏ₁, σ₄ˏ₁, σ₁ˏ₂, σ₂ˏ₂, σ₃ˏ₂, σ₄ˏ₂, σ₁ˏ₃, σ₂ˏ₃, σ₃ˏ₃, σ₄ˏ₃, σ₁ˏ₄, σ₂ˏ₄, σ₃ˏ₄, σ₄ˏ₄, t) = (var"##MTKArg#359"[1], var"##MTKArg#359"[2], var"##MTKArg#359"[3], var"##MTKArg#359"[4], var"##MTKArg#359"[5], var"##MTKArg#359"[6], var"##MTKArg#359"[7], var"##MTKArg#359"[8], var"##MTKArg#359"[9], var"##MTKArg#359"[10], var"##MTKArg#359"[11], var"##MTKArg#359"[12], var"##MTKArg#359"[13], var"##MTKArg#359"[14], var"##MTKArg#359"[15], var"##MTKArg#359"[16], var"##MTKArg#361")
begin
(var"##MTIIPVar#363").nzval[1] = 1
(var"##MTIIPVar#363").nzval[2] = 1
(var"##MTIIPVar#363").nzval[3] = 1
(var"##MTIIPVar#363").nzval[4] = 1
end
end
begin
(var"##MTIIPVar#363").nzval[5] = 1
(var"##MTIIPVar#363").nzval[6] = 1
(var"##MTIIPVar#363").nzval[7] = 1
(var"##MTIIPVar#363").nzval[8] = 1
end
end
begin
(var"##MTIIPVar#363").nzval[9] = 1
(var"##MTIIPVar#363").nzval[10] = 1
(var"##MTIIPVar#363").nzval[11] = 1
(var"##MTIIPVar#363").nzval[12] = 1
end
end
begin
(var"##MTIIPVar#363").nzval[13] = 1
(var"##MTIIPVar#363").nzval[14] = 1
(var"##MTIIPVar#363").nzval[15] = 1
(var"##MTIIPVar#363").nzval[16] = 1
end
end
end
end
nothing
end)


Just to finalize the discussion, you can use the ODEProblem constructor on such a system to generate and solve the ODE. Note that instead of the normal DifferentialEquations.jl syntax where you give an array for the initial condition, here, to convert from symbolic to numerical, you give an array of pairs to tell it how to match symbols to initial conditions.

using OrdinaryDiffEq
u0 = [σ[i,j]=>rand() for i in 1:4, j in 1:4]
p = nothing
tspan = (0.0,1.0)
prob = ODEProblem(sys,u0,tspan,p)
solve(prob,Tsit5())


While this route is not feature-complete with SymPy yet, it can round-trip through SymPy if missing functionality is needed.

Edit: An edit to the question was made after I answered which added "(specifically, I wish to update the ODEs as they change in time)". I would highly recommend generating a time-dependent function or use callbacks to change parameters around instead of trying to generate a new function at every point in time. Note that you won't be able to avoid overhead of hitting the compiler on each new function, so if you're trying to do this as some kind of optimization, it won't actually make things faster and most likely will slow things down. That said, if that's what you need, you can accomplish this with the ModelingToolkit.build_function part of the interface, and I can give an example if you describe a more about what you're trying to do.

Given the timing of the edits going on here, my guess is that this is probably an X-Y problem and it would be easier to help if we can chat about what you're trying to accomplish.

• First, thanks so much for your detailed answer! I'll try to look into the documentation for ModelingToolkit.jl, it seems like a good start for what I am trying to do. I might have some additional questions, but I can reach out over the Julia slack you linked. I actually edited before seeing your initial answer; my apologies for not making the edit more clear. – Christian Hallas Apr 19 '20 at 6:44