Fix a bug in GPT interface.

Add a package called `pyquda_io`.

Author: SaltyChiang

.gitignore

@@ -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

pyproject.toml

@@ -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__" }

pyquda_io/__init__.py

@@ -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,
+)

pyquda_utils/io/gauge_utils.py renamed to pyquda_io/_field_utils.py

@@ -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
 
@@ -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)
@@ -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))

pyquda_utils/io/mpi_file.py renamed to pyquda_io/_mpi_file.py

@@ -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}

pyquda_utils/io/chroma.py renamed to pyquda_io/chroma.py

@@ -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
 
@@ -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
 
 
@@ -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)

pyquda_utils/io/io_general.py renamed to pyquda_io/io_general.py

@@ -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}

pyquda_utils/io/kyu.py renamed to pyquda_io/kyu.py

@@ -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
 

pyquda_utils/io/lime.py renamed to pyquda_io/lime.py

@@ -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
 

pyquda_utils/io/milc.py renamed to pyquda_io/milc.py

@@ -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
 

pyquda_utils/io/nersc.py renamed to pyquda_io/nersc.py

@@ -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
 
@@ -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)
@@ -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:

pyquda_utils/io/npy.py renamed to pyquda_io/npy.py

@@ -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()