PyQUDA-Utils 0.9.9.dev5__tar.gz → 0.9.9.dev6__tar.gz

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PyQUDA-Utils
-Version: 0.9.9.dev5
+Version: 0.9.9.dev6
 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.dev5 → pyquda_utils-0.9.9.dev6}/PyQUDA_Utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PyQUDA-Utils
-Version: 0.9.9.dev5
+Version: 0.9.9.dev6
 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.dev5 → pyquda_utils-0.9.9.dev6}/PyQUDA_Utils.egg-info/SOURCES.txt

@@ -26,6 +26,7 @@ pyquda_utils/io/io_general.py
 pyquda_utils/io/kyu.py
 pyquda_utils/io/lime.py
 pyquda_utils/io/milc.py
+pyquda_utils/io/mpi_file.py
 pyquda_utils/io/nersc.py
 pyquda_utils/io/npy.py
 pyquda_utils/io/openqcd.py

pyquda_utils-0.9.9.dev6/pyquda_utils/_version.py

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

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/__init__.py

@@ -70,23 +70,26 @@ def rotateToDeGrandRossi(propagator: LatticePropagator):
 
 
 def readChromaQIOGauge(filename: str, checksum: bool = True):
+    from pyquda import getGridSize
     from .chroma import readQIOGauge as read
 
-    latt_size, gauge_raw = read(filename, checksum)
+    latt_size, gauge_raw = read(filename, getGridSize(), checksum)
     return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4]))
 
 
 def readILDGBinGauge(filename: str, dtype: str, latt_size: List[int]):
+    from pyquda import getGridSize
     from .chroma import readILDGBinGauge as read
 
-    gauge_raw = read(filename, dtype, latt_size)
+    gauge_raw = read(filename, dtype, latt_size, getGridSize())
     return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4]))
 
 
 def readChromaQIOPropagator(filename: str, checksum: bool = True):
+    from pyquda import getGridSize
     from .chroma import readQIOPropagator as read
 
-    latt_size, staggered, propagator_raw = read(filename, checksum)
+    latt_size, staggered, propagator_raw = read(filename, getGridSize(), checksum)
     if not staggered:
         return LatticePropagator(LatticeInfo(latt_size), cb2(propagator_raw, [0, 1, 2, 3]))
     else:
@@ -94,22 +97,24 @@ def readChromaQIOPropagator(filename: str, checksum: bool = True):
 
 
 def readMILCGauge(filename: str, checksum: bool = True):
+    from pyquda import getGridSize
     from .milc import readGauge as read
 
-    latt_size, gauge_raw = read(filename, checksum)
+    latt_size, gauge_raw = read(filename, getGridSize(), checksum)
     return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4]))
 
 
 def writeMILCGauge(filename: str, gauge: LatticeGauge):
     from .milc import writeGauge as write
 
-    write(filename, gauge.latt_info.global_size, gauge.lexico())
+    write(filename, gauge.latt_info.global_size, gauge.latt_info.grid_size, gauge.lexico())
 
 
 def readMILCQIOPropagator(filename: str):
+    from pyquda import getGridSize
     from .milc import readQIOPropagator as read
 
-    latt_size, staggered, propagator_raw = read(filename)
+    latt_size, staggered, propagator_raw = read(filename, getGridSize())
     if not staggered:
         return LatticePropagator(LatticeInfo(latt_size), cb2(propagator_raw, [0, 1, 2, 3]))
     else:
@@ -117,35 +122,42 @@ def readMILCQIOPropagator(filename: str):
 
 
 def readKYUGauge(filename: str, latt_size: List[int]):
+    from pyquda import getGridSize
     from .kyu import readGauge as read
 
-    gauge_raw = read(filename, latt_size)
+    gauge_raw = read(filename, latt_size, getGridSize())
     return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4]))
 
 
 def writeKYUGauge(filename: str, gauge: LatticeGauge):
     from .kyu import writeGauge as write
 
-    write(filename, gauge.latt_info.global_size, gauge.lexico())
+    write(filename, gauge.latt_info.global_size, gauge.latt_info.grid_size, gauge.lexico())
 
 
 def readKYUPropagator(filename: str, latt_size: List[int]):
+    from pyquda import getGridSize
     from .kyu import readPropagator as read
 
-    propagator_raw = read(filename, latt_size)
-    return rotateToDeGrandRossi(LatticePropagator(LatticeInfo(latt_size), cb2(propagator_raw, [0, 1, 2, 3])))
+    propagator_raw = read(filename, latt_size, getGridSize())
+    propagator = LatticePropagator(LatticeInfo(latt_size), cb2(propagator_raw, [0, 1, 2, 3]))
+    propagator = rotateToDeGrandRossi(propagator)
+    return propagator
 
 
 def writeKYUPropagator(filename: str, propagator: LatticePropagator):
     from .kyu import writePropagator as write
 
