Week 7: Comprehensive Benchmarking of the Native PETSc HEOM Solver

Building on the successful integration of PETSc into QuTiP from Week 5 and initial tests in Week 6, this week was dedicated to extensive benchmarking. We rigorously tested the newly implemented PETSc backend against QuTiP’s traditional SciPy csr solver to identify where PETSc provides the most value, discover the performance crossover points, and analyze MPI strong scaling.

We utilized a standard 1D Ising model spin chain coupled to Drude-Lorentz baths for our tests. Our benchmark scripts (benchmark_ising.py and benchmark_ising_petsc.py) analyzed performance across three main scaling dimensions:

1. System Size: Scaling the number of spins N.
2. Hierarchy Depth: Scaling the number of Auxiliary Density Operators (ADOs).
3. MPI Ranks: Testing the strong scaling efficiency.

All PETSc tests used the implicit bdf time stepper with adaptive step sizing, while the traditional solver used QuTiP’s default settings.

1. Scaling the System Size (N)

First, we maintained a fixed hierarchy depth (max_depth = 4) and a single bath on the first spin, scaling the number of spins from N=2 to N=10.

Spins (N)	ADOs	Traditional (s)	PETSc 4-Ranks (s)
2	126	0.43	1.18
4	126	7.41	43.14
6	126	177.31	1136.67
8	OOM	-	-

Analysis: For small ADO counts, the traditional solver is significantly faster. The PETSc implicit solvers have a large per-step initialization overhead that isn’t justified for small matrices. Furthermore, scaling the system size exponentially increases the Hilbert space (2^N). By N=8, the traditional solver runs out of memory (OOM).

2. Scaling the Hierarchy Depth (The Crossover Point)

Next, we fixed the system size to N=2 with 1 bath and scaled max_depth from 6 up to 22. This drastically increased the number of ADOs from 462 to 80,730, highlighting PETSc’s ability to handle massive hierarchies.

Depth	ADOs	Traditional (s)	PETSc 4-Ranks (s)	Speedup
6	462	1.55	3.39	0.45x
10	3003	29.71	19.65	1.51x
14	11628	174.89	70.46	2.48x
18	33649	644.50	181.89	3.54x
22	80730	1985.08	397.87	4.98x

Evolution Time vs. Number of ADOs (N=2, 1 Bath)

Analysis: Here we see the true power of PETSc’s implicit solver and distributed matrices. While the traditional solver is faster for small hierarchies, PETSc overtakes it at the crossover point of ~3,000 ADOs. As the hierarchy grows, PETSc maintains excellent efficiency, becoming almost 5x faster at 80k ADOs. This is largely due to the bdf solver taking much larger time steps (dt ~ 0.03) for the stiff HEOM ODEs compared to explicit methods.

3. Multiple Baths and Extreme Hierarchies

We also tested a scenario with 2 baths at a depth of 10, generating an enormous 184,756 ADOs. - Traditional Solver: 2820.9 s - PETSc (4 ranks): 1160.0 s (~2.4x faster)

When pushing to depth 14 with 2 baths, both solvers hit Out-Of-Memory limits on the local machine. This clearly defines the boundary where we must transition from local development to HPC clusters.

4. MPI Strong Scaling

To understand how effectively PETSc distributes the workload, we tested the 80,730 ADO system (Depth 22, 1 Bath) across 1, 2, 4, and 8 MPI ranks.

MPI Ranks	Evolution Time (s)
1 (Trad)	1985.08
1 (PETSc)	677.71
2 (PETSc)	424.15
4 (PETSc)	397.87
8 (PETSc)	486.92

MPI Strong Scaling (N=2, 1 Bath, Depth=22, 80730 ADOs)

Analysis: - 1 to 4 Ranks: We see a clear reduction in evolution time. Going from 1 rank to 2 provides a great speedup, though diminishing returns appear by rank 4. - 8 Ranks: Performance actually degrades (486s). This is a classic example of hitting the communication limit for a specific problem size. On a local machine, the time spent communicating matrix boundary data between 8 ranks outweighs the computation time saved.

Conclusion and Next Steps

The benchmarks are a resounding success, demonstrating that the PETSc backend is a powerful tool for solving deep hierarchies and large ADO spaces that bring the traditional solver to a crawl. The implicit bdf solver is particularly effective at handling the stiffness of the HEOM ODEs.

Next Steps for Week 8: - Cluster Benchmarking: Move these tests to an HPC cluster. Testing on a proper distributed environment will eliminate local resource contention and reveal the true scaling limit beyond 4 ranks.