added block_cg, general scalar action for RMHMC; next: stencil versio…

…n of laplace

lib/gpt/algorithms/inverter/__init__.py

@@ -21,6 +21,7 @@
 from gpt.algorithms.inverter.deflate import deflate
 from gpt.algorithms.inverter.coarse_deflate import coarse_deflate
 from gpt.algorithms.inverter.cg import cg
+from gpt.algorithms.inverter.block_cg import block_cg
 from gpt.algorithms.inverter.bicgstab import bicgstab
 from gpt.algorithms.inverter.cagcr import cagcr
 from gpt.algorithms.inverter.fom import fom

lib/gpt/algorithms/inverter/block_cg.py

@@ -0,0 +1,145 @@
+#
+#    GPT - Grid Python Toolkit
+#    Copyright (C) 2024  Christoph Lehner ([email protected], https://github.com/lehner/gpt)
+#                        Adopted from Grid's BlockConjugateGradient
+#
+#    This program is free software; you can redistribute it and/or modify
+#    it under the terms of the GNU General Public License as published by
+#    the Free Software Foundation; either version 2 of the License, or
+#    (at your option) any later version.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU General Public License for more details.
+#
+#    You should have received a copy of the GNU General Public License along
+#    with this program; if not, write to the Free Software Foundation, Inc.,
+#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+#
+import gpt as g
+import numpy as np
+from gpt.algorithms import base_iterative
+
+
+def QC_is_R(Q, R):
+    # checked QC_is_R(Q, R) = C  such that QC = R
+    m_rr = g.inner_product(R, R)
+    m_rr = 0.5 * (m_rr + m_rr.transpose().conjugate())
+    L = np.linalg.cholesky(m_rr)
+    C = L.transpose().conjugate()
+    Cinv = np.linalg.inv(C)
+    g.linear_combination(Q, R, np.ascontiguousarray(Cinv.transpose()))
+    return C
+
+
+class block_cg(base_iterative):
+    @g.params_convention(eps=1e-15, maxiter=1000000, eps_abs=None, miniter=0)
+    def __init__(self, params):
+        super().__init__()
+        self.params = params
+        self.eps = params["eps"]
+        self.eps_abs = params["eps_abs"]
+        self.maxiter = params["maxiter"]
+        self.miniter = params["miniter"]
+
+    def modified(self, **params):
+        return block_cg({**self.params, **params})
+
+    def __call__(self, mat):
+        vector_space = None
+        if isinstance(mat, g.matrix_operator):
+            vector_space = mat.vector_space
+            mat = mat.specialized_list_callable()
+
+        @self.timed_function
+        def inv(X, B, t):
+            nblock = len(B)
+
+            t("reductions")
+            ssq = g.norm2(B)
+
+            AD = [g.lattice(x) for x in B]
+            Q = [g.lattice(x) for x in B]
+            Z = [g.lattice(x) for x in B]
+
+            # QC = R = B-AX, D = Q
+            t("mat")
+            mat(AD, X)
+
+            t("linear")
+            tmp = g(g.expr(B) - g.expr(AD))
+
+            t("QR")
+            m_C = QC_is_R(Q, tmp)
+            D = g.copy(Q)
+
+            for k in range(self.maxiter):
+                # Z  = AD
+                t("mat")
+                mat(Z, D)
+
+                # M  = [D^dag Z]^{-1}
+                t("reduction")
+                m_DZ = g.inner_product(D, Z)
+
+                t("inverse")
+                m_M = np.linalg.inv(m_DZ)
+
+                # X += D MC
+                t("linear")
+                m_tmp = m_M @ m_C
+                g.linear_combination(tmp, D, np.ascontiguousarray(m_tmp.transpose()))
+                for i in range(nblock):
+                    X[i] += tmp[i]
+
+                # QS = Q - ZM = Q - tmp
+                g.linear_combination(tmp, Z, np.ascontiguousarray(m_M.transpose()))
+                for i in range(nblock):
+                    tmp[i] @= Q[i] - tmp[i]
+
+                t("QR")
+                m_S = QC_is_R(Q, tmp)
+
+                # D = Q + D S^dag
+                t("linear")
+                m_tmp = m_S.transpose().conjugate()
+                g.linear_combination(tmp, D, np.ascontiguousarray(m_tmp.transpose()))
+                for i in range(nblock):
+                    D[i] @= Q[i] + tmp[i]
+
+                # C = S C
+                m_C = m_S @ m_C
+
+                m_rr = m_C.transpose().conjugate() @ m_C
+
+                max_resid_rel = 0
+                max_resid_abs = 0
+                for b in range(nblock):
+                    rr = m_rr[b, b].real
+                    if rr > max_resid_abs:
+                        max_resid_abs = rr
+                    rr /= ssq[b]
+                    if rr > max_resid_rel:
+                        max_resid_rel = rr
+
+                self.log_convergence(k, max_resid_rel, self.eps**2.0)
+                if k + 1 >= self.miniter:
+                    if self.eps_abs is not None and max_resid_abs <= self.eps_abs**2.0:
+                        self.log(f"converged in {k+1} iterations (absolute criterion)")
+                        return
+                    if max_resid_rel <= self.eps**2.0:
+                        self.log(f"converged in {k+1} iterations")
+                        return
+
+            self.log(
+                f"NOT converged in {k+1} iterations;  squared resudial relative {max_resid_rel} and absolute {max_resid_abs}"
+            )
+
+        return g.matrix_operator(
+            mat=inv,
+            inv_mat=mat,
+            accept_guess=(True, False),
+            vector_space=vector_space,
+            accept_list=True,
+        )

