PyQUDA-Utils 0.10.1.dev0__tar.gz → 0.10.1.dev2__tar.gz

{pyquda_utils-0.10.1.dev0 → pyquda_utils-0.10.1.dev2}/PKG-INFO

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

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

@@ -26,6 +26,7 @@ pyquda_utils/convert.py
 pyquda_utils/core.py
 pyquda_utils/deprecated.py
 pyquda_utils/gamma.py
+pyquda_utils/gauge_nd_sun.py
 pyquda_utils/gpt.py
 pyquda_utils/hmc_param.py
 pyquda_utils/milc_rhmc_param.py

{pyquda_utils-0.10.1.dev0 → pyquda_utils-0.10.1.dev2}/pyquda_io/_field_utils.py

@@ -173,7 +173,7 @@ def gaugeProject(gauge: numpy.ndarray):
     pass
 
 
-def gaugeReunitarize(gauge: numpy.ndarray, reunitarize_sigma: bool = True):
+def gaugeReunitarize(gauge: numpy.ndarray, reunitarize_sigma: float):
     gauge = numpy.ascontiguousarray(gauge.transpose(5, 6, 0, 1, 2, 3, 4))
     row0_abs = numpy.linalg.norm(gauge[0], axis=0)
     gauge[0] /= row0_abs
@@ -182,7 +182,7 @@ def gaugeReunitarize(gauge: numpy.ndarray, reunitarize_sigma: bool = True):
     row1_abs = numpy.linalg.norm(gauge[1], axis=0)
     gauge[1] /= row1_abs
     row2 = numpy.cross(gauge[0], gauge[1], axis=0).conjugate()
-    if reunitarize_sigma:
+    if reunitarize_sigma > 0:
         assert (
             MPI.COMM_WORLD.allreduce(
                 numpy.sqrt(
@@ -193,13 +193,13 @@ def gaugeReunitarize(gauge: numpy.ndarray, reunitarize_sigma: bool = True):
                 ).max(),
                 MPI.MAX,
             )
-            < 2e-7  # sqrt(Nc) * fp32 machine epsilon
+            < reunitarize_sigma
         )
     gauge[2] = row2
     return gauge.transpose(2, 3, 4, 5, 6, 0, 1)
 
 
-def gaugeReunitarizeReconstruct12(gauge: numpy.ndarray, reunitarize_sigma: bool = True):
+def gaugeReunitarizeReconstruct12(gauge: numpy.ndarray, reunitarize_sigma: float):
     """gauge shape (Nd, Lt, Lz, Ly, Lx, Nc - 1, Nc)"""
     gauge_ = gauge.transpose(5, 6, 0, 1, 2, 3, 4)
     gauge = numpy.empty((Nc, *gauge_.shape[1:]), "<c16")
@@ -211,13 +211,13 @@ def gaugeReunitarizeReconstruct12(gauge: numpy.ndarray, reunitarize_sigma: bool
     row1_abs = numpy.linalg.norm(gauge[1], axis=0)
     gauge[1] /= row1_abs
     row2 = numpy.cross(gauge[0], gauge[1], axis=0).conjugate()
-    if reunitarize_sigma:
+    if reunitarize_sigma > 0:
         assert (
             MPI.COMM_WORLD.allreduce(
                 numpy.sqrt((1 - row0_abs) ** 2 + numpy.abs(row0_row1) ** 2 + (1 - row1_abs) ** 2).max(),
                 MPI.MAX,
             )
-            < 2e-7  # sqrt(Nc) * fp32 machine epsilon
+            < reunitarize_sigma
         )
     gauge[2] = row2
     return gauge.transpose(2, 3, 4, 5, 6, 0, 1)

{pyquda_utils-0.10.1.dev0 → pyquda_utils-0.10.1.dev2}/pyquda_io/chroma.py

@@ -34,7 +34,7 @@ def checksum_qio(latt_size: List[int], grid_size: List[int], data):
     return sum29, sum31
 
 
-def readQIOGauge(filename: str, grid_size: List[int], checksum: bool = True):
+def readQIOGauge(filename: str, grid_size: List[int], checksum: bool = True, reunitarize_sigma: float = 2e-7):
     from .lime import Lime
 
     lime = Lime(filename)
@@ -68,7 +68,7 @@ def readQIOGauge(filename: str, grid_size: List[int], checksum: bool = True):
         ), f"Bad checksum for {filename}"
     gauge = gauge.transpose(4, 0, 1, 2, 3, 5, 6).astype("<c16")
     if precision == 4:
-        gauge = gaugeReunitarize(gauge)
+        gauge = gaugeReunitarize(gauge, reunitarize_sigma)  # 2e-7: Nc**0.5 * 1.1920929e-07
     return latt_size, gauge
 
 

{pyquda_utils-0.10.1.dev0 → pyquda_utils-0.10.1.dev2}/pyquda_io/milc.py

@@ -55,7 +55,7 @@ def checksum_qio(latt_size: List[int], grid_size: List[int], data):
     return sum29, sum31
 
 
-def readGauge(filename: str, grid_size: List[int], checksum: bool = True, reunitarize_sigma: bool = True):
+def readGauge(filename: str, grid_size: List[int], checksum: bool = True, reunitarize_sigma: float = 2e-7):
     filename = path.expanduser(path.expandvars(filename))
     with open(filename, "rb") as f:
         magic = f.read(4)
@@ -79,7 +79,7 @@ def readGauge(filename: str, grid_size: List[int], checksum: bool = True, reunit
             sum31,
         ), f"Bad checksum for {filename}"
     gauge = gauge.transpose(4, 0, 1, 2, 3, 5, 6).astype("<c16")
-    gauge = gaugeReunitarize(gauge, reunitarize_sigma)
+    gauge = gaugeReunitarize(gauge, reunitarize_sigma)  # 2e-7: Nc**0.5 * 1.1920929e-07
     return latt_size, gauge
 
 

{pyquda_utils-0.10.1.dev0 → pyquda_utils-0.10.1.dev2}/pyquda_io/nersc.py

@@ -27,7 +27,7 @@ def readGauge(
     checksum: bool = True,
     plaquette: bool = True,
     link_trace: bool = True,
-    reunitarize_sigma: bool = True,
+    reunitarize_sigma: float = 2e-7,
 ):
     filename = path.expanduser(path.expandvars(filename))
     header: Dict[str, str] = {}
