PyQUDA-Utils 0.10.8.dev5__tar.gz → 0.10.9.dev1__tar.gz

{pyquda_utils-0.10.8.dev5 → pyquda_utils-0.10.9.dev1}/PKG-INFO

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

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

@@ -23,6 +23,7 @@ pyquda_io/openqcd.py
 pyquda_io/xqcd.py
 pyquda_plugins/__main__.py
 pyquda_plugins/plugin_pyx.py
+pyquda_plugins/pycontract/__init__.py
 pyquda_plugins/pygwu/__init__.py
 pyquda_utils/__init__.py
 pyquda_utils/_version.py

{pyquda_utils-0.10.8.dev5 → pyquda_utils-0.10.9.dev1}/pyquda_plugins/plugin_pyx.py

@@ -5,8 +5,13 @@ _C_TO_PYTHON: Dict[str, str] = {
     "void *": "Pointer",
     "void": "None",
     "short": "int",
+    "unsigned short": "int",
     "int": "int",
+    "unsigned int": "int",
     "long": "int",
+    "unsigned long": "int",
+    "long long": "int",
+    "unsigned long long": "int",
     "float": "float",
     "double": "float",
     "long double": "float",
@@ -125,10 +130,29 @@ def parseHeader(header, include_path):
     except ImportError or ModuleNotFoundError:
         from pyquda_core.pycparser.pycparser import parse_file, c_ast  # This is for the language server
 
+    def evaluate(node):
+        if node is None:
+            return
+        elif type(node) is c_ast.UnaryOp:
+            if node.op == "+":
+                return evaluate(node.expr)
+            elif node.op == "-":
+                return -evaluate(node.expr)
+        elif type(node) is c_ast.BinaryOp:
+            if node.op == "+":
+                return evaluate(node.left) + evaluate(node.right)
+            elif node.op == "-":
+                return evaluate(node.left) - evaluate(node.right)
+        elif type(node) is c_ast.Constant:
+            return int(node.value, 0)
+        else:
+            raise ValueError(f"Unknown node {node}")
+
     ast = parse_file(
         os.path.join(include_path, header), use_cpp=True, cpp_path="cc", cpp_args=["-E", Rf"-I{include_path}"]
     )
     funcs: List[FunctionMeta] = []
+    enums: Dict[str, List[Tuple[str, int]]] = {}
     for node in ast:
         if isinstance(node.type, c_ast.FuncDecl):
             a = []
@@ -179,24 +203,53 @@ def parseHeader(header, include_path):
                 funcs.append(FunctionMeta((" ".join(nttt.type.names), "*"), node.name, a))
             else:
                 raise ValueError(f"Unexpected node {node}")
-    return funcs
+        elif type(node.type.type) is c_ast.Enum:
+            print(node.name)
+            current_value = -1
+            enums[node.name] = []
+            for item in node.type.type.values:
+                item_value = evaluate(item.value)
+                if item_value is not None:
+                    current_value = item_value
+                else:
+                    current_value += 1
+                enums[node.name].append((item.name, current_value))
+    return funcs, enums
 
 
 def build_plugin_pyx(plugins_root, lib: str, header: str, include_path: str):
     lib_pxd = f'cdef extern from "{header}":\n'
     pylib_pyx = (
+        "from enum import IntEnum\n"
         "from libcpp cimport bool\n"
         "from numpy cimport ndarray\n"
         "from pyquda_comm.pointer cimport Pointer, _NDArray\n"
         f"cimport {lib}\n"
         "\n"
     )
-    pylib_pyi = "from numpy import int32, float64, complex128\n" "from numpy.typing import NDArray\n" "\n"
-    c_funcs = parseHeader(header, include_path)
+    pylib_pyi = (
+        "from enum import IntEnum\n"
+        "from numpy import int32, float64, complex128\n"
+        "from numpy.typing import NDArray\n"
+        "\n"
+    )
+    c_funcs, c_enums = parseHeader(header, include_path)
+    for c_enum in c_enums.keys():
+        lib_pxd += f"    ctypedef enum {c_enum}:\n        pass\n\n"
+        pylib_pyx += f"class {c_enum}(IntEnum):\n"
+        pylib_pyi += f"class {c_enum}(IntEnum):\n"
+        for key, value in c_enums[c_enum]:
+            pylib_pyx += f"    {key} = {value}\n"
+            pylib_pyi += f"    {key} = {value}\n"
+        pylib_pyx += "\n"
+        pylib_pyi += "\n"
+        _C_TO_PYTHON.update({c_enum: c_enum})
     for c_func in c_funcs:
         lib_pxd += f"    {c_func}\n"
         pylib_pyx += "\n" + c_func.pyx(lib)
         pylib_pyi += c_func.pyi()
+    for c_enum in c_enums.keys():
+        pylib_pyx = pylib_pyx.replace(f"{c_enum} ", f"{lib}.{c_enum} ")
 
     os.makedirs(os.path.join(plugins_root, f"py{lib}", "src"), exist_ok=True)
     with open(os.path.join(plugins_root, f"py{lib}", "src", f"{lib}.pxd"), "w") as f:

pyquda_utils-0.10.9.dev1/pyquda_plugins/pycontract/__init__.py

@@ -0,0 +1,41 @@
+from typing import Union
+
+from pyquda_comm import getLogger
+from pyquda_comm.field import LatticeComplex, LatticePropagator
+from pyquda_utils.gamma import Gamma, Polarize
+
+from . import _pycontract as contract
+from ._pycontract import BaryonContractType
+
+
+def baryonTwoPoint(
+    propag_i: LatticePropagator,
+    propag_j: LatticePropagator,
+    propag_m: LatticePropagator,
+    contract_type: BaryonContractType,
+    gamma_ij: Gamma,
+    gamma_kl: Gamma,
+    gamma_mn: Union[Gamma, Polarize],
+):
+    latt_info = propag_i.latt_info
+    if isinstance(gamma_mn, Gamma):
+        correl = LatticeComplex(latt_info)
+        contract.baryon_two_point(
+            correl.data_ptr,
+            propag_i.data_ptr,
+            propag_j.data_ptr,
+            propag_m.data_ptr,
+            contract_type,
+            latt_info.volume,
+            gamma_ij.index,
+            gamma_kl.index,
+            gamma_mn.index,
+        )
+        correl.data *= gamma_ij.sign * gamma_kl.sign * gamma_mn.sign
+        return correl
+    elif isinstance(gamma_mn, Polarize):
+        correl_left = baryonTwoPoint(propag_i, propag_j, propag_m, contract_type, gamma_ij, gamma_kl, gamma_mn.left)
+        correl_right = baryonTwoPoint(propag_i, propag_j, propag_m, contract_type, gamma_ij, gamma_kl, gamma_mn.right)
+        return correl_left + correl_right
+    else:
+        raise getLogger().critical("gamma_mn should be Gamma or Polarize", ValueError)

pyquda_utils-0.10.9.dev1/pyquda_utils/_version.py

@@ -0,0 +1 @@
+__version__ = "0.10.9.dev1"

{pyquda_utils-0.10.8.dev5 → pyquda_utils-0.10.9.dev1}/pyquda_utils/convert.py

@@ -12,7 +12,7 @@ from pyquda.field import (
 
 
 def linkToFermion(link: LatticeLink):
-    fermion = LatticeLink(link.latt_info)
+    fermion = LatticeFermion(link.latt_info)
     for color in range(link.latt_info.Nc):
         fermion.data[:, :, :, :, :, color, :] = link.data[:, :, :, :, :, :, color]
     return fermion

{pyquda_utils-0.10.8.dev5 → pyquda_utils-0.10.9.dev1}/pyquda_utils/gamma.py

@@ -54,7 +54,7 @@ class GammaMatrix:
     )
 
     @classmethod
-    def matrix(cls, index) -> numpy.ndarray:
+    def matrix(cls, index):
         return (
             (cls.gamma_1 if index & 0b0001 else cls.gamma_0)
             @ (cls.gamma_2 if index & 0b0010 else cls.gamma_0)
@@ -115,14 +115,14 @@ class Gamma:
     popcnt = [0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]
     popsign = [1, -1, -1, 1, -1, 1, 1, -1, -1, 1, 1, -1, 1, -1, -1, 1]
 
-    def __init__(self, index: int, sign: Literal[1, -1]) -> None:
+    def __init__(self, index: int, sign: Literal[1, -1] = 1) -> None:
         assert isinstance(index, int) and 0b0000 <= index <= 0b1111, "index should be int from 0 to 15"
         self.index = index
         self.sign = sign
 
     def __repr__(self) -> str:
         return (
-            f"{'-' if self.sign == -1 else ''}"
+            f"{'-' if self.sign == -1 else '' if self.sign == 1 else self.sign}"
             f"{'γ₀' if not self.index else ''}"
             f"{'γ₁' if self.index & 0b0001 else ''}"
             f"{'γ₂' if self.index & 0b0010 else ''}"
@@ -130,6 +130,35 @@ class Gamma:
             f"{'γ₄' if self.index & 0b1000 else ''}"
         )
 
+    def __neg__(self) -> "Gamma":
+        return Gamma(self.index, -self.sign)
+
+    def __mul__(self, rhs: Union[int, float, complex]) -> "Gamma":
+        assert isinstance(rhs, (int, float, complex)), "rhs should be int, float or complex"
+        return Gamma(self.index, self.sign * rhs)
+
+    def __rmul__(self, lhs: Union[int, float, complex]) -> "Gamma":
+        assert isinstance(lhs, (int, float, complex)), "lhs should be int, float or complex"
+        return Gamma(self.index, lhs * self.sign)
+
+    def __truediv__(self, rhs: Union[int, float, complex]) -> "Gamma":
+        assert isinstance(rhs, (int, float, complex)), "rhs should be int, float or complex"
+        return Gamma(self.index, self.sign / rhs)
+
+    def __add__(self, rhs: "Gamma") -> Union["Gamma", "Polarize"]:
+        assert isinstance(rhs, Gamma), "rhs should be Gamma"
+        if self.index == rhs.index:
+            return Gamma(self.index, self.sign + rhs.sign)
+        else:
+            return Polarize(self, rhs)
+
+    def __sub__(self, rhs: "Gamma") -> Union["Gamma", "Polarize"]:
+        assert isinstance(rhs, Gamma), "rhs should be Gamma"
+        if self.index == rhs.index:
+            return Gamma(self.index, self.sign - rhs.sign)
+        else:
+            return Polarize(self, -rhs)
+
     def __matmul__(self, rhs: "Gamma") -> "Gamma":
         index = self.index ^ rhs.index
         sign = self.sign * rhs.sign
@@ -150,7 +179,7 @@ class Gamma:
         return Gamma(index, sign)
 
     @property
-    def matrix(self) -> numpy.ndarray:
+    def matrix(self):
         return self.sign * GammaMatrix.matrix(self.index)
 
     @property
@@ -158,6 +187,49 @@ class Gamma:
         return GammaSparse.csr_matrix(self.index)
 
 
+class Polarize:
+    def __init__(self, left: Gamma, right: Gamma) -> None:
+        assert isinstance(left, Gamma), "left should be Gamma"
+        assert isinstance(right, Gamma), "right should be Gamma"
+        assert left.index != right.index, "left and right should be different"
+        self.left = left
+        self.right = right
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.left} + {self.right}"
+            if self.right.sign > 0
+            else f"{self.left} - {Gamma(self.right.index, -self.right.sign)}"
+        )
+
+    def __neg__(self) -> "Polarize":
+        return Polarize(-self.left, -self.right)
+
+    def __mul__(self, rhs: Union[int, float, complex]) -> "Gamma":
+        assert isinstance(rhs, (int, float, complex)), "rhs should be int, float or complex"
+        return Polarize(self.left * rhs, self.right * rhs)
+
+    def __rmul__(self, lhs: Union[int, float, complex]) -> "Gamma":
+        assert isinstance(lhs, (int, float, complex)), "lhs should be int, float or complex"
+        return Polarize(lhs * self.left, lhs * self.right)
+
+    def __truediv__(self, rhs: Union[int, float, complex]) -> "Gamma":
+        assert isinstance(rhs, (int, float, complex)), "rhs should be int, float or complex"
+        return Polarize(self.left / rhs, self.right / rhs)
+
+    def __matmul__(self, rhs: Gamma) -> "Polarize":
+        assert isinstance(rhs, Gamma), "rhs should be Gamma"
+        return Polarize(self.left @ rhs, self.right @ rhs)
+
+    def __rmatmul__(self, lhs: Gamma) -> "Polarize":
+        assert isinstance(lhs, Gamma), "rhs should be Gamma"
+        return Polarize(lhs @ self.left, lhs @ self.right)
+
+    @property
+    def matrix(self):
+        return self.left.matrix + self.right.matrix
+
+
 def gamma(n: int):
     assert isinstance(n, int) and 0 <= n <= 15, "n should be int from 0 to 15"
     backend = getCUDABackend()

{pyquda_utils-0.10.8.dev5 → pyquda_utils-0.10.9.dev1}/pyquda_utils/hmc_param.py

@@ -307,7 +307,7 @@ def staggeredFermionRationalParam(
             fermion_action = remez.getIPFE()
             rational_param = RationalParam(*molecular_dynamics, *pseudo_fermion, *fermion_action)
             cache[key] = rational_param
-        _dumpCache(cache)
+            _dumpCache(cache)
     else:
         rational_param = None
     return getMPIComm().bcast(rational_param)

pyquda_utils-0.10.8.dev5/pyquda_utils/_version.py

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