Symbolic Implementation of Extensions of the $\texttt{PyCosmo}$ Boltzmann Solver

$\texttt{PyCosmo}$ is a Python-based framework for the fast computation of cosmological model predictions. One of its core features is the symbolic representation of the Einstein-Boltzmann system of equations. Efficient $\texttt{C/C++}$ code is generated from the $\texttt{SymPy}$ symbolic expressions making use of the $\texttt{sympy2c}$ package. This enables easy extensions of the equation system for the implementation of new cosmological models. We illustrate this with three extensions of the $\texttt{PyCosmo}$ Boltzmann solver to include a dark energy component with a constant equation of state, massive neutrinos and a radiation streaming approximation. We describe the $\texttt{PyCosmo}$ framework, highlighting new features, and the symbolic implementation of the new models. We compare the $\texttt{PyCosmo}$ predictions for the $\Lambda$CDM model extensions with $\texttt{CLASS}$, both in terms of accuracy and computational speed. We find a good agreement, to better than 0.1% when using high-precision settings and a comparable computational speed. Links to the Python Package Index (PyPI) page of the code release and to the PyCosmo Hub, an online platform where the package is installed, are available at: https://cosmology.ethz.ch/research/software-lab/PyCosmo.html.


Introduction
Our understanding of the Universe relies on the possibility to predict cosmological observables from theoretical principles. One of the key theoretical predictions is the evolution over time of the linear order perturbations of the constituents of the Universe, captured by the Einstein-Boltzmann equations (see, e.g., Ma & Bertschinger (1995), Dodelson (2003)). The system of ordinary differential equations, due to its complexity and the coupling of the fields, needs to be solved numerically (see, e.g., Nadkarni-Ghosh & Refregier (2017)). For this purpose, several codes have been developed since the release of the pivotal Boltzmann code COSMICS (Bertschinger, 1995), closely followed in time by CMBFAST (Seljak & Zaldarriaga, 1996), later ported to C++ with the name CMBEASY (Doran, 2005a). The currently maintained Boltzmann solvers are CAMB (Lewis et al., 2000), CLASS (Lesgourgues, 2011a) and PyCosmo (Refregier et al., 2018). Both for CAMB and CLASS several codes have been written to include extensions beyond ΛCDM, for example hi class (Zumalacárregui et al., 2017) and EFTCAMB (Hu et al., 2014) for modified gravity theories, CLASS EDE (Hill et al., 2020) for early dark energy and CLASSgal (Dio et al., 2013) to include general relativistic effects in the computation of galaxy number counts. PyCosmo was introduced by Refregier et al. (2018)

as a novel
Python library that uses symbolic representation of equations for generating efficient C/C++ code.
The framework includes both a Boltzmann solver as well as prediction tools for the computation of cosmological observables with several different fitting functions and approximations (Tarsitano et al., 2021). With these tools, PyCosmo offers similar utilities as, e.g., the Core Cosmology Library CCL (Chisari et al., 2019) developed by the Dark Energy Science Collaboration (DESC).
An important feature of PyCosmo is the possibility to easily implement model extensions in symbolic form in the code, while taking advantage of the computational speed of the generated C/C++ code. This feature has been improved by rewriting and refactoring the related code as a new sympy2c package presented in Schmitt et al. (2022), which expands the idea of generating fast C/C++ code from symbolic representations of equations. Both sympy2c and PyCosmo are publicly available in the Python Package Index (PyPI). In this work, we illustrate how the PyCosmo Boltzmann solver can be extended thanks to the symbolic framework by implementing several extensions. We 3 introduce two extensions of the Standard Model of Cosmology: a constant dark energy equation of state and massive neutrinos. We also include a radiation streaming approximation (RSA) for photons and massless relics, which approximates the evolution of radiation at late times and includes reionisation, following the treatment in Blas et al. (2011) implemented in CLASS. This approximation speeds up the code by reducing the number of equations in the ODE system and avoiding the reflection of power caused by the truncation of the multipole expansion of the radiation equations.
We begin by giving an overview of the new features of the PyCosmo framework in Section 2, where we discuss also the usage of the code and the precision settings used to compare with CLASS.
We then present the equations of the models, implementation details and code comparisons with CLASS, both in terms of agreement and performance. In Section 3 we present the implementation of the constant dark energy equation of state, since it is a minimal modification of the Boltzmann system of equations for ΛCDM. We describe the inclusion of massive neutrinos, treated as noninteracting and non relativistic relics, in Section 4. In Section 5 we present the Radiation Streaming Approximation, which requires sympy2c to handle a switch between two different equation systems and is then applied to all models. In Section 6 we discuss the results obtained by benchmarking the speed of the computations and comparing the numerical results to CLASS. We conclude in Section 7. This work heavily relies on the Boltzmann equations presented in Ma & Bertschinger (1995). To translate the PyCosmo notation to the Ma-Bertschinger and CLASS notation, we refer the reader to Appendix A. Appendix B presents the Einstein-Boltzmann system of ODE in PyCosmo notation, using ln a as the independent variable. The adiabatic initial conditions for ΛCDM and all the other implemented models are shown in Appendix C. We report in Appendix D the parameters of PyCosmo and CLASS that have been kept constant throughout the paper. In Appendix E, we provide a self contained summary of the computation of the total matter power spectrum, including general relativistic corrections, and using the A s normalisation parameter.

