Solving the Dirac equation on a GPU for strong-field processes in multidimensional background fields

In this paper, we show how to solve the Dirac equation, $(i{\gamma }^{\mu }[{\mathrm{\partial }}_{\mu }+ie{A}_{\mu }(t,\mathbf{x})]-m)\psi =0$, on a GPU. This is orders of magnitude faster than solving it on CPU and allows us to consider background fields, $A_{\mu }(t,\mathbf{x})$, that depend on $2+1$ or even $3+1$ coordinates. Our approach is conveniently implemented using the computational library JAX. We show how to obtain the probabilities of Schwinger and nonlinear Breit-Wheeler pair production from these solutions using a scattered-wave-function approach and compare the results with the worldline-instanton approximations.