-    write(filename, propagator.latt_info.global_size, rotateToDiracPauli(propagator).lexico())
+    propagator = rotateToDiracPauli(propagator)
+    propagator_raw = propagator.lexico()
+    write(filename, propagator.latt_info.global_size, propagator.latt_info.grid_size, propagator_raw)
 
 
 def readXQCDPropagator(filename: str, latt_size: List[int], staggered: bool):
+    from pyquda import getGridSize
     from .xqcd import readPropagator as read
 
-    propagator_raw = read(filename, latt_size, staggered)
+    propagator_raw = read(filename, latt_size, getGridSize(), staggered)
     if not staggered:
         return rotateToDeGrandRossi(LatticePropagator(LatticeInfo(latt_size), cb2(propagator_raw, [0, 1, 2, 3])))
     else:
@@ -156,19 +168,21 @@ def writeXQCDPropagator(filename: str, propagator: Union[LatticePropagator, Latt
     from .xqcd import writePropagator as write
 
     latt_size = propagator.latt_info.global_size
+    grid_size = propagator.latt_info.grid_size
     staggered = isinstance(propagator, LatticeStaggeredPropagator)
     if not staggered:
-        write(filename, latt_size, rotateToDiracPauli(propagator).lexico(), staggered)
+        write(filename, latt_size, grid_size, rotateToDiracPauli(propagator).lexico(), staggered)
     else:
-        write(filename, latt_size, propagator.lexico(), staggered)
+        write(filename, latt_size, grid_size, propagator.lexico(), staggered)
 
 
 def readXQCDPropagatorFast(filename: str, latt_size: List[int]):
+    from pyquda import getGridSize
     from .xqcd import readPropagatorFast as read
 
     latt_info = LatticeInfo(latt_size)
     Lx, Ly, Lz, Lt = latt_info.size
-    propagator_raw = read(filename, latt_size)
+    propagator_raw = read(filename, getGridSize(), latt_size)
     propagator = LatticePropagator(latt_info, cb2(propagator_raw, [2, 3, 4, 5]))
     propagator.data = propagator.data.reshape(Ns, Nc, 2, Lt, Lz, Ly, Lx // 2, Ns, Nc)
     propagator.toDevice()
@@ -187,14 +201,15 @@ def writeXQCDPropagatorFast(filename: str, propagator: LatticePropagator):
     propagator.toHost()
     propagator.data = propagator.data.reshape(Ns, Nc, 2, Lt, Lz, Ly, Lx // 2, Ns, Nc)
     propagator_raw = lexico(propagator.data, [2, 3, 4, 5, 6])
-    write(filename, latt_info.global_size, propagator_raw)
+    write(filename, latt_info.global_size, latt_info.grid_size, propagator_raw)
 
 
 def readNPYGauge(filename: str):
+    from pyquda import getGridSize
     from .npy import readGauge as read
 
     filename = filename if filename.endswith(".npy") else filename + ".npy"
-    latt_size, gauge_raw = read(filename)
+    latt_size, gauge_raw = read(filename, getGridSize())
     return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4]))
 
 
@@ -202,27 +217,29 @@ def writeNPYGauge(filename: str, gauge: LatticeGauge):
     from .npy import writeGauge as write
 
     filename = filename if filename.endswith(".npy") else filename + ".npy"
-    write(filename, gauge.latt_info.global_size, gauge.lexico())
+    write(filename, gauge.latt_info.global_size, gauge.latt_info.grid_size, gauge.lexico())
 
 
 def readNPYPropagator(filename: str):
+    from pyquda import getGridSize
     from .npy import readPropagator as read
 
-    latt_size, propagator_raw = read(filename)
+    latt_size, propagator_raw = read(filename, getGridSize())
     return LatticePropagator(LatticeInfo(latt_size), cb2(propagator_raw, [0, 1, 2, 3]))
 
 
 def writeNPYPropagator(filename: str, propagator: LatticePropagator):
     from .npy import writePropagator as write
 
-    write(filename, propagator.latt_info.global_size, propagator.lexico())
+    write(filename, propagator.latt_info.global_size, propagator.latt_info.grid_size, propagator.lexico())
 
 
 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)
+    latt_size, plaquette, gauge_ = read(filename, getGridSize())
     gauge_ = LatticeGauge(LatticeInfo(latt_size), gauge_)
     gauge_.toDevice()
     gauge = gauge_.shift([X, Y, Z, T])
@@ -239,13 +256,14 @@ def writeOpenQCDGauge(filename: str, gauge: LatticeGauge):
     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())
+    write(filename, gauge.latt_info.global_size, gauge.latt_info.grid_size, plaquette, gauge_.getHost())
 
 
 def readNERSCGauge(filename: str, return_plaquette: bool = False, link_trace: bool = True, checksum: bool = True):
+    from pyquda import getGridSize
     from .nersc import readGauge as read
 
-    latt_size, plaquette, gauge_raw = read(filename, link_trace, checksum)
+    latt_size, plaquette, gauge_raw = read(filename, getGridSize(), link_trace, checksum)
     if not return_plaquette:
         return LatticeGauge(LatticeInfo(latt_size), cb2(gauge_raw, [1, 2, 3, 4]))
     else:

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/chroma.py