C/C++ code generation
We reimplemented and improved the C/C++ code generation related parts of previous versions of PyCosmo (Refregier et al., 2018) as a separate Python package named sympy2c which we describe in detail in Schmitt et al. (2022). sympy2c translates symbolic representations of expressions and ordinary differential equations to C/C++ code and compiles this code as a Python extension module.
sympy2c replaces the Backward-Differentiation-Formula (BDF) solver from the previous version of PyCosmo by the established and robust Livermore Solver for Ordinary Differential Equations (LSODA) solver (Petzold, 1983) for improved step-size control and error diagnostics. LSODA detects stiff and nonstiff time domains automatically and switches between the nonstiff Adams method and the stiff BDF method. The BDF method solves a linear system derived from the Jacobian matrix of the differential equations at each time step. This affects runtime significantly for large systems. sympy2c leverages the symbolic form of the ODE and generates code to solve such systems efficiently by avoiding unnecessary computations based on the known sparsity structure of the involved Jacobian matrix.
To solve such linear systems sympy2c unrolls loops occurring in used LU factorization with partial pivoting (LUP) algorithm during code generation. This procedure depends on predetermined row permutations of the system, and the generated code includes checks for whether the considered permutation is appropriate for ensuring numerical accuracy. When solving the ODE, a new, not yet considered, permutation might arise. In this case, the solver delegates to a fall-back general LUP solver and records the new permutation. The result is a valid result but with a sub-optimal computation time. In this case, the warning message "there are new permutations pending, you might want to recompile" will be displayed together with the command necessary to recompile.
Running the code-generator again will then also create optimized code for the newly recorded permutation(s), so that future runs of the solver will benefit from this. This approach starts with the identity permutation and could require several steps of solving the Boltzmann equations followed by code generation and compilation to achieve optimal performance. In our experiments not more than one such iteration is needed.
A large ODE systems can result in C/C++ functions with millions of lines of code which challenge the compiler and can cause long compilation times and high memory consumption, especially during the optimization phase of the compiler. To mitigate this, sympy2c can split the original matrix into smaller blocks and then generate code to implement blocked Gaussian elimination using Schurcomplements. This affects the generated C/C++ code by creating more but significantly shorter functions and thus supports the optimization step of the compiler. Another benefit of this approach is that runtime is improved by reducing the number of cache misses on the CPU.
Depending on the size and sparsity structure of the system, this code generation and compilation step can take seconds up to 30 minutes or even more. PyCosmo and sympy2c use caching strategies 5 that consider previously generated code so that cached solvers are available within fractions of a second.

Usage
The equations for ΛCDM and the extended models are implemented symbolically in PyCosmo, both with and without radiation streaming approximation, in the CosmologyCore model .py and CosmologyCore model rsa.py files. Supported models are currently "lcdm", "wcdm" or "mnulcdm".
The method PyCosmo.build initializes an instance of the Cosmo class for subsequent computations. Cosmo is the class that manages most of the functionalities of PyCosmo and on which all the other classes rely. PyCosmo.build requires the name of the model as well as all parameters which influence the code generation and compilation step. The argument rsa enables or disables the RSA and l max specifies the maximum moment for truncating the photons and massless neutrinos hierarchies. The "mnulcdm" model also accepts parameters l max mnu and mnu relerr which we describe later. Furthermore, parameters controlling the compiler, such as the optimization flag -On and the splits to use to reduce memory consumption and compilation time, can be specified when calling PyCosmo.build. All the parameters that can be passed to build are specified in Table 1.
Parameters which do not affect code generation, such as cosmological parameters, precision settings, parameters specific for approximations and physical constants can be set or modified using the Cosmo.set method. Each cosmological model is equipped with a default set of such parameters, contained in a default model .ini file. Listing 1 demonstrates how to create a cosmology and change parameters. One parameter that is particularly relevant is pk type, since it allows the user to switch between the Boltzmann solver (pk type = "boltz") and the approximations (pk type = "EH" for the fitting function by Eisenstein and Hu (Eisenstein & Hu, 1998) and pk type = "BBKS" for the 6 BBKS polynomial fitting function (Peacock, 1997)). In this work we will always set pk type = "boltz". Other cosmological and precision parameters which are kept fixed throughout the paper are reported in D. Tutorials for the computation of cosmological observables and the usage of the Boltzmann solver can be found on the PyCosmo Hub (see Tarsitano et al. (2021)), a public platform hosting the current version of PyCosmo, along with CLASS, CCL, and iCosmo (Refregier et al., 2011), an IDL predecessor of PyCosmo. The link to the PyCosmo Hub can be found at https://cosmology.ethz.ch/research/software-lab/PyCosmo.html.

Code comparisons setup
In order to validate the newly introduced models, we carry out detailed comparisons with CLASS 1 .
We evaluate the accuracy in terms of relative difference: where X is the cosmological observable we want to compare, for example the total matter power spectrum. Since this is a function of the wavenumber k, we compare it visually by plotting as a function of k or, when specified, we look at the maximum relative difference in a k interval.
Comparing the two codes implies carefully setting the cosmological parameters and the precision settings. In the case of CLASS, we use the precision files shipped with the code: cl permille.pre for fast and accurate computation and pk ref.pre for high precision. cl permille.pre guarantees a precision of 0.1% up to k = 1hMpc −1 for the matter power spectrum, and pk ref.pre a precision of 0.001% on scales k < 50hMpc −1 (Lesgourgues, 2011b). Both CLASS precision settings use the Radiation Streaming Approximation, described in detail in Section 5. They also include the Tight Coupling Approximation and Ultra-relativistic Fluid Approximation (see Blas et al. (2011)), whereas only cl permille.pre uses the fluid approximation for massive neutrinos (presented in ). The main parameters controlling precision in PyCosmo are the l max parameter which defines the truncation of the multipole hierarchy for radiation fields (same for photons and massless neutrinos) and the boltzmann rtol and boltzmann atol parameters defining the relative and absolute tolerance of the LSODA ODE solver. Massive neutrinos also add two important precision parameters which will be described more in detail in section 4: mnu relerr 1 We use CLASS v3.1.0 throughout the paper, through the Python wrapper classy.
When using the RSA in PyCosmo, we set the RSA trigger parameters (detailed in Section 5) to:  Table 1. We do not attempt to exactly match the precision parameters in the two packages, since CLASS includes a number of approximations that are not available in PyCosmo. It is possible to switch off most of the approximations but this would imply losing the precision guarantees of the default precision files.
The cosmological parameters that remain constant throughout the paper are listed in Appendix D, whereas the matter energy density Ω m , the number of massless neutrinos N ν , the number of massive neutrinos N ν,m , the total sum of the massive neutrinos Σm ν and the dark energy equation For a constant dark energy equation of state w de , the dark energy density is given by 2 We immediately see that ρ = const. for a cosmological constant with w de = −1. The Hubble parameter is given by In this equation, Ω r is the radiation density which includes photons and massless neutrinos, Ω m the matter density, Ω κ the curvature density (listed for completeness, even though the Boltzmann solver in PyCosmo currently only supports flat models) and Ω Λ the dark energy density. In all cases, 2 In the code, w de is written as the parameter w0.
Ω i is defined as the fraction of the energy density of the corresponding component today and the critical energy density ρ crit of the Universe.
In the case of a cosmological constant, there are no dark energy perturbations. For w de = −1, we can write down the dark energy equations following Ma & Bertschinger (1995); Dodelson (2003) and obtainδ where all derivatives are with respect to the conformal time η and we use the conformal Newtonian gauge as in Refregier et al. (2018). The anisotropic stress, σ de , vanishes, which deletes one term in the second perturbation equation.
In general, the sound speedc 2 s,de = δp de δρ de is a Gauge-dependent variable, which can be expressed in terms of the rest frame sound speed c 2 s,de and the adiabatic sound speed c 2 a . The latter is equal to w de for a constant dark energy equation of state (Ballesteros & Lesgourgues, 2010). Here, we use the expression (Ballesteros & Lesgourgues, 2010) which is valid for a constant dark energy equation of state.
Inserting this expression in Eq. 3 we obtaiṅ Compared to the system of equations in Refregier et al. (2018), the Einstein equations are modified to In order to evolve the perturbation equations, we also need to define the initial conditions for these. We choose the adiabatic initial conditions from CLASS, outlined in Ballesteros & Lesgourgues (2010) in the synchronous gauge, and then transform them into the conformal Newtonian gauge.
We report the initial conditions for all the fields in Appendix C.

