Skip to content

Conversation

@hersle
Copy link
Contributor

@hersle hersle commented Nov 29, 2025

I noticed two things when profiling with sparse+analytical Jacobians generated by MTK in my package:

using SymBoltz
M = ΛCDM(K = nothing) # creates a System with 219 equations
pars = parameters_Planck18(M)
ks = 1.0:1.0:100.0

@profview prob = CosmologyProblem(M, pars; jac = true, sparse = true) # creates ODEProblem with sparse 54x54 Jacobian
@profview sol = solve(prob, ks; thread = false) # solves 100x ODEProblems with sparse Jacobians
  1. Almost all time in ODEProblem(...; jac = true, sparse = true) is spent inside this iszero check:
bilde
  1. The solve is slowed by this sparsity pattern check every time the analytically generated Jacobian function is called:
bilde

I suggest two changes:

  1. Circumvent that iszero altogether. It is used when nonzeros of W are explicitly stored in J. But this can be done equivalently by recreating the sparse Jacobian in (I, J, V)-form with explicit zeros added at all of W's nonzero entries.
  2. Make it possible to disable the checking of the sparsity pattern at runtime. My first suggestion is for this checking to be done only if checkbounds = true. I am sure there are also other ways to solve this.

With these changes, the respective sections of the flamegraph disappear and the timings improve:

@time prob = CosmologyProblem(M, pars; jac = true, sparse = true)
@time sol = solve(prob, ks; thread = false)
 31.711084 seconds (262.74 M allocations: 8.367 GiB, 1.73% gc time, 0.78% compilation time: <1% of which was recompilation)
  0.594824 seconds (1.32 M allocations: 346.698 MiB, 34.96% gc time)

to

  4.061829 seconds (40.72 M allocations: 1.243 GiB, 2.49% gc time, 0.10% compilation time: <1% of which was recompilation)
  0.396266 seconds (1.10 M allocations: 154.179 MiB, 7.32% gc time)

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the
    contributor guidelines, in particular the SciML Style Guide and
    COLPRAC.
  • Any new documentation only uses public API

@hersle
Copy link
Contributor Author

hersle commented Nov 29, 2025

Alternatively, if you think the bounds check should be preserved for safety, I suggest to change it to check only the number of nonzero entries, and not the correctness of every index. Something like length(out.nzval) == length(jac_precomputed).

@hersle hersle force-pushed the sparse_performance branch from fe2296b to e14cc64 Compare November 29, 2025 19:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant