A Collection of Jacobian Sparsity Acceleration Tools for Julia
October 6 2019 in Differential Equations, Julia, Programming, Scientific ML | Tags: | Author: Christopher Rackauckas
Over the summer there have been a whole suite of sparsity acceleration tools for Julia. These are encoded in the packages:
The toolchain is showcased in the following blog post by Pankaj Mishra, the student who build a lot of the Jacobian coloring and decompression framework. Langwen Huang setup the fast paths for structured matrices (tridiagonal, banded, and block-banded matrices) and also integrated these tools with DifferentialEquations.jl. Shashi Gowda then setup a mechanism for automatically detecting the sparsity of Julia programs (!!!).
A tutorial using this workflow together is described in the SparseDiffTools.jl README. In summary, to use the tools you have the following flow:
- Find your sparsity pattern, Jacobian structure (i.e. Jacobian type), or automatically detect it with SparsityDetection.jl.
- Call `matrix_colors(A)` from SparseDiffTools.jl to get the `colorvec` for A. This is the vector that the differentiation tools need to have to exploit sparsity and reduce the total cost of generating the Jacobian.
- When calling `forwarddiff_color_jacobian` from SparseDiffTools.jl for sparse AD or `finite_difference_jacobian` from DiffEqDiffTools.jl for sparse finite differencing, pass the `colorvec` and `sparsity` (the sparsity pattern by either passing the sparse matrix or the structured matrix), then the differentiation tools will automatically accelerate to be fast for that kind of matrix
- When building the ODEFunction for the DifferentialEquations.jl ODE solver, pass the `colorvec` and `jac_prototype` and all internal functions will automatically specialize on the sparsity pattern and accelerate. If you pass a structured matrix, like a BandedMatrix, the color vector will be determined automatically, making those accelerations free.
Thus together the chain is: get the sparsity, get the colorvec, pass it to packages and boom you’re faster!
The Math of Sparsity
If you’re interested in how this all works, please take a look at the lecture notes for my course 18.337:
- Forward-Mode AD via High Dimensional Algebras (necessary backstory)
- Solving Stiff Ordinary Differential Equations (which explains the sparse AD story)
Additionally, take a look at this paper for an explanation of how you can do automatic sparsity detection of Julia packages.
Conclusion
It will be interesting to see how having an integrated platform for acceleration via sparsity effects a high level language, especially the automatic sparsity detection. It is fitting for Julia to have these tools since, given the focus on performance, this is the piece of math that is required to make your performance work really matter! This is a pervasive mechanism which lets you accelerate differentiation on your own, or directly give the differential equation solvers these to utilize (and it works with ODEs, SDEs, DAEs, DDEs, hybrid equations, etc.). We hope to integrate this with NLsolve.jl, and get the Hessian tools finished for Optim.jl and JuMP.jl. Also, JuMP.jl is getting a new and improved NLP interface which will utilize a lot of this behind the scenes automatically. Stay tuned.