Numerical implementation
In order to implement the new equations outlined in the previous section, we generated a new file for the symbolic Boltzmann equations, called CosmologyCore wcdm.py. This can be used instead of the default equations file CosmologyCore.py for the ΛCDM model, by setting the model to "wcdm" as shown in Listing 1.
In PyCosmo, all derivatives are written with respect to ln a (in the code lna) instead of η.
The perturbation equations for a ΛCDM model in this notation are detailed in B. In the previous section, we presented the dark energy perturbation equations with respect to η. Using the conversion d ln a dη = aH, we can rewrite the two dark energy perturbation equations from Eq. 5 as Then the two dark energy perturbation equations in Eq. 9 can be expressed with SymPy as CosmologyCore wcdm.py also contains the background equations and the initial conditions for the linear perturbations.

Code comparisons
In Figure 1 we show the dark energy perturbations δ de and u de as a function of scale factor a for three wavenumbers k = 0.005, 0.05 and 5 Mpc −1 , plotted both with PyCosmo as well as CLASS.  We also display the relative difference between the evolution of the perturbations obtained by the two codes. The cosmology we consider is with all other parameters as specified in D and using the precision settings from 2.3. In general, we find good agreement between the codes. When the fields are highly oscillating around zero, we observe a degradation of the agreement, as expected, given the impact of step-size control and numerical precision of the solver in that regime. We also notice a discrepancy at initial time for small values of k, which is caused by the tight coupling approximation in CLASS. In Figure 2 we show the wCDM total matter power spectrum computed with the two codes for 200 log-spaced k values between 10 −4 and 10 Mpc −1 at redshifts z = 0, z = 1 and z = 5. In general, we observe that the results for wCDM show the same level of agreement with CLASS as ΛCDM. The discrepancies tend to grow on large scales for z = 5 but remain below the 10 −3 level. The same is observed for higher redshifts. In Section 6 we summarize the comparisons in terms of computing time and power spectrum relative difference for different k ranges at z = 0.

Equations
Oscillation experiments provide evidence that neutrinos have mass (see e.g. Fukuda et al. (1998); Ahmad et al. (2002); Eguchi et al. (2003) and also de Salas et al. (2021) for a recent global fit of neutrino oscillation data) and since their masses are imprinted onto cosmological observables, we need to include the evolution of light massive relics in the system of equations. This allows cosmological probes to constrain the properties of neutrinos, in particular the sum of the neutrino masses (see e.g. Lesgourgues & Pastor (2006); Lesgourgues et al. (2013); Lattanzi & Gerbino (2018) for reviews on neutrino cosmology).
In this section, we present the implementation of the Einstein-Boltzmann equations for massive neutrinos into the PyCosmo Boltzmann solver. Massive neutrinos modify both the background evolution and the linear order perturbations. Qualitatively, massive neutrinos undergo a phase transition: they behave like radiation at early times, when they are fully relativistic, and shift to a matter-like behaviour at late times (see, e.g., Lattanzi & Gerbino (2018); Ichikawa et al. (2005)).