@@ -2,19 +2,20 @@ from os import path
 from typing import List
 from xml.etree import ElementTree as ET
 
-from pyquda import getSublatticeSize, getMPIComm, getGridCoord, readMPIFile
+from mpi4py import MPI
+
+from .mpi_file import getSublatticeSize, getGridCoord, readMPIFile
 
 Nd, Ns, Nc = 4, 4, 3
 _precision_map = {"D": 8, "F": 4, "S": 4}
 
 
-def checksum_qio(latt_size: List[int], data):
+def checksum_qio(latt_size: List[int], grid_size: List[int], data):
     import zlib
     import numpy
-    from mpi4py import MPI
 
-    gx, gy, gz, gt = getGridCoord()
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    gx, gy, gz, gt = getGridCoord(grid_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     gLx, gLy, gLz, gLt = gx * Lx, gy * Ly, gz * Lz, gt * Lt
     GLx, GLy, GLz, GLt = latt_size
     work = numpy.empty((Lt * Lz * Ly * Lx), "<u4")
@@ -27,12 +28,12 @@ def checksum_qio(latt_size: List[int], data):
     )
     rank29 = (rank % 29).astype("<u4")
     rank31 = (rank % 31).astype("<u4")
-    sum29 = getMPIComm().allreduce(numpy.bitwise_xor.reduce(work << rank29 | work >> (32 - rank29)), MPI.BXOR)
-    sum31 = getMPIComm().allreduce(numpy.bitwise_xor.reduce(work << rank31 | work >> (32 - rank31)), MPI.BXOR)
+    sum29 = MPI.COMM_WORLD.allreduce(numpy.bitwise_xor.reduce(work << rank29 | work >> (32 - rank29)), MPI.BXOR)
+    sum31 = MPI.COMM_WORLD.allreduce(numpy.bitwise_xor.reduce(work << rank31 | work >> (32 - rank31)), MPI.BXOR)
     return sum29, sum31
 
 
-def readQIOGauge(filename: str, checksum: bool = True):
+def readQIOGauge(filename: str, grid_size: List[int], checksum: bool = True):
     from .lime import Lime
 
     lime = Lime(filename)
@@ -55,12 +56,12 @@ def readQIOGauge(filename: str, checksum: bool = True):
     assert int(scidac_private_record_xml.find("datacount").text) == Nd
     assert int(scidac_private_file_xml.find("spacetime").text) == Nd
     latt_size = [int(L) for L in scidac_private_file_xml.find("dims").text.split()]
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype = f">c{2 * precision}"
 
-    gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0))
+    gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0), grid_size)
     if checksum:
-        assert checksum_qio(latt_size, gauge.reshape(Lt * Lz * Ly * Lx, Nd * Nc * Nc)) == (
+        assert checksum_qio(latt_size, grid_size, gauge.reshape(Lt * Lz * Ly * Lx, Nd * Nc * Nc)) == (
             int(scidac_checksum_xml.find("suma").text, 16),
             int(scidac_checksum_xml.find("sumb").text, 16),
         ), f"Bad checksum for {filename}"
@@ -68,17 +69,17 @@ def readQIOGauge(filename: str, checksum: bool = True):
     return latt_size, gauge
 
 
-def readILDGBinGauge(filename: str, dtype: str, latt_size: List[int]):
+def readILDGBinGauge(filename: str, dtype: str, latt_size: List[int], grid_size: List[int]):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     offset = 0
 
-    gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0))
+    gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0), grid_size)
     gauge = gauge.transpose(4, 0, 1, 2, 3, 5, 6).astype("<c16")
     return gauge
 
 
-def readQIOPropagator(filename: str, checksum: bool = True):
+def readQIOPropagator(filename: str, grid_size: List[int], checksum: bool = True):
     from .lime import Lime
 
     lime = Lime(filename)
@@ -107,21 +108,21 @@ def readQIOPropagator(filename: str, checksum: bool = True):
     assert int(scidac_private_record_xml.find("datacount").text) == 1
     assert int(scidac_private_file_xml.find("spacetime").text) == Nd
     latt_size = [int(L) for L in scidac_private_file_xml.find("dims").text.split()]
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype = f">c{2 * precision}"
 
     if not staggered:
-        propagator = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc), (3, 2, 1, 0))
+        propagator = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc), (3, 2, 1, 0), grid_size)
         if checksum:
-            assert checksum_qio(latt_size, propagator.reshape(Lt * Lz * Ly * Lx, Ns * Ns * Nc * Nc)) == (
+            assert checksum_qio(latt_size, grid_size, propagator.reshape(Lt * Lz * Ly * Lx, Ns * Ns * Nc * Nc)) == (
                 int(scidac_checksum_xml.find("suma").text, 16),
                 int(scidac_checksum_xml.find("sumb").text, 16),
             ), f"Bad checksum for {filename}"
         propagator = propagator.astype("<c16")
     else:
-        propagator = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nc, Nc), (3, 2, 1, 0))
+        propagator = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nc, Nc), (3, 2, 1, 0), grid_size)
         if checksum:
-            assert checksum_qio(latt_size, propagator.reshape(Lt * Lz * Ly * Lx, Nc * Nc)) == (
+            assert checksum_qio(latt_size, grid_size, propagator.reshape(Lt * Lz * Ly * Lx, Nc * Nc)) == (
                 int(scidac_checksum_xml.find("suma").text, 16),
                 int(scidac_checksum_xml.find("sumb").text, 16),
             ), f"Bad checksum for {filename}"

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/kyu.py

@@ -3,17 +3,17 @@ from typing import List
 
 import numpy
 
-from pyquda import getSublatticeSize, readMPIFile, writeMPIFile
+from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
 
 Nd, Ns, Nc = 4, 4, 3
 
 
-def readGauge(filename: str, latt_size: List[int]):
+def readGauge(filename: str, latt_size: List[int], grid_size: List[int]):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = ">f8", 0
 
-    gauge = readMPIFile(filename, dtype, offset, (Nd, Nc, Nc, 2, Lt, Lz, Ly, Lx), (7, 6, 5, 4))
+    gauge = readMPIFile(filename, dtype, offset, (Nd, Nc, Nc, 2, Lt, Lz, Ly, Lx), (7, 6, 5, 4), grid_size)
     gauge = (
         gauge.transpose(0, 4, 5, 6, 7, 2, 1, 3)
         .astype("<f8")
@@ -24,23 +24,23 @@ def readGauge(filename: str, latt_size: List[int]):
     return gauge
 
 
-def writeGauge(filename: str, latt_size: List[int], gauge: numpy.ndarray):
+def writeGauge(filename: str, latt_size: List[int], grid_size: List[int], gauge: numpy.ndarray):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = ">f8", 0
 
     gauge = (
         gauge.view("<f8").reshape(Nd, Lt, Lz, Ly, Lx, Nc, Nc, 2).astype(dtype).transpose(0, 6, 5, 7, 1, 2, 3, 4).copy()
     )
-    writeMPIFile(filename, dtype, offset, (Nd, Nc, Nc, 2, Lt, Lz, Ly, Lx), (7, 6, 5, 4), gauge)
+    writeMPIFile(filename, dtype, offset, (Nd, Nc, Nc, 2, Lt, Lz, Ly, Lx), (7, 6, 5, 4), grid_size, gauge)
 
 
-def readPropagator(filename: str, latt_size: List[int]):
+def readPropagator(filename: str, latt_size: List[int], grid_size: List[int]):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = ">f8", 0
 
-    propagator = readMPIFile(filename, dtype, offset, (Ns, Nc, 2, Ns, Nc, Lt, Lz, Ly, Lx), (8, 7, 6, 5))
+    propagator = readMPIFile(filename, dtype, offset, (Ns, Nc, 2, Ns, Nc, Lt, Lz, Ly, Lx), (8, 7, 6, 5), grid_size)
     propagator = (
         propagator.transpose(5, 6, 7, 8, 3, 0, 4, 1, 2)
         .astype("<f8")
@@ -51,9 +51,9 @@ def readPropagator(filename: str, latt_size: List[int]):
     return propagator
 
 
-def writePropagator(filename: str, latt_size: List[int], propagator: numpy.ndarray):
+def writePropagator(filename: str, latt_size: List[int], grid_size: List[int], propagator: numpy.ndarray):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = ">f8", 0
 
     propagator = (
@@ -63,4 +63,4 @@ def writePropagator(filename: str, latt_size: List[int], propagator: numpy.ndarr
         .transpose(5, 7, 8, 4, 6, 0, 1, 2, 3)
         .copy()
     )
-    writeMPIFile(filename, dtype, offset, (Ns, Nc, 2, Ns, Nc, Lt, Lz, Ly, Lx), (8, 7, 6, 5), propagator)
+    writeMPIFile(filename, dtype, offset, (Ns, Nc, 2, Ns, Nc, Lt, Lz, Ly, Lx), (8, 7, 6, 5), grid_size, propagator)

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/lime.py

@@ -24,6 +24,9 @@ class Lime:
                 f.seek((length + 7) // 8 * 8, io.SEEK_CUR)
                 buffer = f.read(8)
 
+    def keys(self):
+        return [record.name for record in self._records]
+
     def records(self, key: str):
         return [record for record in self._records if record.name == key]
 

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/milc.py

@@ -5,40 +5,38 @@ from typing import List
 from xml.etree import ElementTree as ET
 
 import numpy
+from mpi4py import MPI
 
-from pyquda import getSublatticeSize, getMPIComm, getMPISize, getMPIRank, getGridCoord, readMPIFile, writeMPIFile
+from .mpi_file import getGridCoord, getSublatticeSize, readMPIFile, writeMPIFile
 
 Nd, Ns, Nc = 4, 4, 3
 _precision_map = {"D": 8, "F": 4, "S": 4}
 
 
-def checksum_milc(latt_size: List[int], data):
-    from mpi4py import MPI
-
-    gx, gy, gz, gt = getGridCoord()
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+def checksum_milc(latt_size: List[int], grid_size: List[int], data):
+    gx, gy, gz, gt = getGridCoord(grid_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     gLx, gLy, gLz, gLt = gx * Lx, gy * Ly, gz * Lz, gt * Lt
     GLx, GLy, GLz, GLt = latt_size
 
     work = data.view("<u4")
     rank = (
-        numpy.arange(getMPISize() * work.size, dtype="<u8")
+        numpy.arange(MPI.COMM_WORLD.Get_size() * work.size, dtype="<u8")
         .reshape(GLt, GLz, GLy, GLx, -1)[gLt : gLt + Lt, gLz : gLz + Lz, gLy : gLy + Ly, gLx : gLx + Lx]
         .reshape(-1)
     )
     rank29 = (rank % 29).astype("<u4")
     rank31 = (rank % 31).astype("<u4")
-    sum29 = getMPIComm().allreduce(numpy.bitwise_xor.reduce(work << rank29 | work >> (32 - rank29)), MPI.BXOR)
-    sum31 = getMPIComm().allreduce(numpy.bitwise_xor.reduce(work << rank31 | work >> (32 - rank31)), MPI.BXOR)
+    sum29 = MPI.COMM_WORLD.allreduce(numpy.bitwise_xor.reduce(work << rank29 | work >> (32 - rank29)), MPI.BXOR)
+    sum31 = MPI.COMM_WORLD.allreduce(numpy.bitwise_xor.reduce(work << rank31 | work >> (32 - rank31)), MPI.BXOR)
     return sum29, sum31
 
 
-def checksum_qio(latt_size: List[int], data):
+def checksum_qio(latt_size: List[int], grid_size: List[int], data):
     import zlib
-    from mpi4py import MPI
 
-    gx, gy, gz, gt = getGridCoord()
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    gx, gy, gz, gt = getGridCoord(grid_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     gLx, gLy, gLz, gLt = gx * Lx, gy * Ly, gz * Lz, gt * Lt
     GLx, GLy, GLz, GLt = latt_size
     work = numpy.empty((Lt * Lz * Ly * Lx), "<u4")
@@ -51,12 +49,12 @@ def checksum_qio(latt_size: List[int], data):
     )
     rank29 = rank % 29
     rank31 = rank % 31
-    sum29 = getMPIComm().allreduce(numpy.bitwise_xor.reduce(work << rank29 | work >> (32 - rank29)), MPI.BXOR)
-    sum31 = getMPIComm().allreduce(numpy.bitwise_xor.reduce(work << rank31 | work >> (32 - rank31)), MPI.BXOR)
+    sum29 = MPI.COMM_WORLD.allreduce(numpy.bitwise_xor.reduce(work << rank29 | work >> (32 - rank29)), MPI.BXOR)
+    sum31 = MPI.COMM_WORLD.allreduce(numpy.bitwise_xor.reduce(work << rank31 | work >> (32 - rank31)), MPI.BXOR)
     return sum29, sum31
 
 
-def readGauge(filename: str, checksum: bool = True):
+def readGauge(filename: str, grid_size: List[int], checksum: bool = True):
     filename = path.expanduser(path.expandvars(filename))
     with open(filename, "rb") as f:
         magic = f.read(4)
@@ -70,12 +68,12 @@ def readGauge(filename: str, checksum: bool = True):
         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()
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype = f"{endian}c8"
 
-    gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0))
+    gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0), grid_size)
     if checksum:
-        assert checksum_milc(latt_size, gauge.astype("<c8").reshape(-1)) == (
+        assert checksum_milc(latt_size, grid_size, gauge.astype("<c8").reshape(-1)) == (
             sum29,
             sum31,
         ), f"Bad checksum for {filename}"
@@ -83,14 +81,14 @@ def readGauge(filename: str, checksum: bool = True):
     return latt_size, gauge
 
 
-def writeGauge(filename: str, latt_size: List[int], gauge: numpy.ndarray):
+def writeGauge(filename: str, latt_size: List[int], grid_size: List[int], gauge: numpy.ndarray):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = "<c8", None
 
     gauge = numpy.ascontiguousarray(gauge.transpose(1, 2, 3, 4, 0, 5, 6).astype(dtype))
-    sum29, sum31 = checksum_milc(latt_size, gauge.reshape(-1))
-    if getMPIRank() == 0:
+    sum29, sum31 = checksum_milc(latt_size, grid_size, gauge.reshape(-1))
+    if MPI.COMM_WORLD.Get_rank() == 0:
         with open(filename, "wb") as f:
             f.write(struct.pack("<i", 20103))
             f.write(struct.pack("<iiii", *latt_size))
@@ -98,12 +96,12 @@ def writeGauge(filename: str, latt_size: List[int], gauge: numpy.ndarray):
             f.write(struct.pack("<i", 0))  # order
             f.write(struct.pack("<II", sum29, sum31))
             offset = f.tell()
-    offset = getMPIComm().bcast(offset)
+    offset = MPI.COMM_WORLD.bcast(offset)
 
-    writeMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0), gauge)
+    writeMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0), grid_size, gauge)
 
 
-def readQIOPropagator(filename: str):
+def readQIOPropagator(filename: str, grid_size: List[int]):
     from .lime import Lime
 
     lime = Lime(filename)
