ad-ft-hmc framework

lib/gpt/ad/forward/foundation.py → lib/gpt/ad/forward/foundation/__init__.py

@@ -17,7 +17,7 @@
 #    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 #
 import gpt as g
-
+import gpt.ad.forward.foundation.matrix
 
 python_sum = sum
 
@@ -124,7 +124,7 @@ def df(x, dx, maxn):
             tr_adjx2 = g.trace(adjx2)
             v += v0 * (tr_adjx * tr_adjx - tr_adjx2) / 2.0
         if maxn >= 4:
-            raise Exception(f"{maxn-1}-derivative of g.matrix.det not yet implemented")
+            raise Exception(f"{maxn - 1}-derivative of g.matrix.det not yet implemented")
         v.otype = v0.otype
         return v
 

lib/gpt/ad/forward/foundation/matrix/__init__.py

@@ -0,0 +1,19 @@
+#
+#    GPT - Grid Python Toolkit
+#    Copyright (C) 2023  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.ad.forward.foundation.matrix.exp

lib/gpt/ad/forward/foundation/matrix/exp.py

@@ -0,0 +1,43 @@
+#
+#    GPT - Grid Python Toolkit
+#    Copyright (C) 2023  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
+
+
+def function(x):
+    fac = 1.0
+    base = 128.0
+    nbase = 7
+    x = x / base
+    c = x
+    r = g.identity(x)
+    for i in range(1, 10):
+        fac /= float(i)
+        r = r + fac * c
+        c = c * x
+    for i in range(nbase):
+        r = r * r
+    return r
+
+
+# gives 1e-14 / 1e-15 errors up to at least |x| < 10
+
+
+def function_and_gradient(x, dx):
+    assert False

lib/gpt/qcd/gauge/smear/differentiable.py

@@ -34,8 +34,7 @@ def __call__(self, fields):
     def jacobian(self, fields, fields_prime, dfields):
         N = len(fields_prime)
         assert len(fields) == N
-        assert len(fields) == N
-        assert len(fields) == N
+        assert len(dfields) == N
         aU_prime = [g(2j * dfields[mu] * fields_prime[mu]) for mu in range(N)]
         for mu in range(N):
             self.aU[mu].value = fields[mu]
@@ -50,55 +49,111 @@ def jacobian(self, fields, fields_prime, dfields):
 
         return gradient
 
+    def adjoint_jacobian(self, fields, dfields):
+        N = len(fields)
+        assert len(dfields) == N
+
+        fad = g.ad.forward
+        eps = fad.infinitesimal("eps")
+
+        On = fad.landau(eps**2)
+        aU = [fad.series(fields[mu], On) for mu in range(N)]
+        for mu in range(N):
+            aU[mu][eps] = g(1j * dfields[mu] * fields[mu])
+
+        aU_prime = self.ft(aU)
+
+        gradient = [g(aU_prime[mu][eps] * g.adj(aU_prime[mu][1]) / 1j) for mu in range(N)]
+
+        return gradient
+
 
 class dft_action_log_det_jacobian(differentiable_functional):
-    def __init__(self, U, ft, dfm, inverter):
+    def __init__(self, U, ft, dfm, inverter_force, inverter_action):
         self.dfm = dfm
-        self.inverter = inverter
+        self.inverter_force = inverter_force
+        self.inverter_action = inverter_action
         self.N = len(U)
         mom = [g.group.cartesian(u) for u in U]
         rad = g.ad.reverse
 
         _U = [rad.node(g.copy(u)) for u in U]
-        _mom = [rad.node(g.copy(u)) for u in mom]
+        _left = [rad.node(g.copy(u), with_gradient=False) for u in mom]
+        _right = [rad.node(g.copy(u), with_gradient=False) for u in mom]
         _Up = dfm(_U)
-        momp = dfm.jacobian(_U, _Up, _mom)
+        J_right = dfm.jacobian(_U, _Up, _right)
 
         act = None
         for mu in range(self.N):
             if mu == 0:
-                act = g.norm2(momp[mu])
+                act = g.inner_product(_left[mu], J_right[mu])
             else:
-                act = g(act + g.norm2(momp[mu]))
+                act = g(act + g.inner_product(_left[mu], J_right[mu]))
+
+        self.left_J_right = act.functional(*(_U + _left + _right))
 
-        self.action = act.functional(*(_U + _mom))
+    def mat_J(self, U, U_prime):
+        def _mat(dst_5d, src_5d):
+            src = g.separate(src_5d, dimension=0)
+            dst = self.dfm.jacobian(U, U_prime, src)
+            dst_5d @= g.merge(dst, dimension=0)
+
+        return _mat
+
+    def mat_Jdag(self, U):
+        def _mat(dst_5d, src_5d):
+            src = g.separate(src_5d, dimension=0)
+            dst = self.dfm.adjoint_jacobian(U, src)
+            dst_5d @= g.merge(dst, dimension=0)
+
+        return _mat
 
     def __call__(self, fields):
-        return self.action(fields)
+        U = fields[0 : self.N]
+        mom = fields[self.N :]
+        mom_xd = g.merge(mom, dimension=0)
+        U_prime = [g(x) for x in self.dfm.ft(U)]
+        mom_prime_xd = self.inverter_action(self.mat_J(U, U_prime))(mom_xd)
+        mom_prime2_xd = self.inverter_action(self.mat_Jdag(U))(mom_prime_xd)
+        return g.inner_product(mom_xd, mom_prime2_xd).real
 
     def gradient(self, fields, dfields):
