Skip to content

Documentation

Jump to bottom
SaltyChiang edited this page Aug 14, 2024 · 5 revisions

Module pyquda

This module handles the MPI configuration to split the lattice, the backend, and the device information.

Initialize MPI environment and QUDA

The init() function initializes the QUDA library to perform lattice QCD calculations on GPUs. You should set the grid to let QUDA know how to split the lattice. The following code defines a grid with size Gx, Gy, Gz, Gt = 1, 1, 1, 2, which means we use 2 GPUs to split the lattice into 2 parts in the t direction.

from pyquda import init

init([1, 1, 1, 2])

mpiexec -n 2 is needed to run the code like this. The number of MPI processes should be equal to the grid volume. The default grid is [1, 1, 1, 1], and you can suppress it when you are using only 1 GPU to perform the calculation.

CAUTION: Initialization should be performed before any operation in PyQUDA.

Get the MPI configuration

There are some functions to get the MPI configuration.

  • getMPIComm()
  • getMPISize()
  • getMPIRank()
  • getGridSize()
  • getGridCoord()
  • getCUDABackend()
  • isHIP()
  • getGPUID()
  • getCUDAComputeCapability()

Print information, warnings, and errors with only one process by logger

pyquda provides a logger to make sure output or error will only be printed once. The logging level is the same as in logging package.

from pyquda import getLogger

logger = getLogger()
logger.debug("This is DEBUG")
logger.info("This is INFO")
logger.warning("This is WARNING", RuntimeWarning)
logger.error("This is ERROR", RuntimeError)
logger.critical("This is CRITICAL", RuntimeError)

Use QUDA's functions directly or get QUDA's parameter struct

QUDA's functions and parameter struct are stored in submodule pyquda.pyquda, which is the main file to wrap up the QUDA functions. You should not call functions from this module directly unless you know what you are doing. QUDA parameter structs are also defined here, and it's not recommended to make such a struct by yourself.

from pyquda import QudaInvertParam

inv_param = QudaInvertParam()

Use QUDA's enum

QUDA's enum are stored in submodule pyquda.enum_quda, which is basically a translation of QUDA's enum_quda.h.

from pyquda.enum_quda import QudaInverterType

inv_param.inv_type = QudaInverterType.QUDA_CG_INVERTER

Manually set the correspondence between GPUID and MPI rank

pyquda.setGPUID(gpuid: int) is the function to manually bind GPUs with MPI ranks. If you need such a thing you should call this before init().

A default lattice can be initialized within init(). See Class LatticeInfo section for more details.

Module pyquda.field

This module defines the classes to store the process-specific data on GPUs. Many pure gauge functions are bound to the LatticeGauge class for convenience.

Defining a lattice with class LatticeInfo

This class handles the information to construct a lattice with multiple processes. It defines the size of the lattice and the grid to split the lattice. It also handles the extra information to define the lattice (the t-boundary and the anisotropy).

For example, the code below defines a lattice with the global size Lx, Ly, Lz, Lt = 4, 4, 4, 8. The lattice is anti-periodic on the t-boundary and isotropic.

from pyquda.field import LatticeInfo

latt_info = LatticeInfo([4, 4, 4, 8], t_boundary=-1, anisotropy=1.0)

A default LatticeInfo can be initialized in init(). And then we can get this LatticeInfo by pyquda.getDefaultLattice().

from pyquda import getDefaultLattice

init([1, 1, 1, 2], latt_size=[4, 4, 4, 8], t_boundary=-1, anisotropy=1.0)
latt_info = getDefaultLattice()

Also, we can set this default lattice by pyquda.setDefaultLattice().

from pyquda import setDefaultLattice, getDefaultLattice

setDefaultLattice([4, 4, 4, 8], t_boundary=-1, anisotropy=1.0)
latt_info = getDefaultLattice()

Once we get a LatticeInfo instance, we can get a lattice field class by instantiating the corresponding lattice field class. The LatticeInfo instance is saved in latt_info class member.