@@ -129,7 +127,7 @@ def readQIOPropagator(filename: str):
     assert int(scidac_private_record_xml.find("datacount").text) == 1
     assert int(scidac_private_file_xml.find("spacetime").text) == Nd
     latt_size = [int(L) for L in scidac_private_file_xml.find("dims").text.split()]
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype = f">c{2 * precision}"
 
     if not staggered:
@@ -137,12 +135,14 @@ def readQIOPropagator(filename: str):
         for spin in range(Ns):
             for color in range(Nc):
                 propagator[spin, color] = readMPIFile(
-                    filename, dtype, offset[spin * Nc + color], (Lt, Lz, Ly, Lx, Ns, Nc), (3, 2, 1, 0)
+                    filename, dtype, offset[spin * Nc + color], (Lt, Lz, Ly, Lx, Ns, Nc), (3, 2, 1, 0), grid_size
                 )
         propagator = propagator.transpose(2, 3, 4, 5, 6, 0, 7, 1).astype("<c16")
     else:
         propagator = numpy.empty((Nc, Lt, Lz, Ly, Lx, Nc), dtype)
         for color in range(Nc):
-            propagator[color] = readMPIFile(filename, dtype, offset[color], (Lt, Lz, Ly, Lx, Nc), (3, 2, 1, 0))
+            propagator[color] = readMPIFile(
+                filename, dtype, offset[color], (Lt, Lz, Ly, Lx, Nc), (3, 2, 1, 0), grid_size
+            )
         propagator = propagator.transpose(1, 2, 3, 4, 5, 0).astype("<c16")
     return latt_size, staggered, propagator

