A Python GPU-accelerated solver for the Gross-Pitaevskii equation and applications to many-body cavity QED

TL;DR

TorchGPE is a Python package that efficiently solves the Gross-Pitaevskii equation using GPU acceleration, enabling advanced simulations in many-body cavity QED with flexible potential modeling.

Contribution

The paper introduces TorchGPE, a GPU-accelerated Python solver for the GPE with a modular design for complex potentials, significantly improving computational efficiency over CPU methods.

Findings

Achieved substantial speed-up using GPU acceleration.

Demonstrated flexibility in modeling various potentials.

Enabled advanced simulations in cavity QED research.

Abstract

TorchGPE is a general-purpose Python package developed for solving the Gross-Pitaevskii equation (GPE). This solver is designed to integrate wave functions across a spectrum of linear and non-linear potentials. A distinctive aspect of TorchGPE is its modular approach, which allows the incorporation of arbitrary self-consistent and time-dependent potentials, e.g., those relevant in many-body cavity QED models. The package employs a symmetric split-step Fourier propagation method, effective in both real and imaginary time. In our work, we demonstrate a significant improvement in computational efficiency by leveraging GPU computing capabilities. With the integration of the latter technology, TorchGPE achieves a substantial speed-up with respect to conventional CPU-based methods, greatly expanding the scope and potential of research in this field.