13
The transition happens smoothly through cosmic time and the dependence on mass, scale factor and momentum in the evolution of the distribution function prevents from integrating out the momentum dependence. This can be done only when considering approximations. For this reason, the inclusion of massive neutrinos in the Boltzmann equations is highly non trivial and has a strong impact on the size of the ODE system.
In this section we write the equations for N ν,m massive neutrinos with degenerate masses and total neutrino mass sum Σm ν , where the sum goes over the three neutrino mass eigenstates. A generalization to neutrinos with different masses is simply achieved by suppressing the N ν,m in front of the equations and writing separate equations for each neutrino species with mass m ν,i . In PyCosmo we introduce degenerate massive neutrinos. The introduction of neutrino hierarchies is left as future development.
The massive neutrino density can be written as where q = ap, with p the proper momentum, related to the 4-momentum by P i = a(1 − Φ)p i in the Newtonian conformal gauge. Then = q 2 + a 2 m 2 is the proper energy measured by a comoving observer multiplied by the scale factor and f 0 (q) is the Fermi-Dirac distribution where g s is the spin degeneracy factor that equals 2N ν,m in the case of degenerate neutrinos.
where Ω r still includes massless neutrinos if present (denoted with the subscript ν such that Ω r = Ω γ + Ω ν ), whereas the massive neutrino energy density is Ω ν,m (a). Note that we factor out the a −4 term from Ω ν,m (a) for similarity with the other energy densities, but Ω ν,m (a) still contains a dependency on the scale factor a, differently from the other species, since the proper energy contains a factor of a that we cannot integrate out (see equation 10, Ω ν,m (a)a −4 = ρ ν,m (a)/ρ crit ).
The massive neutrino perturbations arise from a linear expansion of the distribution function f (x, q,q, η) = f 0 (q) [1 + M(x, q,q, η)] around the Fermi-Dirac distribution f 0 (q). The function M(x, q,q, η) is Fourier transformed and expanded in a Legendre series as with µ =k ·q, P l the Legendre polynomials and M l defined as The massive neutrino Boltzmann equations are then derived similarly to those of the ultra-relativistic fields, setting the collision term to 0, since neutrinos are only weakly interacting. Using the definition of M l to express the Boltzmann equations as a hierarchy of moments, we obtaiṅ The hierarchy is truncated at a multipole l max (l max mnu in the code) when solving the system of equations numerically with a hierarchy truncation from Ma & Bertschinger (1995), which is analogous to that for photons and massless neutrinoṡ Massive neutrinos also modify the Einstein equations due to the extra terms in the stress-energy tensor. These now read where Ω m δ m is a shortcut for Ω dm δ + Ω b δ b , Ω r Θ r0 = Ω γ Θ 0 + Ω ν N 0 and ρ crit is the critical energy density of the Universe. The massive neutrino quantities we are interested in are the density fluctuation, the fluid velocity and the shear stress which are computed by integrating over the 15 moments M l : Note that in the PyCosmo implementation, we substitute q with q = q k b T0 for convenience.

Numerical implementation
The implementation of the massive neutrino equations uses sympy2c similarly to the PyCosmo implementation of ΛCDM and wCDM. The background integral over momentum q is computed using indefinite numerical integration from sympy2c , whereas the integrals at perturbation level use a Gauss-Laguerre quadrature integration scheme in our symbolic representation of the ODE system, since we need to evolve a finite number of equations. This follows the approach used in . The number of discrete q values is governed by the parameter mnu relerr, which sets the relative difference between the Gauss-Laguerre integration of a test function ( n=4 n=2 q n f 0 (q)) with respect to its analytical result. This parameter is specified using PyCosmo.build since it affects the size of the ODE system and thus also code generation. We introduce an additional parameter influencing code generation in addition to the truncation parameter l max of the photon and massless neutrino hierarchies: l max mnu that truncates the Legendre series M l as described above.
The implementation of massive neutrino cosmologies results in a large systems of equations and thus C functions with millions of lines of generated code, challenging the optimizer of the used C compiler. To mitigate significantly compilation time and memory requirements, we enable sympy2c to use the matrix splitting feature (enabled by specifying the splits parameter and reorder=True in PyCosmo.build) described in Section 2. The user can also pass a compilation flags parameter, that enables or disables compiler optimizations of the C code and has a diametrical effect on compilation vs. runtime.
We implement initial conditions for massive neutrinos that match the adiabatic initial conditions in CLASS and Cosmics and can be triggered using the initial conditions parameter. We report the equations in Appendix C.