-        return self.action.gradient(fields, dfields)
-
-    def draw(self, fields, rng):
         U = fields[0 : self.N]
         mom = fields[self.N :]
-        assert len(mom) == self.N
-        assert len(U) == self.N
 
-        rng.normal_element(mom, scale=1.0)
+        # all eigenvalues of J need to be positive because for trivial FT they are all 1 and we
+        # stay invertible so none of them can have crossed a zero
 
-        U_prime = self.dfm(U)
+        # M = J J^dag
+        # ->  S = pi^dag J^{-1}^dag J^{-1} pi
+        # -> dS = -pi^dag J^{-1}^dag dJ^dag J^{-1}^dag J^{-1} pi
+        #         -pi^dag J^{-1}^dag J^{-1} dJ J^{-1} pi
+        #       = -pi^dag J^{-1}^dag J^{-1} dJ J^{-1} pi + C.C.  ->  need two inverse of J
+        assert dfields == U
 
-        def _mat(dst_5d, src_5d):
-            src = g.separate(src_5d, dimension=0)
-            dst = self.dfm.jacobian(U, U_prime, src)
-            dst_5d @= g.merge(dst, dimension=0)
+        U_prime = [g(x) for x in self.dfm.ft(U)]
 
         mom_xd = g.merge(mom, dimension=0)
+        mom_prime_xd = self.inverter_force(self.mat_J(U, U_prime))(mom_xd)
+        mom_prime2_xd = self.inverter_force(self.mat_Jdag(U))(mom_prime_xd)
 
-        mom_prime_xd = self.inverter(_mat)(mom_xd)
+        mom_prime2 = g.separate(mom_prime2_xd, dimension=0)
         mom_prime = g.separate(mom_prime_xd, dimension=0)
 
+        gradients = self.left_J_right.gradient(U + mom_prime2 + mom_prime, U)
+        return [g(-x - g.adj(x)) for x in gradients]
+
+    def draw(self, fields, rng):
+        U = fields[0 : self.N]
+        mom = fields[self.N :]
+        assert len(mom) == self.N
+        assert len(U) == self.N
+
+        rng.normal_element(mom, scale=1.0)
+        U_prime = [g(x) for x in self.dfm.ft(U)]
+
+        mom_prime = self.dfm.jacobian(U, U_prime, mom)
+
         act = 0.0
         for mu in range(self.N):
             act += g.norm2(mom[mu])
@@ -108,11 +163,12 @@ def _mat(dst_5d, src_5d):
 
 
 class differentiable_field_transformation:
-    def __init__(self, U, ft, inverter, optimizer):
+    def __init__(self, U, ft, inverter_force, inverter_action, optimizer):
         self.ft = ft
         self.U = U
         self.dfm = dft_diffeomorphism(self.U, self.ft)
-        self.inverter = inverter
+        self.inverter_force = inverter_force
+        self.inverter_action = inverter_action
         self.optimizer = optimizer
 
     def diffeomorphism(self):
@@ -129,4 +185,6 @@ def inverse(self, Uft):
         return U
 
     def action_log_det_jacobian(self):
-        return dft_action_log_det_jacobian(self.U, self.ft, self.dfm, self.inverter)
+        return dft_action_log_det_jacobian(
+            self.U, self.ft, self.dfm, self.inverter_force, self.inverter_action
+        )

tests/qcd/gauge.py

@@ -304,7 +304,7 @@
 
 
 # test general differentiable field transformation framework
-ft_stout = g.qcd.gauge.smear.differentiable_stout(rho=0.05)
+ft_stout = g.qcd.gauge.smear.differentiable_stout(rho=0.01)
 
 fr = g.algorithms.optimize.fletcher_reeves
 ls2 = g.algorithms.optimize.line_search_quadratic
@@ -313,7 +313,9 @@
     U,
     ft_stout,
     # g.algorithms.inverter.fgmres(eps=1e-15, maxiter=1000, restartlen=60),
-    g.algorithms.inverter.fgcr(eps=1e-15, maxiter=1000, restartlen=60),
+    # g.algorithms.inverter.fgmres(eps=1e-15, maxiter=1000, restartlen=60),
+    g.algorithms.inverter.fgcr(eps=1e-13, maxiter=1000, restartlen=60),
+    g.algorithms.inverter.fgcr(eps=1e-13, maxiter=1000, restartlen=60),
     g.algorithms.optimize.non_linear_cg(
         maxiter=1000, eps=1e-15, step=1e-1, line_search=ls2, beta=fr
     ),
@@ -323,7 +325,7 @@
 ald = dft.action_log_det_jacobian()
 
 # test diffeomorphism of stout against reference implementation
-dfm_ref = g.qcd.gauge.smear.stout(rho=0.05)
+dfm_ref = g.qcd.gauge.smear.stout(rho=0.01)
 Uft = dfm(U)
 Uft_ref = dfm_ref(U)
 for mu in range(4):
@@ -350,7 +352,7 @@
 mom2 = g.copy(mom)
 g.message("Action log det jac:", ald(U + mom2))
 
-ald.assert_gradient_error(rng, U + mom2, U + mom2, 1e-3, 1e-7)
+ald.assert_gradient_error(rng, U + mom2, U, 1e-3, 1e-7)
 
 act = ald.draw(U + mom2, rng)
 act2 = ald(U + mom2)