pyquda_utils-0.9.9.dev6/pyquda_utils/io/mpi_file.py

@@ -0,0 +1,77 @@
+from typing import Sequence
+
+import numpy
+from mpi4py import MPI
+from mpi4py.util import dtlib
+
+
+def getSublatticeSize(latt_size: Sequence[int], grid_size: Sequence[int]):
+    GLx, GLy, GLz, GLt = latt_size
+    Gx, Gy, Gz, Gt = grid_size
+    assert GLx % Gx == 0 and GLy % Gy == 0 and GLz % Gz == 0 and GLt % Gt == 0
+    return [GLx // Gx, GLy // Gy, GLz // Gz, GLt // Gt]
+
+
+def getGridCoord(grid_size: Sequence[int]):
+    rank = MPI.COMM_WORLD.Get_rank()
+    Gx, Gy, Gz, Gt = grid_size
+    gx, gy, gz, gt = rank // Gt // Gz // Gy % Gx, rank // Gt // Gz % Gy, rank // Gt % Gz, rank % Gt
+    return [gx, gy, gz, gt]
+
+
+def getSubarray(shape: Sequence[int], axes: Sequence[int], grid: Sequence[int]):
+    sizes = [d for d in shape]
+    subsizes = [d for d in shape]
+    starts = [d if i in axes else 0 for i, d in enumerate(shape)]
+    coord = getGridCoord(grid)
+    for j, i in enumerate(axes):
+        sizes[i] *= grid[j]
+        starts[i] *= coord[j]
+    return sizes, subsizes, starts
+
+
+def readMPIFile(
+    filename: str,
+    dtype: str,
+    offset: int,
+    shape: Sequence[int],
+    axes: Sequence[int],
+    grid: Sequence[int],
+):
+    assert MPI.COMM_WORLD.Get_size() == int(numpy.prod(grid))
+    sizes, subsizes, starts = getSubarray(shape, axes, grid)
+    native_dtype = dtype if not dtype.startswith(">") else dtype.replace(">", "<")
+    buf = numpy.empty(subsizes, native_dtype)
+
+    fh = MPI.File.Open(MPI.COMM_WORLD, filename, MPI.MODE_RDONLY)
+    filetype = dtlib.from_numpy_dtype(native_dtype).Create_subarray(sizes, subsizes, starts)
+    filetype.Commit()
+    fh.Set_view(disp=offset, filetype=filetype)
+    fh.Read_all(buf)
+    filetype.Free()
+    fh.Close()
+
+    return buf.view(dtype)
+
+
+def writeMPIFile(
+    filename: str,
+    dtype: str,
+    offset: int,
+    shape: Sequence[int],
+    axes: Sequence[int],
+    grid: Sequence[int],
+    buf: numpy.ndarray,
+):
+    assert MPI.COMM_WORLD.Get_size() == int(numpy.prod(grid))
+    sizes, subsizes, starts = getSubarray(shape, axes, grid)
+    native_dtype = dtype if not dtype.startswith(">") else dtype.replace(">", "<")
+    buf = buf.view(native_dtype)
+
+    fh = MPI.File.Open(MPI.COMM_WORLD, filename, MPI.MODE_WRONLY | MPI.MODE_CREATE)
+    filetype = dtlib.from_numpy_dtype(native_dtype).Create_subarray(sizes, subsizes, starts)
+    filetype.Commit()
+    fh.Set_view(disp=offset, filetype=filetype)
+    fh.Write_all(buf)
+    filetype.Free()
+    fh.Close()

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/nersc.py

@@ -1,20 +1,19 @@
 from os import path
-from typing import Dict
+from typing import Dict, List
 
 import numpy
+from mpi4py import MPI
 
-from pyquda import getSublatticeSize, readMPIFile, writeMPIFile, getMPIComm, getMPIRank
+from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
 
 Nd, Ns, Nc = 4, 4, 3
 
 
 def checksum_nersc(data: numpy.ndarray) -> int:
-    from mpi4py import MPI
+    return MPI.COMM_WORLD.allreduce(numpy.sum(data.view("<u4"), dtype="<u4"), MPI.SUM)
 
-    return getMPIComm().allreduce(numpy.sum(data.view("<u4"), dtype="<u4"), MPI.SUM)
 
-
-def readGauge(filename: str, link_trace: bool = True, checksum: bool = True):
+def readGauge(filename: str, grid_size: List[int], link_trace: bool = True, checksum: bool = True):
     filename = path.expanduser(path.expandvars(filename))
     header: Dict[str, str] = {}
     with open(filename, "rb") as f:
@@ -31,7 +30,7 @@ def readGauge(filename: str, link_trace: bool = True, checksum: bool = True):
         int(header["DIMENSION_3"]),
         int(header["DIMENSION_4"]),
     ]
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     assert header["FLOATING_POINT"].startswith("IEEE")
     if header["FLOATING_POINT"][6:] == "BIG":
         endian = ">"
