Beyond the Runge-Kutta-Wentzel-Kramers-Brillouin method
Jamie Bamber Will Handley
physics.comp-ph astro-ph.CO astro-ph.IM


APS/123-QED





[]jarb2@alumni.cam.ac.uk
Astrophysics Group, Cavendish Laboratory, J.J.Thomson Avenue, Cambridge, CB3 0HE, UK
Pembroke College, Pembroke Street, Cambridge, CB2 1RF, UK


[]wh260@mrao.cam.ac.uk
Astrophysics Group, Cavendish Laboratory, J.J.Thomson Avenue, Cambridge, CB3 0HE, UK
Kavli Institute for Cosmology, Madingley Road, Cambridge, CB3 0HA, UK
Gonville & Caius College, Trinity Street, Cambridge, CB2 1TA, UK




             


We explore higher-dimensional generalizations of the Runge-Kutta-Wentzel-Kramers-Brillouin method for integrating coupled systems of first-order ordinary differential equations with highly oscillatory solutions. Such methods could improve the performance and adaptability of the codes which are used to compute numerical solutions to the Einstein-Boltzmann equations. We test Magnus expansion-based methods on the Einstein-Boltzmann equations for a simple universe model dominated by photons with a small amount of cold dark matter. The Magnus expansion methods achieve an increase in run speed of about 50% compared to a standard Runge-Kutta integration method. A comparison of approximate solutions derived from the Magnus expansion and the Wentzel-Kramers-Brillouin (WKB) method implies the two are distinct mathematical approaches. Simple Magnus expansion solutions show inferior long range accuracy compared to WKB. However we also demonstrate how one can improve on the standard Magnus approach to obtain a new “Jordan-Magnus” method. This has a WKB-like performance on simple two-dimensional systems, although its higher-dimensional generalization remains elusive.   



                              
Beyond the Runge-Kutta-Wentzel-Kramers-Brillouin method
    Will Handley
    [10pt]0pt0pt December 30, 2023
=======================================================







§ INTRODUCTION


In the new era of high precision cosmology, one of the most important tools in the cosmologist's toolbox is the suite of Boltzmann codes. These methods, including class <cit.>, camb <cit.>, and pycosmo <cit.> provide numerical solutions to the Einstein-Boltzmann equations, which describe the evolution of linear perturbations in the Universe. The dominant computational cost for these codes is integrating coupled systems of linear first order ordinary differential equations (ODEs) with highly oscillatory solutions <cit.>. Currently the codes use several numerical and semianalytic approximations to get around these bottlenecks. However these approximations are designed for specific cosmological theories. A more generalized theory-independent method would make it easier to switch between different models and test new extensions to the current theory. New numerical integration methods could also reduce the run-time of the code. While the time needed to run a single simulation is currently only a few seconds, the many thousands of simulation runs needed for a typical Monte Carlo analysis add up to a substantial computational cost. 

Handley, Lasenby and Hobson proposed in <cit.>, and developed with Agocs in <cit.>, a new numerical method based on a combination of the standard Runge-Kutta (RK) approach and the Wentzel-Kramers-Brillouin (WKB) approximation, dubbed the “Runge-Kutta-Wentzel-Kramers-Brillouin” (RKWKB) method. In highly oscillatory regions of the solution using the WKB approximation allows integration steps which span many oscillations, giving faster performance than standard RK methods. However RKWKB is limited to essentially one-dimensional systems <cit.>.

We can write a multidimensional coupled system of  ODEs in the form of (<ref>),

    x(t) = A(t) x,

    
where x is a vector of dynamical variables, A is a time-dependent matrix, and an overdot denotes differentiation with respect to time. 

In the mathematics field of numerical analysis there has been substantial work examining the formal solution of (<ref>), the Magnus expansion <cit.>. A number of authors have also used the Magnus expansion and related techniques to develop numerical methods to integrate highly oscillatory differential equations <cit.>, for example Iserles <cit.>.

These Magnus-based numerical integrators have been extensively applied to problems in quantum and atomic physics <cit.>, as the time dependent Schrodinger equation can be written in the form of (1.1) if we use a basis of wave functions <cit.>. Some alternative methods for solving coupled differential equations have also been explored in the context of primordial cosmology <cit.>. However, Magnus-based methods remain unexplored with respect to the Einstein-Boltzmann equations and cosmological systems in general. Furthermore, to our knowledge, no previous authors have tested these methods with adaptive step size control. 

In this paper we examine the application of Magnus-based techniques to the Einstein-Boltzmann equations. However, we also achieve results which apply to coupled systems of first order ODEs more generally, and hence have relevance to quantum systems, along with problems in many other areas of physics and applied mathematics <cit.>.



§ BACKGROUND




 §.§ The WKB approximation


The WKB method is an example of multiple scale analysis <cit.> applied to oscillatory systems. It is most commonly applied to equations of the form of (<ref>) <cit.>,


    x = -ω^2(t)x.


The derivation begins by considering a general solution of the form


    x = exp(∑_n=0^∞S_n(t)).


Using the method of separation of scales we can obtain expressions for S_n <cit.>


    S_0    = ± i ∫_0^tω(t_1)  t_1, 
    
        S_1    =  - 1/2ln(ω(t)), 
    
        2S_0 S_n + S_n-1 +    ∑_j=1^n-1S_j S_n-j  = 0      n ≥ 2.


Truncating with only the S_0, S_1 terms gives the classic WKB approximation:


    x(t) = A/√(ω)e^i ∫_0^t ω(t_1)  t_1 + B/√(ω)e^-i ∫_0^t ω(t_1)  t_1.


If we include only the S_0 term we obtain 


    x + ω^2 x = ± i ω(t) x,


implying we require |ω(t)| ≪ |ω^2(t)| for the approximation to be reasonable.



 §.§ The Magnus expansion


The Magnus expansion <cit.> provides a general solution to equations of the type (<ref>). As with WKB, it starts by by considering a solution in the form of the exponential of an infinite series (<ref>) <cit.>,


    x(t) = exp(∑_n=1^∞_n(t)) x(0),


