Skip to content

Commit

Permalink
Fix a bug in GPT interface.
Browse files Browse the repository at this point in the history
Add a package called `pyquda_io`.
  • Loading branch information
SaltyChiang committed Dec 12, 2024
1 parent a8fccc3 commit 634a8af
Show file tree
Hide file tree
Showing 17 changed files with 135 additions and 74 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ pyquda_core/pyquda/enum_quda.py
pyquda_core/pyquda/src/*.c
pyquda_core/pyquda/src/*.cpp
pyquda_utils/_version.py
pyquda_io/_version.py

# symbolic link
/pyquda
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
requires = ["setuptools", "wheel"]

[tool.setuptools.packages.find]
include = ["pyquda_utils*"]
include = ["pyquda_utils*", "pyquda_io*"]

[tool.setuptools.dynamic]
version = { attr = "pyquda_utils._version.__version__" }
Expand Down
46 changes: 46 additions & 0 deletions pyquda_io/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
# flake8: noqa

from .io_general import (
IOGeneral,
read as readIOGeneral,
write as writeIOGeneral,
)
from .lime import (
Lime,
)
from .npy import (
readGauge as readNPYGauge,
writeGauge as writeNPYGauge,
readPropagator as readNPYPropagator,
writePropagator as writeNPYPropagator,
)
from .chroma import (
readQIOGauge as readChromaQIOGauge,
readQIOPropagator as readChromaQIOPropagator,
readILDGBinGauge,
)
from .milc import (
readGauge as readMILCGauge,
writeGauge as writeMILCGauge,
readQIOPropagator as readMILCQIOPropagator,
)
from .kyu import (
readGauge as readKYUGauge,
writeGauge as writeKYUGauge,
readPropagator as readKYUPropagator,
writePropagator as writeKYUPropagator,
)
from .xqcd import (
readPropagator as readXQCDPropagator,
writePropagator as writeXQCDPropagator,
readPropagatorFast as readXQCDPropagatorFast,
writePropagatorFast as writeXQCDPropagatorFast,
)
from .nersc import (
readGauge as readNERSCGauge,
writeGauge as writeNERSCGauge,
)
from .openqcd import (
readGauge as readOpenQCDGauge,
writeGauge as writeOpenQCDGauge,
)
68 changes: 43 additions & 25 deletions pyquda_utils/io/gauge_utils.py → pyquda_io/_field_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from mpi4py import MPI
import numpy

from .mpi_file import getSublatticeSize, getNeighbourRank
from ._mpi_file import getSublatticeSize, getNeighbourRank

Nd, Nc = 4, 3

Expand Down Expand Up @@ -72,30 +72,12 @@ def gaugeLexicoPlaquette(latt_size: List[int], grid_size: List[int], gauge: nump
extended[:, -1, :-1, :-1, :-1] = buf

plaq = numpy.zeros((6))
plaq[0] = numpy.vdot(
numpy.linalg.matmul(gauge[0], extended[1, :-1, :-1, :-1, 1:]),
numpy.linalg.matmul(gauge[1], extended[0, :-1, :-1, 1:, :-1]),
).real
plaq[1] = numpy.vdot(
numpy.linalg.matmul(gauge[0], extended[2, :-1, :-1, :-1, 1:]),
numpy.linalg.matmul(gauge[2], extended[0, :-1, 1:, :-1, :-1]),
).real
plaq[2] = numpy.vdot(
numpy.linalg.matmul(gauge[1], extended[2, :-1, :-1, 1:, :-1]),
numpy.linalg.matmul(gauge[2], extended[1, :-1, 1:, :-1, :-1]),
).real
plaq[3] = numpy.vdot(
numpy.linalg.matmul(gauge[0], extended[3, :-1, :-1, :-1, 1:]),
numpy.linalg.matmul(gauge[3], extended[0, 1:, :-1, :-1, :-1]),
).real
plaq[4] = numpy.vdot(
numpy.linalg.matmul(gauge[1], extended[3, :-1, :-1, 1:, :-1]),
numpy.linalg.matmul(gauge[3], extended[1, 1:, :-1, :-1, :-1]),
).real
plaq[5] = numpy.vdot(
numpy.linalg.matmul(gauge[2], extended[3, :-1, 1:, :-1, :-1]),
numpy.linalg.matmul(gauge[3], extended[2, 1:, :-1, :-1, :-1]),
).real
plaq[0] = numpy.vdot(gauge[0] @ extended[1, :-1, :-1, :-1, 1:], gauge[1] @ extended[0, :-1, :-1, 1:, :-1]).real
plaq[1] = numpy.vdot(gauge[0] @ extended[2, :-1, :-1, :-1, 1:], gauge[2] @ extended[0, :-1, 1:, :-1, :-1]).real
plaq[2] = numpy.vdot(gauge[1] @ extended[2, :-1, :-1, 1:, :-1], gauge[2] @ extended[1, :-1, 1:, :-1, :-1]).real
plaq[3] = numpy.vdot(gauge[0] @ extended[3, :-1, :-1, :-1, 1:], gauge[3] @ extended[0, 1:, :-1, :-1, :-1]).real
plaq[4] = numpy.vdot(gauge[1] @ extended[3, :-1, :-1, 1:, :-1], gauge[3] @ extended[1, 1:, :-1, :-1, :-1]).real
plaq[5] = numpy.vdot(gauge[2] @ extended[3, :-1, 1:, :-1, :-1], gauge[3] @ extended[2, 1:, :-1, :-1, :-1]).real

plaq /= int(numpy.prod(latt_size)) * Nc
plaq = MPI.COMM_WORLD.allreduce(plaq, MPI.SUM)
Expand Down Expand Up @@ -186,3 +168,39 @@ def gaugeEvenShiftBackward(latt_size: List[int], grid_size: List[int], gauge: nu
MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[3], source=neighbour_rank[7])
gauge_shift[3, 1, 0, :, :, :] = buf
return gauge_shift


# matrices to convert gamma basis bewteen DeGrand-Rossi and Dirac-Pauli
# DP for Dirac-Pauli, DR for DeGrand-Rossi
# \psi(DP) = _DR_TO_DP \psi(DR)
# \psi(DR) = _DP_TO_DR \psi(DP)
_DP_TO_DR = numpy.array(
[
[0, 1, 0, -1],
[-1, 0, 1, 0],
[0, 1, 0, 1],
[-1, 0, -1, 0],
]
)
_DR_TO_DP = numpy.array(
[
[0, -1, 0, -1],
[1, 0, 1, 0],
[0, 1, 0, -1],
[-1, 0, 1, 0],
]
)


def propagatorDeGrandRossiToDiracPauli(propagator: numpy.ndarray):
P = _DR_TO_DP
Pinv = _DP_TO_DR / 2

return numpy.ascontiguousarray(numpy.einsum("ij,tzyxjkab,kl->tzyxilab", P, propagator.data, Pinv, optimize=True))


def propagatorDiracPauliToDeGrandRossi(propagator: numpy.ndarray):
P = _DP_TO_DR
Pinv = _DR_TO_DP / 2

return numpy.ascontiguousarray(numpy.einsum("ij,tzyxjkab,kl->tzyxilab", P, propagator.data, Pinv, optimize=True))
File renamed without changes.
2 changes: 1 addition & 1 deletion pyquda_utils/io/chroma.py → pyquda_io/chroma.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from mpi4py import MPI

from .mpi_file import getSublatticeSize, getGridCoord, readMPIFile
from ._mpi_file import getSublatticeSize, getGridCoord, readMPIFile

Nd, Ns, Nc = 4, 4, 3
_precision_map = {"D": 8, "F": 4, "S": 4}
Expand Down
File renamed without changes.
5 changes: 4 additions & 1 deletion pyquda_utils/io/kyu.py → pyquda_io/kyu.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@

import numpy

from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from ._mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from ._field_utils import propagatorDiracPauliToDeGrandRossi, propagatorDeGrandRossiToDiracPauli

Nd, Ns, Nc = 4, 4, 3

Expand Down Expand Up @@ -48,6 +49,7 @@ def readPropagator(filename: str, latt_size: List[int], grid_size: List[int]):
.reshape(Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc * 2)
.view("<c16")
)
propagator = propagatorDiracPauliToDeGrandRossi(propagator)
return propagator


Expand All @@ -56,6 +58,7 @@ def writePropagator(filename: str, latt_size: List[int], grid_size: List[int], p
Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
dtype, offset = ">f8", 0

propagator = propagatorDeGrandRossiToDiracPauli(propagator)
propagator = (
propagator.view("<f8")
.reshape(Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc, 2)
Expand Down
File renamed without changes.
2 changes: 1 addition & 1 deletion pyquda_utils/io/milc.py → pyquda_io/milc.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import numpy
from mpi4py import MPI

from .mpi_file import getGridCoord, getSublatticeSize, readMPIFile, writeMPIFile
from ._mpi_file import getGridCoord, getSublatticeSize, readMPIFile, writeMPIFile

Nd, Ns, Nc = 4, 4, 3
_precision_map = {"D": 8, "F": 4, "S": 4}
Expand Down
4 changes: 2 additions & 2 deletions pyquda_utils/io/nersc.py → pyquda_io/nersc.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import numpy
from mpi4py import MPI

from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from .gauge_utils import gaugeLexicoPlaquette
from ._mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from ._field_utils import gaugeLexicoPlaquette

Nd, Ns, Nc = 4, 4, 3

Expand Down
2 changes: 1 addition & 1 deletion pyquda_utils/io/npy.py → pyquda_io/npy.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from numpy.lib.format import dtype_to_descr, read_magic, read_array_header_1_0, write_array_header_1_0
from mpi4py import MPI

from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from ._mpi_file import getSublatticeSize, readMPIFile, writeMPIFile

Nd, Ns, Nc = 4, 4, 3

Expand Down
4 changes: 2 additions & 2 deletions pyquda_utils/io/openqcd.py → pyquda_io/openqcd.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import numpy
from mpi4py import MPI

from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from .gauge_utils import gaugeEvenOdd, gaugeLexico, gaugeLexicoPlaquette, gaugeOddShiftForward, gaugeEvenShiftBackward
from ._mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from ._field_utils import gaugeEvenOdd, gaugeLexico, gaugeLexicoPlaquette, gaugeOddShiftForward, gaugeEvenShiftBackward

Nd, Ns, Nc = 4, 4, 3

Expand Down
5 changes: 4 additions & 1 deletion pyquda_utils/io/xqcd.py → pyquda_io/xqcd.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@

import numpy

from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from ._mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
from ._field_utils import propagatorDiracPauliToDeGrandRossi, propagatorDeGrandRossiToDiracPauli

Ns, Nc = 4, 3

Expand All @@ -16,6 +17,7 @@ def readPropagator(filename: str, latt_size: List[int], grid_size: List[int], st
if not staggered:
propagator = readMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2), grid_size)
propagator = propagator.transpose(2, 3, 4, 5, 6, 0, 7, 1).astype("<c16")
propagator = propagatorDiracPauliToDeGrandRossi(propagator)
else:
# QDP_ALIGN16 makes the last Nc to be aligned with 16 Bytes.
propagator_align16 = readMPIFile(filename, dtype, offset, (Nc, Lt, Lz, Ly, Lx, 4), (4, 3, 2, 1), grid_size)
Expand All @@ -32,6 +34,7 @@ def writePropagator(
dtype, offset = "<c8", 0

if not staggered:
propagator = propagatorDeGrandRossiToDiracPauli(propagator)
propagator = propagator.astype(dtype).transpose(5, 7, 0, 1, 2, 3, 4, 6).copy()
writeMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2), grid_size, propagator)
else:
Expand Down
2 changes: 1 addition & 1 deletion pyquda_utils/gpt.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def LatticeGaugeGPT(lattice: List[g.lattice], gen_simd_width: int, gauge: Lattic
else:
assert latt_info.size == gauge.latt_info.size
for index in range(latt_info.Nd):
gpt_shape = [i for sl in zip(gpt_simd, gpt_latt) for i in sl]
gpt_shape = [i for sl in zip(gpt_simd[::-1], gpt_latt[::-1]) for i in sl]
lattice[index].mview()[0][:] = (
gauge[index]
.lexico()
Expand Down
Loading

0 comments on commit 634a8af

Please sign in to comment.