PyQUDA-Utils 0.9.9.dev8__tar.gz → 0.9.10.dev0__tar.gz

{pyquda_utils-0.9.9.dev8 → pyquda_utils-0.9.10.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PyQUDA-Utils
-Version: 0.9.9.dev8
+Version: 0.9.10.dev0
 Summary: Utility scripts based on PyQUDA
 Author-email: SaltyChiang <SaltyChiang@users.noreply.github.com>
 Maintainer-email: SaltyChiang <SaltyChiang@users.noreply.github.com>

{pyquda_utils-0.9.9.dev8 → pyquda_utils-0.9.10.dev0}/PyQUDA_Utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PyQUDA-Utils
-Version: 0.9.9.dev8
+Version: 0.9.10.dev0
 Summary: Utility scripts based on PyQUDA
 Author-email: SaltyChiang <SaltyChiang@users.noreply.github.com>
 Maintainer-email: SaltyChiang <SaltyChiang@users.noreply.github.com>

pyquda_utils-0.9.10.dev0/pyquda_utils/_version.py

@@ -0,0 +1 @@
+__version__ = "0.9.10.dev0"

{pyquda_utils-0.9.9.dev8 → pyquda_utils-0.9.10.dev0}/pyquda_utils/io/__init__.py

@@ -236,27 +236,19 @@ def writeNPYPropagator(filename: str, propagator: LatticePropagator):
 
 def readOpenQCDGauge(filename: str):
     from pyquda import getGridSize
-    from pyquda.field import X, Y, Z, T
     from .openqcd import readGauge as read
 
-    latt_size, plaquette, gauge_ = read(filename, getGridSize())
-    gauge_ = LatticeGauge(LatticeInfo(latt_size), gauge_)
-    gauge_.toDevice()
-    gauge = gauge_.shift([X, Y, Z, T])
-    gauge.data[:, 1] = gauge_.data[:, 0]
+    latt_size, plaquette, gauge = read(filename, getGridSize(), False)
+    gauge = LatticeGauge(LatticeInfo(latt_size), gauge)
     assert numpy.isclose(gauge.plaquette()[0], plaquette)
     return gauge
 
 
 def writeOpenQCDGauge(filename: str, gauge: LatticeGauge):
-    from pyquda.field import X, Y, Z, T
     from .openqcd import writeGauge as write
 
-    gauge.toDevice()
     plaquette = gauge.plaquette()[0]
-    gauge_ = gauge.shift([-X, -Y, -Z, -T])
-    gauge_.data[:, 0] = gauge.data[:, 1]
-    write(filename, gauge.latt_info.global_size, gauge.latt_info.grid_size, plaquette, gauge_.getHost())
+    write(filename, gauge.latt_info.global_size, gauge.latt_info.grid_size, plaquette, gauge.getHost(), False)
 
 
 def readNERSCGauge(filename: str, link_trace: bool = True, checksum: bool = True):

{pyquda_utils-0.9.9.dev8 → pyquda_utils-0.9.10.dev0}/pyquda_utils/io/milc.py

@@ -64,7 +64,7 @@ def readGauge(filename: str, grid_size: List[int], checksum: bool = True):
         else:
             raise ValueError(f"Broken magic {magic} in MILC gauge")
         latt_size = struct.unpack(f"{endian}iiii", f.read(16))
-        time_stamp = f.read(64).decode()  # noqa: F841
+        timestamp = f.read(64).decode()  # noqa: F841
         assert struct.unpack(f"{endian}i", f.read(4))[0] == 0  # order
         sum29, sum31 = struct.unpack(f"{endian}II", f.read(8))
         offset = f.tell()

{pyquda_utils-0.9.9.dev8 → pyquda_utils-0.9.10.dev0}/pyquda_utils/io/mpi_file.py

@@ -19,6 +19,21 @@ def getGridCoord(grid_size: Sequence[int]):
     return [gx, gy, gz, gt]
 
 
+def getNeighbourRank(grid_size: Sequence[int]):
+    Gx, Gy, Gz, Gt = grid_size
+    gx, gy, gz, gt = getGridCoord(grid_size)
+    return [
+        (((gx + 1) % Gx * Gy + gy) * Gz + gz) * Gt + gt,
+        ((gx * Gy + (gy + 1) % Gy) * Gz + gz) * Gt + gt,
+        ((gx * Gy + gy) * Gz + (gz + 1) % Gz) * Gt + gt,
+        ((gx * Gy + gy) * Gz + gz) * Gt + (gt + 1) % Gt,
+        (((gx - 1) % Gx * Gy + gy) * Gz + gz) * Gt + gt,
+        ((gx * Gy + (gy - 1) % Gy) * Gz + gz) * Gt + gt,
+        ((gx * Gy + gy) * Gz + (gz - 1) % Gz) * Gt + gt,
+        ((gx * Gy + gy) * Gz + gz) * Gt + (gt - 1) % Gt,
+    ]
+
+
 def getSubarray(shape: Sequence[int], axes: Sequence[int], grid: Sequence[int]):
     sizes = [d for d in shape]
     subsizes = [d for d in shape]

{pyquda_utils-0.9.9.dev8 → pyquda_utils-0.9.10.dev0}/pyquda_utils/io/nersc.py

@@ -77,13 +77,15 @@ def writeGauge(
     grid_size: List[int],
     plaquette: float,
     gauge: numpy.ndarray,
-    float_nbytes: int = 8,
+    use_fp32: bool = False,
 ):
     filename = path.expanduser(path.expandvars(filename))
+    float_nbytes = 4 if use_fp32 else 8
     dtype, offset = f"<c{2 * float_nbytes}", None
     gauge = numpy.ascontiguousarray(gauge.transpose(1, 2, 3, 4, 0, 5, 6).astype(dtype))
     link_trace = link_trace_nersc(gauge)
     checksum = checksum_nersc(gauge.reshape(-1))
+    timestamp = datetime.now().astimezone().strftime(R"%a %b %d %H:%M:%S %Y %Z")
     Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     header: Dict[str, str] = {
         "HDR_VERSION": "1.0",
@@ -107,8 +109,8 @@ def writeGauge(
         "SEQUENCE_NUMBER": "1",
         "CREATOR": "pyquda",
         "CREATOR_HARDWARE": f"{uname().nodename}-{uname().machine}-{uname().sysname}-{uname().release}",
-        "CREATION_DATE": datetime.now().strftime(R"%a %b %d %H:%M:%S %Y %z"),
-        "ARCHIVE_DATE": datetime.now().strftime(R"%a %b %d %H:%M:%S %Y %z"),
+        "CREATION_DATE": timestamp,
+        "ARCHIVE_DATE": timestamp,
         "FLOATING_POINT": f"IEEE{float_nbytes * 8}LITTLE",
     }
     if MPI.COMM_WORLD.Get_rank() == 0:

pyquda_utils-0.9.10.dev0/pyquda_utils/io/openqcd.py

@@ -0,0 +1,189 @@
+from os import path
+import struct
+from typing import List
+
+import numpy
+from mpi4py import MPI
+
+from .mpi_file import getSublatticeSize, getNeighbourRank, readMPIFile, writeMPIFile
+
+Nd, Ns, Nc = 4, 4, 3
+
+
+def cb2Gauge(latt_size: List[int], gauge: numpy.ndarray):
+    Lx, Ly, Lz, Lt = latt_size
+    gauge_cb2 = numpy.zeros_like(gauge).reshape(Nd, 2, Lt, Lz, Ly, Lx // 2, Nc, Nc)
+    for t in range(Lt):
+        for z in range(Lz):
+            for y in range(Ly):
+                eo = (t + z + y) % 2
+                if eo == 0:
+                    gauge_cb2[:, 0, t, z, y, :] = gauge[:, t, z, y, 0::2]
+                    gauge_cb2[:, 1, t, z, y, :] = gauge[:, t, z, y, 1::2]
+                else:
+                    gauge_cb2[:, 0, t, z, y, :] = gauge[:, t, z, y, 1::2]
+                    gauge_cb2[:, 1, t, z, y, :] = gauge[:, t, z, y, 0::2]
+    return gauge_cb2
+
+
+def lexicoGauge(latt_size: List[int], gauge: numpy.ndarray):
+    Lx, Ly, Lz, Lt = latt_size
+    gauge_lexico = numpy.empty_like(gauge).reshape(Nd, Lt, Lz, Ly, Lx, Nc, Nc)
+    for t in range(Lt):
+        for z in range(Lz):
+            for y in range(Ly):
+                eo = (t + z + y) % 2
+                if eo == 0:
+                    gauge_lexico[:, t, z, y, 0::2] = gauge[:, 0, t, z, y, :]
+                    gauge_lexico[:, t, z, y, 1::2] = gauge[:, 1, t, z, y, :]
+                else:
+                    gauge_lexico[:, t, z, y, 1::2] = gauge[:, 0, t, z, y, :]
+                    gauge_lexico[:, t, z, y, 0::2] = gauge[:, 1, t, z, y, :]
+    return gauge_lexico
+
+
+def shiftGaugeOddForward(latt_size: List[int], grid_size: List[int], gauge: numpy.ndarray):
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
+    gauge_shift = numpy.empty_like(gauge)
+    gauge_shift[:, 1] = gauge[:, 0]
+    rank = MPI.COMM_WORLD.Get_rank()
+    neighbour_rank = getNeighbourRank(grid_size)
+    for t in range(Lt):
+        for z in range(Lz):
+            for y in range(Ly):
+                if (t + z + y) % 2 == 0:
+                    gauge_shift[0, 0, t, z, y, :] = gauge[0, 1, t, z, y, :]
+                else:
+                    gauge_shift[0, 0, t, z, y, :-1] = gauge[0, 1, t, z, y, 1:]
+                    if rank == neighbour_rank[0] and rank == neighbour_rank[4]:
+                        gauge_shift[0, 0, t, z, y, -1] = gauge[0, 1, t, z, y, 0]
+                    else:
+                        buf = gauge[0, 1, t, z, y, 0].copy()
+                        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[4], source=neighbour_rank[0])
+                        gauge_shift[0, 0, t, z, y, -1] = buf
+    gauge_shift[1, 0, :, :, :-1, :] = gauge[1, 1, :, :, 1:, :]
+    if rank == neighbour_rank[1] and rank == neighbour_rank[5]:
+        gauge_shift[1, 0, :, :, -1, :] = gauge[1, 1, :, :, 0, :]
+    else:
+        buf = gauge[1, 1, :, :, 0, :].copy()
+        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[5], source=neighbour_rank[1])
+        gauge_shift[1, 0, :, :, -1, :] = buf
+    gauge_shift[2, 0, :, :-1, :, :] = gauge[2, 1, :, 1:, :, :]
+    if rank == neighbour_rank[2] and rank == neighbour_rank[6]:
+        gauge_shift[2, 0, :, -1, :, :] = gauge[2, 1, :, 0, :, :]
+    else:
+        buf = gauge[2, 1, :, 0, :, :].copy()
+        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[6], source=neighbour_rank[2])
+        gauge_shift[2, 0, :, -1, :, :] = buf
+    gauge_shift[3, 0, :-1, :, :, :] = gauge[3, 1, 1:, :, :, :]
+    if rank == neighbour_rank[3] and rank == neighbour_rank[7]:
+        gauge_shift[3, 0, -1, :, :, :] = gauge[3, 1, 0, :, :, :]
+    else:
+        buf = gauge[3, 1, 0, :, :, :].copy()
+        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[7], source=neighbour_rank[3])
+        gauge_shift[3, 0, -1, :, :, :] = buf
+    return gauge_shift
+
+
+def shiftGaugeEvenBackward(latt_size: List[int], grid_size: List[int], gauge: numpy.ndarray):
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
+    gauge_shift = numpy.empty_like(gauge)
+    gauge_shift[:, 0] = gauge[:, 1]
+    rank = MPI.COMM_WORLD.Get_rank()
+    neighbour_rank = getNeighbourRank(grid_size)
+    for t in range(Lt):
+        for z in range(Lz):
+            for y in range(Ly):
+                if (t + z + y) % 2 == 0:
+                    gauge_shift[0, 1, t, z, y, :] = gauge[0, 0, t, z, y, :]
+                else:
+                    gauge_shift[0, 1, t, z, y, 1:] = gauge[0, 0, t, z, y, :-1]
+                    if rank == neighbour_rank[0] and rank == neighbour_rank[4]:
+                        gauge_shift[0, 1, t, z, y, 0] = gauge[0, 0, t, z, y, -1]
+                    else:
+                        buf = gauge[0, 0, t, z, y, -1].copy()
+                        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[0], source=neighbour_rank[4])
+                        gauge_shift[0, 1, t, z, y, 0] = buf
+    gauge_shift[1, 1, :, :, 1:, :] = gauge[1, 0, :, :, :-1, :]
+    if rank == neighbour_rank[1] and rank == neighbour_rank[5]:
+        gauge_shift[1, 1, :, :, 0, :] = gauge[1, 0, :, :, -1, :]
+    else:
+        buf = gauge[1, 0, :, :, -1, :].copy()
+        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[1], source=neighbour_rank[5])
+        gauge_shift[1, 1, :, :, 0, :] = buf
+    gauge_shift[2, 1, :, 1:, :, :] = gauge[2, 0, :, :-1, :, :]
+    if rank == neighbour_rank[2] and rank == neighbour_rank[6]:
+        gauge_shift[2, 1, :, 0, :, :] = gauge[2, 0, :, -1, :, :]
+    else:
+        buf = gauge[2, 0, :, -1, :, :].copy()
+        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[2], source=neighbour_rank[6])
+        gauge_shift[2, 1, :, 0, :, :] = buf
+    gauge_shift[3, 1, 1:, :, :, :] = gauge[3, 0, :-1, :, :, :]
+    if rank == neighbour_rank[3] and rank == neighbour_rank[7]:
+        gauge_shift[3, 1, 0, :, :, :] = gauge[3, 0, -1, :, :, :]
+    else:
+        buf = gauge[3, 0, -1, :, :, :].copy()
+        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=neighbour_rank[3], source=neighbour_rank[7])
+        gauge_shift[3, 1, 0, :, :, :] = buf
+    return gauge_shift
+
+
+def readGauge(filename: str, grid_size: List[int], lexico: bool = True):
+    filename = path.expanduser(path.expandvars(filename))
+    with open(filename, "rb") as f:
+        latt_size = struct.unpack("<iiii", f.read(16))[::-1]
+        plaquette = struct.unpack("<d", f.read(8))[0] / Nc
+        offset = f.tell()
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
+    dtype = "<c16"
+
+    gauge_reorder = readMPIFile(filename, dtype, offset, (Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), (1, 2, 3, 0), grid_size)
+
+    gauge = numpy.zeros((Nd, 2, Lt, Lz, Ly, Lx // 2, Nc, Nc), dtype)
+    for t in range(Lt):
+        for y in range(Ly):
+            for z in range(Lz):
+                for x in range(Lx // 2):
+                    x_ = 2 * x + (1 - (t + z + y) % 2)
+                    z_ = z // 2
+                    gauge[[3, 0, 1, 2], :, t, z, y, x, :, :] = gauge_reorder[t, x_, y, z_]
+
+    gauge = shiftGaugeOddForward(latt_size, grid_size, gauge)
+    if lexico:
+        gauge = lexicoGauge([Lx, Ly, Lz, Lt], gauge).astype("<c16")
+    return latt_size, plaquette, gauge
+
+
+def writeGauge(
+    filename: str,
+    latt_size: List[int],
+    grid_size: List[int],
+    plaquette: float,
+    gauge: numpy.ndarray,
+    lexico: bool = True,
+):
+    filename = path.expanduser(path.expandvars(filename))
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
+    dtype, offset = "<c16", None
+
+    if lexico:
+        gauge = cb2Gauge([Lx, Ly, Lz, Lt], gauge.astype(dtype))
+    gauge = shiftGaugeEvenBackward(latt_size, grid_size, gauge)
+    gauge_reorder = numpy.zeros((Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), dtype)
+    for t in range(Lt):
+        for y in range(Ly):
+            for z in range(Lz):
+                for x in range(Lx // 2):
+                    x_ = 2 * x + (1 - (t + z + y) % 2)
+                    z_ = z // 2
+                    gauge_reorder[t, x_, y, z_] = gauge[[3, 0, 1, 2], :, t, z, y, x, :, :]
+
+    gauge = gauge_reorder.astype(dtype)
+    if MPI.COMM_WORLD.Get_rank() == 0:
+        with open(filename, "wb") as f:
+            f.write(struct.pack("<iiii", *latt_size[::-1]))
+            f.write(struct.pack("<d", plaquette * Nc))
+            offset = f.tell()
+    offset = MPI.COMM_WORLD.bcast(offset)
+
+    writeMPIFile(filename, dtype, offset, (Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), (1, 2, 3, 0), grid_size, gauge)

pyquda_utils-0.9.9.dev8/pyquda_utils/_version.py

@@ -1 +0,0 @@
-__version__ = "0.9.9.dev8"

pyquda_utils-0.9.9.dev8/pyquda_utils/io/openqcd.py

@@ -1,59 +0,0 @@
-from os import path
-import struct
-from typing import List
-
-import numpy
-from mpi4py import MPI
-
-from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
-
-Nd, Ns, Nc = 4, 4, 3
-
-
-def readGauge(filename: str, grid_size: List[int]):
-    filename = path.expanduser(path.expandvars(filename))
-    with open(filename, "rb") as f:
-        latt_size = struct.unpack("<iiii", f.read(16))[::-1]
-        plaquette = struct.unpack("<d", f.read(8))[0] / Nc
-        offset = f.tell()
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
-    dtype = "<c16"
-
-    gauge_reorder = readMPIFile(filename, dtype, offset, (Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), (1, 2, 3, 0), grid_size)
-
-    gauge = numpy.zeros((Nd, 2, Lt, Lz, Ly, Lx // 2, Nc, Nc), dtype="<c16")
-    for t in range(Lt):
-        for y in range(Ly):
-            for z in range(Lz):
-                for x in range(Lx // 2):
-                    x_ = 2 * x + (1 - (t + z + y) % 2)
-                    z_ = z // 2
-                    gauge[[3, 0, 1, 2], :, t, z, y, x, :, :] = gauge_reorder[t, x_, y, z_]
-
-    gauge = gauge.astype("<c16")
-    return latt_size, plaquette, gauge
-
-
-def writeGauge(filename: str, latt_size: List[int], grid_size: List[int], plaquette: float, gauge: numpy.ndarray):
-    filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
-    dtype, offset = "<c16", None
-
-    gauge_reorder = numpy.zeros((Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), dtype="<c16")
-    for t in range(Lt):
-        for y in range(Ly):
-            for z in range(Lz):
-                for x in range(Lx // 2):
-                    x_ = 2 * x + (1 - (t + z + y) % 2)
-                    z_ = z // 2
-                    gauge_reorder[t, x_, y, z_] = gauge[[3, 0, 1, 2], :, t, z, y, x, :, :]
-
-    gauge = gauge_reorder.astype(dtype)
-    if MPI.COMM_WORLD.Get_rank() == 0:
-        with open(filename, "wb") as f:
-            f.write(struct.pack("<iiii", *latt_size[::-1]))
-            f.write(struct.pack("<d", plaquette * Nc))
-            offset = f.tell()
-    offset = MPI.COMM_WORLD.bcast(offset)
-
-    writeMPIFile(filename, dtype, offset, (Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), (1, 2, 3, 0), grid_size, gauge)