Code comparisons
One of the key effects of massive neutrinos on cosmological observables is the suppression of small scale matter overdensities due to neutrino free streaming (Lesgourgues & Pastor, 2006). In 2010; Bonvin & Durrer, 2011;Challinor & Lewis, 2011;Dio et al., 2013). In general, we observe a very good agreement, with a maximum relative discrepancy of 5 × 10 −4 on intermediate scales.
The redshift evolution does not impact the agreement. We also show the suppression of the total matter power spectrum with respect to the power spectrum with massless neutrinos in Figure 4 both for Σm ν = 60 meV and 120 meV (Ω m =0.29737), when keeping Ω m,tot = Ω m + Ω ν,m fixed to Ω m,tot = 0.3 and N eff = 3.044. The cosmological parameters for the ΛCDM model are set to: 5. An approximation scheme: Radiation Streaming Approximation

Equations
After decoupling, photons and massless neutrinos behave approximately like test particles freestreaming in the gravitational field determined by the massive components, making it possible to derive a non-oscillatory solution of the inhomogeneous Boltzmann equations inside the Hubble radius. This approximation, called the Radiation Streaming Approximation (RSA), was introduced in the Newtonian gauge by Doran (2005b) and in the synchronous gauge by Blas et al. (2011). This treatment allows both to avoid unphysical oscillations resulting from the hierarchy truncation and to speed up the integration, which is slowed down by fast late time oscillations of the radiation fields especially on small scales. At late times, the approximation does not need to be precise, since the contribution of the radiation energy density to the overall energy density is negligible. This approximation only impacts the linear perturbations. The evolution of the relativistic fields can be In the right plot, we display the same quantities when setting l max = 50 for all the radiation fields and massive neutrinos and switching off the UFA in CLASS.
written as 19 with all the higher order multipoles set to 0. This approximation is switched on when two conditions are satisfied, following the CLASS (Blas et al., 2011) scheme: corresponding to decoupled radiation within the horizon.

Numerical implementation
In order to implement the Radiation Streaming Approximation within PyCosmo we use a functionality from sympy2c to switch between two different ODEs at a dynamically computed time point.
This requires: • Two ODE systems specified as sympy2c OdeFast objects. In our case these will be the symbolic representations of the model of interest (ΛCDM, ΛCDM with massive neutrinos or wCDM) and its equivalent RSA system. Note that the two systems can have different dimensions.
• A switch time function which determines at which time to switch from the first system of equations to the second. In the RSA implementation we use the switching conditions described above.
• A switch function, computing the initial conditions for the second system of equations from the state of the first system before and at the switching time. In the RSA implementation, this function just discards the matrix entries for the fields Θ i , Θ P i and N i , i ≥ 0 since Θ 0 , Θ 1 , N 0 and N 1 have analytical expressions and thus do not need initial conditions.
• A merge function which specifies how to combine the matrix valued results from both numerical solutions into a final matrix. Our RSA implementation keeps the full matrix from the solution of the full system of equations and extends the matrix from the RSA solution using the analytical formula for Θ i and N i for i ≤ 1. All other entries are set to 0 for i ≥ 2 and for all the polarization terms Θ P i .

20
In order to use the RSA in PyCosmo one needs to pass the rsa = True flag to PyCosmo.build for any of the models. This will switch the CosmologyCore model .py equation file to a CosmologyCore model rsa.py file containing the RSA equations for the radiation fields and all the other equations of the system.

Internal code consistency
The main achievement obtained by implementing the RSA is a significant reduction of the computation time for the fields, especially for high values of k, as we will discuss in detail in the next section. This is especially true when choosing the ΛCDM and wCDM models, where most of the system of ODEs consists of radiation perturbations. In Figure 5, 10 −3 and 10 Mpc −1 . We observe that the oscillations around zero correspond to a relative difference of the order of 0.001%, but the computation time is reduced by approximately 80%. In Figure 5 we also display the relative difference between the power spectrum computed using PyCosmo speed (panel a) and precision (panel b) settings from Section 2.3, both with and without RSA, and the reference power spectrum just described.
Reducing the l max parameter to 50 (precision settings) increases the discrepancy with the reference power spectrum to ∼0.08% for the full system, regardless of whether the RSA is turned on or not. For the speed settings, the discrepancy is ∼0.6% both when using or not using the RSA.
The computation time is again reduced by roughly 80% when using RSA compared to solving the full equation system with the same l max, meaning that the approximation is essential to reduce the computational time of the perturbations for mid to high k values without sacrificing the accuracy.
The unphysical reflection of power caused by the hierarchy truncation dominates on small scales, making the inaccuracy introduced by the approximation completely negligible.  (bottom panel) settings with and without RSA. We also display in both panels in light blue the relative difference between the reference matter power spectrum with and without RSA. The grey horizontal lines correspond to a 10 −3 (dash-dotted) and 10 −4 (dotted) precision level.