lib/gpt/core/operator/matrix_operator.py

@@ -76,6 +76,10 @@ def specialized_singlet_callable(self):
         # removing possible overhead for specialized call
         return self.mat if not self.accept_list else self
 
+    def specialized_list_callable(self):
+        # removing possible overhead for specialized call
+        return self.mat if self.accept_list else self
+
     def inv(self):
         return matrix_operator(
             mat=self.inv_mat,

lib/gpt/qcd/scalar/action/__init__.py

@@ -18,4 +18,4 @@
 #    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 #
 from gpt.qcd.scalar.action.phi4 import phi4
-from gpt.qcd.scalar.action.mass_term import mass_term, fourier_mass_term
+from gpt.qcd.scalar.action.mass_term import mass_term, fourier_mass_term, general_mass_term

lib/gpt/qcd/scalar/action/mass_term.py

@@ -94,14 +94,18 @@ def __call__(self, pi):
     def draw(self, pi, rng):
         # P(pi) = e^{-sum(trace(pi^dag scale^2 fft^dag mass fft pi)) * 2 / 2}    # fft^dag fft = 1/scale^2 ; inv(fft) = fft^dag / scale^2
         rng.normal_element(pi)
+
+        value = g.group.inner_product(pi, pi) * 0.5
+
         pi_mom = [g(self.scale_unitary * self.fft * p) for p in pi]
         for mu in range(self.n):
             r = g.lattice(pi_mom[mu])
             r[:] = 0
             for nu in range(self.n):
                 r += self.fourier_inv_sqrt_mass_field[mu][nu] * pi_mom[nu]
             pi[mu] @= g.inv(self.fft) * r / self.scale_unitary
-        return self.__call__(pi)
+
+        return value
 
     @differentiable_functional.multi_field_gradient
     def gradient(self, pi, dpi):
@@ -121,3 +125,54 @@ def gradient(self, pi, dpi):
             dS.append(ret)
 
         return dS
+
+
+class general_mass_term(differentiable_functional):
+    def __init__(self, M, sqrt_M):
+        self.M = M
+        self.sqrt_M = sqrt_M
+        self.inv_sqrt_M = sqrt_M.inv()
+        # Need:
+        # - sqrt_M^dag = sqrt_M
+        # - sqrt_M^2 = M
+
+    def __call__(self, pi):
+        pi_prime = self.M(pi)
+        n = len(pi)
+        A = 0.0
+        for mu in range(n):
+            A += g.inner_product(pi[mu], pi_prime[mu])
+        return A.real
+
+    def draw(self, pi, rng):
+        # P(pi) = e^{-pi^dag sqrt_M^dag sqrt_M pi * 2 / 2}
+        rng.normal_element(pi)  # pi = sqrt(2) sqrt_M pi_desired -> pi_desired = inv_sqrt_M pi
+
+        n = len(pi)
+
+        value = g.group.inner_product(pi, pi) * 0.5
+
+        pi_prime = self.inv_sqrt_M(pi)
+        for mu in range(n):
+            pi[mu] @= pi_prime[mu]
+
+        return value
+
+    # U -> V(0) U V(1)^-1
+    # e^{i A} U -> V(0) e^{i A} U V(1)^{-1} -> A -> V(0) A V(0)^-1
+
+    @differentiable_functional.multi_field_gradient
+    def gradient(self, pi, dpi):
+        dS = []
+
+        pi_prime = self.M(pi)
+
+        for _pi in dpi:
+            mu = pi.index(_pi)
+
+            ret = pi_prime[mu]
+
+            ret = g(g.qcd.gauge.project.traceless_hermitian(ret))
+            dS.append(ret)
+
+        return dS

tests/qcd/scalar.py

@@ -35,7 +35,7 @@
 g.message(f"Regress fourier mass term: {eps}")
 assert eps < 1e-10
 
-# now test with general Hermitian mass matrix
+# now test with general Hermitian Fourier mass matrix
 L = grid.gdimensions
 scale_unitary = float(np.prod(L)) ** 0.5
 
@@ -55,7 +55,7 @@
 
 # test distribution of A0 draw
 r = A0.draw(U_mom, rng)
-assert abs(A0(U_mom) - r) < 1e-10
+assert abs(A0(U_mom) / r - 1) < 1e-10
 eps = g.group.defect(U_mom[0])
 g.message("Group defect:", eps)
 assert eps < 1e-10
@@ -69,7 +69,7 @@
 
 # test distribution of A1 draw P ~ e^{-momdag inv(FT) mass FT mom}
 r = A1.draw(U_mom, rng)
-assert abs(A1(U_mom) - r) < 1e-10
+assert abs(A1(U_mom) / r - 1) < 1e-10
 
 # group defect would be triggered if sqrt_mass does not have sqrt_mass[k] = sqrt_mass[-k]
 eps = g.group.defect(U_mom[0])
@@ -142,3 +142,51 @@
     eps = abs(x - 1)
     g.message("Compare variance:", x, eps)
     assert eps < 0.01
+
+# Test general mass term
+U = g.qcd.gauge.random(U_mom[0].grid, rng)
+
+
+def __slap(dst, src):
+    assert len(src) == 4
+    for nu in range(len(dst)):
+        dst[nu][:] = 0
+        for mu in range(4):
+            src_p = g.cshift(src[nu], mu, 1)
+            src_m = g.cshift(src[nu], mu, -1)
+            U_m = g.cshift(U[mu], mu, -1)
+            dst[nu] += (
+                1 / 16 * (U[mu] * src_p * g.adj(U[mu]) + g.adj(U_m) * src_m * U_m - 2 * src[nu])
+            )
+
+
+cg = g.algorithms.inverter.block_cg({"eps": 1e-12, "maxiter": 100})
+_slap = g.matrix_operator(__slap, accept_list=True)
+slap = g.matrix_operator(mat=_slap, inv_mat=cg(_slap), accept_list=True)
+slap2 = slap * slap
+
+# TODO: need stencil version of __slap
+
+A2 = g.qcd.scalar.action.general_mass_term(M=slap2, sqrt_M=slap)
+
+A2.assert_gradient_error(rng, U_mom, U_mom, 1e-3, 1e-8)
+
+# test distribution
+r = A2.draw(U_mom, rng)
+assert abs(A2(U_mom) / r - 1) < 1e-10
+eps = g.group.defect(U_mom[0])
+g.message("Group defect:", eps)
+assert eps < 1e-10
+
+ngen = len(U_mom[0].otype.generators(np.complex128))
+x = 0.0
+U_mom_prime = g(slap2 * U_mom)
+for mu in range(4):
+    x += (
+        g.sum(g(g.trace(g.adj(U_mom[mu]) * U_mom_prime[mu]) * 2.0 / ngen)).real
+        / U_mom[0].grid.gsites
+    )
+x /= 4
+eps = abs(x - 1)
+g.message("Compare variance:", x, eps)
+assert eps < 0.01