Skip to content

Exploit linearity of perturbation ODEs for solver performance #34

Description

@hersle

Consider the perturbation ODEs (without neutrinos for simplicity):

using SymBoltz
M = ΛCDM(h = nothing, ν = nothing)
pars = parameters_Planck18(M)
prob = CosmologyProblem(M, pars)
equations(prob.pt.f.sys)
19-element Vector{Equation}:
 Differential(τ)((γ₊G(τ, k))[6]) ~ (-7(γ₊G(τ, k))[6]) / τ + k*(γ₊G(τ, k))[5] + γ₊κ̇(τ)*(γ₊G(τ, k))[6]
 Differential(τ)((γ₊G(τ, k))[5]) ~ (1//11)*k*(5(γ₊G(τ, k))[4] - 6(γ₊G(τ, k))[6]) + γ₊κ̇(τ)*(γ₊G(τ, k))[5]
 Differential(τ)((γ₊G(τ, k))[4]) ~ (1//9)*k*(4(γ₊G(τ, k))[3] - 5(γ₊G(τ, k))[5]) + γ₊κ̇(τ)*(γ₊G(τ, k))[4]
 Differential(τ)((γ₊G(τ, k))[3]) ~ (1//7)*k*(3(γ₊G(τ, k))[2] - 4(γ₊G(τ, k))[4]) + γ₊κ̇(τ)*(γ₊G(τ, k))[3]
 Differential(τ)((γ₊G(τ, k))[2]) ~ (1//5)*k*(2(γ₊G(τ, k))[1] - 3(γ₊G(τ, k))[3]) + (-(1//10)*γ₊Π(τ, k) + (γ₊G(τ, k))[2])*γ₊κ̇(τ)
 Differential(τ)((γ₊G(τ, k))[1]) ~ (1//3)*k*(γ₊G0(τ, k) - 2(γ₊G(τ, k))[2]) + γ₊κ̇(τ)*(γ₊G(τ, k))[1]
 Differential(τ)(γ₊G0(τ, k)) ~ -k*(γ₊G(τ, k))[1] + (γ₊G0(τ, k) - (1//2)*γ₊Π(τ, k))*γ₊κ̇(τ)
 Differential(τ)((γ₊F(τ, k))[6]) ~ (-7(γ₊F(τ, k))[6]) / τ + k*(γ₊F(τ, k))[5] + γ₊κ̇(τ)*(γ₊F(τ, k))[6]
 Differential(τ)((γ₊F(τ, k))[5]) ~ (1//11)*k*(5(γ₊F(τ, k))[4] - 6(γ₊F(τ, k))[6]) + γ₊κ̇(τ)*(γ₊F(τ, k))[5]
 Differential(τ)((γ₊F(τ, k))[4]) ~ (1//9)*k*(4(γ₊F(τ, k))[3] - 5(γ₊F(τ, k))[5]) + γ₊κ̇(τ)*(γ₊F(τ, k))[4]
 Differential(τ)((γ₊F(τ, k))[3]) ~ (1//7)*k*(3(γ₊F(τ, k))[2] - 4(γ₊F(τ, k))[4]) + γ₊κ̇(τ)*(γ₊F(τ, k))[3]
 Differential(τ)((γ₊F(τ, k))[2]) ~ (1//5)*k*(2(γ₊F(τ, k))[1] - 3(γ₊F(τ, k))[3]) + (-(1//10)*γ₊Π(τ, k) + (γ₊F(τ, k))[2])*γ₊κ̇(τ)
 Differential(τ)((γ₊F(τ, k))[1]) ~ ((-4//3)*(-γ₊θ(τ, k) + γ₊θb(τ, k))*γ₊κ̇(τ)) / k + (1//3)*k*(γ₊F0(τ, k) + 4Ψ(τ, k) - 2(γ₊F(τ, k))[2])
 Differential(τ)(Φ(τ, k)) ~ (-4.1887902047863905(a(τ)^2)*G₊δρ(τ, k)) / ℋ(τ) + (-(k^2)*Φ(τ, k)) / (3ℋ(τ)) - ℋ(τ)*Ψ(τ, k)
 Differential(τ)(γ₊F0(τ, k)) ~ -k*(γ₊F(τ, k))[1] + 4((-4.1887902047863905(a(τ)^2)*G₊δρ(τ, k)) / ℋ(τ) + (-(k^2)*Φ(τ, k)) / (3ℋ(τ)) - ℋ(τ)*Ψ(τ, k))
 Differential(τ)(c₊θ(τ, k)) ~ ((k^2)*c₊cₛ²(τ)*c₊δ(τ, k)) / (1 + c₊w(τ)) + c₊θinteraction(τ, k) - (k^2)*c₊σ(τ, k) + (k^2)*Ψ(τ, k) - ℋ(τ)*c₊θ(τ, k)*(1 - 3c₊w(τ))
 Differential(τ)(c₊δ(τ, k)) ~ (c₊θ(τ, k) - 3((-4.1887902047863905(a(τ)^2)*G₊δρ(τ, k)) / ℋ(τ) + (-(k^2)*Φ(τ, k)) / (3ℋ(τ)) - ℋ(τ)*Ψ(τ, k)))*(-1 - c₊w(τ)) - 3ℋ(τ)*(c₊cₛ²(τ) - c₊w(τ))*c₊δ(τ, k)
 Differential(τ)(b₊θ(τ, k)) ~ b₊θinteraction(τ, k) + ((k^2)*b₊δ(τ, k)*b₊cₛ²(τ)) / (1 + b₊w(τ)) + (k^2)*Ψ(τ, k) - (k^2)*b₊σ(τ, k) - ℋ(τ)*b₊θ(τ, k)*(1 - 3b₊w(τ))
 Differential(τ)(b₊δ(τ, k)) ~ (b₊θ(τ, k) - 3((-4.1887902047863905(a(τ)^2)*G₊δρ(τ, k)) / ℋ(τ) + (-(k^2)*Φ(τ, k)) / (3ℋ(τ)) - ℋ(τ)*Ψ(τ, k)))*(-1 - b₊w(τ)) - 3ℋ(τ)*b₊δ(τ, k)*(-b₊w(τ) + b₊cₛ²(τ))

Unless I am missing something, this is in the linear form

$$\boldsymbol{u}^\prime(t) = A(t) \boldsymbol{u}(t),$$

where $A(t) = J(t)$ is also the Jacobian of the system and only depends on the background. It should be possible to exploit this linearity for performance gains. Look at this profiling:

using ProfileCanvas
ks = 10 .^ range(0, 4, length = 128)
@profview sol = solve(prob, ks; thread = false)
Image

And the solver stats:

SymBoltz.statspt(sol) # summarized perturbation stats
SciMLBase.DEStats
Number of function 1 evaluations:                  672101
Number of function 2 evaluations:                  0
Number of W matrix evaluations:                    44928
Number of linear solves:                           627062
Number of Jacobians created:                       44523
Number of nonlinear solver iterations:             582671
Number of nonlinear solver convergence failures:   159
Number of fixed-point solver iterations:           0
Number of fixed-point solver convergence failures: 0
Number of rootfind condition calls:                0
Number of accepted steps:                          44145
Number of rejected steps:                          246

This computes $\boldsymbol{f}$ and $J$ separately. But one could instead first compute $J$, then $\boldsymbol{f} = J \boldsymbol{u}$ with a simple matrix multiplication. Maybe the nonlinear solve at every time step can also be simplified with a single linear solve when the $f(\boldsymbol{u},t) = J(t) \boldsymbol{u}$?

Resources:

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions