Arxiv:1907.11638V2 [Physics.Comp-Ph] 12 Feb 2020

Beyond the Runge-Kutta-Wentzel-Kramers-Brillouin method

Jamie Bamber1, 2, ∗ and Will Handley1, 3, 4, † 1Astrophysics Group, Cavendish Laboratory, J.J.Thomson Avenue, Cambridge, CB3 0HE, UK 2Pembroke College, Pembroke Street, Cambridge, CB2 1RF, UK 3Kavli Institute for Cosmology, Madingley Road, Cambridge, CB3 0HA, UK 4Gonville & Caius College, Trinity Street, Cambridge, CB2 1TA, UK (Dated: February 13, 2020) We explore higher-dimensional generalizations of the Runge-Kutta-Wentzel-Kramers-Brillouin method for integrating coupled systems of ﬁrst-order ordinary diﬀerential 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 Mag- nus 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. How- ever 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.

I. INTRODUCTION many oscillations, giving faster performance than standard RK methods. However RKWKB is limited to es- In the new era of high precision cosmology, one of sentially one-dimensional systems [4, 5]. the most important tools in the cosmologist’s toolbox is We can write a multidimensional coupled system of the suite of Boltzmann codes. These methods, including linear ODEs in the form of (1.1), class [1], camb [2], and pycosmo [3] provide numeri- x˙ (t) = A(t)x, (1.1) cal solutions to the Einstein-Boltzmann equations, which describe the evolution of linear perturbations in the Uni- where x is a vector of dynamical variables, A is a time- verse. The dominant computational cost for these codes dependent matrix, and an overdot denotes differentiation is integrating coupled systems of linear first order ordi- with respect to time. nary differential equations (ODEs) with highly oscilla- In the mathematics field of numerical analysis there tory solutions [1, 2]. Currently the codes use several nu- has been substantial work examining the formal solution merical and semianalytic approximations to get around of (1.1), the Magnus expansion [6]. A number of authors these bottlenecks. However these approximations are de- have also used the Magnus expansion and related tech- signed for specific cosmological theories. A more general- niques to develop numerical methods to integrate highly ized theory-independent method would make it easier to oscillatory differential equations [6], for example Iserles switch between different models and test new extensions [7, 8]. to the current theory. New numerical integration meth- These Magnus-based numerical integrators have been ods could also reduce the run-time of the code. While extensively applied to problems in quantum and atomic the time needed to run a single simulation is currently physics [9, 10], as the time dependent Schrodinger equa- only a few seconds, the many thousands of simulation tion can be written in the form of (1.1) if we use a basis of runs needed for a typical Monte Carlo analysis add up to wave functions [11]. Some alternative methods for solv- a substantial computational cost. ing coupled differential equations have also been explored Handley, Lasenby and Hobson proposed in [4], in the context of primordial cosmology [12]. However, and developed with Agocs in [5], a new numeri- Magnus-based methods remain unexplored with respect cal method based on a combination of the standard arXiv:1907.11638v2 [physics.comp-ph] 12 Feb 2020 to the Einstein-Boltzmann equations and cosmological Runge-Kutta (RK) approach and the Wentzel-Kramers- systems in general. Furthermore, to our knowledge, no Brillouin (WKB) approximation, dubbed the “Runge- previous authors have tested these methods with adap- Kutta-Wentzel-Kramers-Brillouin” (RKWKB) method. tive step size control. In highly oscillatory regions of the solution using the In this paper we examine the application of Magnus- WKB approximation allows integration steps which span 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 ∗ [email protected] in many other areas of physics and applied mathematics † [email protected] [4]. 2

II. BACKGROUND Further terms can be generated by a recursive procedure [6] and involve an increasing number of nested integrals, A. The WKB approximation commutators, and powers of A. For the Ω series to converge, A must be suﬃciently The WKB method is an example of multiple scale anal- small in some sense. It can be shown [6] that the Mag- ysis [13] applied to oscillatory systems. It is most com- nus series is absolutely convergent for a general complex A ≤ ≤ monly applied to equations of the form of (2.1) [14], matrix , for t0 t T , if x¨ = −ω2(t)x. (2.1) Z T kA(s)k2 ds < π, (2.12) The derivation begins by considering a general solution t0 of the form where k · k denotes the two-norm of a matrix, which is ∞ 2 X [15] x = exp Sn(t) . (2.2) n=0 q kA(s)k = max λ (AT A), (2.13) Using the method of separation of scales we can obtain 2 i=1,...,n i expressions for Sn [14] where λi(Q) denotes the ith eigenvalue of matrix Q. Z t However, the actual convergence domain can be larger S0 = ±i ω(t1)dt1, (2.3) than that given by (2.12) [6]. 0 1 S1 = const. − ln(ω(t)), (2.4) 2 C. Stepping numerical integration methods n−1 X 2S˙0S˙n + S¨n−1+ S˙jS˙n−j = 0 for n ≥ 2. (2.5) j=1 A general stepping numerical integration method proceeds by using an approximate solution to the equation Truncating with only the S ,S terms gives the classic 0 1 to advance from one integration point xn at time tn to WKB approximation: the next xn+1 at time tn+1. For a system of ﬁrst order

A i R t ω(t )dt B −i R t ω(t )dt linear ODEs such as (1.1) we can write the approximate x(t) = √ e 0 1 1 + √ e 0 1 1 . (2.6) ω ω solution for initial conditions x(t0) in terms of a solution matrix M(t0, t) [16], If we include only the S0 term we obtain x¨ + ω2x = ±iω˙ (t)x, (2.7) x(t) = M(t0, t)x(t0). (2.14) implying we require |ω˙ (t)| |ω2(t)| for the approximation to be reasonable. Then the general stepping procedure is

B. The Magnus expansion xn+1 = M(tn, tn+1)xn. (2.15)

