Beyond the Runge-Kutta-Wentzel-Kramers-Brillouin method Jamie Bamber1, 2, ∗ and Will Handley1, 3, 4, y 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 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 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 meth- ods achieve an increase in run speed of about 50% compared to a standard Runge-Kutta integra- tion 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 stan- dard 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 y [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 sufficiently 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 1 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 pro- ceeds 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 first 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) = p e 0 1 1 + p 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 j!_ (t)j j!2(t)j for the approxima- tion 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 infinite series (2.8) [6], 1 X M(t0; t) = 1 + (t − t0)A(t0): (2.16) x(t) = exp Wn(t) x(0); (2.8) n=1 The RKWKB method (to second order in the WKB the first 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 Wi then formal time. We denote the nonconformal cosmic Hub- we can use those directly. In that case the W series trun- ble rate as H = H=a. Theτ _ = −neσT a is the confor- cated at W = Ω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 Wi cannot be computed analytically, we tion. Mean densities are described using parameters of can use an approach from Blanes et al.