from pyquda.field import LatticeGauge, LatticePropagator

gauge = LatticeGauge(latt_info)
propagator = LatticePropagator(latt_info)
gauge_latt_info = gauge.latt_info
print(gauge_latt_info == latt_info)  # True

Data stored in these fields can be accessed by data class member. If you want to get a copy of the data, use backup() class method. copy() class method will return a deep copy of the field but with the same latt_info.

gauge_data = gauge.data
gauge_data_backup = gauge.backup()
gauge_copy = gauge.copy()

backend is the package to store the data on the device memory. location can be used to identify where the data is stored now. toDevice and toHost class methods will transfer the data between the host memory and the device memory.

print(gauge.backend)  # cupy
print(gauge.location)  # cupy
gauge.toHost()
print(gauge.location)  # numpy
gauge.toDevice()
print(gauge.location)  # cupy

All the data are stored in even-odd preconditioned format. lexico() is the method to get a numpy.ndarray copy of data but without even-odd preconditioning.

print(gauge.data)  # (Nd, 2, Lt, Lz, Ly, Lx // 2, Nc, Nc)
print(gauge.lexico().shape)  # (Nd, Lt, Lz, Ly, Lx, Nc, Nc)
print(propagator.data)  # (2, Lt, Lz, Ly, Lx // 2, Ns, Ns Nc, Nc)
print(propagator.lexico().shape)  # (Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc)

Data format in the lattice field

Data in data class member could be a numpy.ndarray, cupy.ndarray, or torch.Tensor depending on the backend you defined in the init() progress. The dtype should be complex128 in most field types and float64 in momentum and clover fields. All Dirac indices are calculated under the DeGrand-Rossi basis of Dirac matrices.

