Sparsity Notes

[xzackli]

Bolt solves the photon-baryon tight coupling problem using standard "brute force" methods -- no physics, just numerics. We employ implicit RK (ESDIRK) instead of the commonly-used TCA with explicit RK (i.e. DVERK). Implicit solvers (standard for stiff problems) need to repeatedly construct a Jacobian and perform a linear solve using the matrix W = I - J where I is the identity (this implicit calculation basically integrates over the fast timescale solutions within a timestep). The Jacobian doesn't need to be constructed every timestep, and such solvers also rely on taking fewer timesteps at fixed accuracy in order to be competitive.

There are two things that suck with implicit methods:

1. Constructing the Jacobian (done with AD).
2. Linear solve.

Explicit methods require neither of these, but they tend to take many more steps -- which is fine for CAMB/CLASS because Boltzman steps are cheap (i.e. the cost of hierarchy!). It thus remains to be seen whether the implicit approach taken in Bolt can match the speed of CAMB or CLASS, but the upside is that we can write down some truly wacky cosmologies without having to come up with new TCA expressions.

However, it remains the case that Jacobian construction and the ensuing linear solve dominates the perturbation solve, so we should do something about that.

Sparsity Handling

If you know which elements of the Jacobian you want to compute with AD, then you can speed up the calculation a lot. The SciML packages make accelerating this part trivial, and we should probably make this automatic at some point. One needs to assess the sparsity of the Jacobian, and then pass that to ODEFunction.

Setup

using Bolt
using BenchmarkTools

par = CosmoParams()
bg = Background(par)
ih = IonizationHistory(Peebles(), par, bg)

k_grid = quadratic_k(0.1bg.H₀, 1000bg.H₀, 100)
k_i = 10
hierarchy = Hierarchy(BasicNewtonian(), par, bg, ih, k_grid[k_i])

xᵢ = log(1e-8)
u₀ = Bolt.initial_conditions(xᵢ, hierarchy)
du_dummy = deepcopy(u₀)

Sparsity Prototype

using SparsityDetection, SparseArrays, OrdinaryDiffEq
hf(du, u) = Bolt.hierarchy!(du, u, hierarchy, xᵢ)
sparsity_pattern = jacobian_sparsity(hf, du_dummy, u₀)
jac_sparsity = Float64.(sparse(sparsity_pattern))
f = ODEFunction(Bolt.hierarchy!;jac_prototype=jac_sparsity)
prob = ODEProblem{true}(Bolt.hierarchy!, u₀, (xᵢ , zero(Float64)), hierarchy)
sparseprob = ODEProblem{true}(f, u₀, (xᵢ , zero(Float64)), hierarchy)

@btime sol = solve(sparseprob, KenCarp4(), reltol=1e-6,
            # saveat=hierarchy.bg.x_grid, 
            dense=false);
@btime sol = solve(prob, KenCarp4(), reltol=1e-6,
            # saveat=hierarchy.bg.x_grid, 
            dense=false);

I get a 250% speed boost from doing this! The SparsityDetection package is remarkable but we definitely need some precompile statements for it. It's in the process of being replaced by ModelingToolkit (the issue to watch is SciML/ModelingToolkit.jl#361).