@@ -44,7 +43,7 @@ def readGauge(filename: str, link_trace: bool = True, checksum: bool = True):
     plaquette = float(header["PLAQUETTE"])
 
     if header["DATATYPE"] == "4D_SU3_GAUGE_3x3":
-        gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0))
+        gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc, Nc), (3, 2, 1, 0), grid_size)
         if link_trace:
             assert numpy.isclose(
                 numpy.einsum("tzyxdaa->", gauge.real) / (gauge.size // Nc), float(header["LINK_TRACE"])

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/npy.py

@@ -3,8 +3,9 @@ from typing import List
 
 import numpy
 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 pyquda import getSublatticeSize, getMPIComm, getMPIRank, readMPIFile, writeMPIFile
+from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
 
 Nd, Ns, Nc = 4, 4, 3
 
@@ -19,49 +20,49 @@ def _readHeader(filename: str):
     return shape, dtype, offset
 
 
-def readGauge(filename: str):
+def readGauge(filename: str, grid_size: List[int]):
     filename = path.expanduser(path.expandvars(filename))
     shape, dtype, offset = _readHeader(filename)
     latt_size = [shape[i] for i in [4, 3, 2, 1]]
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
 
-    gauge = readMPIFile(filename, dtype, offset, (Nd, Lt, Lz, Ly, Lx, Nc, Nc), (4, 3, 2, 1))
+    gauge = readMPIFile(filename, dtype, offset, (Nd, Lt, Lz, Ly, Lx, Nc, Nc), (4, 3, 2, 1), grid_size)
     return latt_size, gauge
 
 
-def writeGauge(filename: str, latt_size: List[int], gauge: numpy.ndarray):
+def writeGauge(filename: str, latt_size: List[int], grid_size: List[int], gauge: numpy.ndarray):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
-    if getMPIRank() == 0:
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
+    if MPI.COMM_WORLD.Get_rank() == 0:
         with open(filename, "wb") as f:
             write_array_header_1_0(
                 f, {"shape": (Nd, *latt_size[::-1], Nc, Nc), "fortran_order": False, "descr": "<c16"}
             )
-    getMPIComm().Barrier()
+    MPI.COMM_WORLD.Barrier()
     shape, dtype, offset = _readHeader(filename)
 
-    writeMPIFile(filename, dtype, offset, (Nd, Lt, Lz, Ly, Lx, Nc, Nc), (4, 3, 2, 1), gauge)
+    writeMPIFile(filename, dtype, offset, (Nd, Lt, Lz, Ly, Lx, Nc, Nc), (4, 3, 2, 1), grid_size, gauge)
 
 
-def readPropagator(filename: str):
+def readPropagator(filename: str, grid_size: List[int]):
     filename = path.expanduser(path.expandvars(filename))
     shape, dtype, offset = _readHeader(filename)
     latt_size = [shape[i] for i in [3, 2, 1, 0]]
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
 
-    propagator = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc), (3, 2, 1, 0))
+    propagator = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc), (3, 2, 1, 0), grid_size)
     return latt_size, propagator
 
 
-def writePropagator(filename: str, latt_size: List[int], propagator: numpy.ndarray):
+def writePropagator(filename: str, latt_size: List[int], grid_size: List[int], propagator: numpy.ndarray):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
-    if getMPIRank() == 0:
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
+    if MPI.COMM_WORLD.Get_rank() == 0:
         with open(filename, "wb") as f:
             write_array_header_1_0(
                 f, {"shape": (*latt_size[::-1], Ns, Ns, Nc, Nc), "fortran_order": False, "descr": "<c16"}
             )
