A python package simulating the quasi-2D pseudospin-1/2 Gross-Pitaevskii equation with NVIDIA GPU acceleration.

spinor-gpe is high-level, object-oriented Python package for numerically solving the quasi-2D, psuedospinor (two component) Gross-Piteavskii equation (GPE), for both ground state solutions and real-time dynamics. This project grew out of a desire to make high-performance simulations of the GPE more accessible to the entering researcher.

While this package is primarily built on NumPy, the main computational heavy-lifting is performed using PyTorch, a deep neural network library commonly used in machine learning applications. PyTorch has a NumPy-like interface, but a backend that can run either on a conventional processor or a CUDA-enabled NVIDIA(R) graphics card. Accessing a CUDA device will provide a significant hardware acceleration of the simulations.

This package has been tested on Windows, Mac, and Linux systems.

View the documentation on ReadTheDocs

1. PyTorch >= 1.8.0
2. cudatoolkit >= 11.1
3. NumPy

4. matplotlib (visualizing results)
5. tqdm (progress messages)
6. scikit-image (matrix signal processing)
7. ffmpeg = 4.3.1 (animation generation)

The dependencies for spinor-gpe can be installed directly into the new conda virtual environment spinor using the environment.yml file included with the package: :

conda env create --file environment.yml

This installation may take a while.

The above dependencies can also be installed manually using conda into a virtual environment: :

conda activate <new_virt_env_name>
conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge
conda install numpy matplotlib tqdm scikit-image ffmpeg spyder

To verify that Pytorch was installed correctly, you should be able to import it:

>>> import torch
>>> x = torch.rand(5, 3)
>>> print(x)
tensor([[0.2757, 0.3957, 0.9074],
        [0.6304, 0.1279, 0.7565],
        [0.0946, 0.7667, 0.2934],
        [0.9395, 0.4782, 0.9530],
        [0.2400, 0.0020, 0.9569]])

Also, if you have an NVIDIA GPU, you can test that it is available for GPU computing:

>>> torch.cuda.is_available()
True

CUDA is the API that interfaces with the computing resources on NVIDIA graphics cards, and it can be accessed through the PyTorch package. If your computer has an NVIDIA graphics card, start by verifying that it is CUDA-compatible. This page lists out the compute capability of many NVIDIA devices. (Note: yours may still be CUDA-compatible even if it is not listed here.)

Given that your graphics card can run CUDA, the following are the steps to install CUDA on a Windows computer:

1. Install the NVIDIA CUDA Toolkit. Go to the CUDA download page for the most recent version. Select the operating system options and installer type. Download the installer and install it via the wizard on the screen. This may take a while. For reference, here is the Windows CUDA Toolkit installation guide.

   To test that CUDA is installed, run which nvcc, and, if instlled correctly, will return the installation path. Also run nvcc --version to verify that the version of CUDA matches the PyTorch CUDA toolkit version (>=11.1).

2. Download the correct drivers for your NVIDIA device. Once the driver is installed, you will have the NVIDIA Control Panel installed on your computer.

1. Clone the repository.
2. Navigate to the spinor_gpe/examples/ directory, and start to experiment with the examples there.

This package has a simple, object-oriented interface for imaginary- and real-time propagations of the pseudospinor-GPE. While there are other operations and features to this package, all simulations will have the following basic structure:

>>> import pspinor as spin
>>> DATA_PATH = '<project_name>/Trial_###'
>>> ps = spin.PSpinor(DATA_PATH)

The program will create a new directory DATA_PATH, in which the data and results from this simulation trial will be saved. If DATA_PATH is a relative path, as shown above, then the trial data will be located in the /data/ folder. When working with multiple simulation projects, it can be helpful to specify a <project_name> directory; furthermore, the form Trial_### is convenient, but not strictly required.

The example below demonstrates imaginary-time propagation. The method PSpinor.imaginary performs the propagation loop and returns a PropResult object. This object contains the results, including the final wavefunctions and populations, and analysis and plotting methods (described below).

>>> DT = 1/50
>>> N_STEPS = 1000
>>> DEVICE = 'cuda'  
>>> res = ps.imaginary(DT, N_STEPS, DEVICE, is_sampling=True, n_samples=50)

For real-time propagation, use the method PSpinor.real.

PropResult provides several methods for viewing and understanding the final results. The code block below demonstrates several of them:

>>> res.plot_spins()  # Plots the spin-dependent densities and phases.
>>> res.plot_total()  # Plots the total densities and phases.
>>> res.plot_pops()   # Plots the spin populations throughout the propagation. 
>>> res.make_movie()  # Generates a movie from the sampled wavefunctions.

Note that PSpinor also exposes methods to plot the spin and total densities. These can be used independent of PropResult:

>>> ps.plot_spins()

Likely you will want to repeat or chain together different segments of this structure. Demonstrations of this are shown in the Examples gallery.