Computational Graphs for High Performance Simulation of the Dynamics of Open Quantum Systems

[hodgestar]

Project

QuTiP

Summary

Describing computations that will be performed in a graph, and then compiling that graph optimally for the hardware the computation will be performed on has proved an effective approach to obtain good performance (e.g. TensorFlow, PyTorch, Jax). This project will attempt to apply a similar technique to calculating the evolution of open quantum systems in QuTiP.

submitter

Simon Cross

project lead

@ajgpitch

Community benefit

QuTiP users will be able to solve for the dynamics of many bigger quantum systems more quickly and using much less memory:

• Students with cheap laptops will be able to simulate more realistic systems.
• Educators will be able to show more challenging examples.
• Engineers modeling quantum hardware will be able to do so more precisely.
• HPC users will be able to make more effective use of their supercomputers.

The computational graph will allow generating optimal code and algorithms
for both CPUs and GPUs.

Once the technique is established, it will likely be ported to other
libraries and languages, for example, QuTiP's
Quantum Toolbox in Julia.

This feature will be a core component of QuTiP 6.

Amount requested

10000

Execution plan

This grant will fund Simon Cross (me),
to work on this project part-time for six months.

I will be working alongside QuTiP's other part-time developer, Eric Giguère, who is
already funded separately to work on QuTiP more broadly.

Me and Eric were the QuTiP developers responsible for delivering many of the
new core features in QuTiP 5, including the components most relevant to this project (QuTiP 5 data layer, QObjEvo, solvers and integrators).

Month 1 & 2: Building the graph

During the first two months I will pull together concepts from existing
QuTiP components and experiments into a framework for constructing and
recording the system Hamiltonian as a computational graph.

Initial features of the new graph will include:

• Time-dependent scalar coefficients (Coefficients in QuTiP 5)
• Time-independent Hilbert space operators (Qobj in QuTiP 5)
• A variety of sparse representations of operationrs (QuTiP 5 data layer)
• Sums of products scalar coefficients and operators (QobjEvo in QuTiP 5)
• Tensor products of these (no corresponding feature in QuTiP 5)
• Tracking the system dimensions (Dimensions in QuTiP 5)

The initial graph will be convertible to existing QuTiP 5 objects, allowing
the graph to be constructed and tested.

Month 3 & 4: Building a computational engine

The second two months will focus on producing a proper computational engine
to allow solving the ODE (ordinary differential equation) for the quantum
dynamics on a CPU.

Two possible approaches will be explored:

• Handwrite a Cython extension for executing the computational graph.
  This fits well with QuTiP's existing code base, and allows maximum flexibility
  and re-use of existing QuTiP Cython code.

• Write the computational engine on top of Jax. This leverages an existing
  library with great performance, support for GPUs, and support for
  auto-differentiation. However, it runs the risk of having to extend a
  complex library if limitations are encountered.

Month 5 & 6: A first application

QuTiP includes a Hierarchical Equations of Motion (HEOM) solver. This is an
advancted algorithm for solving for the dynamics of a quantum system
that interacts closely with its environemnt.

The Hierarchical Equations of Motion, as the name suggests, are highly
structured. QuTiP 5 cannot capture this structure well, but the new
computational graph should do so well.

The last two months will focus on extending the graph to support the HEOM.
This should reduce both the memory and computation required to solve for the
dynamics, greatly extending the classes of systems that can be modelled
on easily available hardware.