TorchGPE is a python package for solving the Gross-Pitaevskii equation (GPE). The numerical solver is designed to integrate a wave function in a variety of linear and non-linear potentials. The code is based on symmetric Fourier split-step propagation, both in real and imaginary time. It has a modular approach that allows the integration of arbitrary self-consistent and time-dependent potentials.
The code uses PyTorch tensors to speed up the calculation (which heavily relies on FFTs). You can check if cuda is available on your system with
import torch
torch.cuda.is_available()If not, the default device is set to CPU.
The package can be installed via pip
pip install git+https://github.com/qo-eth/TorchGPE.gitOr, if you already downloaded the repository, run
pip install .from the locally cloned directory.
Full documentation for this repository is available here.
The code was created within the Quantum Optics research group at ETH Zurich.