GitHub

Using automatic differentiation

This tutorial shows how to compute the power spectrum $P(k; \theta)$ and its (logarithmic) derivatives

\[\frac{\partial \lg P}{\partial \lg \theta_i}\]

using automatic differentiation with ForwardDiff.jl. This technique can also differentiate any other quantity.

1. Wrap the evaluation

First, we must decide which parameters $\theta$ the power spectrum $P(k; \theta)$ should be considered a function of. To do so, let us write a small wrapper function that calculates the power spectrum as a function of the parameters $(T_{\gamma 0}, \Omega_{c0}, \Omega_{b0}, N_\textrm{eff}, h, Y_p)$, following the Getting started tutorial:

using SymBoltz
M = ΛCDM(K = nothing)
pars = [M.γ.T₀, M.c.Ω₀, M.b.Ω₀, M.ν.Neff, M.g.h, M.b.YHe, M.h.m_eV, M.I.ln_As1e10, M.I.ns]
prob0 = CosmologyProblem(M, Dict(pars .=> NaN))

probgen = parameter_updater(prob0, pars)
P(k, θ) = spectrum_matter(probgen(θ), k; verbose = true, ptopts = (reltol = 1e-3,))
P (generic function with 1 method)

It is now easy to evaluate the power spectrum:

using Unitful, UnitfulAstro
θ = [2.7, 0.27, 0.05, 3.0, 0.7, 0.25, 0.02, 3.0, 0.95]
ks = 10 .^ range(-3, 0, length=100) / u"Mpc"
Ps = P(ks, θ)
100-element Vector{Unitful.Quantity{Float64, 𝐋^3, Unitful.FreeUnits{(Mpc^3,), 𝐋^3, nothing}}}:
 15091.382819797198 Mpc^3
 16046.459687146058 Mpc^3
 17053.403792107027 Mpc^3
 18113.476476359803 Mpc^3
 19228.362225165445 Mpc^3
 20398.74955632339 Mpc^3
 21625.64662800446 Mpc^3
 22909.56976534061 Mpc^3
 24250.627294118287 Mpc^3
 25648.75626151622 Mpc^3
     ⋮
   234.87525494631112 Mpc^3
   200.7802647832025 Mpc^3
   171.16642641100168 Mpc^3
   145.98788115974426 Mpc^3
   124.34218437061641 Mpc^3
   105.5788986722152 Mpc^3
    89.81464786855148 Mpc^3
    76.33371583241178 Mpc^3
    64.76077116958713 Mpc^3

This can be plotted with

using Plots
plot(log10.(ks/u"1/Mpc"), log10.(Ps/u"Mpc^3"); xlabel = "lg(k/Mpc⁻¹)", ylabel = "lg(P/Mpc³)", label = nothing)
Example block output

2. Calculate the derivatives

To get $\partial \lg P / \partial \lg \theta$, we can simply pass the wrapper function P(k, θ) through ForwardDiff.jacobian:

using ForwardDiff
lgP(lgθ) = log10.(P(ks, 10 .^ lgθ) / u"Mpc^3") # in log-space
dlgP_dlgθs = ForwardDiff.jacobian(lgP, log10.(θ))
100×9 Matrix{Float64}:
 -0.148128  -1.23994  -0.229643  -0.00985095  …  -0.0045604   3.0  -3.71642
 -0.162161  -1.23398  -0.228539  -0.0108777      -0.00452997  3.0  -3.65014
 -0.177475  -1.22758  -0.227353  -0.0119988      -0.00452758  3.0  -3.58385
 -0.194093  -1.22058  -0.226053  -0.0132165      -0.00455557  3.0  -3.51756
 -0.21219   -1.21307  -0.224661  -0.0145465      -0.00461821  3.0  -3.45128
 -0.231787  -1.20508  -0.223178  -0.0159914   …  -0.00471995  3.0  -3.38499
 -0.252868  -1.19617  -0.22152   -0.0175526      -0.00486003  3.0  -3.3187
 -0.275642  -1.18659  -0.219747  -0.0192399      -0.00503854  3.0  -3.25242
 -0.300196  -1.17643  -0.217871  -0.0210666      -0.0052547   3.0  -3.18613
 -0.326577  -1.16545  -0.215846  -0.0230388      -0.00550595  3.0  -3.11984
  ⋮                                           ⋱                    
 -5.97727    1.88842  -0.36711   -0.382855       -0.0253137   3.0   2.31565
 -6.02361    1.89996  -0.371122  -0.386789       -0.0253184   3.0   2.38194
 -6.0566     1.92271  -0.363939  -0.384746       -0.0253186   3.0   2.44823
 -6.08032    1.93603  -0.367462  -0.386885       -0.0253198   3.0   2.51451
 -6.11367    1.94875  -0.365832  -0.388269    …  -0.0253207   3.0   2.5808
 -6.14018    1.9635   -0.366959  -0.389298       -0.0253215   3.0   2.64709
 -6.16293    1.97795  -0.368343  -0.389868       -0.0253223   3.0   2.71337
 -6.19015    1.99125  -0.368893  -0.391279       -0.0253221   3.0   2.77966
 -6.21864    2.00429  -0.368066  -0.392242       -0.0253234   3.0   2.84595

The matrix element dlgP_dlgθs[i, j] now contains $\partial \lg P(k_i) / \partial \lg \theta_j$. We can plot them all at once:

plot(
    log10.(ks/u"1/Mpc"), dlgP_dlgθs;
    xlabel = "lg(k/Mpc⁻¹)", ylabel = "∂ lg(P) / ∂ lg(θᵢ)",
    labels = "θᵢ=" .* ["Tγ0" "Ωc0" "Ωb0" "Neff" "h" "YHe" "mh" "ln(10¹⁰As)" "ns"]
)
Example block output

Get values and derivatives together

The above example showed how to calculate the power spectrum values and their derivatives through two separate calls. If you need both, it is faster to calculate them simultaneously with the package DiffResults.jl:

using DiffResults

# Following DiffResults documentation:
Pres = DiffResults.JacobianResult(ks/u"1/Mpc", θ) # allocate buffer for values+derivatives for a function with θ-sized input and ks-sized output
Pres = ForwardDiff.jacobian!(Pres, lgP, log10.(θ)) # evaluate values+derivatives of lgP(log10.(θ)) and store the results in Pres
lgPs = DiffResults.value(Pres) # extract values
dlgP_dlgθs = DiffResults.jacobian(Pres) # extract derivatives

p1 = plot(
   log10.(ks/u"1/Mpc"), lgPs;
   ylabel = "lg(P/Mpc³)", label = nothing
)
p2 = plot(
   log10.(ks/u"1/Mpc"), dlgP_dlgθs;
   xlabel = "lg(k/Mpc⁻¹)", ylabel = "∂ lg(P) / ∂ lg(θᵢ)",
   labels = "θᵢ=" .* ["Tγ0" "Ωc0" "Ωb0" "Neff" "h" "YHe" "mh" "ln(10¹⁰As)" "ns"]
)
plot(p1, p2, layout=(2, 1), size = (600, 600))
Example block output

General approach

The technique shown here can be used to calculate the derivative of any SymBoltz.jl output quantity:

  1. Write a wrapper function output(input) that calculates the desired output quantities from the desired input quantities.
  2. Use ForwardDiff.derivative(output, input) (scalar-to-scalar), ForwardDiff.gradient(output, input) (vector-to-scalar) or ForwardDiff.jacobian(output, input) (vector-to-vector) to evaluate the derivative of output at the values input. Or use the similar functions in DiffResults to calculate the value and derivatives simultaneously.