groundwork for phmc

lib/gpt/algorithms/group/__init__.py

@@ -19,3 +19,5 @@
 from gpt.algorithms.group.symmetric_functional import symmetric_functional
 from gpt.algorithms.group.locally_coherent_functional import locally_coherent_functional
 from gpt.algorithms.group.repeat_arguments_functional import repeat_arguments_functional
+from gpt.algorithms.group.polar_decomposition_functional import polar_decomposition_functional
+from gpt.algorithms.group.polar_regulator import polar_regulator

lib/gpt/algorithms/group/polar_decomposition_functional.py

@@ -0,0 +1,137 @@
+#
+#    GPT - Grid Python Toolkit
+#    Copyright (C) 2024  Christoph Lehner ([email protected], https://github.com/lehner/gpt)
+#
+#    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.core.group import differentiable_functional
+
+
+def decompose(w, phase0=None):
+    if phase0 is None:
+        phase0 = g.complex(w.grid)
+        phase0[:] = 0
+
+    h, u = g.matrix.polar.decompose(w)
+    rel_det = g(g.matrix.det(g.component.exp(-1j * phase0) * u))
+    rel_phase = g(g.component.log(rel_det) / 1j / u.otype.Ndim)
+    phase = g(phase0 + rel_phase)
+    su = g(g.component.exp(-1j * phase) * u)
+    su.otype = g.ot_matrix_su_n_fundamental_group(u.otype.Ndim)
+    return h, phase, su
+
+
+class polar_decomposition_functional(differentiable_functional):
+    def __init__(self, u_functional, h_functional):
+        self.u_functional = u_functional
+        self.h_functional = h_functional
+        self.reference_phases = None
+
+    def reduce(self, fields):
+        h = []
+        p = []
+        u = []
+        p0 = self.reference_phases if self.reference_phases is not None else [None] * len(fields)
+        for i in range(len(fields)):
+            hf, pf, uf = decompose(fields[i], p0[i])
+            h.append(hf)
+            p.append(pf)
+            u.append(uf)
+        self.reference_phases = p
+        return h, p, u
+
+    def __call__(self, fields):
+        h, p, u = self.reduce(fields)
+        return self.u_functional(u) + self.h_functional(h)
+
+    def gradient(self, fields, dfields):
+        indices = [fields.index(f) for f in dfields]
+
+        h, p, u = self.reduce(fields)
+
+        u_gradient = self.u_functional.gradient(u, [u[mu] for mu in indices])
+        h_gradient = self.h_functional.gradient(h, [h[mu] for mu in indices])
+
+        dSdA = []
+
+        grid = fields[0].grid
+        lsites = grid.gsites // grid.Nprocessors
+        d_cartesian_space = g.group.cartesian(fields[0])
+        cartesian_space = g.group.cartesian(u[0])
+        gen = cartesian_space.otype.generators(grid.precision.complex_dtype)
+        for nu, mu in enumerate(indices):
+            # fill jacobian
+            Nc = fields[0].otype.Ndim
+            N = Nc**2 * 2
+            Na = len(gen)
+            jac = np.ndarray(shape=(lsites, N, N), dtype=np.float64)
+            for a in range(Na):
+                ta = gen[a]
+                ta_u = g(ta * u[mu])
+                h_ta_u = g(h[mu] * ta_u)
+                eitheta_h_ta_u = g(g.component.exp(1j * p[mu]) * h_ta_u)
+                eitheta_ta_u = g(g.component.exp(1j * p[mu]) * ta_u)
+
+                v_eitheta_h_ta_u = eitheta_h_ta_u[:]
+                v_eitheta_ta_u = eitheta_ta_u[:]
+                for i in range(Nc):
+                    for j in range(Nc):
+                        jac[:, 0 * Na + a, 0 * Nc * Nc + i * Nc + j] = -v_eitheta_h_ta_u[
+                            :, i, j
+                        ].imag
+                        jac[:, 0 * Na + a, 1 * Nc * Nc + i * Nc + j] = v_eitheta_h_ta_u[
+                            :, i, j
+                        ].real
+                        jac[:, 1 * Na + a, 0 * Nc * Nc + i * Nc + j] = v_eitheta_ta_u[:, i, j].real
+                        jac[:, 1 * Na + a, 1 * Nc * Nc + i * Nc + j] = v_eitheta_ta_u[:, i, j].imag
+
+            v_w = fields[mu][:]
+            eitheta_u = g(g.component.exp(1j * p[mu]) * u[mu])
+            v_eitheta_u = eitheta_u[:]
+            for i in range(Nc):
+                for j in range(Nc):
+                    jac[:, 2 * Na + 0, 0 * Nc * Nc + i * Nc + j] = -v_w[:, i, j].imag
+                    jac[:, 2 * Na + 0, 1 * Nc * Nc + i * Nc + j] = v_w[:, i, j].real
+                    jac[:, 2 * Na + 1, 0 * Nc * Nc + i * Nc + j] = v_eitheta_u[:, i, j].real
+                    jac[:, 2 * Na + 1, 1 * Nc * Nc + i * Nc + j] = v_eitheta_u[:, i, j].imag
+
+            inv_jac = np.linalg.inv(jac)
+
+            # next, project out each a
+            gr_w = np.zeros(shape=(lsites, 2 * Nc * Nc), dtype=np.complex128)
+            for a in range(Na):
+                u_gradient_a = g(g.trace(gen[a] * u_gradient[nu]))
+                v_u_gradient_a = u_gradient_a[:]
+
+                h_gradient_a = g(g.trace(gen[a] * h_gradient[nu]))
+                v_h_gradient_a = h_gradient_a[:]
+
+                gr_w += inv_jac[:, :, 0 * Na + a] * v_u_gradient_a.real
+                gr_w += inv_jac[:, :, 1 * Na + a] * v_h_gradient_a.real / 2.0
+
+            h_gradient_a = g(g.trace(h_gradient[nu]))
+            v_h_gradient_a = h_gradient_a[:]
+            gr_w += inv_jac[:, :, 2 * Na + 1] * v_h_gradient_a.real / 2.0
+
+            x = gr_w[:, 0 : Nc * Nc].reshape(lsites, Nc, Nc)
+            x += 1j * gr_w[:, Nc * Nc : 2 * Nc * Nc].reshape(lsites, Nc, Nc)
+            y = g.lattice(d_cartesian_space)
+            y[:] = np.ascontiguousarray(x)
+            dSdA.append(g(2 * y))
+
+        return dSdA