The shapes of data in different fields are listed below.

  • LatticeGauge.data
    • (Nd, 2, Lt, Lz, Ly, Lx // 2, Nc, Nc), row-column order, complex.
  • LatticeMom.data
    • (2, Lt, Lz, Ly, Lx // 2, 10), lower triangle (6) + diagnol (3) + preserved (1), real.
  • LatticeClover.data
    • (2, Lt, Lz, Ly, Lx // 2, 2, ((Ns // 2) * Nc) ** 2), upper-left + lower-right, diagnol (6) + lower triangle (30), real.
  • LatticeFermion.data
    • (2, Lt, Lz, Ly, Lx // 2, Ns, Nc), complex.
  • LatticePropagator.data
    • (2, Lt, Lz, Ly, Lx // 2, Ns, Ns, Nc, Nc), sink-source order, complex.
  • LatticeStaggeredFermion.data
    • (2, Lt, Lz, Ly, Lx // 2, Nc), complex.
  • LatticeStaggeredPropagator.data
    • (2, Lt, Lz, Ly, Lx // 2, Nc, Nc), sink-source order, complex.

Read/write filed with field I/O functions

The I/O functions to read/write fields from/to the hard drive are defined in the pyquda.utils.io module.

For example, we can read a gauge configuration generated by Chroma in .lime format, and convert it to LatticeGauge.

from pyquda.utils.io import readChromaQIOGauge

gauge = readChromaQIOGauge("weak_field.lime")

But we cannot write a LatticeGauge to the .lime file because we don't have the QIO support.

The supported operations are listed below.

  • readChromaQIOGauge()
  • readChromaQIOPropagator()
  • readMILCGauge()
  • readKYUGauge()
  • writeKYUGauge()
  • readKYUPropagator()
  • writeKYUPropagator()
  • readXQCDPropagator()
  • writeXQCDPropagator()
  • readNPYGauge()
  • writeNPYGauge()
  • readNPYPropagator()
  • writeNPYPropagator()

Here we defined a new format "NPY". This is the default file format of numpy.ndarray, and can be loaded and saved easily by numpy.load() and numpy.save(). The data format is inherited from the lexico() result. For example

import numpy
from pyquda.utils.io import writeNPYGauge

writeNPYGauge("weak_field.npy", gauge)
print(numpy.load("weak_field.npy").shape)  # [Nd, Lt, Lz, Ly, Lx, Nc, Nc]

Pure gauge functions bound to LatticeGauge

  • covDev(x, covdev_mu)

    • Applies the covariant derivative on x in direction covdev_mu. x should be LatticeFermion. 0/1/2/3 represent +x/+y/+z/+t and 4/5/6/7 represent -x/-y/-z/-t. The covariant derivative is defined as $\psi'(x)=U_\mu(x)\psi(x+\hat{\mu})$.

  • LatticeGauge.laplace(x, laplace3D)

    • Applies the Laplacian operator on x, and laplace3D takes 3 or 4 to apply Laplacian on spacial or all directions. x should be LatticeStaggeredFermion. The Laplacian operator is defined as $\psi'(x)=\frac{1}{N_\mathrm{Lap}}\sum_\mu\psi(x)-\dfrac{1}{2}\left[U_\mu(x)\psi(x+\hat{\mu})+U_\mu^\dagger(x-\hat{\mu})\psi(x-\hat{\mu})\right]$

  • staggeredPhase()

    • Applies the staggered phase to the gauge field. The convention is controlled by LatticeGauge.pure_gauge.gauge_param.staggered_phase_type, which is QudaStaggeredPhase.QUDA_STAGGERED_PHASE_MILC by default.

  • projectSU3(tol)

    • Projects the gauge field onto SU(3) matrix. tol is the tolerance of how the matrix deviates from SU(3). 2e-15 (which is 10x the epsilon of fp64) should be a good choice.

  • path(paths, coeff)

    • paths is a list of length 4, which is [paths_x, paths_y, paths_z, paths_t]. All paths_* should have the same shape. paths_x is a list of any length.

  • loopTrace(paths)

    • paths is similar to paths_x in LatticeGauge.path, but the function returns the traces of all loops. The traces is a numpy.ndarray of dtype complex128, and every element is defined as $\sum_x\mathrm{Tr}W(x)$.

  • apeSmear(n_steps, alpha, dir_ignore)

    • Applies the APE smearing to the gauge field. alpha is the smearing strength.

  • smearAPE(n_steps, factor, dir_ignore)

    • Similar to LatticeGauge.apeSmear() but factor matches with Chroma.

  • stoutSmear(n_steps, rho, dir_ignore

    • Applies the stout smearing to the gauge field. rho is the smearing strength.

  • smearSTOUT(n_steps, rho, dir_ignore)

    • Similar to LatticeGauge.stoutSmear().

  • hypSmear(n_steps, alpha1, alpha2, alpha3, dir_ignore)

    • Applies the stout smearing to the gauge field. alpha1/alpha2/alpha3 is the smearing strength on level 3/2/1.

  • smearHYP(n_steps, alpha1, alpha2, alpha3, dir_ignore) Similar to LatticeGauge.hypSmear().

  • wilsonFlow(n_steps, epsilon)

    • Applies the Wilson flow to the gauge field. Returns the energy (all, spatial, temporal) for every step.

  • wilsonFlowScale(max_steps, epsilon)

    • Returns $t_0$ and $w_0$ with up to max_steps Wilson flow step.

  • symanzikFlow(n_steps, epsilon)

    • Applies the Symanzik flow to the gauge field. Returns the energy (all, spatial, temporal) for every step.

  • symanzikFlowScale(max_steps, epsilon)

    • Returns $t_0$ and $w_0$ with up to max_steps Symanzik flow step.

  • plaquette()

    • Returns the plaquette (all, spatial, temporal) of the gauge field.

  • polyakovLoop()

    • Returns the Polyakov loop (real, image) of the gauge field.

  • energy()

    • Returns the energy (all, spatial, temporal) of the gauge field.

  • qcharge()

    • Returns the topological charge of the gauge field.

  • qchargeDensity()

    • Returns the topological charge density (with shape (Lt, Lz, Ly, Lx))of the gauge field.

  • gauss(seed, sigma)

    • Fills the gauge field with random SU(3) matrices. sigma=1 corresponds to the standard normal distribution.

  • fixingOVR(gauge_dir, Nsteps, verbose_interval, relax_boost, tolerance, reunit_interval, stopWtheta)

    • Applies the gauge fixing to the gauge field and over-relaxation is used to speed up the operation.

      gauge_dir: {3, 4}
    3 for Coulomb gauge fixing, 4 for Landau gauge fixing
Nsteps: int
    maximum number of steps to perform gauge fixing
verbose_interval: int
    print gauge fixing info when iteration count is a multiple of this
relax_boost: float
    gauge fixing parameter of the overrelaxation method, most common value is 1.5 or 1.7.
tolerance: float
    torelance value to stop the method, if this value is zero then the method stops when
    iteration reachs the maximum number of steps defined by Nsteps
reunit_interval: int
    reunitarize gauge field when iteration count is a multiple of this
stopWtheta: int
    0 for MILC criterion and 1 to use the theta value

  • fixingFFT(gauge_dir, Nsteps, verbose_interval, alpha, autotune, tolerance, stopWtheta)

    • Applies the gauge fixing to the gauge field and a fast Fourier transform is used to speed up the operation.

      gauge_dir: {3, 4}
    3 for Coulomb gauge fixing, 4 for Landau gauge fixing
Nsteps: int
    maximum number of steps to perform gauge fixing
verbose_interval: int
    print gauge fixing info when iteration count is a multiple of this
alpha: float
    gauge fixing parameter of the method, most common value is 0.08
autotune: int
    1 to autotune the method, i.e., if the Fg inverts its tendency we decrease the alpha value
tolerance: float
    torelance value to stop the method, if this value is zero then the method stops when
    iteration reachs the maximum number of steps defined by Nsteps
stopWtheta: int
    0 for MILC criterion and 1 to use the theta value

Module pyquda.core

Most useful functions in lattice QCD calculation are defined here.

Get a Wilson fermion action

You can generate a clover fermion action by calling getWilson().

from pyquda.core import getWilson

xi_0, nu = 2.464, 0.95
kappa = 0.125
mass = 1 / (2 * kappa) - 4

latt_info = LatticeInfo([4, 4, 4, 8], t_boundary=-1, anisotropy=xi_0 / nu)
wilson_dirac = getWilson(latt_info, mass=mass, tol=1e-12, maxiter=1000)

The code above defines a Wilson fermion action with a specific mass on an anisotropic lattice. tol and maxiter are used to control the iterative solver. The default solver for Wilson action is BiCGStab.

Multigrid will be enabled if the multigrid parameter is set to the blocking dimensions for multigrid algorithm.

wilson_dirac = getWilson(latt_info, mass=mass, tol=1e-12, maxiter=1000, multigrid=[[4, 4, 4, 4], [4, 4, 4, 4]])

The code above creates a 2-level multigrid preconditioner for the Wilson action.

Get a clover fermion action

You can generate a clover fermion action by calling getClover().

from pyquda.core import getClover

xi_0, nu = 2.464, 0.95
kappa = 0.115
mass = 1 / (2 * kappa) - 4
coeff = 1.17
coeff_r, coeff_t = 0.91, 1.07

latt_info = LatticeInfo([4, 4, 4, 8], t_boundary=-1, anisotropy=xi_0 / nu)
clover_dirac = getClover(latt_info, mass=mass, tol=1e-12, maxiter=1000, xi_0=xi_0, clover_csw_t=coeff_t, clover_csw_r=coeff_r)

The code above defines a clover fermion action with specific mass and clover coefficient on an anisotropic lattice. The definitions of clover_csw_t and clover_csw_r are the same as clovCoeffT and clovCoeffR in Chroma. tol and maxiter are used to control the iterative solver. The default solver for clover action is BiCGStab.

Multigrid will be enabled if the multigrid parameter is set to the blocking dimensions for multigrid algorithm.

clover_dirac = getClover(latt_info, mass=mass, tol=1e-12, maxiter=1000, xi_0=xi_0, clover_csw_t=coeff_t, clover_csw_r=coeff_r, multigrid=[[4, 4, 4, 4], [4, 4, 4, 4]])

The code above creates a 2-level multigrid preconditioner for the Clover action.

Get a HISQ fermion action

You can generate an HISQ fermion action by calling getHISQ().

from pyquda.core import getHISQ

mass = 0.0102

hisq_dirac = getHISQ(latt_info, mass=mass, tol=1e-12, maxiter=1000, tadpole_coeff=1.0, naik_epsilon=0.0)

The code above defines a HISQ fermion action with specific mass and tadpole/naik improvement on an isotropic lattice. tol and maxiter are used to control the iterative solver. The default solver for HISQ action is CG. The staggered phase is defined as the MILC convention.

Multigrid for HISQ action is unavailable for now. We are testing the parameters to get a good preconditioner.

Get a source

Utilities to generate specific sources are defined in the submodule pyquda.utils.source.

from pyquda.utils.source import source

point_source = source12(latt_info, source_type="point", t_srce=[0, 0, 0, 0])

We generated a point source at x, y, z, t = 0, 0, 0, 0 through the source12() function (12 means $N_s\times Nc=12$). The source_type could be "point", "wall" or "volume". The returned value is a LatticePropagator. If you want to get a source for staggered fermion action, you should call source3() instead.

The form required by the t_srce parameter varies depending on the source_type:

  • "point": A list of int [x, y, z, t]
  • "wall": A int t
  • "volume": Will not use it

Add a phase to the source

Phases are often used in contraction and source. pyquda.utils.phase defines the momentum phase and the grid phase.

from pyquda.utils.phase import MomentumPhase, GridPhase

momentum_phase = MomentumPhase(latt_info).getPhase([1, 2, 3])
grid_phase = GridPhase(latt_info, stride=[2, 2, 2, 2]).getPhase([1, 1, 1, 1])

Here we get a momentum phase with three-momentum nx, ny, nz = 1, 2, 3. The number is the momentum mode defined on the lattice: $p_i=\frac{2\pi n_i}{L_i}$. The MomentumPhase accepts the four-momentum. We also get a grid phase which fills the lattice every 2 lattice sites starting from coord 1 in all four directions.

If we want to get a momentum source of momentum mode [1, 2, 3] at time 4, we can use "wall" as the source type and then pass the momentum phase to the source_phase parameter.

momentum_source = source(latt_info, source_type="wall", t_srce=4, source_phase=momentum_phase)

If we want to get a grid source with the momentum phase at every all-odd coordinates, we can use "volume" as the source type and then pass the product of the momentum phase and the grid phase to the source_phase parameter.

momentum_source = source(latt_info, source_type="volume", t_srce=None, source_phase=momentum_phase * grid_phase)

Get the propagator by inversion

invert() and invertStaggered() are used to perform the inversion and return the propagator.

wilson_dirac.loadGauge(gauge)
wilson_propag = core.invert(wilson_dirac, source_type="point", [0, 0, 0, 0])

hisq_dirac.loadGauge(gauge)
hisq_propag = core.invertStaggered(hisq_dirac, source_type="point", t_srce=[0, 0, 0, 0])

Here we perform the inversion on a point source with predefined fermion actions. invert() and invertStaggered() accept extra arguments to generate a source, see Get a source for details.

Also, you can get the propagator based on an existing propagator, which could be the source generated before or a sequential source.

wilson_seq_propag = core.invertPropagator(wilson_dirac, point_source)

hisq_seq_propag = core.invertStaggeredPropagator(hisq_dirac, hisq_propag.timeslice(4))

We perform the inversion on the point source with the Wilson action and on the previous propagators at time 4 with the HISQ action.

Getting started

Clone this wiki locally