QuTiPv5 Paper Example: QuTiP-JAX with mesolve and auto-differnetiation¶

Authors: Maximilian Meyer-Mölleringhof (m.meyermoelleringhof@gmail.com), Rochisha Agarwal (rochisha.agarwal2302@gmail.com), Neill Lambert (nwlambert@gmail.com)

For many years now, GPUs have been a fundamental tool for accelerating numerical tasks. Today, many libraries enable off-the-shelf methods to leverage GPUs' potential to speed up costly calculations. QuTiP’s flexible data layer can directly be used with many such libraries and thereby drastically reduce computation time. Despite a big variety of frameworks, in connection to QuTiP, development has centered on the QuTiP-JAX integration [1] due to JAX's robust auto-differentiation features and widespread adoption in machine learning.

In these examples we illustrate how JAX naturally integrates into QuTiP v5 [2] using the QuTiP-JAX package. As a simple first example, we look at a one-dimensional spin chain and how we might employ mesolve() and JAX to solve the related master equation. In the second part, we focus on the auto-differentiation capabilities. For this we first consider the counting statistics of an open quantum system connected to an environment. Lastly, we look at a driven qubit system and computing the gradient of its state population in respect to the drive's frequency using mcsolve().

Introduction¶

In addition to the standard QuTiP package, in order to use QuTiP-JAX, the JAX package needs to be installed. This package also comes with jax.numpy which mirrors all of numpys functionality for seamless integration with JAX.

An immediate effect of importing qutip_jax is the availability of the data layer formats jax and jax_dia. They allow for dense (jax) and custom sparse (jax_dia) formats.

In order to use the JAX data layer also within the master equation solver, there are two settings we can choose. First, is simply adding method: diffrax to the options parameter of the solver. Second, is to use the qutip_jax.set_as_default() method. It automatically switches all data data types to JAX compatible versions and sets the default solver method to diffrax.

To revert this setting, we can set the function parameter revert = True.

Using the GPU with JAX and Diffrax - 1D Ising Spin Chain¶

Before diving into the example, it is worth noting here that GPU acceleration depends heavily on the type of problem. GPUs are good at parallelizing many small matrix-vector operations, such as integrating small systems across multiple parameters or simulating quantum circuits with repeated small matrix operations. For a single ODE involving large matrices, the advantages are less straightforward since ODE solvers are inherently sequential. However, as it is illustrated in the QuTiP v5 paper [2], there is a cross-over point at which using JAX becomes beneficial.

1D Ising Spin Chain¶

To illustrate the usage of QuTiP-JAX, we look at the one-dimensional spin chain with the Hamiltonian

$H = \sum_{i=1}^N g_0 \sigma_z^{(n)} - \sum_{n=1}^{N-1} J_0 \sigma_x^{(n)} \sigma_x^{(n+1)}$.

We hereby consider $N$ spins that share an energy splitting of $g_0$ and have a coupling strength $J_0$. The end of the chain connects to an environment described by a Lindbladian dissipator which we model with the collapse operator $\sigma_x^{(N-1)}$ and coupling rate $\gamma$.

As part of the QuTiPv5 paper, we see an extensive study on the computation time depending on the dimensionality $N$. In this example we cannot replicate the performance of a supercomputer of course, so we rather focus on the correct implementation to solve the Lindblad equation for this system using JAX and mesolve().

Auto-Differentiation¶

We have seen in the previous example how the new JAX data-layer in QuTiP works. On top of that, JAX adds the features of auto-differentiation. To compute derivatives, it is often numerical approximations (e.g., finite difference method) that need to be employed. Especially for higher order derivatives, these methods can turn into costly and inaccurate calculations.

Auto-differentiation, on the other hand, exploits the chain rule to compute such derivatives. The idea is that any numerical function can be expressed by elementary analytical functions and operations. Consequently, using the chain rule, the derivatives of almost any higher-level function become accessible.

Although there are many applications for this technique, in this chapter we want to focus on two examples where auto-differentiation becomes relevant.

Statistics of Excitations between Quantum System and Environment¶

We consider an open quantum system that is in contact with an evironment via a single jump operator. Additionally, we have a measurement device that tracks the flow of excitations between the system and the environment. The probability distribution that describes the number of such exchanged excitations $n$ in a certain time $t$ is called the full counting statistics and denoted by $P_n(t)$. This statistics is a defining property that allows to derive many experimental observables like shot noise or current.

For the example here, we can calculate this statistics by using a modified version of the density operator and Lindblad master equation. We introduce the tilted density operator $G(z,t) = \sum_n e^{zn} \rho^n (t)$ with $\rho^n(t)$ being the density operator of the system conditioned on $n$ exchanges by time $t$, so $\text{Tr}[\rho^n(t)] = P_n(t)$. The master equation for this operator, including the jump operator $C$, is then given as

$\dot{G}(z,t) = -\dfrac{i}{\hbar} [H(t), G(z,t)] + \dfrac{1}{2} [2 e^z C \rho(t)C^\dagger - \rho C^\dagger C - C^\dagger C \rho(t)]$.

We see that for $z = 0$, this master equation becomes the regular Lindblad equation and $G(0,t) = \rho(t)$. However, it also allows us to describe the counting statistics through its derivatives

$\langle n^m \rangle (t) = \sum_n n^m \text{Tr} [\rho^n (t)] = \dfrac{d^m}{dz^m} \text{Tr} [G(z,t)]|_{z=0}$.

These derivatives are precisely where the auto-differention by JAX finds its application for us.

When working with JAX you can choose the type of device / processor to be used. In our case, we will resort to the CPU since this is a simple Jupyter Notebook. However, when running this on your machine, you can opt for using your GPU by simpy changing the argument below.

Driven One Qubit System & Frequency Optimization¶

As a second example for auto differentiation, we consider the driven Rabi model, which is given by the time-dependent Hamiltonian

$H(t) = \dfrac{\hbar \omega_0}{2} \sigma_z + \dfrac{\hbar \Omega}{2} \cos (\omega t) \sigma_x$

with the energy splitting $\omega_0$, $\Omega$ as the Rabi frequency, the drive frequency $\omega$ and $\sigma_{x/z}$ are Pauli matrices. When we add dissipation to the system, the dynamics is given by the Lindblad master equation, which introduces collapse operator $C = \sqrt{\gamma} \sigma_-$ to describe energy relaxation.

For this example, we are interested in the population of the excited state of the qubit

$P_e(t) = \bra{e} \rho(t) \ket{e}$

and its gradient with respect to the frequency $\omega$.

We want to optimize this quantity by adjusting the drive frequency $\omega$. To achieve this, we compute the gradient of $P_e(t)$ in respect to $\omega$ by using JAX's auto-differentiation tools and QuTiP's mcsolve().

References¶

[1] QuTiP-JAX

[2] QuTiP 5: The Quantum Toolbox in Python

About¶

Testing¶