lib/gpt/algorithms/group/polar_regulator.py

@@ -0,0 +1,46 @@
+#
+#    GPT - Grid Python Toolkit
+#    Copyright (C) 2024  Christoph Lehner ([email protected], https://github.com/lehner/gpt)
+#
+#    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
+from gpt.core.group import differentiable_functional
+
+
+class polar_regulator(differentiable_functional):
+    def __init__(self, lam, kap):
+        self.lam = lam
+        self.kap = kap
+
+    def __call__(self, fields):
+        I = g.identity(fields[0])
+        Nc = fields[0].otype.Ndim
+        r = 0.0
+        for mu in range(len(fields)):
+            r += (self.lam / 2 / Nc) * g.sum(g.trace((fields[mu] - I) * (fields[mu] - I))).real
+            r -= (self.kap / Nc) * g.sum(g.component.log(g.matrix.det(fields[mu]))).real
+        return r
+
+    def gradient(self, fields, dfields):
+        # log(det(A + dA)) = log(det(A(1+inv(A)dA))) = log(det(A)) + log(1+tr(invA dA))
+        #                  = log(det(A)) + tr(invA dA)
+        dAdS = []
+        I = g.identity(dfields[0])
+        Nc = fields[0].otype.Ndim
+        for df in dfields:
+            x = g(2.0 * (self.lam / 2 / Nc) * (df - I) - (self.kap / Nc) * g.matrix.inv(df))
+            dAdS.append(x)
+        return dAdS

lib/gpt/core/matrix/__init__.py

@@ -21,3 +21,4 @@
 from gpt.core.matrix.exp import exp
 from gpt.core.matrix.log import log
 from gpt.core.matrix.det import det
+import gpt.core.matrix.polar

lib/gpt/core/matrix/polar.py

@@ -0,0 +1,40 @@
+#
+#    GPT - Grid Python Toolkit
+#    Copyright (C) 2024  Christoph Lehner ([email protected], https://github.com/lehner/gpt)
+#
+#    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
+
+
+def angle(w):
+    # Heron's method
+    uk = g(w / g.norm2(w) ** 0.5 * 1e3)
+    I = g.identity(w)
+    nrm = g.norm2(I)
+    for i in range(20):
+        uk = g(0.5 * uk + g.matrix.inv(g.adj(uk)) * 0.5)
+        err2 = g.norm2(uk * g.adj(uk) - I) / nrm
+        if err2 < w.grid.precision.eps**2 * 10:
+            return uk
+    raise Exception("angle did not converge")
+
+
+def decompose(w):
+    u = angle(w)
+    h = g(w * g.adj(u))
+    err2 = g.norm2(h - g.adj(h)) / g.norm2(h)
+    assert err2 < w.grid.precision.eps**2 * 100
+    return h, u

lib/gpt/core/object_type/__init__.py

@@ -92,6 +92,10 @@ def matrix_complex_additive(grid, n):
     return gpt_object(grid, ot_matrix_complex_additive_group(n))
 
 
+def matrix_color_complex_additive(grid, n):
+    return gpt_object(grid, ot_matrix_color_complex_additive_group(n))
+
+
 def real_additive(grid):
     return gpt_object(grid, ot_real_additive_group())
 
@@ -162,6 +166,7 @@ def str_to_otype(s):
             "ot_matrix_real_additive_group",
             "ot_vector_complex_additive_group",
             "ot_matrix_complex_additive_group",