Agreement and performance comparison with CLASS
In this section, we present benchmarks of the PyCosmo Boltzmann solver for all the models described in the previous sections, both in terms of relative difference to CLASS and in terms of computing times. The cosmological parameters that are modified in each model are specified in the ΛCDM, wCDM and degenerate Σm ν = 60 meV parameter settings in the previous sections, while the fixed parameters are reported in D. In order to perform a comparison purely on the Boltzmann solver, we read in the CLASS recombination files for each model and set the same initial  conditions 4 . All the computations are carried out on a single core on a full node of the ETH Zurich Euler cluster 5 , by disabling parallel execution in CLASS and not enabling parallel computation for the PyCosmo power spectrum. We run on a full node on the cluster, instead of on a laptop, in order to only run the Boltzmann solver and not get impacted by other processes being executed by the operating system. In the case of PyCosmo, we only report the time necessary for the power spectrum computation, which does not include the compilation time for the first time the C/C++ code is generated for each model. Most models also require a second compilation (recompilation) that applies permutations to the existing equations in order to enable the use of specialized solvers instead of the standard solver as explained in Section 2.1. The runtime we report is the best time obtained in three executions.
We begin by comparing the runtime between the computation of the matter power spectrum in PyCosmo, with and without turning on the RSA, and CLASS, both for the speed and precision settings (defined in Section 2.3). In Table 2, we display the time necessary to compute the power spectrum with PyCosmo and CLASS for a 100 log-spaced k values between k min and k max . We fix k min = 10 −4 Mpc −1 , and set k max first to 1 Mpc −1 and then to 10 Mpc −1 , where the effects of the Radiation Streaming Approximation are most evident.
We observe that, while PyCosmo achieves a slower runtime than CLASS before introducing the RSA, the RSA reverts the situation for all the models, except for massive neutrinos. This happens despite the presence of further approximations in the CLASS implementation. Previous versions of PyCosmo achieved a comparable execution time with CLASS without physical approximations (Refregier et al., 2018). This was due to the reduction of the dynamic range of the time step based on a consistency relation of the Einstein equations. The adaptive control of the time step was highly optimized for ΛCDM and proved difficult to extend to more general models. The new approach has also the advantage of reducing the number of permutations necessary to create the optimized code.
In the case of massive neutrinos, the size of the system determines a considerable decline in the performance of both codes. PyCosmo is significantly slower than CLASS when using the speed settings, mainly due to the presence of the fluid approximation for non cold dark matter in CLASS. For the precision settings, the time needed for the computation is comparable for CLASS and PyCosmo. The ODE system used by CLASS is larger in this case, since the tol ncdm parameters are set to 10 −10 in pk ref.pre, and the three neutrinos are treated independently despite having the same mass. We do not decrease mnu relerr, equivalent to tol ncdm, in PyCosmo since we do not deem it necessary to achieve the desired precision. The error caused by sampling less q values is always subdominant compared to the hierarchy truncation, due to the small contribution of massive neutrinos to the overall matter density.
We compare the numerical results of PyCosmo and CLASS in Table 3. We report the maximum relative difference between the matter power spectra in PyCosmo and CLASS in five k ranges for the different models and precision settings. The RSA is not relevant here because it leads to a 10 −5 relative difference, when compared to the full computation with the same l max, as shown for ΛCDM in the previous section.
We start by noting that the size of the relative differences is comparable in all models when using the same k ranges, with the exception of the model with massive neutrinos for k > 0.1 Mpc −1 .
We also see that the difference in precision between speed and precision settings is dominant on scales k > 0.1 Mpc −1 , where the truncation effects have the largest impact. The precision settings lead to a relative difference to CLASS that is better than 0.1% on all scales k < 10 Mpc −1 , while the speed settings lead to a relative difference of order 0.5 − 1% beyond k = 0.1 Mpc −1 . This is acceptable, especially since cl permille.pre comes with no guarantees for scales k > 1 hMpc −1 .