@@ -63,7 +63,7 @@ def readGauge(
             assert checksum_nersc(gauge.reshape(-1)) == int(header["CHECKSUM"], 16), f"Bad checksum for {filename}"
         gauge = gauge.transpose(4, 0, 1, 2, 3, 5, 6).astype("<c16")
         if float_nbytes == 4:
-            gauge = gaugeReunitarize(gauge, reunitarize_sigma)
+            gauge = gaugeReunitarize(gauge, reunitarize_sigma)  # 2e-7: Nc**0.5 * 1.1920929e-07
     elif header["DATATYPE"] == "4D_SU3_GAUGE":
         gauge = readMPIFile(filename, dtype, offset, (Lt, Lz, Ly, Lx, Nd, Nc - 1, Nc), (3, 2, 1, 0), grid_size)
         gauge = gauge.astype(f"<c{2 * float_nbytes}")
@@ -71,7 +71,7 @@ def readGauge(
             assert checksum_nersc(gauge.reshape(-1)) == int(header["CHECKSUM"], 16), f"Bad checksum for {filename}"
         gauge = gauge.transpose(4, 0, 1, 2, 3, 5, 6).astype("<c16")
         if float_nbytes == 4:
-            gauge = gaugeReunitarizeReconstruct12(gauge, reunitarize_sigma)
+            gauge = gaugeReunitarizeReconstruct12(gauge, reunitarize_sigma)  # 2e-7: Nc**0.5 * 1.1920929e-07
         elif float_nbytes == 8:
             gauge = gaugeReconstruct12(gauge)
     else:

pyquda_utils-0.10.1.dev2/pyquda_utils/_version.py

@@ -0,0 +1 @@
+__version__ = "0.10.1.dev2"

pyquda_utils-0.10.1.dev2/pyquda_utils/gauge_nd_sun.py

@@ -0,0 +1,331 @@
+from os import environ
+from typing import Sequence
+
+import numpy
+from mpi4py import MPI
+
+_RANK = MPI.COMM_WORLD.Get_rank()
+_GPUID: int = -1
+
+
+def initGPU(gpuid: int = -1):
+    from platform import node as gethostname
+    import cupy
+
+    rank = _RANK
+
+    # quda/include/communicator_quda.h
+    # determine which GPU this rank will use
+    hostname = gethostname()
+    hostname_recv_buf = MPI.COMM_WORLD.allgather(hostname)
+
+    if gpuid < 0:
+        device_count = cupy.cuda.runtime.getDeviceCount()
+        if device_count == 0:
+            raise RuntimeError("No devices found")
+
+        # We initialize gpuid if it's still negative.
+        gpuid = 0
+        for i in range(rank):
+            if hostname == hostname_recv_buf[i]:
+                gpuid += 1
+
+        if gpuid >= device_count:
+            if "QUDA_ENABLE_MPS" in environ and environ["QUDA_ENABLE_MPS"] == "1":
+                gpuid %= device_count
+                print(f"MPS enabled, rank={rank} -> gpu={gpuid}")
+            else:
+                raise RuntimeError(f"Too few GPUs available on {hostname}")
+
+    cupy.cuda.Device(gpuid).use()
+    print(f"Rank {rank} uses GPU {gpuid}")
+
+    global _GPUID
+    _GPUID = gpuid
+
+
+def getSublatticeSize(latt_size: Sequence[int], grid_size: Sequence[int]):
+    assert len(latt_size) == len(grid_size)
+    for GL, G in zip(latt_size, grid_size):
+        assert GL % G == 0
+    return [GL // G for GL, G in zip(latt_size, grid_size)]
+
+
+def getGridCoord(grid_size: Sequence[int]):
+    rank = _RANK
+    grid_coord = []
+    for G in grid_size[::-1]:
+        grid_coord.append(rank % G)
+        rank //= G
+    return grid_coord[::-1]
+
+
+def getShiftedRank(grid_coord: Sequence[int], grid_size: Sequence[int], delta: Sequence[int]):
+    Nd = len(grid_size)
+    grid_coord = [(g + d) % G for g, G, d in zip(grid_coord, grid_size, delta)]
+    rank = grid_coord[0]
+    for mu in range(1, Nd):
+        rank = rank * grid_size[mu] + grid_coord[mu]
+    return rank
+
+
+def gaugeSendRecv(extended, gauge, dest, source):
+    rank = _RANK
+    if rank == dest and rank == source:
+        extended[:] = gauge
+    else:
+        buf = gauge.copy()
+        MPI.COMM_WORLD.Sendrecv_replace(buf, dest=dest, source=source)
+        extended[:] = buf
+
+
+class LatticeLink:
+    def __init__(
+        self,
+        latt_size: Sequence[int],
+        grid_size: Sequence[int],
+        color: int,
+        matrix: numpy.ndarray = None,
+        mu: int = None,
+    ):
+        assert len(latt_size) == len(grid_size)
+        self.Nd = len(latt_size)
+        self.Nc = color
+        self.latt_size = tuple(latt_size)
+        self.grid_size = tuple(grid_size)
+        self.grid_coord = getGridCoord(grid_size)
+        self.sublatt_size = getSublatticeSize(latt_size, grid_size)
+        if matrix is None:
+            self.matrix = numpy.empty((*self.sublatt_size[::-1], self.Nc, self.Nc), numpy.complex128)
+        else:
+            self.matrix = matrix.reshape(*self.sublatt_size[::-1], self.Nc, self.Nc)
+        self.mu = mu
+
+    def __getitem__(self, key):
+        return self.matrix[key]
+
+    def __matmul__(self, other: "LatticeLink"):
+        return self.matrix @ other.matrix
+
+    @property
+    def backend(self):
+        if type(self.matrix).__module__ == "numpy":
+            return numpy
+        elif type(self.matrix).__module__ == "cupy":
+            import cupy
+
+            return cupy
+        else:
+            raise RuntimeError(f"Unknown array type {type(self.matrix)}")
+
+    def shift(self, mu: int, dagger: bool = False):
+        assert 0 <= mu < 2 * self.Nd
+        backend = self.backend
+        Nd = self.Nd
+        dir = 1 - 2 * (mu // self.Nd)
+        mu = mu % self.Nd
+        left_slice = [slice(None, None) for nu in range(self.Nd)]
+        right_slice = [slice(None, None) for nu in range(self.Nd)]
+        result = backend.empty_like(self.matrix)
+        right = self.matrix
+        rank = _RANK
+        dest = getShiftedRank(self.grid_coord, self.grid_size, [0 if nu != mu else -dir for nu in range(Nd)])
+        source = getShiftedRank(self.grid_coord, self.grid_size, [0 if nu != mu else dir for nu in range(Nd)])
+
+        left_slice[mu] = slice(-1, None) if dir == 1 else slice(None, 1)
+        right_slice[mu] = slice(None, 1) if dir == 1 else slice(-1, None)
+
+        # gaugeSendRecv(shifted[*shift_slice[::-1]], matrix[*matrix_slice[::-1]], dest, source)
+        sendbuf = right[*right_slice[::-1]] if not dagger else right[*right_slice[::-1]].swapaxes(-2, -1).conjugate()
+        if rank == source and rank == dest:
+            pass
+        else:
+            sendbuf = backend.ascontiguousarray(sendbuf)
+            request = MPI.COMM_WORLD.Isend(sendbuf, dest)
+
+        left_slice[mu] = slice(None, -1) if dir == 1 else slice(1, None)
+        right_slice[mu] = slice(1, None) if dir == 1 else slice(None, -1)
+        result[*left_slice[::-1]] = (
+            right[*right_slice[::-1]] if not dagger else right[*right_slice[::-1]].swapaxes(-2, -1).conjugate()
+        )
+        left_slice[mu] = slice(-1, None) if dir == 1 else slice(None, 1)
+        right_slice[mu] = slice(None, 1) if dir == 1 else slice(-1, None)
+
+        if rank == source and rank == dest:
+            recvbuf = sendbuf
+        else:
+            recvbuf = backend.empty_like(sendbuf)
+            MPI.COMM_WORLD.Recv(recvbuf, source)
+            request.Wait()
+        result[*left_slice[::-1]] = recvbuf
+
+        return LatticeLink(self.latt_size, self.grid_size, self.Nc, result, self.mu)
+
+    def link(self, right: "LatticeLink"):
+        assert self.mu is not None, "Ambiguous dimension and direction"
+        backend = self.backend
+        Nd = self.Nd
+        dir = 1 - 2 * (self.mu // self.Nd)
+        mu = self.mu % self.Nd
+        left_slice = [slice(None, None) for nu in range(self.Nd)]
+        right_slice = [slice(None, None) for nu in range(self.Nd)]
+        result = backend.empty_like(self.matrix)
+        left = self.matrix
+        right = right.matrix
+        rank = _RANK
+        dest = getShiftedRank(self.grid_coord, self.grid_size, [0 if nu != mu else -dir for nu in range(Nd)])
+        source = getShiftedRank(self.grid_coord, self.grid_size, [0 if nu != mu else dir for nu in range(Nd)])
+
+        left_slice[mu] = slice(-1, None) if dir == 1 else slice(None, 1)
+        right_slice[mu] = slice(None, 1) if dir == 1 else slice(-1, None)
+
+        sendbuf = right[*right_slice[::-1]]
+        if rank == source and rank == dest:
+            pass
+        else:
+            sendbuf = backend.ascontiguousarray(sendbuf)
+            request = MPI.COMM_WORLD.Isend(sendbuf, dest)
+
+        left_slice[mu] = slice(None, -1) if dir == 1 else slice(1, None)
+        right_slice[mu] = slice(1, None) if dir == 1 else slice(None, -1)
+        result[*left_slice[::-1]] = left[*left_slice[::-1]] @ right[*right_slice[::-1]]
+        left_slice[mu] = slice(-1, None) if dir == 1 else slice(None, 1)
+        right_slice[mu] = slice(None, 1) if dir == 1 else slice(-1, None)
+
+        if rank == source and rank == dest:
+            recvbuf = sendbuf
+        else:
+            recvbuf = backend.empty_like(sendbuf)
+            MPI.COMM_WORLD.Recv(recvbuf, source)
+            request.Wait()
+        result[*left_slice[::-1]] = left[*left_slice[::-1]] @ recvbuf
+
+        return LatticeLink(self.latt_size, self.grid_size, self.Nc, result)
+
+    def dagger(self):
+        return LatticeLink(
+            self.latt_size,
+            self.grid_size,
+            self.Nc,
+            self.matrix.swapaxes(-2, -1).conjugate(),
+        )
+
+    def toDevice(self):
+        import cupy
+
+        if _GPUID < 0:
+            initGPU()
+        self.matrix = cupy.asarray(self.matrix)
+
+    def toHost(self):
+        self.matrix = self.matrix.get()
+
+
+class LatticeGauge:
+    def __init__(
+        self,
+        latt_size: Sequence[int],
+        grid_size: Sequence[int],
+        color: int,
+        border: int = 0,
+        gauge: numpy.ndarray = None,
+        extended: numpy.ndarray = None,
+    ):
+        assert len(latt_size) == len(grid_size)
+        self.Nd = len(latt_size)
+        self.Nc = color
+        self.latt_size = tuple(latt_size)
+        self.grid_size = tuple(grid_size)
+        self.grid_coord = getGridCoord(grid_size)
+        self.sublatt_size = getSublatticeSize(latt_size, grid_size)
+        shape = (self.Nd, *self.sublatt_size[::-1], self.Nc, self.Nc)
+        if gauge is None:
+            self.gauge = numpy.empty(shape, numpy.complex128)
+        else:
+            self.gauge = gauge.reshape(shape)
+        self.extend(border, extended)
+
+    def __getitem__(self, mu):
+        assert 0 <= mu < 2 * self.Nd
+        gauge_mu = LatticeLink(self.latt_size, self.grid_size, self.Nc, self.gauge[mu % self.Nd], mu)
+        return gauge_mu if mu < self.Nd else gauge_mu.shift(mu, True)
+
+    @property
+    def backend(self):
+        if type(self.gauge).__module__ == "numpy":
+            return numpy
+        elif type(self.gauge).__module__ == "cupy":
+            import cupy
+
+            return cupy
+        else:
+            raise RuntimeError(f"Unknown array type {type(self.gauge)}")
+
+    def extend(self, border: int, extended: numpy.ndarray = None):
+        if border <= 0:
+            self.Lb = 0
+            self.extlatt_size = self.sublatt_size
+            self.extended = None
+        else:
+            self.Lb = border
+            self.extlatt_size = [L + 2 * border for L in self.sublatt_size]
+            shape = (self.Nd, *self.extlatt_size[::-1], self.Nc, self.Nc)
+            if extended is None:
+                self.extended = self.backend.empty(shape, self.gauge.dtype)
+            else:
+                self.extended = extended.reshape(shape)
+            self.exchange()
+
+    def exchange(self):
+        assert self.extended is not None
+        Nd = self.Nd
+        Lb = self.Lb
+        extended_slice = [slice(Lb, -Lb) for mu in range(Nd)]
+        gauge_slice = [slice(None, None) for mu in range(Nd)]
+        stride = [3 ** (Nd - 1 - mu) for mu in range(Nd)]
+        for tag in range(3**Nd):
+            delta = [(tag // stride[mu] % 3 - 1) for mu in range(Nd)]
+            for mu in range(Nd):
+                if delta[mu] == -1:
+                    extended_slice[mu] = slice(-Lb, None)
+                    gauge_slice[mu] = slice(None, Lb)
+                elif delta[mu] == 1:
+                    extended_slice[mu] = slice(None, Lb)
+                    gauge_slice[mu] = slice(-Lb, None)
+            gaugeSendRecv(
+                self.extended[:, *extended_slice[::-1]],
+                self.gauge[:, *gauge_slice[::-1]],
+                getShiftedRank(self.grid_coord, self.grid_size, delta),
+                getShiftedRank(self.grid_coord, self.grid_size, [-d for d in delta]),
+            )
+            for mu in range(Nd):
+                if delta[mu] != 0:
+                    extended_slice[mu] = slice(Lb, -Lb)
+                    gauge_slice[mu] = slice(None, None)
+
+    def shift(self, delta: Sequence[int]):
+        assert numpy.abs(delta).max() <= self.Lb
+        Lb = self.Lb
+        extended_slice = [slice(Lb + d, None if Lb == d else -(Lb - d)) for d in delta[::-1]]
+        return LatticeGauge(self.latt_size, self.grid_size, self.Nc, 0, self.extended[:, *extended_slice], None)
+
+    def toDevice(self):
+        import cupy
+
+        if _GPUID < 0:
+            initGPU()
+        self.gauge = cupy.asarray(self.gauge)
+        if self.extended is not None:
+            self.extended = cupy.asarray(self.extended)
+
+    def toHost(self):
+        self.gauge = self.gauge.get()
+        if self.extended is not None:
+            self.extended = self.extended.get()
+
+
+def link(*color_matrices: LatticeLink):
+    linked = color_matrices[-1]
+    for color_matrix in color_matrices[::-1][1:]:
+        linked = color_matrix.link(linked)
+    return linked

{pyquda_utils-0.10.1.dev0 → pyquda_utils-0.10.1.dev2}/pyquda_utils/gpt.py

@@ -82,7 +82,7 @@ def LatticePropagatorGPT(lattice: g.lattice, gen_simd_width: int, propagator: La
         return propagator
     else:
         assert latt_info.size == propagator.latt_info.size
-        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.mview()[0][:] = (
             propagator.lexico()
             .astype(f"<c{2 * gpt_prec}")

{pyquda_utils-0.10.1.dev0 → pyquda_utils-0.10.1.dev2}/pyquda_utils/io/__init__.py

@@ -69,11 +69,11 @@ def rotateToDeGrandRossi(propagator: LatticePropagator):
     )
 
 
-def readChromaQIOGauge(filename: str, checksum: bool = True):
+def readChromaQIOGauge(filename: str, checksum: bool = True, reunitarize_sigma: float = 2e-7):
     from pyquda import getGridSize
     from pyquda_io.chroma import readQIOGauge as read
 
-    latt_size, gauge_raw = read(filename, getGridSize(), checksum)
+    latt_size, gauge_raw = read(filename, getGridSize(), checksum, reunitarize_sigma)
     return LatticeGauge(LatticeInfo(latt_size), evenodd(gauge_raw, [1, 2, 3, 4]))
 
 
@@ -96,7 +96,7 @@ def readChromaQIOPropagator(filename: str, checksum: bool = True):
         return LatticeStaggeredPropagator(LatticeInfo(latt_size), evenodd(propagator_raw, [0, 1, 2, 3]))
 
 
-def readMILCGauge(filename: str, checksum: bool = True, reunitarize_sigma: bool = True):
+def readMILCGauge(filename: str, checksum: bool = True, reunitarize_sigma: float = 2e-7):
     from pyquda import getGridSize
     from pyquda_io.milc import readGauge as read
 
@@ -244,7 +244,7 @@ def readNERSCGauge(
     checksum: bool = True,
     plaquette: bool = True,
     link_trace: bool = True,
-    reunitarize_sigma: bool = True,
+    reunitarize_sigma: float = 2e-7,
 ):
     from pyquda import getGridSize
     from pyquda_io.nersc import readGauge as read

pyquda_utils-0.10.1.dev0/pyquda_utils/_version.py

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