+            "ot_matrix_color_complex_additive_group",
         ]
     )
 

lib/gpt/core/object_type/complex_additive_group.py

@@ -20,7 +20,13 @@
 import numpy
 
 # need a basic container
-from gpt.core.object_type import ot_singlet, ot_matrix_singlet, ot_vector_singlet
+from gpt.core.object_type import (
+    ot_singlet,
+    ot_matrix_singlet,
+    ot_vector_singlet,
+    ot_matrix_color,
+    ot_vector_color,
+)
 
 
 ###
@@ -205,3 +211,72 @@ def defect(self, U):
 
     def project(self, U, method):
         return None
+
+
+class ot_matrix_color_complex_additive_group(ot_matrix_color):
+    def __init__(self, n):
+        self.Ndim = n
+        super().__init__(n)
+        self.__name__ = f"ot_matrix_color_complex_additive_group({n})"
+        self.data_alias = lambda: ot_matrix_color(n)
+        self.vector_type = ot_vector_color(n)
+        self.mtab = {
+            self.__name__: (lambda: self, (1, 0)),
+            f"ot_vector_color({n})": (
+                lambda: ot_vector_color(n),
+                (1, 0),
+            ),
+            "ot_singlet": (lambda: self, None),
+            "ot_complex_additive_group": (lambda: self, None),
+        }
+        self.rmtab = {
+            "ot_singlet": (lambda: self, None),
+            "ot_complex_additive_group": (lambda: self, None),
+        }
+
+    def compose(self, a, b):
+        return a + b
+
+    def cartesian(self):
+        return self
+
+    def generators(self, dt):
+        n = self.shape[0]
+
+        def basis_real(i, j):
+            m = numpy.zeros(self.shape, dtype=dt)
+            m[i, j] = 1.0
+            return gpt.matrix_color(m, n)
+
+        def basis_imag(i, j):
+            m = numpy.zeros(self.shape, dtype=dt)
+            m[i, j] = 1.0j
+            return gpt.matrix_color(m, n)
+
+        return [basis_real(i, j) for i in range(n) for j in range(n)] + [
+            basis_imag(i, j) for i in range(n) for j in range(n)
+        ]
+
+    def inner_product(self, left, right):
+        return gpt.sum(gpt.trace(gpt.adj(left) * right)).real
+
+    def coordinates(self, l, c=None):
+        assert l.otype.__name__ == self.__name__
+        gen = self.generators(l.grid.precision.complex_dtype)
+        if c is None:
+            nhalf = len(gen) // 2
+            l_real = gpt.component.real(l)
+            l_imag = gpt.component.imag(l)
+            return [gpt.eval(gpt.trace(gpt.adj(l_real) * Ta)) for Ta in gen[0:nhalf]] + [
+                gpt.eval(gpt.trace(gpt.adj(l_imag) * Ta)) for Ta in gen[0:nhalf]
+            ]
+        else:
+            l[:] = 0
+            for ca, Ta in zip(c, gen):
+                l += ca * Ta
+
+    def defect(self, U):
+        return 0.0
+
+    def project(self, U, method):
+        return None

tests/algorithms/optimize.py

@@ -91,3 +91,13 @@ def gradient(self, fields, dfields):
 V1 = g.lattice(cgrid, V0.otype)
 rng.element([U1, V1])
 lc.assert_gradient_error(rng, [U0, V0, U1, V1], [U0, U1], 1e-4, 1e-10)
+
+# test polar decomposition functional
+r = g.algorithms.group.polar_regulator(1.3, 0.6)
+f = g.qcd.gauge.action.wilson(5.3)
+s = g.algorithms.group.polar_decomposition_functional(f, r)
+W = [g.matrix_color_complex_additive(grid, 3) for _ in range(4)]
+rng.element(W)
+for w in W:
+    w += g.identity(w) * 2
+s.assert_gradient_error(rng, W, W, 1e-4, 1e-8)

tests/core/matrix.py

@@ -67,6 +67,26 @@ def mod0p1(x):
 
 # test inv
 for grid, eps in [(grid_dp, 1e-14), (grid_sp, 1e-6)]:
+    g.message(
+        f"""
+
+    Test polar.decomposition for {grid.precision.__name__}
+
+"""
+    )
+    rng = g.random("test")
+    W = rng.normal_element(g.matrix_color_complex_additive(grid, 3))
+    H, U = g.matrix.polar.decompose(W)
+    err2 = g.norm2(H * U - W) / g.norm2(W)
+    g.message(f"Polar decomposition closure: {err2}")
+    assert err2 < eps**2
+    err2 = g.norm2(H - g.adj(H)) / g.norm2(H)
+    g.message(f"Polar decomposition H: {err2}")
+    assert err2 < eps**2
+    err2 = g.norm2(U * g.adj(U) - g.identity(U)) / g.norm2(U)
+    g.message(f"Polar decomposition U: {err2}")
+    assert err2 < eps**2
+
     g.message(
         f"""