jacobian_sparsity speed

[sjdaines]

Some benchmarks for Symbolics.jacobian_sparsity demonstrating that speed can be marginal (and ~40x slower than SparsityTracing https://github.com/PALEOmodel/SparsityTracing.jl which is a minimalist tracing approach). Simplified example here with 50000 state variables takes ~20s with Symbolics - actual code is 10x this or more.

import SparsityTracing  # formerly ADelemtree
import Symbolics
using BenchmarkTools

# define a random chemical reaction network with nreactions among nspecies in each of ncells
ncells = 1000
nspecies = 50
nreactions = 200

r1idx = rand(1:nspecies, nreactions);  # reactant 1 species
r2idx = rand(1:nspecies, nreactions);  # reactant 2 species
pidx = rand(1:nspecies, nreactions);   # product species

p = (ncells, nspecies, r1idx, r2idx, pidx);

# time derivative of species concentrations for idealized reaction network 
function react!(du_vec, u_vec, p)
    ncells, nspecies, r1idx, r2idx, pidx = p
    u = reshape(u_vec, nspecies, ncells)
    du = reshape(du_vec, nspecies, ncells)
    du .= 0.0
    for i in 1:ncells
        for (r1i, r2i, pi) in zip(r1idx, r2idx, pidx)
            rate = u[r1i, i]*u[r2i, i]
            du[r1i, i] -= 0.1234*rate  # implausible stoichiometries for test purposes
            du[r2i, i] -= 0.2345*rate
            du[pi, i] += 0.3456*rate
        end
    end
    
    return nothing
end

u = ones(nspecies*ncells);
du = zeros(nspecies*ncells);

println("benchmark function call:")
@time react!(du, u, p)
@time react!(du, u, p)
@time react!(du, u, p)

benchmark function call:
0.070629 seconds (240.95 k allocations: 14.479 MiB, 98.67% compilation time)
0.000889 seconds (4 allocations: 192 bytes)
0.000905 seconds (4 allocations: 192 bytes)

println("benchmark SparsityTracing")
u_ad = SparsityTracing.create_advec(u);
du_ad = similar(u_ad);
println("    function call:")
@time react!(du_ad, u_ad, p)
@time react!(du_ad, u_ad, p)
@time react!(du_ad, u_ad, p)
println("    jacobian:")
@time Jst = SparsityTracing.jacobian(du_ad, length(u_ad));
@time Jst = SparsityTracing.jacobian(du_ad, length(u_ad));
@time Jst = SparsityTracing.jacobian(du_ad, length(u_ad));

benchmark SparsityTracing
function call:
0.381003 seconds (2.55 M allocations: 88.460 MiB, 16.16% gc time, 35.22% compilation time)
0.204951 seconds (2.40 M allocations: 79.346 MiB, 11.08% gc time)
0.223499 seconds (2.40 M allocations: 79.346 MiB, 17.41% gc time)
jacobian:
1.391276 seconds (4.15 M allocations: 300.907 MiB, 34.48% gc time, 56.99% compilation time)
0.252109 seconds (2.85 M allocations: 226.638 MiB, 9.84% gc time)
0.258635 seconds (2.85 M allocations: 226.638 MiB, 16.48% gc time)

println("benchmark Symbolics")
Symbolics.@variables u_sym[1:length(u)];
u_sym_s = Symbolics.scalarize(u_sym);
du_sym_s = similar(u_sym_s);
println("    function call:")
@time react!(du_sym_s, u_sym_s, p)
@time react!(du_sym_s, u_sym_s, p)
@time react!(du_sym_s, u_sym_s, p)
println("    jacobian_sparsity:")
@time Jsym = Float64.(Symbolics.jacobian_sparsity(du_sym_s, u_sym_s));
@time Jsym = Float64.(Symbolics.jacobian_sparsity(du_sym_s, u_sym_s));
@time Jsym = Float64.(Symbolics.jacobian_sparsity(du_sym_s, u_sym_s));

benchmark Symbolics
function call:
10.942955 seconds (62.87 M allocations: 2.829 GiB, 22.33% gc time, 12.15% compilation time)
9.523435 seconds (59.88 M allocations: 2.661 GiB, 22.74% gc time)
9.695355 seconds (59.88 M allocations: 2.661 GiB, 22.44% gc time)
jacobian_sparsity:
35.813055 seconds (234.05 M allocations: 8.488 GiB, 12.64% gc time, 8.15% compilation time)
13.600838 seconds (88.15 M allocations: 3.666 GiB, 8.12% gc time)
13.384328 seconds (88.15 M allocations: 3.666 GiB, 6.25% gc time)

# check 
Jst.nzval .= 1.0
println("Jacobians equal: ", Jst == Jsym)

Jacobians equal: true

┆Issue is synchronized with this Trello card by Unito