PyQCD

English | 中文

PyQCD is a lattice-QCD toolkit in Python with optional native acceleration. It provides building blocks for propagator and operator contractions, gradient flow, renormalisation and gluon TMD-PDF extraction, plus a self-contained contraction engine with an opt_einsum-compatible API.

Install

pip install pyqcd

The native kernels are optional: without a C++ compiler (or without CUDA) the package installs anyway and the contraction engine runs its pure-Python planner — same results, only slower.

Set PYQCD_EINSUM_REQUIRE_NATIVE=1 to make a missing toolchain a hard error.

Check what you actually got (no imports needed):

python -m pyqcd            # version, native backend, optional deps, subpackages
python -m pyqcd --json     # machine-readable

Documentation

• English — https://py-qcd.readthedocs.io/en/latest/usage.html
• 中文 — https://py-qcd.readthedocs.io/zh-cn/latest/usage.html

Repositories

Quick start

import numpy as np
import pyqcd.einsum as pe

a = np.arange(6.0).reshape(2, 3)
b = np.arange(12.0).reshape(3, 4)

pe.contract("ij,jk->ik", a, b)                       # like a @ b
path, info = pe.contract_path("ij,jk->ik", a, b)     # (path, PathInfo)
int(info.opt_cost)

expr = pe.contract_expression("ij,jk->ik", (2, 3), (3, 4))
expr(a, b)                                           # reusable plan

The public surface mirrors opt_einsum: contract, contract_path, contract_expression, shared_intermediates, PathInfo, paths, backends, helpers, parser, blas.

What is inside

Contraction engine

• Pure-Python planning (greedy / optimal / branch / dp) is always available.
• Two optional native backends: KUPL (CPU) and CUDA (GPU). Both export the same C ABI and share one planning core.
• Fallback is a hard invariant: a missing library, a failed native call or a rejected shape silently degrades to the Python path — results never change.
• A small set of lattice-QCD equation shapes gets a specialised execution path; it is enabled by default and can be turned off per call (contract(..., specialize=False)) or per process (PYQCD_EINSUM_SPECIALIZE=0).

Introspection:

import pyqcd.einsum as pe

pe.build_info()        # backend, mpi, mpi_label, native_library
pe.mpi_enabled()       # whether the KUPL library has real MPI support
pe.specializations()   # registered equation shapes

MPI sharding exists only in the KUPL backend and only shards a single output axis; the CUDA backend has no MPI implementation.

Performance

Numbers you can reproduce on the published package, measured with pyqcd==0.1.12 (build_info(): backend=kupl, array_backend=numpy), opt_einsum==3.4.0 with the dp planner, complex128, median of 15 calls after 3 warm-up calls:

Those three run through the shipped kernels/plan cache with no extra setup. Dedicated kernels in this release cover the nucleon 2pt and 3pt shapes only; other lattice shapes (γ sandwich, perambulator × gauge, the sush variants) get the plan cache and the general execution layer.

import numpy as np
import pyqcd.einsum as pe

expr = "gahb,mcnd,oepf,ac,og,Mbdf,Mnph,em->M"   # nucleon 2pt
fast = pe.contract(expr, *operands)                    # specialised path (default)
base = pe.contract(expr, *operands, specialize=False)  # general path
assert np.allclose(fast, base)                         # same numbers, different speed

Per-call specialize=False and the process-wide PYQCD_EINSUM_SPECIALIZE=0 switch the optimisation off; both paths must agree numerically.

• Memory. The Python execution path peaks at roughly 1.4–2× the RSS of opt_einsum on the same shape (measured 18→41 MB at 1200², 73→104 MB at 2000²), because one extra output-sized intermediate stays alive per step. Size long jobs accordingly.

The gain comes from recognising a handful of lattice-QCD equation shapes and routing them to purpose-built kernels (fusion, buffer reuse) instead of a generic BLAS chain. It is not universal: on the GPU nucleon 2pt/3pt against opt_einsum both sides sit at the hardware limit (~1x), small tensors lose to per-call overhead, and shapes without a specialised kernel run the generic path. Full protocol and limitations: https://py-qcd.readthedocs.io/en/latest/performance.html

Tests

These suites live in pyqcd/testing/, which is not part of the PyPI distribution; run them from a source checkout:

python -m pyqcd.testing                       # integration suite
bash pyqcd/testing/simple/run-all.sh          # per-function API parity vs opt_einsum
bash pyqcd/testing/einsum/simple/run-all.sh   # einsum sub-system suite

License

MIT