Comparison to CLASS

using SymBoltz
using CLASS
using DelimitedFiles
using DataInterpolations
using Unitful, UnitfulAstro
using CairoMakie

lmax = 6
reionization = true
Hswitch = 1
Heswitch = 6
M = w0waCDM(; lmax, reionization, Hswitch, Heswitch)
pars = merge(parameters_Planck18(M), Dict(
    M.X.w0 => -0.9,
    M.X.wa => 0.1,
    M.X.cₛ² => 0.9
))
h = pars[M.g.h] # needed later
prob = CosmologyProblem(M, pars)

function solve_class(pars, k = nothing)
    prob = CLASSProblem(
        "write_background" => "yes",
        "write_thermodynamics" => "yes",
        "background_verbose" => 2,
        "output" => "mPk, tCl, pCl", # need one to evolve perturbations

        "k_output_values" => isnothing(k) ? "" : NoUnits(k / u"1/Mpc"),
        "ic" => "ad",
        "modes" => "s",
        "gauge" => "newtonian",

        # metric
        "h" => pars[M.g.h],

        # photons
        "T_cmb" => pars[M.γ.T₀],
        "l_max_g" => lmax,
        "l_max_pol_g" => lmax,

        # baryons
        "Omega_b" => pars[M.b.Ω₀],
        "YHe" => pars[M.b.YHe],
        "recombination" => "recfast", # TODO: HyREC
        "recfast_Hswitch" => Hswitch,
        "recfast_Heswitch" => Heswitch,
        "reio_parametrization" => reionization ? "reio_camb" : "reio_none",

        # cold dark matter
        "Omega_cdm" => pars[M.c.Ω₀],

        # neutrinos
        "N_ur" => SymBoltz.have(M, :ν) ? pars[M.ν.Neff] : 0.0,
        "N_ncdm" => SymBoltz.have(M, :h) ? 1 : 0,
        "m_ncdm" => SymBoltz.have(M, :h) ? pars[M.h.m_eV] : 0.0, # in eV/c^2
        "T_ncdm" => SymBoltz.have(M, :h) ? (4/11)^(1/3) : 0.0,
        "l_max_ur" => lmax,
        "l_max_ncdm" => lmax,

        # primordial power spectrum
        "ln_A_s_1e10" => pars[M.I.ln_As1e10],
        "n_s" => pars[M.I.ns],

        # w0wa dark energy
        "Omega_Lambda" => 0.0, # unspecified
        "w0_fld" => SymBoltz.have(M, :X) ? pars[M.X.w0] : -1.0,
        "wa_fld" => SymBoltz.have(M, :X) ? pars[M.X.wa] : 0.0,
        "cs2_fld" => SymBoltz.have(M, :X) ? pars[M.X.cₛ²] : 1.0,
        "use_ppf" => SymBoltz.have(M, :X) ? "no" : "yes", # use full equations (must be yes with CC to prevent crash)

        # other stuff
        "Omega_k" => SymBoltz.have(M, :K) ? pars[M.K.Ω₀] : 0.0, # curvature
        "Omega_scf" => 0.0,
        "Omega_dcdmdr" => 0.0,

        # approximations (see include/precisions.h)
        "tight_coupling_trigger_tau_c_over_tau_h" => 1e-2, # cannot turn off?
        "tight_coupling_trigger_tau_c_over_tau_k" => 1e-3, # cannot turn off
        "radiation_streaming_approximation" => 3, # turns off RSA
        "ur_fluid_approximation" => 3, # turns off UFA
        "ncdm_fluid_approximation" => 3, # turns off NCDM fluid approximation

        #"l_max_scalars" => 1500,
        #"temperature_contributions" => "pol",
    )
    return solve(prob)
end

k = 1e1 / u"Mpc" # 1/Mpc
sol1 = solve_class(pars, k)
sol2 = solve(prob, k; ptopts = (alg = SymBoltz.Rodas4P(),)) # looks like lower-precision KenCarp4 and Kvaerno5 "emulate" radiation streaming, while higher-precision Rodas5P continues in an exact way

