xzackli / Bolt.jl

differentiable boltzmann code
MIT License
43 stars 5 forks source link

Background speed #83

Open jmsull opened 1 year ago

jmsull commented 1 year ago

The background is clocking in at ~200ms in btime - that's too long. Most of the time is in the FD background integral, so we should think about how to optimize this better.

xzackli commented 1 year ago

Yeah, it should be straightforward. Here's a 4x faster version with just

function f0(q,Tν)
    gs =  2 #should be 2 for EACH neutrino family (mass eigenstate)
    return gs / (2π)^3 / ( exp(q/Tν) +1)
end

#This is just copied from perturbations.jl for now - but take out Pressure - maybe later restore for FD tests?
function ρP_0(a,par::AbstractCosmoParams,quad_pts,quad_wts)
    #Do q integrals to get the massive neutrino metric perturbations
    #MB eqn (55)
    Tν =  (par.N_ν/3)^(1/4) *(4/11)^(1/3) * (15/ π^2 *ρ_crit(par) *par.Ω_r)^(1/4)
    #Not allowed to set Neff=0 o.w. breaks this #FIXME add an error message
    logqmin,logqmax=log10(Tν/30),log10(Tν*30)
    #FIXME: avoid repeating code? and maybe put general integrals in utils?
    m = par.Σm_ν
    xq,wq =quad_pts,quad_wts
    ρ, P = zero(a), zero(a)

    am² = a * m
    for i in eachindex(xq)
        xq2qx = xq2q(xq[i],logqmin,logqmax)
        xq2qx² = xq2qx^2
        ϵxx = √(xq2qx² + am²)
        f0_over_dxdq = f0(xq2qx,Tν) / dxdq(xq2qx,logqmin,logqmax)

        term = xq2qx² * f0_over_dxdq * wq[i]
        ρ += term * ϵxx 
        P += term * (xq2qx²/ϵxx)
    end
    ρ *= 4π * a^(-4) 
    P *= 4π/3 * a^(-4)

    return ρ,P
end

I need to write some tests for this, but it just reuses the various values it computes.