Spinsim: A GPU optimized python package for simulating spin-half and spin-one quantum systems

Abstract

Spinsim simulates the quantum dynamics of unentangled spin-1/2 and spin-1 systems evolving under time-dependent control. While other solvers for the time-dependent Schrödinger equation optimize for larger state spaces but less temporally-rich control, spinsim is optimized for intricate time evolution of a minimalist system. Efficient simulation of individual or ensemble quanta driven by adiabatic sweeps, elaborate pulse sequences, complex signals and non-Gaussian noise is the primary target application. We achieve fast and robust evolution using a geometric integrator to bound errors over many steps, and split the calculation parallel-in-time on a GPU using the numba just-in-time compiler. Speed-up is three orders of magnitude over QuTip's sesolve and Mathematica's NDSolve, and four orders over SciPy's ivp_solve for equal accuracy. Interfaced through python, spinsim should be useful for simulating robust state preparation, inversion and dynamical decoupling sequences in NMR and MRI, and in quantum control, memory and sensing applications with two- and three-level quanta. Program summaryProgram Title: SpinsimCPC Library link to program files:https://doi.org/10.17632/f6rdk4gyxr.1Developer's repository link:https://github.com/alexander-tritt-monash/spinsimLicensing provisions: BSD 3-clauseProgramming language: Python (3.7 or greater)Nature of problem: Quantum sensing is a domain of quantum technology where the dynamics of quantum systems are used to infer properties of the systems' environments. The development of quantum sensing protocols is greatly sped-up by software simulations of the new techniques. Quantum sensing simulation benefits from temporally-rich control of individual quanta. However, current specialized time-dependent Schrödinger equation solvers are instead optimized only for simple pulses in large Hilbert spaces. Thus, there is a need for efficient simulation of individual or ensemble quanta driven by adiabatic sweeps, elaborate pulse sequences, complex signals and non-Gaussian noise.Solution method: Spinsim simulates the quantum dynamics of spin-1/2 and spin-1 systems evolving under time-dependent control. We first speed up the integration of the time-dependent Schrödinger equation by splitting the calculation parallel-in-time on a GPU using the numba [1] just-in-time compiler. We achieve fast and robust evolution using a geometric integrator to bound errors over many steps. A dynamic rotating frame transformation and Lie-Trotter decomposition are used to decrease effective step sizes on long and short time scales, respectively. Hence, each individual step is more accurate. Spinsim is interfaced via a python package, meaning it can be used by researchers inexperienced with geometric integrators and GPU parallelism.Additional comments including restrictions and unusual features: To use the (default) Nvidia cuda GPU parallelization, one needs to have a cuda compatible Nvidia GPU [2]. For cuda mode to function, one also needs to install the Nvidia cuda toolkit [3]. If cuda is not available on the system, the simulator will automatically parallelize over multicore CPUs instead. Developed and tested on Windows 10; tested on Windows 11. CPU functionality tested on MacOS 10.16 Big Sur (note that MacOS 10.14 Mojave and higher are not compatible with cuda hardware and software). The package (including cuda functionality) is in principle compatible with Linux, but functionality has not been tested.