Prepare for arbitrary Ns/Nc field.

Add loading and saving to a gpt.lattice object.

pyquda_core/pyquda/__init__.py

@@ -67,14 +67,22 @@ class _ComputeCapability(NamedTuple):
 _COMPUTE_CAPABILITY: _ComputeCapability = _ComputeCapability(0, 0)
 
 
-def getRankFromCoord(coord: List[int], grid: List[int]) -> int:
-    x, y, z, t = grid
-    return ((coord[0] * y + coord[1]) * z + coord[2]) * t + coord[3]
+def getRankFromCoord(grid_coord: List[int], grid_size: List[int] = None) -> int:
+    Gx, Gy, Gz, Gt = _GRID_SIZE if grid_size is None else grid_size
+    gx, gy, gz, gt = grid_coord
+    return ((gx * Gy + gy) * Gz + gz) * Gt + gt
 
 
-def getCoordFromRank(rank: int, grid: List[int]) -> List[int]:
-    x, y, z, t = grid
-    return [rank // t // z // y, rank // t // z % y, rank // t % z, rank % t]
+def getCoordFromRank(mpi_rank: int, grid_size: List[int] = None) -> List[int]:
+    Gx, Gy, Gz, Gt = _GRID_SIZE if grid_size is None else grid_size
+    return [mpi_rank // Gt // Gz // Gy, mpi_rank // Gt // Gz % Gy, mpi_rank // Gt % Gz, mpi_rank % Gt]
+
+
+def getSublatticeSize(latt_size: List[int], grid_size: List[int] = None) -> List[int]:
+    Gx, Gy, Gz, Gt = _GRID_SIZE if grid_size is None else grid_size
+    Lx, Ly, Lz, Lt = latt_size
+    assert Lx % Gx == 0 and Ly % Gy == 0 and Lz % Gz == 0 and Lt % Gt == 0
+    return [Lx // Gx, Ly // Gy, Lz // Gz, Lt // Gt]
 
 
 def _composition4(n):
@@ -405,13 +413,6 @@ def getCUDAComputeCapability():
     return _COMPUTE_CAPABILITY
 
 
-def getSublatticeSize(latt_size: List[int]):
-    Lx, Ly, Lz, Lt = latt_size
-    Gx, Gy, Gz, Gt = _GRID_SIZE
-    assert Lx % Gx == 0 and Ly % Gy == 0 and Lz % Gz == 0 and Lt % Gt == 0
-    return [Lx // Gx, Ly // Gy, Lz // Gz, Lt // Gt]
-
-
 def _getSubarray(shape: Sequence[int], axes: Sequence[int]):
     sizes = [d for d in shape]
     subsizes = [d for d in shape]

pyquda_core/pyquda/_version.py

@@ -1 +1 @@
-__version__ = "0.9.5"
+__version__ = "0.9.6"

pyquda_core/pyquda/action/clover_wilson.py

@@ -12,7 +12,7 @@
     QudaSolveType,
     QudaVerbosity,
 )
-from ..field import Nd, Nc, Ns, LatticeInfo, LatticeFermion, MultiLatticeFermion
+from ..field import LatticeInfo, LatticeFermion, MultiLatticeFermion
 from ..dirac import CloverWilsonDirac
 
 nullptr = Pointers("void", 0)
@@ -38,7 +38,7 @@ def __init__(
             getLogger().critical("anisotropy != 1.0 not implemented", NotImplementedError)
         super().__init__(latt_info, CloverWilsonDirac(latt_info, mass, tol, maxiter, clover_csw, 1, None))
 
-        kappa = 1 / (2 * (mass + Nd))
+        kappa = 1 / (2 * (mass + latt_info.Nd))
         self.setForceParam(rational_param, kappa, clover_csw, n_flavor)
         self.quark = MultiLatticeFermion(self.latt_info, self.max_num_offset)
         self.phi = LatticeFermion(latt_info)
@@ -83,7 +83,7 @@ def action(self, new_gauge: bool) -> float:
         self.invert_param.compute_action = 0
         return (
             self.invert_param.action[0]
-            - self.latt_info.volume // 2 * Ns * Nc  # volume_cb2 here
+            - self.latt_info.volume // 2 * self.latt_info.Ns * self.latt_info.Nc  # volume_cb2 here
             - self.multiplicity * self.invert_param.trlogA[1]
         )
 

pyquda_core/pyquda/action/gauge.py

@@ -5,7 +5,7 @@
 
 from ..pointer import Pointers
 from ..pyquda import computeGaugeLoopTraceQuda, computeGaugeForceQuda
-from ..field import Nc, LatticeInfo
+from ..field import LatticeInfo
 from ..dirac import GaugeDirac
 
 nullptr = Pointers("void", 0)
@@ -94,7 +94,7 @@ def __init__(self, latt_info: LatticeInfo, loop_param: LoopParam, beta: float, u
         super().__init__(latt_info, GaugeDirac(latt_info))
 
         # S=\frac{\beta}{N_c}\sum_{i}c_i\mathrm{ReTr}(I-W_i)
-        self.action_path = actionPath(loop_param.path, [-beta / Nc * coeff for coeff in loop_param.coeff])
+        self.action_path = actionPath(loop_param.path, [-beta / latt_info.Nc * coeff for coeff in loop_param.coeff])
         self.force_path = forcePath(self.action_path)
 
     def action(self) -> float:

pyquda_core/pyquda/dirac/clover_wilson.py

@@ -1,6 +1,6 @@
 from typing import List, Union
 
-from ..field import Nd, LatticeInfo, LatticeGauge, LatticeClover
+from ..field import LatticeInfo, LatticeGauge, LatticeClover
 from ..enum_quda import QudaDslashType, QudaPrecision
 
 from . import general
@@ -18,7 +18,7 @@ def __init__(
         clover_xi: float = 1.0,
         multigrid: Union[List[List[int]], Multigrid] = None,
     ) -> None:
-        kappa = 1 / (2 * (mass + 1 + (Nd - 1) / latt_info.anisotropy))
+        kappa = 1 / (2 * (mass + 1 + (latt_info.Nd - 1) / latt_info.anisotropy))
         super().__init__(latt_info)
         self.clover: LatticeClover = None
         self.clover_inv: LatticeClover = None

pyquda_core/pyquda/dirac/wilson.py

@@ -1,6 +1,6 @@
 from typing import List, Union
 
-from ..field import Nd, LatticeInfo, LatticeGauge
+from ..field import LatticeInfo, LatticeGauge
 from ..enum_quda import QudaDslashType, QudaPrecision
 
 from . import general
@@ -16,7 +16,7 @@ def __init__(
         maxiter: int,
         multigrid: Union[List[List[int]], Multigrid] = None,
     ) -> None:
-        kappa = 1 / (2 * (mass + 1 + (Nd - 1) / latt_info.anisotropy))
+        kappa = 1 / (2 * (mass + 1 + (latt_info.Nd - 1) / latt_info.anisotropy))
         super().__init__(latt_info)
         self.newQudaGaugeParam()
         self.newQudaMultigridParam(multigrid, mass, kappa, 0.25, 16, 1e-6, 1000, 0, 8)