Conclusion
In this paper, we demonstrated how the PyCosmo Boltzmann solver can be easily modified to include extensions of the ΛCDM cosmological model and approximation schemes, by taking advantage of the SymPy symbolic implementation of equations. The symbolic expressions are translated into optimized C/C++ code by the sympy2c package presented in Schmitt et al. (2022). In this way, PyCosmo combines the speed of C/C++ with the user-friendliness of symbolic Python.
We first presented two cosmological model extensions: dark energy with a constant equation of state, which is a minimal modification of ΛCDM, and massive neutrinos, which enlarge considerably the ODE system and comprise numerical integrations, constituting a more complex extension. The inclusion of these models makes PyCosmo more widely applicable for constraining cosmology. We also implemented an approximation scheme, the radiation streaming approximation. In order to trigger the approximation, sympy2c includes a functionality to switch between two different ODE systems when a condition is verified. The radiation streaming approximation makes the solution of the ODE system considerably faster, up to an 80% speed-up, since it suppresses oscillations of the radiation fields for large k values. The errors introduced by the approximation are largely sub-dominant compared to the artificial power reflection induced by the hierarchy truncation. For convenience, we presented a conversion table between common conventions for cosmological perturbations (A) and a clarification of the computation of the total matter power spectrum with A s normalization (E).
We compared the numerical results obtained by computing the total matter power spectrum 25 with CLASS and found an agreement better than 0.1% with high precision settings. With more relaxed precision settings, we found an agreement of 0.5% for scales k < 1Mpc −1 for all models.
The PyCosmo Boltzmann solver achieves precision and speed that is comparable to CLASS, while not relying on physical approximations (such as tight coupling and ultra relativistic fluid approximation) other than the radiation streaming approximation introduced in this work. In the future, we plan to include more beyond ΛCDM models in PyCosmo. Possible extensions include time-varying dark energy, early dark energy (Poulin et al., 2019;Hill et al., 2020), curvature (Pitrou et al., 2020), dark matter models (Hui, 2021), such as axions and fuzzy dark matter, and extensions of the neutrino sector. We believe that our symbolic implementation will be applicable for most model extensions, with some refactoring needed when a model introduces an ODE system already at background level (for example in the case of scalar field models).

Acknowledgements
This work was supported in part by grant 200021 192243 from the Swiss National Science Foundation. The authors thank Thomas Tram for answering questions about the CLASS implementation.
We thank Joel Mayor for his comments that led to some corrections and clean up of the code, and Silvan Fischbacher for helpful discussions. The authors also thank Joel Akeret, Adam Amara, Lukas Gamper, Jörg Herbel, Tomasz Kacprzak and Andrina Nicola for discussions and contributions to earlier versions of PyCosmo. In this work, we rely on the Python packages numpy (van der Walt et al., 2011), scipy (Virtanen et al., 2020), matplotlib (Hunter, 2007) and SymPy (Meurer et al., 2017 Table 4: Relations between the notation in Ma-Bertschinger and PyCosmo. Ma-Bertschinger notation is used in both the Boltzmann solvers COSMICS and CLASS. Some exceptions include θur that is denoted θr in COSMICS and all the massive neutrino quantities that are denoted δ ncdm , θ ncdm and σ ncdm in CLASS (ncdm = non cold dark matter) and δn, θn and σn in COSMICS.
• Photons polarization: For l > 2 dΘ P,l dlna = k aH(2l + 1) (lΘ P,l−1 − (l + 1)Θ P,l+1 ) +τ aH Θ P,l • Massless neutrinos: The hierarchy truncations for relativistic species read: The equations for the extended models are easily obtained with the same change of variables from the equations in η, reported in the corresponding sections of the paper.
The fields where we did not specify a conversion are gauge invariant, including the massive neutrinos' perturbations of the distribution function. The initial conditions imposed to the multipoles of M(µ) N 3 e −q/k b T0 + 1 and are set after transforming the other fields to the conformal Newtonian gauge. We can then relate the fields from CLASS to those of PyCosmo with the following conversions (already introduced in Appendix A): u = θ c /k u b = θ b /k u de = θ de /k Θ 0 = δ g /4 N 0 = δ ur /4 Θ 1 = θ g /3k N 1 = θ ur /3k N 2 = σ ur /2 N 3 = l3 ur /4.

Appendix D Fixed cosmological parameters
We report all the parameters of the configuration file of PyCosmo that are kept fixed throughout the paper: Note that some parameters (for instance T ν,m or c 2 s,de ) are specified only when necessary.