Add readOpenQCDGauge() and writeOpenQCDGauge().

pyquda_utils/io/__init__.py

@@ -1,4 +1,4 @@
-from typing import List, Literal, Union
+from typing import List, Union
 
 import numpy
 
@@ -218,20 +218,28 @@ def writeNPYPropagator(filename: str, propagator: LatticePropagator):
     write(filename, propagator.latt_info.global_size, propagator.lexico())
 
 
-def readOpenQCDGauge(filename: str, return_plaquette: bool = False):
+def readOpenQCDGauge(filename: str):
+    from pyquda.field import X, Y, Z, T
     from .openqcd import readGauge as read
 
-    latt_size, plaquette, gauge_raw = read(filename)
-    if not return_plaquette:
-        return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4]))
-    else:
-        return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4])), plaquette
+    latt_size, plaquette, gauge_ = read(filename)
+    gauge_ = LatticeGauge(LatticeInfo(latt_size), gauge_)
+    gauge_.toDevice()
+    gauge = gauge_.shift([X, Y, Z, T])
+    gauge.data[:, 1] = gauge_.data[:, 0]
+    assert numpy.isclose(gauge.plaquette()[0], plaquette)
+    return gauge
 
 
-def writeOpenQCDGauge(filename: str, gauge: LatticeGauge, plaquette: float):
+def writeOpenQCDGauge(filename: str, gauge: LatticeGauge):
+    from pyquda.field import X, Y, Z, T
     from .openqcd import writeGauge as write
 
-    write(filename, gauge.latt_info.global_size, plaquette, gauge.lexico())
+    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, plaquette, gauge_.getHost())
 
 
 def readNERSCGauge(filename: str, return_plaquette: bool = False, link_trace: bool = True, checksum: bool = True):

pyquda_utils/io/openqcd.py

@@ -18,22 +18,41 @@ def readGauge(filename: str):
     Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
     dtype = "<c16"
 
-    gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0))
-    gauge = gauge.transpose(4, 0, 1, 2, 3, 5, 6).astype("<c16")
+    gauge_reorder = readMPIFile(filename, dtype, offset, (Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), (1, 2, 3, 0))
+
+    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], plaquette: float, gauge: numpy.ndarray):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = latt_size
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
     dtype, offset = "<c16", None
 
-    gauge = numpy.ascontiguousarray(gauge.transpose(1, 2, 3, 4, 0, 5, 6).astype(dtype))
+    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 getMPIRank() == 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 = getMPIComm().bcast(offset)
 
-    writeMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0), gauge)
+    writeMPIFile(filename, dtype, offset, (Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), (1, 2, 3, 0), gauge)