Use QUDA for lattice QCD calculation with Python

Xiangyu Jiang,Chunjiang Shi,Ying Chen,Ming Gong,Yi-Bo Yang

2024-11-13

Abstract:We developed PyQUDA, a Python wrapper for QUDA written in Cython, designed to facilitate lattice QCD calculations using the Python programming language. PyQUDA leverages the optimized linear algebra capabilities of NumPy/CuPy/PyTorch, along with the highly optimized lattice QCD operations provided by QUDA to accelerate research. This integration simplifies the process of writing calculation codes, enabling researchers to build more complex Python packages like EasyDistillation for specific physics objectives. PyQUDA supports a range of lattice QCD operations, including hybrid Monte Carlo (HMC) with N-flavor clover/HISQ fermions and inversion for the Wilson/clover/HISQ fermion action with the multigrid solver. It also includes utility functions for reading lattice QCD data stored in Chroma, MILC, and $\chi$QCD formats. Type hints are supported by stub files and multi-GPU support is provided through mpi4py.

High Energy Physics - Lattice

What problem does this paper attempt to address?

The problem that this paper attempts to solve is to simplify and accelerate the calculation of Lattice Quantum Chromodynamics (Lattice QCD) in the Python environment. Specifically, the author has developed a Python package named PyQUDA. This package binds the QUDA library through Cython, aiming to provide a user - friendly Python interface, enabling researchers to directly call high - performance QUDA functions from Python for Lattice QCD calculations. PyQUDA takes advantage of the optimized linear algebra capabilities of NumPy/CuPy/PyTorch and combines the efficient Lattice QCD operations provided by QUDA, thereby accelerating the research process. In addition, PyQUDA supports multiple Lattice QCD operations, including the Hybrid Monte Carlo (HMC) algorithm, the reciprocal calculation of multi - flavor Clover/HISQ fermions, etc., and provides utility functions for reading Lattice QCD data stored in Chroma, MILC and χQCD formats.

### Main contributions of the paper
1. **Simplify the workflow**: PyQUDA enables researchers to write complex calculation codes in Python without frequently converting data formats, thus simplifying the workflow.
2. **High - performance computing**: Through the efficient combination with QUDA, PyQUDA can fully utilize the computing power of the GPU and improve the computing efficiency.
3. **User - friendly**: PyQUDA provides type hints and supports code auto - completion, making it more convenient to use.
4. **Multi - GPU support**: Through mpi4py, PyQUDA supports parallel computing in a multi - GPU environment.
5. **Compatibility**: PyQUDA supports multiple Lattice QCD operations and data formats and can be seamlessly integrated with other existing Python scientific computing libraries (such as NumPy, CuPy, PyTorch).

### Specific implementation of the paper
- **Cython binding**: Use Cython to statically bind the C API of QUDA to Python, ensuring performance while providing a user - friendly interface.
- **Data processing**: Data is stored in the form of NumPy or CuPy NDArray, facilitating data processing and analysis using general scientific computing functions.
- **Parallel computing**: Parallel computing in a multi - GPU environment is achieved through MPI for Python (mpi4py).
- **File format support**: Supports reading data files in Chroma, MILC and χQCD formats and provides support for writing in some formats.

### Example applications
The paper gives two specific example applications:
1. **Calculation of quark propagators and correlation functions**: It shows how to use PyQUDA to generate quark propagators and calculate two - point and three - point correlation functions through the einsum function.
2. **Hybrid Monte Carlo algorithm**: It shows how to use PyQUDA to implement the Hybrid Monte Carlo algorithm with Symanzik gauge action and Nf = 2 + 1 Clover fermion action.

### Conclusion
PyQUDA provides a powerful tool for the research of Lattice QCD, which not only simplifies the calculation process but also significantly improves the calculation efficiency. By combining the high - performance computing power of QUDA with the ease of use of Python, PyQUDA is expected to become an important tool for Lattice QCD research.