-    getMPIComm().Barrier()
+    MPI.COMM_WORLD.Barrier()
     shape, dtype, offset = _readHeader(filename)
 
-    writeMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc), (3, 2, 1, 0), propagator)
+    writeMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Ns, Ns, Nc, Nc), (3, 2, 1, 0), grid_size, propagator)

{pyquda_utils-0.9.9.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/openqcd.py

@@ -3,22 +3,23 @@ import struct
 from typing import List
 
 import numpy
+from mpi4py import MPI
 
-from pyquda import getSublatticeSize, getMPIRank, getMPIComm, readMPIFile, writeMPIFile
+from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
 
 Nd, Ns, Nc = 4, 4, 3
 
 
-def readGauge(filename: str):
+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)
+    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))
+    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):
@@ -33,9 +34,9 @@ def readGauge(filename: str):
     return latt_size, plaquette, gauge
 
 
-def writeGauge(filename: str, latt_size: List[int], plaquette: float, gauge: numpy.ndarray):
+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)
+    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")
@@ -48,11 +49,11 @@ def writeGauge(filename: str, latt_size: List[int], plaquette: float, gauge: num
                     gauge_reorder[t, x_, y, z_] = gauge[[3, 0, 1, 2], :, t, z, y, x, :, :]
 
     gauge = gauge_reorder.astype(dtype)
-    if getMPIRank() == 0:
+    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 = getMPIComm().bcast(offset)
+    offset = MPI.COMM_WORLD.bcast(offset)
 
-    writeMPIFile(filename, dtype, offset, (Lt, Lx, Ly, Lz // 2, Nd, 2, Nc, Nc), (1, 2, 3, 0), gauge)
+    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.dev5 → pyquda_utils-0.9.9.dev6}/pyquda_utils/io/xqcd.py

@@ -3,55 +3,57 @@ from typing import List
 
 import numpy
 
-from pyquda import getSublatticeSize, readMPIFile, writeMPIFile
+from .mpi_file import getSublatticeSize, readMPIFile, writeMPIFile
 
 Ns, Nc = 4, 3
 
 
-def readPropagator(filename: str, latt_size: List[int], staggered: bool):
+def readPropagator(filename: str, latt_size: List[int], grid_size: List[int], staggered: bool):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = "<c8", 0
 
     if not staggered:
-        propagator = readMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2))
+        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")
     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))
+        propagator_align16 = readMPIFile(filename, dtype, offset, (Nc, Lt, Lz, Ly, Lx, 4), (4, 3, 2, 1), grid_size)
         propagator = propagator_align16[:, :, :, :, :, :Nc]
         propagator = propagator.transpose(1, 2, 3, 4, 5, 0).astype("<c16")
     return propagator
 
 
-def writePropagator(filename: str, latt_size: List[int], propagator: numpy.ndarray, staggered: bool):
+def writePropagator(
+    filename: str, latt_size: List[int], grid_size: List[int], propagator: numpy.ndarray, staggered: bool
+):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = "<c8", 0
 
     if not staggered:
         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), propagator)
+        writeMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2), grid_size, propagator)
     else:
         # QDP_ALIGN16 makes the last Nc to be aligned with 16 Bytes.
         propagator = propagator.astype(dtype).transpose(5, 0, 1, 2, 3, 4)
         propagator_align16 = numpy.zeros((Nc, Lt, Lz, Ly, Lx, 4), dtype)
         propagator_align16[:, :, :, :, :, :Nc] = propagator
-        writeMPIFile(filename, dtype, offset, (Nc, Lt, Lz, Ly, Lx, 4), (4, 3, 2, 1), propagator_align16)
+        writeMPIFile(filename, dtype, offset, (Nc, Lt, Lz, Ly, Lx, 4), (4, 3, 2, 1), grid_size, propagator_align16)
 
 
-def readPropagatorFast(filename: str, latt_size: List[int]):
+def readPropagatorFast(filename: str, latt_size: List[int], grid_size: List[int]):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = "<c8", 0
 
-    propagator = readMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2))
+    propagator = readMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2), grid_size)
     return propagator
 
 
-def writePropagatorFast(filename: str, latt_size: List[int], propagator: numpy.ndarray):
+def writePropagatorFast(filename: str, latt_size: List[int], grid_size: List[int], propagator: numpy.ndarray):
     filename = path.expanduser(path.expandvars(filename))
-    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size)
+    Lx, Ly, Lz, Lt = getSublatticeSize(latt_size, grid_size)
     dtype, offset = "<c8", 0
 
-    writeMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2), propagator)
+    writeMPIFile(filename, dtype, offset, (Ns, Nc, Lt, Lz, Ly, Lx, Ns, Nc), (5, 4, 3, 2), grid_size, propagator)

pyquda_utils-0.9.9.dev5/pyquda_utils/_version.py

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