the first terms of which are


    Ω_1(t)    = ∫_0^tA(t_1)  t_1, 
    Ω_2(t)    = 1/2∫_0^t t_1 ∫_0^t_1A(t_1)A(t_2) t_2, 
    Ω_3(t)    = 1/6∫_0^t t_1 ∫_0^t_1 t_2 ∫_0^t_2( A(t_1)A(t_2)A(t_3
       + A(t_3)A(t_2)A(t_1)  t_3.


Further terms can be generated by a recursive procedure <cit.> and involve an increasing number of nested integrals, commutators, and powers of A.

For the Ω series to converge, A must be sufficiently small in some sense. It can be shown <cit.> that the Magnus series is absolutely convergent for a general complex matrix A, for t_0 ≤ t ≤ T, if 


    ∫_t_0^T ‖A(s) ‖_2  s < π,


where ‖·‖_2 denotes the two-norm of a matrix, which is <cit.>


    ‖A(s) ‖_2 = _i=1,...,n√(λ_i(A^TA)),


where λ_i(Q) denotes the ith eigenvalue of matrix Q. However, the actual convergence domain can be larger than that given by (<ref>) <cit.>.



 §.§ Stepping numerical integration methods


A general stepping numerical integration method proceeds by using an approximate solution to the equation to advance from one integration point x_n at time t_n to the next x_n+1 at time t_n+1. For a system of first order linear ODEs such as (<ref>) we can write the approximate solution for initial conditions x(t_0) in terms of a solution matrix M(t_0, t) <cit.>,


    x(t) = M(t_0, t)x(t_0).


Then the general stepping procedure is


    x_n+1 = M(t_n, t_n+1)x_n.


Runge-Kutta methods work by using truncated Taylor series solutions for M(t_0, t). For example, Euler's method, the simplest Runge-Kutta method <cit.>, has


    M(t_0,t) = 1 + (t - t_0)A(t_0).


The RKWKB method (to second order in the WKB series) is obtained by replacing the Taylor series solution with the two independent WKB solutions in (<ref>). As the WKB solutions typically remain good approximations to the true solution over many oscillations, unlike truncated Taylor series solutions, the RKWKB method is able to take much larger steps than Runge-Kutta methods, which in general reduces the required run-time <cit.>. Handley et al. <cit.> and Agocs et al. <cit.> also enable the program to switch to a standard Runge-Kutta-Fehlberg method in slowly oscillating regions; however, in this paper we use RKWKB to refer to the simpler nonswitching approach. 

We can also obtain stepping numerical integration methods of the form of (<ref>) from the Magnus expansion. If we can compute the analytic forms of _i then we can use those directly. In that case the  series truncated at = Ω_1 + ⋯ + Ω_i where i = 2s-1 or 2s-2 for integer s gives a method order h^2s+1 in the step size h <cit.>. However, if _i cannot be computed analytically, we can use an approach from Blanes et al. <cit.> Sec. 5.4. The A is Taylor expanded about the midpoint t_1/2 = t_n + h/2 of each step, where h = t_n+1 - t_n is the step size.


    A(t) = ∑_j=0^∞ a_j(t - t_1/2)^j      a_j = 1/j!. [n]At|_t=t_1/2.


One can then use the defining recursion relation for the Magnus expansion to express  to the desired order in h. For example to order four we have 


    ^[4] = α_1 - 1/12α_1α_2 + 𝒪(h^5),


where α_j = h^ja_j-1. The α_j's can then be approximated using various quadrature rules (including Gauss-Legendre, Simpson, Newton-Cotes) and linear combinations of A evaluated at different points in the interval [t_n, t_n+1].



 §.§ The Einstein-Boltzmann equations


In this section we describe the background theory and the Einstein-Boltzmann equations that will be used in Sec. <ref>.



  §.§.§ Overview


We seek to express the Einstein-Boltzmann equations in the form of (<ref>) <cit.>. Working in the conformal Newtonian gauge in flat spacetime, the metric is given by 


    s^2 = a^2(η)[(1 + 2Ψ)η^2 - (1 - 2Φ)δ_ij x^i  x^j].


where the scalar perturbations are Ψ, Φ, and η is conformal time <cit.>. As we work at early times we neglect dark energy and treat the neutrinos as relativistic. We denote the photon and neutrino temperature perturbations as Θ and 𝒩 respectively. The analogous perturbation Θ_P describes the strength of the photon polarization. We treat the cold dark matter (CDM) and baryonic matter as nonrelativistic and describe them with overdensities δ, δ_b and peculiar velocities v, v_b. 

To achieve a linear set of equations we expand the perturbations in spatial Fourier modes, introducing wave number k, and we expand the relativistic perturbations in multipole moments Θ_l(k, η) for l = 0, 1, 2, …, ∞. The resulting 14 × 14 A(k, η) matrix for l ≤ 2 is shown in <ref>.

The conformal Hubble rate is ℋ = a/a, where an overdot is used to denote the derivative with respect to conformal time. We denote the nonconformal cosmic Hubble rate as H = ℋ/a. The τ = -n_e σ_T a is the conformal time derivative of the optical depth, τ, with n_e the mean electron density and σ_T the Thompson cross section. Mean densities are described using parameters of the form Ω = ρ/ρ_ where ρ is the mean density of a component and  ρ_ is the critical density = 3H^2/8π G. These parameters are denoted Ω_γ, Ω_ν, Ω_ and Ω_b for the photons, neutrinos, CDM and baryonic matter respectively.



  §.§.§ Photon and CDM universe




As a test case we consider a simplified universe with only photons and cold dark matter. By neglecting the baryonic matter we can avoid the complexities of the Coulomb interaction between the baryons and the photons. We also neglect the photon polarization and all harmonics l > 1, approximating the photons as a fluid <cit.>. In this limit the gravitational potential Φ = Ψ. By making these approximations we obtain a five-dimensional system with x = [Θ _0,Θ _1,δ,v,Φ]^T. The reduced dimensions makes this much more amenable than the 14-dimensional system, while still preserving some of the key features of the early universe. By testing on this simplified system we provide a proof-of-principle study for application to the full Einstein-Boltzman equations.  

For adiabatic perturbations and radiation domination we obtain a second order evolution equation for the potential <cit.>


    Φ + 4/ηΦ + 1/3k^2Φ = 0.


This has a general solution (<ref>) <cit.> 


    Φ = A_kj_1(x)/x + B_ky_1(x)/x,


where x = kη / √(3) and j_1, y_1 are first order spherical Bessel functions of the first and second kinds respectively. These take the form 

    j_1(x) = sin(x)/x^2 - cos(x)/x, 
    
        y_1(x) = -cos(x)/x^2 - sin(x)/x.


As radiation dominates the energy density we also find


    δ_r = 4Θ_0 = -2/3(kη)^2Φ - 2ηΦ - 2Φ,


where δ_r is the overdensity for radiation. Hence for kη≫ 1 we have Θ_0 ∝ -(kη)^2Φ∼ A_kcos(kη/√(3)) + B_ksin(kη/√(3)). This gives the acoustic oscillations of the photon fluid at the sound speed c_s=1/√(3) that one would expect. 
The solution in (<ref>) suggests the potential also oscillates for kη≳ 1, but with the amplitude decaying as 1/(kη)^2. 







To use the numerical methods we must first obtain the relevant A. Working from the 14 × 14 matrix in <ref> we can derive the reduced A matrix shown in (<ref>),


    [ Θ_0; Θ_1;   δ;   v;   Φ ] = 
    [   -2 ℋΩ _γ         -k   -1/2ℋΩ _          0 -k^2/3 ℋ-ℋ;        k/3          0          0          0        k/3;   -6 ℋΩ _γ          0   -3/2ℋΩ _       -i k  -k^2/ℋ-3ℋ;          0          0          0         -H       -i k;   -2 ℋΩ _γ          0   -1/2ℋΩ _          0 -k^2/3 ℋ-ℋ;            ][ Θ_0; Θ_1;   δ;   v;   Φ ]_.


We set the conformal time at matter-radiation equality to be η_eq = 1, set a(η_eq) = 1 and assume Ω_γ(η_eq) = Ω_(η_eq) = 1/2. This gives


    A(k, η) = 
    [          -1/η            -k          -1/4             0 -η  k^2/3-1/η;           k/3             0             0             0           k/3;          -3/η             0          -3/4          -i k   -η  k^2-3/η;             0             0             0        -1/η^2          -i k;          -1/η             0          -1/4             0 -η  k^2/3-1/η;               ]_.









§ COMPARING WKB AND THE MAGNUS EXPANSION






 §.§ Comparing analytic solutions




To investigate the relationship between the WKB approximation and the Magnus expansion we compared the lowest order approximate analytic solutions to (<ref>).

Equation (<ref>) can be written in the form of (<ref>) by setting


    x = 
        [ x; x ]_,    A = 
        [    0    1; -ω^2    0 ]_.


For the Magnus expansion, the first two  terms are given by 


    _1    = 
        [               0               t; -∫_0^t ω^2  t_1               0 ]_,
    _2    = 
        [  a  0;  0 -a ]_,


where 

    a = 12 t ∫_0^t ω^2  t_1 - ∫_0^t  t_1 ∫_0^t_1ω^2  t_2 = 12∫_0^t t_1^2 t⟨ω^2 ⟩ t_1.



Including only _1 we obtain x(t) = M(t) x(0), where M = exp(_1(t)) given by (<ref>),


    M = 
       [            cos(√(⟨ω^2 ⟩)t) 1/√(⟨ω^2 ⟩)sin(√(⟨ω^2 ⟩)t);   √(⟨ω^2 ⟩)sin(√(⟨ω^2 ⟩)t)            cos(√(⟨ω^2 ⟩)t) ]_.


We can then write x(t) as 


    x(t) = Acos(√(⟨ω^2 ⟩)t) + B1/√(⟨ω^2 ⟩)sin(√(⟨ω^2 ⟩)t).


The frequency is approximated as ω∼√(⟨ω^2 ⟩) = √(1/t∫_0^t ω^2  t_1). We also get some amplitude correction from the B or M_12 term. We denote this solution “first Magnus.” 




Including both _1 and _2 we obtain


    M   = 1cos(ω̂t) + (_1 + _2)1/ω̂tsin(ω̂t) 
       =
       [ 1 0; 0 1 ]cos(ω̂t) + 
        1/ω̂t[               a               t; -∫_0^t ω^2  t_1              -a ]sin(ω̂t)_.


In (<ref>) the frequency is approximated as ω̂, 

    ω̂ = √(⟨ω^2 ⟩ - (a/t)^2) = √(⟨ω^2 ⟩ - 1/4⟨ t^2 ⟨ω^2 ⟩⟩^2),


where ⟨ω^2 ⟩ denotes the time derivative of ⟨ω^2 ⟩. Then we obtain (<ref>) and (<ref>),


    x(t) = A[cos(ω̂t) + a/tω̂sin(ω̂t)] + B1/ω̂sin(ω̂t),



    x(t) = A√(⟨ω^2 ⟩)/ω̂cos(ω̂t + ϕ̂(t)) + B1/ω̂sin(ω̂t),


with corrections to both the phase and the amplitude, where ϕ̂ is a time dependent phase term. We denote this solution “second Magnus.” For our matrix in (<ref>) ‖A(s) ‖_2 = (ω^2, 1). So (<ref>) implies a convergence domain of T-t_0 < π, ∫_t_0^T ω^2  t < π. One can also see from (<ref>) that at sufficiently large times ω̂ becomes imaginary, giving exponentially growing or shrinking solutions. 

For the WKB approximation in Eq. (<ref>) the  term approximates the frequency as 

    ⟨ω⟩ = 1/t∫_0^tω(t_1)  t_1,


while the  term provides a correction to the amplitude.

We can write (<ref>) in the form x(t) = M(t_0, t)x(t_0). Setting t_0=0 for simplicity we obtain 


    M = [                             C + ω(0)2ω(0)^2S                                       1ω(0)S; (ωω(0)2ω(0)^2 - ω2ω)C - (ω + ωω(0)4ωω(0)^2)S                            ωω(0)C - ω2ωω(0)S ]_,


where C = (ω(0)/ω)^1/2cos(⟨ω⟩ t) and S = (ω(0)/ω)^1/2sin(⟨ω⟩ t) are the two independent solutions. Applying the stepping procedure in (<ref>) to this M gives you the RKWKB method (to second order) <cit.>.


We can thus observe that the Magnus expansion and the WKB method take fundamentally different approaches for solving the same problem. The classic WKB approximation approximates the oscillation frequency as ⟨ω⟩, while the first Magnus solution uses the root mean square, √(⟨ω^2 ⟩). The terms of the WKB series involve successively higher order derivatives of ω, while the successive Magnus expansion terms involve successively higher order integrals. 

The WKB approximation is specialized to highly oscillatory solutions, and fails when ω goes to zero. The Magnus expansion can cope with nonoscillatory solutions and is less specialized, which comes at the cost of inferior long range accuracy compared to WKB, as we shall see in Sec. <ref>.



 §.§ Testing on trial equations


These solutions were plotted for two trial equations taken from Handley et al. <cit.>, namely the Airy equation and the “burst” equation, which both have  analytic exact solutions. The Airy equation takes the form


    x(t) = -tx,



and the burst equation is given by


    x = -ω^2(t)x    ,   ω^2(t) = n^2-1/(1+t^2)^2,


where n is an integer.

The analytic WKB approximate solution (<ref>) and the first and second Magnus solutions (<ref>) and (<ref>) for the Airy equation (<ref>) are plotted in <ref>.



For the upper plot, with t_0 = 1, the WKB solution remains a very good approximation to the true solution; however, the first and second Magnus solutions perform poorly. The first Magnus reproduces the increase in oscillation frequency, but does not reproduce the changes in amplitude. The second Magnus solution breaks down completely around t-t_0 ∼π showing the exponential growth predicted in Sec. <ref>.









If we shift the starting point to t=11, as in the bottom plot in <ref>, the first Magnus solution follows the true solution more closely. However unlike the WKB solution it fails to follow the frequency of the true oscillation  at large times. The second Magnus solution still breaks down as predicted.

Next we consider the burst equation (<ref>). Figure <ref> shows the analytic WKB, first Magnus and second Magnus approximate solutions from t_0=-10 up to t=+10 with n=40. While WKB provides an excellent match to the true solution throughout the whole range, both Magnus solutions give very poor approximations. The first Magnus solution fails to follow the changes in amplitude or frequency, and the second Magnus solution once again diverges after a relatively short period.
 


Figure <ref> shows the result if we try to apply the WKB and Magnus solutions to only the small, central, highly oscillatory region. This region has a smaller ω/ω ratio, so it should be easier to model. Even then the WKB solution dramatically outperforms both Magnus solutions. 




  §.§.§ A closer look at the Magnus error


To show why the first Magnus solution compares so poorly in terms of accuracy we can differentiate the Magnus solutions. Let C_M and S_M be the two independent first Magnus solutions,


    C_M    = cos(√(⟨ω^2 ⟩)t), 
    
        S_M    = 1/√(⟨ω^2 ⟩)sin(√(⟨ω^2 ⟩)t).


Let 1 + ϵ=ω^2/⟨ω^2 ⟩ and ξ = ω/ω, and then


    C_M =    -ω^2(1 + ϵ^2/4(1 + ϵ))C_M  …
       … - ω^2ξ S_M - ω^2/t( ϵ - 3ϵ^2/4(1+ϵ)) S_M,



    S_M =    -ω^2(1 + ϵ^2/4(1 + ϵ))S_M  …
       … -[ ξϵ - 1/4t(2ϵ + ϵ^2) ] ( 1/t S_M - C_M ).


One can see that for ϵ = 0, ξ = 0 we get back (<ref>). As the errors can increase with ϵ as well as ξ, they are potentially much larger than for WKB.



§ THE JORDAN-MAGNUS METHOD



The WKB approximation provides good long range accuracy but is limited to one-dimensional systems. In contrast the truncated Magnus expansion provides a matrix-based approach suitable for multidimensional systems, but with inferior long range accuracy as we saw in Sec. <ref>. Ideally we want a way to combine the positive aspects of both WKB and Magnus, that is achieve WKB-like long range accuracy with a solution that works on multidimensional systems. In this section we introduce a new approach based on the Magnus expansion, dubbed the “Jordan-Magnus” method, which goes some way toward this goal.  

If the A matrix were time-independent, we could solve the system by diagonalizing A to obtain the independent oscillating modes. Consider introducing a linear transformation to (<ref>). Let x_P = P^-1x, and let J = P^-1AP for some transformation matrix P(t). Then


    x_P = A_Px_P = [J + (P^-1)P]x_P,


where (P^-1) denotes the time derivative of P^-1. We want to choose a P such that A_P is as diagonal as possible. Let A_P = Λ + K where Λ is diagonal and K has zeros on the diagonal. 

Once again we seek a solution in the form x(t) = M(t)x(0).
We can obtain an approximate solution by neglecting K, 


    M(t) ≈P(t)exp(∫^t_0 t_1)P^-1(0).


This avoids the potentially costly matrix exponential. Alternatively we can include K and obtain


    M(t) ≈P(t)exp(∫^t_0 ( + K)  t_1)P^-1(0),


which is the first Magnus solution to (<ref>). 

An obvious choice of linear transformation is one that diagonalizes A. Now in general A may not have a diagonal form; however, all N × N matrices can be placed in Jordan normal form <cit.> (which reduces to diagonal form for diagonalizable matrices). 

For A = [[    0    1; -ω^2    0 ]], A is indeed diagonalizable, and we obtain 


    J   = [ -iω   0;   0  iω ]_,     P = [  i/ω -i/ω;    1    1 ]_,
       = [   ω/2ω-iω         0;         0 ω/2ω + iω ]_,     K = [     0 -ω/2ω; -ω/2ω     0 ]_.


We then find the Λ only solution (<ref>) returns


    M(t)    = [       (ω(0)/ω(t))^1/2cos(⟨ω⟩) 1/ω(0)(ω(0)/ω(t))^1/2sin(⟨ω⟩); - ω(0)(ω(0)/ω(t))^1/2sin(⟨ω⟩)       (ω(t)/ω(0))^1/2cos(⟨ω⟩) ]
       = 
        [       C 1/ω(0)S; - ω(0)S ω/ω(0)C ]_,


where C and S are the independent WKB solutions from Sec. <ref>. Thus we have obtained a WKB-like framework from a Magnus-like approach. 

The result is not identical to the WKB solution matrix (<ref>). We have no ω terms, and indeed because the Magnus expansion proceeds by integration as opposed to differentiation, extending the series does not introduce any further derivatives. This means that, while similar, the Jordan-Magnus approach remains distinct from WKB. In certain senses this lack of ω terms may be an advantage as derivatives can be difficult to evaluate numerically <cit.>.

If we try plotting the results of the Jordan-Magnus method for the Airy and burst equations (<ref>) we can see it dramatically outperforms the first and second Magnus solutions in accuracy, achieving results similar to the WKB solution. 





However, this method also has some disadvantages over standard Magnus methods. It is ultimately reliant on being able to obtain P, Λ, and K as functions of t. In addition, numerical Jordan decomposition is inherently unstable <cit.>. This implies we need to compute the analytic, symbolic, Jordan normal form of A to apply the method. While this is possible for the two-dimensional (<ref>) system, it quickly becomes prohibitive for more complicated systems with more dimensions.

While there is no stable numerical algorithm for obtaining Jordan-normal form, algorithms do exist for diagonalization <cit.>. Furthermore, although a complex square matrix is not guaranteed to be diagonalizable, a very large proportion of such matrices are diagonalizable <cit.>. Hence diagonalization algorithms could form the basis for a more fully numerical Jordan-Magnus method, although at present efficient methods remain elusive.



§ RESULTS




A python code was developed <cit.> to implement linear numerical integration methods of the form of (<ref>) using the WKB, truncated Magnus, and Jordan-Magnus approximate solutions. The numpy package <cit.> was used to perform the numerical calculations, and the sympy package <cit.> was used to perform the symbolic manipulation. For comparison a RKF4(5) method and a second order (in the WKB series) RKWKB method were also implemented. The method of adaptive step size control used is described in Appendix <ref>.

Figure <ref> shows the result of the numerical code for the different methods on the Airy and burst equations described in Sec. <ref>. For both equations we can see that the RKF4(5) method requires the smallest step size, and the Magnus methods use roughly similar intermediate step sizes. As expected the RKWKB method gives the largest step sizes, quickly reaching the maximum allowed step size h_, although it also gives the largest relative error. The Jordan-Magnus method gives smaller step sizes than RKWKB, but larger step sizes compared to the other methods.

Each plot includes a note of T, the time taken for the integration for each method. This of course varies substantially depending on the machine and the code implementation. However the ratio between the run-times for the different methods should give a reasonable indication of the relative computational cost. <ref> shows the averaged results for <ref>. We can see that the methods with the largest step size give the smallest time ratio and thus the smallest computational cost. The RKWKB method performs best and Jordan-Magnus second best. This result also supports the approach taken in the previous sections of seeking an approximate solution with good long range accuracy, as this allows large integration step sizes. 

We note that the relative errors shown in <ref> are not representative of what could be achieved by an optimized code (for example see Agocs et al. <cit.> for a practical implementation of RKWKB to fourth order in the WKB series). Nonetheless, the results concerning the relative step sizes and relative run-times achieved should still apply.
































 §.§ Application to the Einstein-Boltzmann equations: Photon and CDM universe



Finally we applied the methods and numerical integration code to the Einstein-Boltzmann equations for a simple radiation dominated universe with only photons and CDM, as described in Sec. <ref>. 

Unfortunately it was not possible to use the new Jordan-Magnus method on this system due to the difficulty of obtaining a symbolic, analytic Jordan normal form. However we can test the analytic first order Magnus solution, fourth and sixth order numerical Magnus solutions, and the RKF4(5) method. The results for Θ_0 and Φ are shown in <ref>. 

Fitting a solution of the form A_0cos(kη/√(3)) + B_0sin(kη/√(3)), where A_0 and B_0 are constants, to the results gives good agreement for <ref> for all but the earliest times as expected. 





For <ref>, fitting the adiabatic solution for Φ (<ref>) to the initial conditions gives good agreement only for very early times up to about η≈ 0.015. However, a separate adiabatic solution (<ref>) fitted to the RKF4(5) data for η > 0.015 gives good agreement for η≳ 0.03 onwards. Hence while the results do not fit a single adiabatic solution they can be well described by two adiabatic regimes with a transition region between 0.015 ≲η≲ 0.03.

The plots of log_10(h) show RKF4(5) gives the smallest step size, the sixth  order GL Magnus method has intermediate step sizes, and the first Magnus and fourth order GL Magnus methods give the largest step sizes.

As in Sec. <ref> we can compare the average ratio of run-times (<ref>). This shows that the sixth order GL Magnus method actually underperforms RKF4(5), with the smaller number of steps not compensating for the larger time per step. However, the analytic first Magnus and fourth order GL Magnus methods both give improved performance compared to RKF4(5). 



This shows that despite their inferior performance compared to WKB in Sec. <ref>, Magnus expansion based methods can still outperform RKF4(5) on a real cosmological higher-dimensional system.











§ CONCLUSION





In this work we sought to explore methods for efficiently integrating coupled systems of first order ODEs with highly oscillatory solutions, which would be a higher-dimensional extension of the RKWKB method, with particular applications to cosmological systems of interest.

The work of previous authors <cit.> suggests methods based on the Magnus expansion could be candidates for a multidimensional equivalent of RKWKB. However comparison of the analytic Magnus expansion and WKB approximate solutions implies the two methods are fundamentally distinct mathematical approaches. For equations of the form of (<ref>) WKB approximates the frequency as ⟨ω⟩, whereas the first Magnus solution approximates it as √(⟨ω^2 ⟩). The WKB solution series involves successively higher derivatives of ω while the Magnus expansion involves successively higher order integrals. Testing on trial equations of the form of (<ref>), the WKB solution give much higher accuracy. This indicates that the Magnus expansion and WKB based methods are orthogonal lines of research; however, the relationship between the two may benefit from more formal mathematical inquiry.

The simplest Magnus expansion based approximate solutions were found to gave inferior long range accuracy compared to WKB. By introducing a linear transformation we can convert the A matrix to a more diagonal form, obtaining a new “Jordan-Magnus” method. This gave excellent long range accuracy on simple two-dimensional problems which closely resembled the WKB solutions. However, currently the method relies on finding the symbolic Jordan normal form of A which is very difficult for high-dimensional cases. It could form the basis for a more fully numerical method which could be applied to arbitrarily high dimensions. 

A python code was created to implement the corresponding numerical integration methods, incorporating an adaptive step size procedure and the RKWKB and RKF4(5) methods for comparison. The code leaves plenty of room for optimization, and one could certainly improve the performance by switching to a compiled language like c. On simple two-dimensional systems, simple Magnus expansion based numerical methods ran significantly faster than RKF4(5), with the Jordan-Magnus method running faster still, although not achieving the speed of RKWKB.

Finally we applied this code and these methods to a simplified set of Einstein-Boltzmann equations corresponding to photon and CDM oscillations in the radiation era. The results again demonstrate that Magnus expansion based methods performed better than RKF4(5) by giving a speed improvement of about 50%. However, the improvement was not as impressive as for the simpler systems, and we could not implement the Jordan-Magnus method. For large times the numerical results fit the expected solutions for adiabatic perturbations. However, we could not fit a single adiabatic solution to the numerical results for the potential Φ. Instead the results suggest a transition between two different adiabatic regimes. This could indicate a breakdown of the adiabatic approximation for that system.

Future work would involve testing these Magnus expansion based methods on other multidimensional systems, including the 14 × 14 Einstein-Boltzmann matrix for l ≤ 2. If successful, one could extend up to high l values and see how one might include these methods in future Boltzmann solver codes. 

These results demonstrate that multidimensional numerical integration methods based on the Magnus expansion show promise for improving the performance of not only cosmological Boltzmann solvers but also many similar problems in other areas of physics.



§ ACKNOWLEDGEMENTS


J.B. thanks the University of Cambridge Physics Part III project scheme, which formed the basis for this work. W.H. was supported by a
Gonville & Caius research fellowship. We are also grateful for the advice of Fruzsina Agocs and Lukas Hergt of the Kavli Institute for Cosmology, Cambridge.











§ PROCEDURE FOR ADAPTIVE STEP SIZE CONTROL




We implemented the adaptive step size procedure laid out in Press and Flannery (1988) <cit.>. To estimate the error, and thus the correct step size for each step, we compute the estimate of x_n+1 from one step of size h, and a separate estimate using two steps size h/2, denoting this x^*_n+1. The error is then Δ_i = |x_i  n+1^* - x_i   n+1| for each component i.  We define a maximum tolerable error as Δ^_i = ϵ(a_ + r_|x_i n+1^*|), where a_ and r_ parametrize the absolute and relative error tolerance, and ϵ is a small scale factor.

The procedure is then as follows:

    
  * Calculate x^*_n+1 and x_n+1 for step size h.
    
    
  * Calculate Δ_i and Δ^_i and the error ratio R = √(1/N∑_i^N (Δ_i/Δ^_i)^2).
    
    
  * If R ≤ 1, compute a new h as h_ = hSR^-1/(ν+1) where S = safety factor ∼ 0.99, ν = order of method. 
    
    Go to 5.
    
    
  * Else if R > 1, compute a new h as h_ = hSR^-1/ν.
    
    If h_ < h_, set h_ = h_ and go to 5. 
    
    Otherwise go to 1.
    
    
  * If h_ < 0.2h, set h_ = 0.2h.
    
    If h_ > h_, set h_ = h_.
    
    If h_ > 10h, set h_ = 10h.
    
    
  * Use current x^*_n+1 as the estimate. Advance time. Set h = h_. Go to next step.
    

The RKWKB method used the adaptive step size procedure outlined above, setting ν = 1. The RKF4(5) was implemented following <cit.>. The approach to step size control is the same, except the error is calculated using the difference between the fourth and fifth order RK methods.

The a_ and r_ values for the RKF4(5) and Magnus methods were chosen using trial and error to achieve the largest step sizes for the target error ϵ when applied to the burst equation as in <ref>. For the RKWKB method the values for the Magnus methods were used.

