Large Scale HEOM Solver Using PETSC - GSoC 2026

Student: Chinmay Tangal
Organization: QuTiP - Quantum Toolbox in Python, NumFOCUS
Mentors: Neill Lambert, Emil Vatai, Alex Pitchford
Project Duration: June 2026 - August 2026

Project Overview

What This Feature Enables

The PETSc backend for QuTiP’s Hierarchical Equations of Motion (HEOM) solver enables the simulation of open quantum systems at unprecedented scales by leveraging high-performance distributed computing. As the hierarchy depth and system size increase, the number of Auxiliary Density Operators (ADOs) in HEOM grows exponentially, often leading to severe memory bottlenecks. By utilizing petsc4py and mpi4py, this backend distributes the massive hierarchy matrix assembly across multiple MPI ranks. Once assembled, it taps into PETSc’s robust suite of scalable solvers—using its Time Stepping (TS) module for stable, parallel time-evolution of stiff ODEs, and its Krylov Subspace (KSP) solvers to directly and efficiently find non-equilibrium steady states. Ultimately, this allows researchers to seamlessly transition from running modest HEOM models on a local laptop to executing massive, computationally demanding open quantum system simulations on large-scale HPC clusters.

Key Achievements

Distributed Matrix Assembly: Implemented MPI-based distributed assembly of the HEOM superoperator matrix, breaking the single-machine memory barrier.
Parallel ODE Integration (TS): Integrated PETSc’s Time Stepping module (e.g., implicit BDF solvers) to handle the inherently stiff ODEs produced by the HEOM hierarchy, allowing for fast, adaptive, and scalable time evolution.
Efficient Steady-State Solvers (KSP): Developed a steady-state solver utilizing PETSc’s Krylov Subspace iterative linear solvers, drastically reducing the time needed to compute non-equilibrium steady states.
Seamless QuTiP Integration: Wrapped the complex PETSc and MPI routines inside standard QuTiP objects, preserving a familiar API for existing QuTiP users.
HPC Benchmarking & Validation: Validated the backend extensively against established models (e.g., Ising spin chains with Drude-Lorentz and Fermionic baths) on high-performance supercomputers.

Integration with QuTiP

The new PETSc backend has been designed to integrate gracefully with the existing QuTiP HEOMSolver infrastructure:

solver = HEOMSolver(H, baths, max_depth=max_depth, backend="petsc", options=options)
result = solver.run(rho0, tlist)

Unified API: Accessed simply by passing backend="petsc" and corresponding integration/solver options directly into the standard HEOMSolver initialization.
State Restoration: Returns standard QuTiP Qobj states and ADO representations.
Modular Integrator: Built a custom IntegratorPETSc subclass that adheres to QuTiP’s Integrator API, ensuring that future enhancements in QuTiP will naturally propagate to the PETSc backend.

Technical Implementation

Code Integration (Final PRs and Merges)

Core Infrastructure: Distributed Assembly and TS/KSP Integrations

PR #2976: Implementation of PETSC Backend into Qutip.
Status: PR under review.

Documentation and Tutorials

Documentation

PR #2976: Documentation in same PR.
Status: PR under review.

Tutorial

PR #161: Tutorial for the PETSc backend of HEOM.
Status: PR under review.

Comprehensive tutorial covering:

Setting up MPI and petsc4py environments.
Running time-evolution for an Ising Spin Chain coupled to Drude-Lorentz baths across multiple nodes.
Direct computation of the steady state utilizing PETSc KSP solvers.
Deploying the scripts on cluster infrastructure using SLURM batch scripts.

Challenges and Solutions

Throughout the project, I encountered several technical hurdles that required careful architectural decisions:

1. Distributed Assembly of the RHS Matrix

Initially, one major bottleneck was that the Right-Hand Side (RHS) matrix for the ODE was being assembled in full on a single MPI rank before being passed to PETSc. This negated the memory benefits of distributed computing for very deep hierarchies. To fix this, I completely redesigned the matrix assembly process. Using PETSc.Mat.Type.MPIAIJ, each MPI rank now queries its local ownership range (getOwnershipRange()) and directly computes and injects only the block rows it is responsible for. This ensures the global matrix never exists on any single machine in its entirety, successfully breaking the single-machine memory barrier.

2. Backend Abstraction and Clean Integration

Integrating the PETSc logic directly into qutip/solver/heom/bofin_solvers.py initially led to ad-hoc branching, making the codebase harder to maintain and read. To make the backend work in a tidy way without messing up the existing solver logic, I implemented a factory pattern directly into the HEOMSolver base class. When a user initializes the solver with backend="petsc", HEOMSolver.__new__ dynamically loads and instantiates the PETScHEOMSolver from a separate backend_petsc.py module.

The PETScHEOMSolver seamlessly inherits from HEOMSolver, overriding _calculate_rhs() to bypass the standard scipy.sparse matrix construction. Instead, it triggers the distributed PETSc matrix assembly. Additionally, a custom IntegratorPETSc is decoupled and registered via add_integrator(). This architecture keeps bofin_solvers.py completely clean of PETSc-specific code, preserving standard QuTiP functionality while establishing a flexible registry (_heom_backends) so that any future backends can be easily added.

Benchmarks and Scaling Performance

The primary goal of the PETSc backend was to allow massive scaling. We validated this by benchmarking on high-performance supercomputers, demonstrating excellent scaling capabilities.

MPI Scaling

To understand how effectively PETSc distributes the workload, we tested the 80,730 ADO system across up to 64 MPI ranks. The implementation scales nearly linearly, distributing the matrix construction and solver effortlessly, proving that communication bottlenecks were successfully minimized.

Steady State Solver Performance

For the Steady State implementation leveraging the Krylov subspace solvers, the backend again demonstrated phenomenal scaling, solving large steady states in fractions of the time taken by traditional solvers.

Project Impact and Conclusion

The integration of the PETSc backend into QuTiP’s HEOM solver represents a significant leap forward for the open quantum systems community. By shattering the memory constraints of single-machine computation and providing access to state-of-the-art scalable solvers, researchers can now tackle complex phenomena such as large spin arrays, intricate light-harvesting complexes, and detailed single-molecule electronics—that were previously completely out of reach.

The implementation proves that complex, distributed HPC workflows can be packaged into highly usable, Pythonic interfaces without sacrificing performance. Through extensive testing, robust documentation, and step-by-step tutorials, this project ensures that both experts in distributed computing and standard researchers can harness the full power of modern supercomputing hardware for quantum physics simulations.

Student Contact: GitHub | Email