function plot_compare(x1s, x2s, y1s, y2s, xlabel, ylabels; lgx=false, lgy=false, common=false, errtype=:auto, errlim=NaN, tol = nothing, kwargs...)
    if !(ylabels isa AbstractArray)
        y1s = [y1s]
        y2s = [y2s]
        ylabels = [ylabels]
    end

    xplot(x) = lgx ? log10.(abs.(x)) : x
    yplot(y) = lgy ? log10.(abs.(y)) : y
    xlab(x) = lgx ? "lg(|$x|)" : x
    ylab(y) = lgy ? "lg(|$y|)" : y

    fig = Figure(width = 800, height = 600)
    ax1 = Axis(fig[2:3, 1]; xticklabelsvisible = false)
    ax2 = Axis(fig[4, 1]; xlabel = xlab(xlabel), yminorticks = IntervalsBetween(10), yminorgridvisible = true)
    maxerr = 0.0
    xlims = (NaN, NaN)
    linewidth = 2
    lines!(ax1, [NaN]; color = :black, linewidth, linestyle = :solid, label = "CLASS") # dummy for legend entry
    lines!(ax1, [NaN]; color = :black, linewidth, linestyle = :dash, label = "SymBoltz") # dummy for legend entry
    for (i, (y1, y2, ylabel)) in enumerate(zip(y1s, y2s, ylabels))
        x1, x2 = x1s, x2s
        x1min, x1max = extrema(x1)
        x2min, x2max = extrema(x2)

        plotx1 = xplot(x1) # https://discourse.julialang.org/t/makie-axis-limits-with-inf/85784
        plotx2 = xplot(x2)
        ploty1 = replace!(yplot(y1), -Inf => NaN, Inf => NaN)
        ploty2 = replace!(yplot(y2), -Inf => NaN, Inf => NaN)
        color = Makie.wong_colors()[i]
        lines!(ax1, plotx1, ploty1; color, alpha = 0.6, linewidth, linestyle = :solid)
        lines!(ax1, plotx2, ploty2; color, alpha = 0.6, linewidth, linestyle = :dash)
        lines!(ax1, [NaN]; color, linewidth, label = ylab(ylabel)) # dummy for legend entry

        if i == 1
            xlims = (x1min, x1max) # initialize so not NaN
        end
        xlims = if common
            (max(xlims[1], x1min, x2min), min(xlims[2], x1max, x2max))
        else
            (min(xlims[1], x1min, x2min), max(xlims[2], x1max, x2max))
        end

        xmin = max(x1min, x2min) # need common times to compare error
        xmax = min(x1max, x2max) # need common times to compare error
        i1s = xmin .≤ x1 .≤ xmax # select x values in common range and exclude points too close in time (probably due to CLASS switching between approximation schemes)
        i2s = xmin .≤ x2 .≤ xmax # select x values in common range
        x1 = x1[i1s]
        x2 = x2[i2s]
        y1 = y1[i1s]
        y2 = y2[i2s]
        x = x2 # compare ratios at SymBoltz' times
        y1 = LinearInterpolation(y1, x1; extrapolation = ExtrapolationType.Linear).(x) # TODO: use built-in CosmoloySolution interpolation
        y2 = LinearInterpolation(y2, x2; extrapolation = ExtrapolationType.Linear).(x)

        !isnothing(tol) && @assert all(isapprox.(y1, y2; atol = tol)) "$ylabel does not match within absolute tolerance $tol. Maximum difference was $(maximum(abs.(y1.-y2)))."

        # Compare absolute error if quantity crosses zero, otherwise relative error (unless overridden)
        abserr = (errtype == :abs) || (errtype == :auto && (any(y1 .<= 0) || any(y2 .<= 0)))
        if abserr
            err = y2 .- y1
            ylabel = "SymBoltz - CLASS"
        else
            err = y2 ./ y1 .- 1
            ylabel = "SymBoltz / CLASS - 1"
        end
        ax2.ylabel = ylabel
        maxerr = max(maxerr, maximum(abs.(err)))
        lines!(ax2, xplot(x), err; color, alpha = 1.0, linewidth)
    end

    # write maximum error like a * 10^b (for integer a and b)
    if isnan(errlim)
        b = floor(log10(maxerr))
        a = ceil(maxerr / 10^b)
        errlim = a * 10^b
    end
    Makie.ylims!(ax2, (-errlim, +errlim))

    xlims = xplot.(xlims) # transform to log?
    Makie.xlims!(ax1, xlims)
    Makie.xlims!(ax2, xlims)

    #axislegend(ax1)
    fig[1, 1] = Legend(fig, ax1; padding = (0, 0, 0, 0), margin = (0, 0, 0, 0), framevisible = false, orientation = :horizontal)

    return fig
end

Running CLASS version v3.3.1
Computing background
 -> non-cold dark matter species with i=1 has m_i = 6.000000e-02 eV (so m_i / omega_i =9.405928e+01 eV)
 -> ncdm species i=1 sampled with 11 (resp. 11) points for purpose of background (resp. perturbation) integration. In the relativistic limit it gives Delta N_eff = 0.999997
Chose ndf15 as generic_evolver
 -> age = 13.524806 Gyr
 -> conformal age = 13964.401666 Mpc
 -> N_eff = 3.99 (summed over all species that are non-relativistic at early times)
 -> radiation/matter equality at z = 3020.944008
    corresponding to conformal time = 119.706325 Mpc
 ---------------------------- Budget equation -----------------------
 ---> Nonrelativistic Species
-> Bayrons                        Omega = 0.0493678       , omega = 0.0224
-> Cold Dark Matter               Omega = 0.26447         , omega = 0.12
 ---> Non-Cold Dark Matter Species (incl. massive neutrinos)
-> Non-Cold Species Nr.      1    Omega = 0.00140587      , omega = 0.000637896
 ---> Relativistic Species
-> Photons                        Omega = 5.45025e-05     , omega = 2.47298e-05
-> Ultra-relativistic relics      Omega = 3.701e-05       , omega = 1.67928e-05
 ---> Other Content
-> Dark Energy Fluid              Omega = 0.684664        , omega = 0.310658
 ---> Total budgets
 Radiation                        Omega = 9.15124e-05     , omega = 4.15226e-05
 Non-relativistic                 Omega = 0.315244        , omega = 0.143038
 - Non-Free-Streaming Matter      Omega = 0.313838        , omega = 0.1424
 - Non-Cold Dark Matter           Omega = 0.00140587      , omega = 0.000637896
 Other Content                    Omega = 0.684664        , omega = 0.310658
 TOTAL                            Omega = 1               , omega = 0.453737
 --------------------------------------------------------------------
┌ Warning: Multi-threading over perturbation modes was requested, but BLAS is running with 2 threads.
│ It is recommended to restrict BLAS to one thread with `using LinearAlgebra: BLAS; BLAS.set_num_threads(1)`.
│ For more information, see https://docs.julialang.org/en/v1/manual/performance-tips/#man-multithreading-linear-algebra.
└ @ SymBoltz ~/work/SymBoltz.jl/SymBoltz.jl/src/solve.jl:414