GPU-accelerated solutions of the nonlinear Schrödinger equation for simulating 2D spinor BECs

Abstract

As a first approximation beyond linearity, the nonlinear Schrödinger equation (NLSE) reliably describes a broad class of physical systems. Though numerical solutions of this model are well-established, these methods can be computationally complex. In this paper, we showcase a code development approach, demonstrating how computational time can be significantly reduced with readily available graphics processing unit (GPU) hardware and a straightforward code migration using open-source libraries. This process shows how CPU computations with power-law scaling in computation time with grid size can be made linear using GPUs. As a specific case study, we investigate the Gross-Pitaevskii equation, a specific version of the nonlinear Schrödinger model, as it describes in two dimensions a trapped, interacting, two-component Bose-Einstein condensate (BEC) subject to a spatially dependent interspin coupling, resulting in an analog to a spin-Hall system. This computational approach lets us probe high-resolution spatial features – revealing an interaction-dependent phase transition – all in a reasonable amount of time. Our computational approach is particularly relevant for research groups looking to easily accelerate straightforward numerical simulation of physical phenomena. Program summaryProgram Title:spinor-gpeCPC Library link to program files:https://doi.org/10.17632/sn6gjp99td.1Developer's repository link:https://github.com/ultracoldYEG/spinor-gpeLicensing provisions: MITProgramming language: PythonNature of problem: Calculate the ground- and time-evolved states to the quasi-2D, pseudospinor (two-component) nonlinear Schrodinger equation, with additional Zeeman interaction and momentum-dependent, interspin coupling.Solution method:spinor-gpe is a high-level, object-oriented Python package built on Numpy and PyTorch. It propagates the pseudospinor NLSE using a time-splitting spectral method. This package accelerates solutions using NVIDIA hardware and PyTorch's cuFFT libraries and tensor functionality.Additional comments including restrictions and unusual features: The NVIDIA CUDA backend of PyTorch is not supported on Mac computers. To run this code on a Mac system, the dependency installation will need to exclude the CUDA toolkit.