The Magnus expansion [6] provides a general solution Runge-Kutta methods work by using truncated Tay- M to equations of the type (1.1). As with WKB, it starts by lor series solutions for (t0, t). For example, Euler’s by considering a solution in the form of the exponential method, the simplest Runge-Kutta method [4], has of an inﬁnite series (2.8) [6], ∞ X M(t0, t) = 1 + (t − t0)A(t0). (2.16) x(t) = exp Ωn(t) x(0), (2.8) n=1 The RKWKB method (to second order in the WKB the ﬁrst terms of which are series) is obtained by replacing the Taylor series solu- Z t tion with the two independent WKB solutions in (2.6). As the WKB solutions typically remain good approxima- Ω1(t) = A(t1)dt1, (2.9) 0 tions to the true solution over many oscillations, unlike 1 Z t Z t1 truncated Taylor series solutions, the RKWKB method is Ω (t) = dt [A(t ), A(t )]dt , (2.10) 2 2 1 1 2 2 able to take much larger steps than Runge-Kutta meth- 0 0 ods, which in general reduces the required run-time [5]. Z t Z t1 Z t2 1 et al. et al. Ω3(t) = dt1 dt2 [A(t1), [A(t2), A(t3]] Handley [4] and Agocs [5] also enable the 6 0 0 0 program to switch to a standard Runge-Kutta-Fehlberg method in slowly oscillating regions; however, in this pa- + [A(t3), [A(t2), A(t1]] dt3. per we use RKWKB to refer to the simpler nonswitching (2.11) approach. 3

We can also obtain stepping numerical integration The conformal Hubble rate is H =a/a ˙ , where an over- methods of the form of (2.15) from the Magnus expan- dot is used to denote the derivative with respect to con- sion. If we can compute the analytic forms of Ωi then formal time. We denote the nonconformal cosmic Hub- we can use those directly. In that case the Ω series trun- ble rate as H = H/a. Theτ ˙ = −neσT a is the confor- cated at Ω = Ω1 + ··· + Ωi where i = 2s − 1 or 2s − 2 for mal time derivative of the optical depth, τ, with ne the 2s+1 integer s gives a method order h in the step size h mean electron density and σT the Thompson cross sec- [6]. However, if Ωi cannot be computed analytically, we tion. Mean densities are described using parameters of can use an approach from Blanes et al. [6] Sec. 5.4. The the form Ω = ρ/ρcrit where ρ is the mean density of a 3H2 A is Taylor expanded about the midpoint t1/2 = tn +h/2 component and ρcrit is the critical density = 8πG . These of each step, where h = tn+1 − tn is the step size. parameters are denoted Ωγ , Ων , Ωdm and Ωb for the photons, neutrinos, CDM and baryonic matter respectively. ∞ n X j 1 d A A(t) = aj(t − t ) where aj = . 1/2 j! dtn j=0 t=t1/2 2. Photon and CDM universe (2.17) One can then use the defining recursion relation for the Magnus expansion to express Ω to the desired order in As a test case we consider a simplified universe with h. For example to order four we have only photons and cold dark matter. By neglecting the baryonic matter we can avoid the complexities of the 1 Ω[4] = α − [α , α ] + O(h5), (2.18) Coulomb interaction between the baryons and the pho- 1 12 1 2 tons. We also neglect the photon polarization and all har- monics l > 1, approximating the photons as a fluid [17]. where α = hja . The α ’s can then be approxi- j j−1 j In this limit the gravitational potential Φ = Ψ. By mak- mated using various quadrature rules (including Gauss- ing these approximations we obtain a five-dimensional Legendre, Simpson, Newton-Cotes) and linear combina- system with x = [Θ , Θ , δ, v, Φ]T . The reduced di- tions of A evaluated at different points in the interval 0 1 mensions makes this much more amenable than the 14- [t , t ]. n n+1 dimensional system, while still preserving some of the key features of the early universe. By testing on this D. The Einstein-Boltzmann equations simplified system we provide a proof-of-principle study for application to the full Einstein-Boltzman equations. For adiabatic perturbations and radiation domination In this section we describe the background theory and we obtain a second order evolution equation for the po- the Einstein-Boltzmann equations that will be used in tential [17, 19] Sec. V. 4 1 Φ¨ + Φ˙ + k2Φ = 0. (2.20) η 3 1. Overview This has a general solution (2.21) [17, 19] j (x) y (x) We seek to express the Einstein-Boltzmann equations Φ = A 1 + B 1 , (2.21) k x k x in the form of (1.1) [17]. Working in the conformal New- √ tonian gauge in flat spacetime, the metric is given by where x = kη/ 3 and j1, y1 are first order spherical Bessel functions of the first and second kinds respectively. These take the form 2 2 2 − − i j ds = a (η)[(1 + 2Ψ)dη (1 2Φ)δijdx dx ]. (2.19) sin(x) cos(x) − j1(x) = 2 , where the scalar perturbations are Ψ, Φ, and η is confor- x x (2.22) cos(x) sin(x) mal time [18]. As we work at early times we neglect dark y (x) = − − . energy and treat the neutrinos as relativistic. We denote 1 x2 x the photon and neutrino temperature perturbations as As radiation dominates the energy density we also find Θ and N respectively. The analogous perturbation ΘP describes the strength of the photon polarization. We 2 2 ˙ treat the cold dark matter (CDM) and baryonic matter δr = 4Θ0 = − (kη) Φ − 2ηΦ − 2Φ, (2.23) 3 as nonrelativistic and describe them with overdensities δ, δ and peculiar velocities v, v . where δr is the overdensity for radiation. Hence√ for b b ∝ − 2 ∼ To achieve a linear set of equations we expand the kη 1 we√ have Θ0 (kη) Φ Ak cos (kη/ 3) + perturbations in spatial Fourier modes, introducing wave Bk sin (kη/ 3). This gives the acoustic oscillations√ of number k, and we expand the relativistic perturbations the photon fluid at the sound speed cs = 1/ 3 that one in multipole moments Θl(k, η) for l = 0, 1, 2,..., ∞. The would expect. The solution in (2.21) suggests the poten- resulting 14 × 14 A(k, η) matrix for l ≤ 2 is shown in tial also oscillates for kη & 1, but with the amplitude 2 Fig. 1. decaying as 1/(kη) . 4

 3 3  12 Ωγ 12 Ω 1 1 k2 H ν −2HΩγ −k k2 0 0 0 −2HΩν 0 Hk2 − 2 HΩdm 0 − 2 HΩb 0 − 3 − H  H   2 2   1 4Ωγ 2 4 Ων 1 1   k τ˙ − H − k 0 0 0 0 0 − H 0 0 0 − iτ˙ k   3 k 3 k 3 3     2 21 1 1   0 k τ˙ τ˙ 0 τ˙ 0 0 0 0 0 0 0 0   5 20 20 20     1 1 1   0 0 − τ˙ τ˙ −k − τ˙ 0 0 0 0 0 0 0 0   2 2 2     0 0 0 1 ˙ − 2 0 0 0 0 0 0 0 0   3 k τ 3 k       0 0 − 1 ˙ − 1 ˙ 2 19 ˙ 0 0 0 0 0 0 0 0   20 τ 20 τ 5 k 20 τ     3 3 2  12 Ωγ 12 Ων 1 1 k −2HΩγ 0 H 0 0 0 −2HΩν −k H − HΩ 0 − HΩb 0 − − H   k2 k2 2 dm 2 3   H  2 2  4 Ωγ 1 4Ων 2 1   0 0 − H 0 0 0 k 0 − H − k 0 0 0 0 k   k 3 k 3 3       0 0 0 0 0 0 0 2 k 0 0 0 0 0 0   5    3 3  36 Ωγ 36 Ω 3 3 k2  H ν −6HΩγ 0 k2 0 0 0 −6HΩν 0 Hk2 − 2 HΩdm 0 − 2 HΩb 0 − − 3H  H   2 2   12i Ωγ 12i Ων   0 0 H 0 0 0 0 0 H 0 −H 0 0 −ik   k k   3 3  36 Ωγ 36 Ω 3 3 k2  H ν  −6HΩγ 0 k2 0 0 0 −6HΩν 0 Hk2 − 2 HΩdm 0 − 2 HΩb −ik − − 3H  H   2 2  4iτ˙ Ωγ 12i Ωγ 12i Ω 4τ ˙ Ωγ  H ν   0 k 0 0 0 0 0 Hk 0 0 0 − H −ik   Ωb 3Ωb   3 3  12 Ωγ 12 Ω 1 1 k2 H ν −2HΩγ 0 k2 0 0 0 −2HΩν 0 Hk2 − 2 HΩdm 0 − 2 HΩb 0 − 3 − H H

T FIG. 1: 14 × 14 A matrix for l ≤ 2. Here x = [Θ0, Θ1, Θ2, ΘP 0, ΘP 1, ΘP 2, N0, N1, N2, δ, v, δb, vb, Φ] .

To use the numerical methods we must ﬁrst obtain the Ωdm(ηeq) = 1/2. This gives relevant A. Working from the 14 × 14 matrix in Fig. 1 we can derive the reduced A matrix shown in (2.24),  1 1 ηk2 1  − η −k − 4 0 − 3 − η    k 0 0 0 k   3 3    A(k, η) = − 3 0 − 3 −ik −ηk2 − 3  (2.25)  η 4 η   1   0 0 0 − η2 −ik   2  − 1 0 − 1 0 − ηk − 1 η 4 3 η .

 ˙   1 k2    Θ0 −2HΩγ −k − 2 HΩdm 0 − 3H − H Θ0       III. COMPARING WKB AND THE MAGNUS Θ˙   k 0 0 0 k  Θ  EXPANSION  1  3 3   1    2     δ˙  = −6HΩ 0 − 3 HΩ −ik − k − 3H  δ     γ 2 dm H          A. Comparing analytic solutions  v˙   0 0 0 −H −ik   v     2    Φ˙ −2HΩ 0 − 1 HΩ 0 − k − H Φ γ 2 dm 3H . To investigate the relationship between the WKB ap- (2.24) proximation and the Magnus expansion we compared the We set the conformal time at matter-radiation equality lowest order approximate analytic solutions to (2.1). to be ηeq = 1, set a(ηeq) = 1 and assume Ωγ (ηeq) = Equation (2.1) can be written in the form of (1.1) by 5 setting with corrections to both the phase and the amplitude, where φˆ is a time dependent phase term. We denote x 0 1 x = A = (3.1) this solution “second Magnus.” For our matrix in (3.1) x˙ −ω2 0 2 , . kA(s)k2 = max(ω , 1). So (2.12) implies a convergence R T 2 domain of T −t0 < π, ω dt < π. One can also see from For the Magnus expansion, the first two Ω terms are t0 given by (3.8) that at sufficiently large timesω ˆ becomes imaginary, giving exponentially growing or shrinking solutions. 0 t For the WKB approximation in Eq. (2.6) the n = 0 Ω = (3.2) 1 − R t 2 term approximates the frequency as 0 ω dt1 0 , t a 0 1 Z hωi = ω(t )dt , (3.11) Ω2 = − (3.3) 1 1 0 a , t 0 where while the n = 1 term provides a correction to the amplitude. We can write (2.6) in the form x(t) = Z t Z t Z t1 Z t d M(t , t)x(t ). Setting t = 0 for simplicity we obtain a = 1 t ω2dt − dt ω2dt = 1 t2 hω2idt . 0 0 0 2 1 1 2 2 1 dt 1 0 0 0 0 " ω˙ (0) 1 # (3.4) C + 2 S S M 2ω(0) ω(0) = ωω˙ (0) ω˙ ω˙ ω˙ (0) ω ω˙ 2ω(0)2 − 2ω C − ω + 4ωω(0)2 S ω(0) C − 2ωω(0) S Including only Ω we obtain x(t) = M(t)x(0), where , 1 (3.12) M = exp(Ω1(t)) given by (3.5), 1 1 where C = ω(0) 2 cos(hωit) and S = ω(0) 2 sin(hωit)  p p  ω ω cos hω2it √ 1 sin hω2it are the two independent solutions. Applying the stepping M hω2i procedure in (2.15) to this M gives you the RKWKB = p p p  hω2i sin hω2it cos hω2it method (to second order) [4]. . (3.5) We can thus observe that the Magnus expansion We can then write x(t) as and the WKB method take fundamentally different approaches for solving the same problem. The classic WKB p 1 p approximation approximates the oscillation frequency as x(t) = A cos hω2it + B sin hω2it . (3.6) ph 2i hωi, while the first Magnus solution uses the root mean ω p square, hω2i. The terms of the WKB series involve suc- p The frequency is approximated as ω ∼ hω2i = cessively higher order derivatives of ω, while the succes- q sive Magnus expansion terms involve successively higher 1 R t ω2dt . We also get some amplitude correction t 0 1 order integrals. from the B or M12 term. We denote this solution “first The WKB approximation is specialized to highly os- Magnus.” cillatory solutions, and fails when ω goes to zero. The Magnus expansion can cope with nonoscillatory solutions Including both Ω1 and Ω2 we obtain and is less specialized, which comes at the cost of inferior 1 long range accuracy compared to WKB, as we shall see M = 1 cos(ˆωt) + (Ω + Ω ) sin(ˆωt) in Sec. III B. 1 2 ωtˆ 1 0 1 a t = cos(ˆωt) + sin(ˆωt) 0 1 R t 2 . ωtˆ − 0 ω dt1 −a B. Testing on trial equations (3.7) These solutions were plotted for two trial equations In (3.7) the frequency is approximated asω ˆ, taken from Handley et al. [4], namely the Airy equation r and the “burst” equation, which both have analytic exact p 1 ωˆ = hω2i − (a/t)2 = hω2i − ht2hω˙2ii2, (3.8) solutions. The Airy equation takes the form 4 x¨(t) = −tx, (3.13) where hω˙2i denotes the time derivative of hω2i. Then we obtain (3.9) and (3.10), and the burst equation is given by

a 1 n2 − 1 x(t) = A cos(ˆωt) + sin(ˆωt) + B sin(ˆωt), (3.9) x¨ = −ω2(t)x , ω2(t) = , (3.14) tωˆ ωˆ (1 + t2)2 where n is an integer. p hω2i 1 The analytic WKB approximate solution (2.6) and the x(t) = A cos ωtˆ + φˆ(t) + B sin(ˆωt), (3.10) ωˆ ωˆ ﬁrst and second Magnus solutions (3.6) and (3.9) for the 6

Burst equation (n = 40) : WKB vs Magnus expansion Airy equation : comparing WKB vs Magnus 1.0 ﬁrst Magnus WKB 0.4 second Magnus true soln. 0.5

0.2 x 0.0

0.5 x 0.0 −

1.0 − 0.2 WKB second Magnus − ﬁrst Magnus true soln. 0.5 0.4 −

x 0.0 10.0 7.5 5.0 2.5 0.0 2.5 5.0 7.5 10.0 − − − − t 0.5 − FIG. 3: WKB, and the Ω1 and Ω1 + Ω2 Magnus solutions to the burst equation, n = 40, analytic solutions plotted for 1.0 − t = −10 to t = +10 0 5 10 15 20 25 30 35 40 45 t FIG. 2: WKB, and the Ω1 and Ω1 + Ω2 Magnus analytic solutions to the Airy equation. TOP: from t0 = 1 to t = 35. Burst equation (n = 40) : WKB vs Magnus expansion 0.100 BOTTOM: from t0 = 11 to t = 45. ﬁrst Magnus WKB 0.075 second Magnus true soln.

0.050 Airy equation (3.13) are plotted in Fig. 2. For the upper plot, with t0 = 1, the WKB solution remains a very good 0.025 approximation to the true solution; however, the ﬁrst

and second Magnus solutions perform poorly. The ﬁrst x 0.000 Magnus reproduces the increase in oscillation frequency, 0.025 but does not reproduce the changes in amplitude. The − second Magnus solution breaks down completely around 0.050 − t − t0 ∼ π showing the exponential growth predicted in Sec. III A. 0.075 − If we shift the starting point to t = 11, as in the bottom 0.100 − plot in Fig. 2, the ﬁrst Magnus solution follows the true 2.0 1.5 1.0 0.5 0.0 0.5 1.0 1.5 2.0 solution more closely. However unlike the WKB solu- − − − − t

tion it fails to follow the frequency of the true oscillation FIG. 4: WKB, and the Ω1 and Ω1 + Ω2 Magnus solutions at large times. The second Magnus solution still breaks to the burst equation, n = 40. Solutions plotted for the down as predicted. central region, t = −2 to t = +2 (note the change in scale). The solutions were again computed analytically. Next we consider the burst equation (3.14). Figure 3 shows the analytic WKB, first Magnus and second Mag- nus approximate solutions from t0 = −10 up to t = +10 with n = 40. While WKB provides an excellent match to 1. A closer look at the Magnus error the true solution throughout the whole range, both Mag- nus solutions give very poor approximations. The first Magnus solution fails to follow the changes in amplitude To show why the first Magnus solution compares so or frequency, and the second Magnus solution once again poorly in terms of accuracy we can differentiate the Mag- diverges after a relatively short period. nus solutions. Let CM and SM be the two independent first Magnus solutions, Figure 4 shows the result if we try to apply the WKB and Magnus solutions to only the small, central, highly p 2 oscillatory region. This region has a smallerω/ω ˙ ratio, CM = cos hω it , so it should be easier to model. Even then the WKB so- (3.15) 1 p 2 lution dramatically outperforms both Magnus solutions. SM = p sin hω it . hω2i 7

Let 1 + = ω2/hω2i and ξ =ω/ω ˙ , and then An obvious choice of linear transformation is one that diagonalizes A. Now in general A may not have a diag- 2 2 onal form; however, all N × N matrices can be placed in C¨M = − ω 1 + CM ... 4(1 + ) Jordan normal form [20] (which reduces to diagonal form (3.16) ω2 32 for diagonalizable matrices). · · · − 2 − − 0 1 ω ξSM SM , A 2 A t 4(1 + ) For = −ω 0 , is indeed diagonalizable, and we obtain

2 − − 2 iω 0 i/ω i/ω S¨M = − ω 1 + SM ... J = P = (4.4) 4(1 + ) 0 iω , 1 1 , (3.17) 1 1 ω˙ − iω 0 0 − ω˙ · · · − ξ − (2 + 2) S − C . Λ = 2ω K = 2ω (4.5) M M ω˙ − ω˙ 4t t 0 2ω + iω , 2ω 0 .

One can see that for = 0, ξ = 0 we get back (2.1). We then find the Λ only solution (4.2) returns As the errors can increase with as well as ξ, they are potentially much larger than for WKB.  1 1  ω(0) 2 1 ω(0) 2 ω(t) cos(hωi) ω(0) ω(t) sin(hωi) M(t) =  1 1   ω(0) 2 ω(t) 2  IV. THE JORDAN-MAGNUS METHOD −ω(0) ω(t) sin(hωi) ω(0) cos(hωi) " # C 1 S The WKB approximation provides good long range ac- = ω(0) −ω(0)S ω C curacy but is limited to one-dimensional systems. In con- ω(0) , trast the truncated Magnus expansion provides a matrix- (4.6) based approach suitable for multidimensional systems, but with inferior long range accuracy as we saw in Sec. where C and S are the independent WKB solutions from III. Ideally we want a way to combine the positive as- Sec. II A. Thus we have obtained a WKB-like framework pects of both WKB and Magnus, that is achieve WKB- from a Magnus-like approach. like long range accuracy with a solution that works on The result is not identical to the WKB solution ma- multidimensional systems. In this section we introduce a trix (3.12). We have noω ˙ terms, and indeed because the new approach based on the Magnus expansion, dubbed Magnus expansion proceeds by integration as opposed to the “Jordan-Magnus” method, which goes some way to- differentiation, extending the series does not introduce ward this goal. any further derivatives. This means that, while simi- If the A matrix were time-independent, we could solve lar, the Jordan-Magnus approach remains distinct from the system by diagonalizing A to obtain the independent WKB. In certain senses this lack ofω ˙ terms may be an oscillating modes. Consider introducing a linear transfor- advantage as derivatives can be difficult to evaluate nu- −1 −1 mation to (1.1). Let xP = P x, and let J = P AP merically [5]. for some transformation matrix P(t). Then If we try plotting the results of the Jordan-Magnus method for the Airy and burst equations (Figs. 5 and 6) −˙ 1 we can see it dramatically outperforms the first and sec- x˙ P = AP xP = J + (P )P xP , (4.1) ond Magnus solutions in accuracy, achieving results similar to the WKB solution. −˙ 1 −1 where (P ) denotes the time derivative of P . We However, this method also has some disadvantages over want to choose a P such that AP is as diagonal as pos- standard Magnus methods. It is ultimately reliant on be- sible. Let AP = Λ + K where Λ is diagonal and K has ing able to obtain P, Λ, and K as functions of t. In ad- zeros on the diagonal. dition, numerical Jordan decomposition is inherently un- Once again we seek a solution in the form x(t) = stable [20]. This implies we need to compute the analytic, M(t)x(0). We can obtain an approximate solution by symbolic, Jordan normal form of A to apply the method. neglecting K, While this is possible for the two-dimensional (1.1) sys- Z t tem, it quickly becomes prohibitive for more complicated −1 M(t) ≈ P(t) exp Λdt1 P (0). (4.2) systems with more dimensions. 0 While there is no stable numerical algorithm for ob- This avoids the potentially costly matrix exponential. Al- taining Jordan-normal form, algorithms do exist for di- ternatively we can include K and obtain agonalization [21]. Furthermore, although a complex square matrix is not guaranteed to be diagonalizable, a Z t −1 very large proportion of such matrices are diagonalizable M(t) ≈ P(t) exp (Λ + K)dt1 P (0), (4.3) [22]. Hence diagonalization algorithms could form the 0 basis for a more fully numerical Jordan-Magnus method, which is the first Magnus solution to (4.1). although at present efficient methods remain elusive. 8

Airy equation : comparing WKB vs Jordan-Magnus proximate solutions. The numpy package [24] was used to perform the numerical calculations, and the sympy WKB package [25] was used to perform the symbolic manipu- 0.6 Jordan-Magnus (Λ only) Jordan-Magnus (Λ and K) lation. For comparison a RKF4(5) method and a second 0.4 true soln. order (in the WKB series) RKWKB method were also implemented. The method of adaptive step size control 0.2 used is described in Appendix A. Figure 7 shows the result of the numerical code for x 0.0 the diﬀerent methods on the Airy and burst equations described in Sec. III B. For both equations we can see 0.2 − that the RKF4(5) method requires the smallest step size, and the Magnus methods use roughly similar intermedi- 0.4 − ate step sizes. As expected the RKWKB method gives 0.6 the largest step sizes, quickly reaching the maximum al- − lowed step size hmax, although it also gives the largest relative error. The Jordan-Magnus method gives smaller 0 5 10 15 20 25 30 35 t step sizes than RKWKB, but larger step sizes compared to the other methods. FIG. 5: WKB, and the Λ only and Λ + K Jordan-Magnus Each plot includes a note of T , the time taken for the analytic solutions to the Airy equation with t0 = 1. Note all three lines lie very close on top of one another and the true integration for each method. This of course varies sub- solution. stantially depending on the machine and the code im- plementation. However the ratio between the run-times for the diﬀerent methods should give a reasonable indi- Burst equation (n = 40) : WKB vs Jordan-Magnus cation of the relative computational cost. Table I shows the averaged results for Fig. 7. We can see that the meth- WKB ods with the largest step size give the smallest time ratio 0.4 Jordan-Magnus (Λ only) Jordan-Magnus (Λ and K) and thus the smallest computational cost. The RKWKB true soln. method performs best and Jordan-Magnus second best. 0.2 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

x 0.0 step sizes. We note that the relative errors shown in Fig. 7 are not representative of what could be achieved by an optimized 0.2 − code (for example see Agocs et al. [5] for a practical im- plementation of RKWKB to fourth order in the WKB 0.4 series). Nonetheless, the results concerning the relative − step sizes and relative run-times achieved should still apply. 10.0 7.5 5.0 2.5 0.0 2.5 5.0 7.5 10.0 − − − − t FIG. 6: WKB, and the Λ only and Λ + K Jordan-Magnus Method Airy eq. Burst eq. (n = 40) solutions to the burst equation with n = 40, analytic solutions plotted for t = −10 to t = +10. Note the First Magnus 0.115 ± 0.004 0.78 ± 0.02 Jordan-Magnus lines lie on top of one another for most the graph. Magnus 4◦ GL 0.244 ± 0.004 0.81 ± 0.01 Second Magnus 0.253 ± 0.003 0.503 ± 0.008 RKWKB 0.0029 ± 0.0003 0.0065 ± 0.0004 V. RESULTS Jordan-Magnus 0.050 ± 0.007 0.246 ± 0.028

A python code was developed [23] to implement linear TABLE I: Ratio of integration times vs. RKF4(5) time for the diﬀerent plots. Averaged over 20 runs. numerical integration methods of the form of (2.15) using the WKB, truncated Magnus, and Jordan-Magnus ap- 9

Airy equation : comparing diﬀerent methods true soln second Magnus, analytic, T=0.16s 1.00 RKF 4(5), T=0.9178s RKWKB, T=0.002284s 0.75 ﬁrst Magnus, analytic, T=0.2054s JM (Λ and K) (scipy quad, maxiter=200), T=0.105s

Magnus 4◦, GL quad, T=0.1797s 0.50

x 0.25

0.00

0.25 − 1 ) h

( 0 10 1 log − 0

(rel. error) 5

10 − log 5 10 15 20 25 30 35 t Burst equation (n = 40) : comparing diﬀerent methods true soln second Magnus, analytic, T=0.2513s 0.4 RKF 4(5), T=0.504s RKWKB, T=0.002604s ﬁrst Magnus, analytic, T=0.4212s JM (Λ and K) (scipy quad, maxiter=200), T=0.1242s 0.2 Magnus 4◦, GL quad, T=0.4079s

x 0.0

0.2 −

0.4 − )

h 0 ( 10

log 2 − 0

(rel. error) 5 10 − log 10.0 7.5 5.0 2.5 0.0 2.5 5.0 7.5 10.0 − − − − t

FIG. 7: Applying numerical integration with adaptive step size control to the Airy equation and burst equation with n = 40. For the Magnus and RKWKB solutions = 0.005, atol = 0.005, rtol = 1. For the RKF solution, RKF = 0.005, atol = 1, rtol = 2. Initial step size h0 = 0.5, hmax = 10, hmin = 0.01 for Magnus methods. For RKF4(5) hmax = 2.5, hmin = 0.005. In each plot the layout is TOP: the integration results. MIDDLE: log10 |h| where h=step size. BOTTOM: log10 |relative error|. The line shows the target error of = 0.005. 10

Radiation dominated photon and CDM system : Θ0, k =200 sinusoid soln. ﬁtted to data ﬁrst Magnus, analytic, T=0.2808s 6 RKF 4(5), T=0.5525s Magnus 4◦, GL quad, T=0.2672s Magnus 6◦, GL quad, T=0.7868s 4

2 x 0

2 − 4 −

) 2 h − ( 10 3 log −

(rel. error) 2 10 − log 0.1 0.2 0.3 0.4 0.5 η

FIG. 8: Applying numerical integration√ methods to the photon & CDM system. The black dashed line is the analytic solution for large kη which is C0 cos (k(η + C1)/ 3) where C0,C1 are coeﬃcients ﬁtted to the RKF45 data. The initial conditions are η0 = 0.01, x0 = [1, 2, 1, 2, 1]. The parameters are k = 200, h0 = 0.025, hmin = 0.00025, hmax = 2.5, = 0.005. For the Magnus methods atol = 0.005, rtol = 1, for the RKF method atol = 4, rtol = 2. The layout is as in Fig. 7.

A. Application to the Einstein-Boltzmann The plots of log10(h) show RKF4(5) gives the smallest equations: Photon and CDM universe step size, the sixth order GL Magnus method has inter- mediate step sizes, and the first Magnus and fourth order Finally we applied the methods and numerical inte- GL Magnus methods give the largest step sizes. gration code to the Einstein-Boltzmann equations for a As in Sec. V we can compare the average ratio of run- simple radiation dominated universe with only photons times (Table II). This shows that the sixth order GL Mag- and CDM, as described in Sec. II D 2. nus method actually underperforms RKF4(5), with the smaller number of steps not compensating for the larger Unfortunately it was not possible to use the new time per step. However, the analytic first Magnus and Jordan-Magnus method on this system due to the dif- fourth order GL Magnus methods both give improved ficulty of obtaining a symbolic, analytic Jordan normal performance compared to RKF4(5). form. However we can test the analytic first order Mag- nus solution, fourth and sixth order numerical Magnus Method Ratio solutions, and the RKF4(5) method. The results for Θ0 and Φ are shown in Figs. 8 and 9. √ Magnus 6◦ GL 1.58 ± 0.06 Fitting√ a solution of the form A0 cos (kη/ 3) + First Magnus 0.71 ± 0.03 B0 sin (kη/ 3), where A0 and B0 are constants, to the results gives good agreement for Fig. 8 for all but the Magnus 4◦ GL 0.62 ± 0.03 earliest times as expected. TABLE II: Ratio of integration times for different methods For Fig. 9, fitting the adiabatic solution for Φ (2.21) to vs. RKF4(5) for the photon & CDM system. Parameter the initial conditions gives good agreement only for very settings and initial conditions are as in Figs. 8 and 9. early times up to about η ≈ 0.015. However, a separate Averaged over 20 runs. adiabatic solution (2.21) fitted to the RKF4(5) data for η > 0.015 gives good agreement for η & 0.03 onwards. This shows that despite their inferior performance Hence while the results do not fit a single adiabatic solu- compared to WKB in Sec. III, Magnus expansion based tion they can be well described by two adiabatic regimes methods can still outperform RKF4(5) on a real cosmo- with a transition region between 0.015 . η . 0.03. logical higher-dimensional system. 11

Radiation dominated photon and CDM system : Φ, k =200 1.0 adiabatic soln. fitted to data for η > 0.015 Magnus 6◦, GL quad, T=0.6741s 0.8 adiabatic soln. fitted to init. cond. first Magnus, analytic, T=0.2767s RKF 4(5), T=0.4353s Magnus 4 , GL quad, T=0.2703s 0.6 ◦

x 0.4

0.2

0.0

0.2 −

) 2 h − ( 10 3 log −

(rel. error) 2 10 − log 0.1 0.2 0.3 0.4 0.5 η FIG. 9: Applying numerical integration methods to the photon and CDM system, showing potential Φ. The black dotted line is the adiabatic solution (2.21) ﬁtted to the initial conditions, and the black dashed line is the adiabatic solution ﬁtted to the RKF data for η ≥ 0.05. The initial conditions and all other parameters are as in Fig. 8

VI. CONCLUSION mation 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- In this work we sought to explore methods for effi- dimensional problems which closely resembled the WKB ciently integrating coupled systems of first order ODEs solutions. However, currently the method relies on find- with highly oscillatory solutions, which would be a ing the symbolic Jordan normal form of A which is very higher-dimensional extension of the RKWKB method, difficult for high-dimensional cases. It could form the with particular applications to cosmological systems of basis for a more fully numerical method which could be interest. applied to arbitrarily high dimensions. The work of previous authors [4, 5, 7, 8] suggests meth- A python code was created to implement the cor- ods based on the Magnus expansion could be candidates responding numerical integration methods, incorporat- for a multidimensional equivalent of RKWKB. However ing an adaptive step size procedure and the RKWKB comparison of the analytic Magnus expansion and WKB and RKF4(5) methods for comparison. The code leaves approximate solutions implies the two methods are fun- plenty of room for optimization, and one could certainly damentally distinct mathematical approaches. For equa- improve the performance by switching to a compiled lan- tions of the form of (2.1) WKB approximates the fre- guage like c. On simple two-dimensional systems, simple quency as hωi, whereas the first Magnus solution approx- p Magnus expansion based numerical methods ran signif- imates it as hω2i. The WKB solution series involves icantly faster than RKF4(5), with the Jordan-Magnus successively higher derivatives of ω while the Magnus ex- method running faster still, although not achieving the pansion involves successively higher order integrals. Test- speed of RKWKB. ing on trial equations of the form of (2.1), the WKB so- Finally we applied this code and these methods to lution give much higher accuracy. This indicates that the a simplified set of Einstein-Boltzmann equations corre- Magnus expansion and WKB based methods are orthog- sponding to photon and CDM oscillations in the radia- onal lines of research; however, the relationship between tion era. The results again demonstrate that Magnus ex- the two may benefit from more formal mathematical in- pansion based methods performed better than RKF4(5) quiry. by giving a speed improvement of about 50%. However, The simplest Magnus expansion based approximate so- the improvement was not as impressive as for the simpler lutions were found to gave inferior long range accuracy systems, and we could not implement the Jordan-Magnus compared to WKB. By introducing a linear transfor- method. For large times the numerical results fit the ex- 12 pected solutions for adiabatic perturbations. However, merical integration methods based on the Magnus ex- we could not fit a single adiabatic solution to the numer- pansion show promise for improving the performance of ical results for the potential Φ. Instead the results sug- not only cosmological Boltzmann solvers but also many gest a transition between two different adiabatic regimes. similar problems in other areas of physics. This could indicate a breakdown of the adiabatic approximation for that system. ACKNOWLEDGEMENTS Future work would involve testing these Magnus expansion based methods on other multidimensional sys- J.B. thanks the University of Cambridge Physics Part tems, including the 14 × 14 Einstein-Boltzmann matrix III project scheme, which formed the basis for this work. for l ≤ 2. If successful, one could extend up to high l W.H. was supported by a Gonville & Caius research fel- values and see how one might include these methods in lowship. We are also grateful for the advice of Fruzsina future Boltzmann solver codes. Agocs and Lukas Hergt of the Kavli Institute for Cos- These results demonstrate that multidimensional nu- mology, Cambridge.

[1] D. Blas, J. Lesgourgues, and T. Tram, CLASS: Cosmic [14] C. M. Benderand S. A. Orszag, Advanced Mathemati- Linear Anisotropy Solving System, Astrophysics Source cal Methods for Scientists and Engineers (Springer, New Code Library arxiv:1104.2932. York, 1999) Chap. 10: WKB theory, pp. 549–568. [2] A. Lewis, A. Challinor, and A. Lasenby, Efficient compu- [15] C. Führe, Lunds Universitet, How to compute tation of cosmic microwave background anisotropies in matrix norms, http://www.maths.lth.se/na/courses/ closed Friedmann-Robertson-Walker models, Astrophys. FMN081/FMN081-06/lecture6.pdf. J. 538, 473 (2000). [16] W. I. J. Haddadin and W. J. Handley, Rapid nu- [3] A. Refregier, L. Gamper, A. Amara, and L. Heisenberg, merical solutions for the Mukhanov-Sazaki equation, PyCosmo: An integrated cosmological Boltzmann solver, arXiv:1809.11095. Astron. Comput. 25, 38 (2018), arxiv:1708.05177. [17] S. Dodelson, Modern Cosmology (Academic Press, Ams- [4] W. J. Handley, A. N. Lasenby, and M. P. Hobson, terdam, 2003. http://www.slac.stanford.edu/spires/ The Runge-Kutta-Wentzel-Kramers-Brillouin method, find/books/www?cl=QB981:D62:2003). arXiv:1612.02288. [18] C. P. Ma and E. Bertschinger, Cosmological perturba- [5] F. J. Agocs, W. J. Handley, A. N. Lasenby, and M. P. tion theory in the synchronous and conformal newtonian Hobson, An efficient method for solving highly oscilla- gauges, Astrophys. J. 455, 7 (1995). tory ordinary differential equations with applications to [19] D. Baumann, Cosmology, Part III Mathematical Tri- physical systems, Phys. Rev. Res. 2 , 013030 (2020), pos (DAMTP, University of Cambridge, 2014 http:// arXiv:1906.01421. cosmology.amsterdam/education/cosmology/). [6] S. Blanes, F. Casas, J. Oteo, and J. Ros, The Magnus [20] J. B. Fraleigh and R. A. Beauregard, Linear Algebra, expansion and some of its applications, Phys. Rep. 470, 3rd ed. (Addison-Wesley, Reading, MA, 1995) Chap. 9.4: 151 (2009). Jordan Canonical Form. [7] A. Iserles, On the global error of discretization methods [21] G. H. Golub and C. F. Van Loan, Matrix Computations, for highly-oscillatory ordinary differential equations, BIT 4th ed., Johns Hopkins Studies in the Mathematical Sci- Numer. Math. 42, 561 (2002). ences (Johns Hopkins University Press, 2012) Chap. 7: [8] A. Iserles, On the numerical analysis of rapid Unsymmetric Eigenvalue Problems. oscillation, www.damtp.cam.ac.uk/user/na/NA_papers/ [22] N. Cai and M. J. Khan, Almost decouplability of any NA2003_06.ps.gz (2004). directed weighted network topology, Physica A 346, 637 [9] O. Chuluunbaatar, V. L. Derbov, A. Galtbayar, A. A. (2015), arXiv:1408.0355. Gusev, M. S. Kaschiev, S. I. Vinitsky, and T. Zhanlav, [23] J. Bamber, https://doi.org/10.5281/zenodo. Explicit Magnus expansions for solving the time- 3347094, https://github.com/jarb2/Beyond_the_ dependent Schrdinger equation, J. Phys. A 41, 295203 RKWKB_method (2019). (2008). [24] T. Oliphant, NumPy: A Guide to NumPy (Trelgol Pub- [10] E. Mananga, J. Moghaddasi, A. Sana, A. Akinmoladun, lishing USA, 2006). and M. Sadoqi, Advances in theory of solid-state nuclear [25] A. Meurer et al., SymPy: Symbolic computing in Python, magnetic resonance, J. Nat. Sci. 1, e109 (2016). PeerJ Comput. Sci. 3, e103 (2017). [11] P. Pechukas and J. Light, On the exponential form of [26] W. H. Press, B. P. Flannery, S. A. Teukolsky, and W. T. time-displacement operators in quantum mechanics, J. Vetterling, Numerical Recipes in C: The Art of Scien- Chem. Phys. 44, 3897 (1966). tific Computing (Cambridge University Press, New York, [12] J. Galvez Ghersiand A. V. Frolov, Two-point correlators 1988) Chap. 17.2: Adaptive Stepsize Control for Runge- revisited: Fast and slow scales in multifield models of Kutta. Inflation, J. Cosmol. Astropart. Phys. 05 (2017) 047. [13] P. Jakobsen, Introduction to the method of multiple scales, arxiv:1312.3651. 13

Appendix A: Procedure for adaptive step size control

We implemented the adaptive step size procedure laid out in Press and Flannery (1988) [26]. To estimate the error, and thus the correct step size for each step, we compute the estimate of xn+1 from one step of size h, and a separate ∗ ∗ estimate using two steps size h/2, denoting this xn+1. The error is then ∆i = |xi n+1 − xi n+1| for each component max ∗ i. We deﬁne a maximum tolerable error as ∆i = (atol + rtol|xi n+1|), where atol and rtol parametrize the absolute and relative error tolerance, and is a small scale factor. The procedure is then as follows:

∗ 1. Calculate xn+1 and xn+1 for step size h. q max 1 PN max 2 2. Calculate ∆i and ∆i and the error ratio R = N i (∆i/∆i ) .

−1/(ν+1) 3. If R ≤ 1, compute a new h as hnew = hSR where S = safety factor ∼ 0.99, ν = order of method. Go to 5.

−1/ν 4. Else if R > 1, compute a new h as hnew = hSR .

If hnew < hmin, set hnew = hmin and go to 5. Otherwise go to 1.

5. If hnew < 0.2h, set hnew = 0.2h.

If hnew > hmax, set hnew = hmax.

If hnew > 10h, set hnew = 10h. ∗ 6. Use current xn+1 as the estimate. Advance time. Set h = hnew. Go to next step. The RKWKB method used the adaptive step size procedure outlined above, setting ν = 1. The RKF4(5) was implemented following [26]. The approach to step size control is the same, except the error is calculated using the diﬀerence between the fourth and ﬁfth order RK methods. The atol and rtol 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 Fig. 7. For the RKWKB method the values for the Magnus methods were used.