You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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:
Constructing the Jacobian (done with AD).
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).
reacted with thumbs up emoji reacted with thumbs down emoji reacted with laugh emoji reacted with hooray emoji reacted with confused emoji reacted with heart emoji reacted with rocket emoji reacted with eyes emoji
-
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
whereI
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:
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
Sparsity Prototype
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).Beta Was this translation helpful? Give feedback.
All reactions