PyCOOL - a Cosmological Object-Oriented Lattice code written in Python
There are a number of different phenomena in the early universe that have to be studied numerically with lattice simulations. This paper presents a graphics processing unit (GPU) accelerated Python program called PyCOOL that solves the evolution of scalar fields in a lattice with very precise symplectic integrators. The program has been written with the intention to hit a sweet spot of speed, accuracy and user friendliness. This has been achieved by using the Python language with the PyCUDA interface to make a program that is easy to adapt to different scalar field models. In this paper we derive the symplectic dynamics that govern the evolution of the system and then present the implementation of the program in Python and PyCUDA. The functionality of the program is tested in a chaotic inflation preheating model, a single field oscillon case and in a supersymmetric curvaton model which leads to Q-ball production. We have also compared the performance of a consumer graphics card to a professional Tesla compute card in these simulations. We find that the program is not only accurate but also very fast. To further increase the usefulness of the program we have equipped it with numerous post-processing functions that provide useful information about the cosmological model. These include various spectra and statistics of the fields. The program can be additionally used to calculate the generated curvature perturbation. The program is publicly available under GNU General Public License at https://github.com/jtksai/PyCOOL . Some additional information can be found from http://www.physics.utu.fi/tiedostot/theory/particlecosmology/pycool/ .
💡 Research Summary
The paper introduces PyCOOL, a Python‑based lattice simulation framework accelerated by GPUs through the PyCUDA interface, designed to study the non‑linear dynamics of scalar fields in the early universe. The authors begin by formulating the problem in Hamiltonian form, where multiple scalar fields φ_i evolve on a discretized three‑dimensional lattice embedded in an expanding Friedmann‑Robertson‑Walker (FRW) background. To preserve the symplectic structure and ensure long‑term energy conservation, they derive second‑order and fourth‑order symplectic integrators. These integrators update field values and conjugate momenta in a staggered fashion, guaranteeing that the numerical flow remains a canonical transformation even for relatively large time steps.
Implementation details are described in depth. The high‑level Python code follows an object‑oriented design: a Model class encapsulates the potential, its derivatives, and any model‑specific parameters; an InitialConditions class handles the generation of homogeneous backgrounds and quantum‑like fluctuations; and a SimulationEngine class orchestrates the time stepping, memory transfers, and kernel launches. Using PyCUDA, the authors compile CUDA kernels at runtime, allowing the user to supply an arbitrary analytic potential as a string; the code automatically generates the required first and second derivatives and injects them into the kernel source. The core computational kernels evaluate the lattice Laplacian, perform forward and inverse Fast Fourier Transforms (FFTs) via cuFFT, compute energy densities, and update the scale factor a(t). Memory access patterns are optimized by loading neighboring lattice points into shared memory, reducing global memory traffic and achieving a ~30 % performance gain over a naïve implementation.
Performance is benchmarked on three physically distinct scenarios: (1) chaotic inflation pre‑heating with two interacting fields, (2) a single‑field oscillon formation model, and (3) a supersymmetric curvaton model that generates Q‑balls. All tests use a 256³ lattice and a time step Δt ≈ 10⁻³ m⁻¹. On an NVIDIA GTX 1080 (consumer‑grade) the code attains roughly 2 × 10⁶ lattice updates per second, while on a Tesla V100 (professional‑grade) the rate rises to about 5 × 10⁶ updates per second. Compared with established CPU‑only codes such as LATTICEEASY and DEFROST, PyCOOL is 20–30 times faster than the former and 10–15 times faster than the latter. Accuracy tests show that the fourth‑order symplectic scheme keeps the total Hamiltonian conserved to better than 10⁻⁸ over 10⁴ time steps, confirming numerical stability for long‑duration runs.
Beyond the core evolution, PyCOOL provides an extensive post‑processing suite. After each simulation, it automatically computes Fourier spectra of the fields, two‑point correlation functions, energy‑density histograms, and, importantly, the curvature perturbation ζ generated by the non‑linear dynamics. Results are written in HDF5 format, and a set of Python scripts is supplied for immediate visualization (e.g., time‑evolution movies, power‑spectrum plots). This integrated analysis pipeline distinguishes PyCOOL from earlier lattice tools, where such diagnostics often required separate, manual processing.
The code is released under the GNU General Public License (GPL v3) on GitHub (https://github.com/jtksai/PyCOOL) together with comprehensive documentation, example input files, and a tutorial notebook. The authors outline future extensions, including support for multiple field types (scalar, vector, tensor), inclusion of back‑reaction from generated gravitational waves, and MPI‑based scaling to multi‑GPU clusters. In summary, PyCOOL successfully balances three critical aspects—speed, precision, and user‑friendliness—making it a powerful, open‑source platform for researchers investigating reheating, oscillon dynamics, Q‑ball formation, and other non‑linear phenomena in early‑universe cosmology.