Speaker
Description
GPUs offer a way to massively parallelise certain types of physical models, allowing us to compare more models for the same compute budget (and hence do more science) whilst avoiding the use of emulators (which must be retrained every time you modify the underlying physics). More efficient forward models also allow us to tackle higher-dimensional problems (such as in field-level inference) and more complex foreground systematics and nuisance parameters.
In particular, the Einstein-Boltzmann equations that underly models of the CMB are linear ordinary differential equations that are highly amenable to parallelisation on a GPU, across both "k-modes" and parameters. Meanwhile, modern cosmological analyses are burdened with the technical debt of low-level codes written in FORTRAN and C, which are difficult to modify when trying to test new models of physics.
We have therefore written a new CMB code in Python and JAX (DISCO-EB) which can take advantage of the massive computational throughput of modern GPUs, but is easy to modify with new physics. As part of this work, we have developed new ODE solvers that are better optimised for solving many small ODEs at once and are 100x faster than the current Python/JAX state-of-the-art (diffrax). The result is precision similar to that of existing codes but at vastly reduced $ inference cost. The ODE solvers themselves are packaged separately and could have broad application in the scientific community.