La^2 for each a

Author: simone-romiti

links/canonical_momenta.py

@@ -76,6 +76,7 @@ def __init__(self, U: GaugeConfiguration):
         """
         self.Nc = U.Nc # number of colors of the group
         self.Ng = U.Ng # number of generators in the algebra
+        self.n_links = U.n_links # number of gauge links on the lattice
         self.batchsize = U.batch_size #  number of configurations
         self.tau = suN.get_generators(Nc=self.Nc, device=U.device, dtype=U.dtype)
 
@@ -209,10 +210,28 @@ def La_chain_rule(self, f: typing.Callable, U: GaugeConfiguration):
         f_U = f(GaugeConfiguration(U + delta)) # f(U+delta)
         # NOTE: I can differentiate the sum over configurations
         # because f(U) acts configuration-wise
-        f_U_flat = f_U.sum() + 0.0*1j # flattened view
+        # f_U_flat = f_U.sum() + 0.0*1j # flattened view
         A = -1j * torch.einsum("aij,B...jk->B...aik", self.tau, U.as_subclass(torch.Tensor)) # d(e^{-i*omega*tau_a})/domega at omega==0
         dRef_dU = my_autograd(y=f_U.real, x=delta, grad_outputs=torch.ones_like(f_U.real), create_graph=True, retain_graph=True).unsqueeze(-3)
         dImf_dU = my_autograd(y=f_U.imag, x=delta, grad_outputs=torch.ones_like(f_U.imag), create_graph=True, retain_graph=True).unsqueeze(-3)
         df_domega = chain_rule_contributions(A=A, dRef_dU=dRef_dU, dImf_dU=dImf_dU).sum(dim=(-2,-1)) # summing over the color components
         La_f = -1j * df_domega
         return La_f
+
+    # def LaLa_chain_rule_EXPERIMENTAL(self, f: typing.Callable, U: GaugeConfiguration):
+    #     delta1 = torch.zeros(size=U.shape, dtype=U.dtype, device=U.device, requires_grad=True) # Note: delta should be complex
+    #     delta2 = torch.zeros(size=U.shape, dtype=U.dtype, device=U.device, requires_grad=True) # Note: delta should be complex
+    #     Up = U + delta1 + delta2
+    #     f_U = f(GaugeConfiguration(Up)) # f(U+delta)
+    #     A = -1j * torch.einsum("aij,B...jk->B...aik", self.tau, Up.as_subclass(torch.Tensor)) # d(e^{-i*omega*tau_a})/domega at omega==0
+    #     dRef_dU = my_autograd(y=f_U.real, x=delta, grad_outputs=torch.ones_like(f_U.real), create_graph=True, retain_graph=True).unsqueeze(-3)
+    #     dImf_dU = my_autograd(y=f_U.imag, x=delta, grad_outputs=torch.ones_like(f_U.imag), create_graph=True, retain_graph=True).unsqueeze(-3)
+    #     df_domega_ab = chain_rule_contributions(A=A, dRef_dU=dRef_dU, dImf_dU=dImf_dU).sum(dim=(-2,-1)) # summing over the color components
+    #     Laf_ab = -1j * df_domega
+    #     # print(dRef_dU.shape, Laf_ab.shape, delta.shape)
+    #     cosd_Laf = torch.einsum("...ij,...aij->...aij", torch.cos(delta.abs()), Laf_ab)
+    #     dReLaf_dU = my_autograd(y=cosd_Laf.real, x=delta, grad_outputs=torch.ones_like(cosd_Laf.real), create_graph=True, retain_graph=True).unsqueeze(-3)
+    #     dImLaf_dU = my_autograd(y=cosd_Laf.imag, x=delta, grad_outputs=torch.ones_like(cosd_Laf.imag), create_graph=True, retain_graph=True).unsqueeze(-3)
+    #     d2f_domega = chain_rule_contributions(A=A, dRef_dU=dReLaf_dU, dImf_dU=dImLaf_dU).sum(dim=(-2,-1)) # summing over the color components
+    #     LaLa_f = -1j * d2f_domega
+    #     return LaLa_f

links/canonical_momenta_squared.py

@@ -55,10 +55,10 @@ def get_La2_per_link(self, f: typing.Callable, U: GaugeConfiguration):
             f(..., e^{-i \\omega \\tau_a } U(x,\\mu), ...)|_{\\omega = 0}
         $$
         """
-        Nc = U.Nc # number of colors
-        Ng = U.Ng # number of generators in the Lie algebra
-        n_links = U.n_links # number of links
-        batchsize = U.batch_size # number of configurations
+        Nc = self.Nc # number of colors
+        Ng = self.Ng # number of generators in the Lie algebra
+        n_links = self.n_links # number of links
+        batchsize = self.batchsize # number of configurations
         Id = torch.eye(Nc).to(device=U.device)
         Id_arr = Id.expand(n_links, Nc, Nc)
         omega = torch.tensor(0.0, requires_grad=True, dtype=U.real.dtype, device=U.device)
@@ -100,7 +100,7 @@ def f_i(i):
             #---
             La2_per_link.append(torch.stack(sum_La_squared, dim=0))
         #---
-        return torch.stack(La2_per_link, dim=1)
+        return -torch.stack(La2_per_link, dim=1)
 
     def with_La_twice(self, f: typing.Callable, U: GaugeConfiguration) -> torch.Tensor:
         """
@@ -210,12 +210,12 @@ def f_i(i):
                 f_minus = make_perturbed_U(Va_minus[a,:,:]) # sum_i f(Va(-eps) . U_i)
 
                 dir_der = (f_plus - f_minus)/(2.0*eps) # 1st derivative found with finite difference
-                laplacian_b = my_autograd(dir_der, omega, create_graph=False, retain_graph=True) # 2nd derivative through autodifferentiation
+                laplacian_b = my_autograd(dir_der, omega, grad_outputs=torch.ones_like(dir_der), create_graph=False, retain_graph=True) # 2nd derivative through autodifferentiation
                 sum_La_squared.append(laplacian_b)
             #---
             La2_per_link.append(torch.stack(sum_La_squared, dim=0))
         #---
-        return torch.stack(La2_per_link, dim=1)
+        return -torch.stack(La2_per_link, dim=1) # `-1` factor coming from the product of the 2 `-i` in front of the L_a
 
 
     def TODO_get_sum_La_squared_per_link_fast(self, a: int, f: typing.Callable, U: GaugeConfiguration, f_is_real: bool, eps: float = 1e-8):

links/lie_derivatives.py

@@ -16,6 +16,9 @@
 import typing
 
 import torch
+from torch._functorch.vmap import _flat_vmap
+from torch.distributions import laplace
+from torch.nn.functional import relu
 
 import lattice_data_tools.links.suN as suN
 from lattice_data_tools.links.configuration import GaugeConfiguration, ColorMatrix
@@ -154,56 +157,65 @@ def L_a(self, a: int, f: typing.Callable, U: GaugeConfiguration, f_is_real: bool
         return -1j*df_domega
 
 
-    def La_squared_per_link(self, a: int, f: typing.Callable, U: GaugeConfiguration, f_is_real: bool):
+    def La_squared_per_link(self, f: typing.Callable, U: GaugeConfiguration, f_is_real: bool):
         """
         Returns ${ \\sum_{x,\\mu} L_a^2(x,\\mu) f(U) }$
         NOTE: the index `a` is NOT summed over.
         
         `f(U)` should return a tensor of shape (batchsize,1)
         """
+        batchsize = U.batch_size # number of configurations
         Nc = U.Nc # number of colors
+        Ng = U.Ng
         n_links = U.n_links # number of links
-        tau_a = self.tau[a,:,:]
-        d, M = torch.linalg.eigh(tau_a)  # diagonalization of the generator tau_a
+        #tau_a = self.tau[a,:,:]
+        d, M = torch.linalg.eigh(self.tau)  # diagonalization of the generator tau_a
         omega = torch.tensor(0.0, requires_grad=True, dtype=U.real.dtype, device=U.device)
-        phase = omega * d
+        phase = omega*d
         exp_iphase = torch.exp(-1j * phase)
         exp_iD = torch.diag_embed(exp_iphase)
-        Va = (M @ exp_iD @ M.adjoint()) #.expand(n_links, Nc, Nc)
-        batchsize = U.batch_size # number of configurations
+        V = (M @ exp_iD @ M.adjoint()) #.expand(n_links, Nc, Nc)
         Id = torch.eye(Nc).to(device=U.device)
         Id_arr = Id.expand(n_links, Nc, Nc)
-        sum_La_squared = []
-        for b in range(batchsize):
-            U_b = U[b,...].reshape(n_links, Nc, Nc)
-            def f_i(i):
-                e_i = torch.nn.functional.one_hot(i, n_links).to(dtype=U.dtype, device=U.device)
-                Va_arr = Id_arr + torch.einsum("ab,i->iab", Va-Id, e_i)
-                VaU_i = (Va_arr @ U_b).reshape(*(U[b,...].shape)).unsqueeze(0)
-                f_VaU = f(GaugeConfiguration(VaU_i))/n_links
-                return f_VaU
-            #---
-            # vectorize over all indices 0..N-1
-            indices = torch.arange(n_links)
-            sum_f = torch.vmap(f_i)(indices).sum()
+        single_conf_shape = (1, *U.shape[1:])
+        def f_bai(Ub, Va, i):
+            Ub_flat = Ub.reshape(n_links, Nc, Nc)
+            e_i = torch.nn.functional.one_hot(i, n_links).to(dtype=U.dtype, device=U.device)
+            Va_arr = Id_arr + torch.einsum("ab,i->iab", Va-Id, e_i)
+            VaU_i = (Va_arr @ Ub_flat).reshape(*(single_conf_shape))
+            f_VaU = f(GaugeConfiguration(VaU_i))/n_links
+            return f_VaU
+        #---
+        f_vmap = torch.func.vmap(
+            torch.func.vmap(
+                torch.func.vmap(
+                    f_bai,
+                    in_dims=(None, None, 0) # over link index
+                ),
+                in_dims=(None, 0, None) # over generators
+            ),
+            in_dims=(0, None, None) # over configurations
+        )
+        f_arr = f_vmap(U.as_subclass(torch.Tensor), V, torch.arange(n_links))
+        sum_f = f_arr.sum()
+        
 
-            # \\sum_i \\partial_{x_i} f : directional derivative along (1,...,1)
-            if f_is_real:
-                dir_der = torch.autograd.grad(sum_f, omega, create_graph=True)[0]
-                laplacian_b = torch.autograd.grad(dir_der, omega, create_graph=True)[0]
-            else:
-                Re_dir_der = torch.autograd.grad(sum_f.real, omega, create_graph=True)[0]
-                Im_dir_der = torch.autograd.grad(sum_f.imag, omega, create_graph=True)[0]
-                Re_laplacian_b = torch.autograd.grad(Re_dir_der, omega, create_graph=True)[0]
-                Im_laplacian_b = torch.autograd.grad(Im_dir_der, omega, create_graph=True)[0]
-                laplacian_b = Re_laplacian_b + 1j*Im_laplacian_b
-            #---
-            sum_La_squared.append(laplacian_b)
+        # \\sum_i \\partial_{x_i} f : directional derivative along (1,...,1)
+        if f_is_real:
+            dir_der = torch.autograd.grad(sum_f, omega, create_graph=True)[0]
+            laplacian_b = torch.autograd.grad(dir_der, omega, create_graph=True)[0]
+        else:
+            Re_dir_der = torch.autograd.grad(sum_f.real, omega, create_graph=True)[0]
+            Im_dir_der = torch.autograd.grad(sum_f.imag, omega, create_graph=True)[0]
+            Re_laplacian_b = torch.autograd.grad(Re_dir_der, omega, create_graph=True)[0]
+            Im_laplacian_b = torch.autograd.grad(Im_dir_der, omega, create_graph=True)[0]
+            laplacian_b = Re_laplacian_b + 1j*Im_laplacian_b
         #---
-        return torch.stack(sum_La_squared, dim=0)
+        La2_tensor = - laplacian_b # accounting for the two factors `i`
+        return La2_tensor
 
 
-    def La_squared_per_link_FD(self, a: int, f: typing.Callable, U: GaugeConfiguration, f_is_real: bool, eps: float = 1e-8):
+    def La_squared_per_link_FD(self, f: typing.Callable, U: GaugeConfiguration, f_is_real: bool, eps: float = 1e-8):
         """
         Returns ${ \\sum_{x,\\mu} L_a^2(x,\\mu) f(U) }$ via Finite Differences.
         NOTE: the index `a` is NOT summed over.
@@ -214,60 +226,65 @@ def La_squared_per_link_FD(self, a: int, f: typing.Callable, U: GaugeConfigurati
         `f(U)` should return a tensor of shape (batchsize,1)
         """
         Nc = U.Nc
+        Ng = U.Ng
         n_links = U.n_links
-        tau_a = self.tau[a, :, :]
-        batchsize = U.batch_size
         Id = torch.eye(Nc, dtype=U.dtype, device=U.device)
         Id_arr = Id.expand(n_links, Nc, Nc)
 
         # Compute Va(+eps) and Va(-eps): the group elements e^{±i eps tau_a}
         # Use matrix exponential via diagonalization: tau_a = M D M†, e^{i w tau_a} = M e^{i w D} M†
-        d, M = torch.linalg.eigh(tau_a)
+        d, M = torch.linalg.eigh(self.tau)
 
-        def make_Va(omega: float):
+        def make_V(omega: float):
             phase = omega * d
             exp_iD = torch.diag_embed(torch.exp(-1j * phase).to(dtype=U.dtype))
             return M @ exp_iD @ M.adjoint()
 
+        single_conf_shape = (1, *U.shape[1:])
+        def f_bai(Ub, Va, i):
+            Ub_flat = Ub.reshape(n_links, Nc, Nc)
+            # Apply Va to link i, leaving all other links unchanged.
+            e_i = torch.nn.functional.one_hot(i, n_links).to(dtype=U.dtype, device=U.device)
+            Va_arr = Id_arr + torch.einsum("ab,i->iab", Va - Id, e_i)
+            VaU_i = (Va_arr @ Ub_flat).reshape(*single_conf_shape)
+            return f(GaugeConfiguration(VaU_i))/n_links/Ng
+        #---
+  
+        f_vmap = torch.func.vmap(
+            torch.func.vmap(
+                torch.func.vmap(
+                    f_bai,
+                    in_dims=(None, None, 0) # over link index
+                ),
+                in_dims=(None, 0, None) # over generators
+            ),
+            in_dims=(0, None, None) # over configurations
+        )
+        
         omega = torch.tensor(0.0, requires_grad=True, dtype=U.real.dtype, device=U.device)
-        Va_plus  = make_Va(+eps+omega)   # e^{+i eps tau_a}
-        Va_minus = make_Va(-eps+omega)   # e^{-i eps tau_a}
-
-        sum_La_squared = []
-
-        for b in range(batchsize):
-            U_b = U[b, ...].reshape(n_links, Nc, Nc)
-            shape_b = U[b, ...].shape
-
-            def make_perturbed_U(Va):
-                """Apply Va to link i, leaving all other links unchanged."""
-                def f_i(i):
-                    e_i = torch.nn.functional.one_hot(i, n_links).to(dtype=U.dtype, device=U.device)
-                    Va_arr = Id_arr + torch.einsum("ab,i->iab", Va - Id, e_i)
-                    VaU_i = (Va_arr @ U_b).reshape(*shape_b).unsqueeze(0)
-                    return f(GaugeConfiguration(VaU_i))/n_links
-                #---
-                indices = torch.arange(n_links)
-                return torch.vmap(f_i)(indices).sum()  # sum over all links
-
-            f_plus  = make_perturbed_U(Va_plus)   # sum_i f(Va(+eps) . U_i)
-            f_minus = make_perturbed_U(Va_minus)  # sum_i f(Va(-eps) . U_i)
-            # f_0     = f(GaugeConfiguration(U[b, ...].unsqueeze(0))) * n_links  # n_links * f(U)
-
-            # Central difference: [f(+eps) - 2f(U) + f(-eps)] / eps^2
-            # laplacian_b = (f_plus - 2 * f_0 + f_minus) / (eps ** 2)
-            dir_der = (f_plus - f_minus)/(2.0*eps)
+        Va_p1 = make_V(+eps +omega)   # e^{+i eps tau_a}
+        Va_m1 = make_V(-eps +omega)   # e^{-i eps tau_a}
+        Va_0 = make_V(+omega)   # e^{-i eps tau_a}
+        f_plus  = f_vmap(U.as_subclass(torch.Tensor), Va_p1 , torch.arange(n_links)).sum()   # sum_i f(Va(+eps) . U_i)
+        f_minus = f_vmap(U.as_subclass(torch.Tensor), Va_m1, torch.arange(n_links)).sum()  # sum_i f(Va(-eps) . U_i)
+        # f_0     = f_vmap(U.as_subclass(torch.Tensor), Va_0, torch.arange(n_links)).sum()  # sum_i f(Va(-eps) . U_i)
 
-            if f_is_real:
-                laplacian_b = torch.autograd.grad(dir_der, omega)[0]
-            else:
-                Re_laplacian_b  = torch.autograd.grad(dir_der.real, omega, create_graph=True)[0]
-                Im_laplacian_b  = torch.autograd.grad(dir_der.imag, omega)[0]
-                laplacian_b  = Re_laplacian_b + 1j*Im_laplacian_b
+        # Central difference: [f(+eps) - 2f(U) + f(-eps)] / eps^2
+        #laplacian_b = (f_plus - 2 * f_0 + f_minus) / (eps ** 2)
+        # return laplacian_b
 
-            sum_La_squared.append(laplacian_b)
+        dir_der = (f_plus - f_minus)/(2.0*eps)
+        # dir_der = (f_plus - f_0)/eps
 
-        return torch.stack(sum_La_squared, dim=0)
+        if f_is_real:
+            laplacian_b = torch.autograd.grad(dir_der, omega)[0]
+        else:
+            Re_laplacian_b  = torch.autograd.grad(dir_der.real, omega, create_graph=True)[0]
+            Im_laplacian_b  = torch.autograd.grad(dir_der.imag, omega)[0]
+            laplacian_b  = Re_laplacian_b + 1j*Im_laplacian_b
+            
+        La2_tensor = - Ng * laplacian_b  # accounting for the two factors `i`
+        return La2_tensor
 
 
     def La_squared_per_link_FD_fast(self, a: int, f: typing.Callable, U: GaugeConfiguration, f_is_real: bool, eps: float = 1e-8):

machine_learning/lcnn.py

@@ -184,6 +184,13 @@ def gen_random_beta(self, N_out: int, seed: int):
         beta =  torch.rand(*(d, N_out)).to(self.dtype_U).to(self.device_U)
         return beta
     #---
+    def ibetaWah(self, W: LocallyGaugeCovariant, beta: torch.Tensor):
+        """ Argument of the exponential in `self.exp_ibetaWah()` """
+        # building the anti-hermitian part of W --> i*W_ah lies in the algebra su(N)
+        W_ah_traceless = ColorMatrix(W).get_ah_traceless()
+        arg_exp = torch.einsum("mi,...iab->...mab", 1j*beta, W_ah_traceless) # NOTE: arg_exp has to be hermitean --> we include the `1j`
+        return arg_exp
+    #---
     def exp_ibetaWah(self, W: LocallyGaugeCovariant, beta: torch.Tensor):
         """
         Eq. 9 of https://arxiv.org/pdf/2012.12901
@@ -193,9 +200,7 @@ def exp_ibetaWah(self, W: LocallyGaugeCovariant, beta: torch.Tensor):
         W: N_ch locally transforming variables (obtained after L_conv()). shape: (batch, L1, ..., Ld, N_ch, Nc, Nc)
         beta: shape: (d,N_ch)
         """
-        # building the anti-hermitian part of W --> i*W_ah lies in the algebra su(N)
-        W_ah_traceless = ColorMatrix(W).get_ah_traceless()
-        arg_exp = torch.einsum("mi,...iab->...mab", 1j*beta, W_ah_traceless) # NOTE: arg_exp has to be hermitean --> we include the ""1j"
+        arg_exp = self.ibetaWah(W=W, beta=beta)
         E = suN.get_exp_iA(arg_exp.to_tensor()) # eq. 9 of https://arxiv.org/pdf/2012.12901
         return LocallyGaugeCovariant(E)
     #---
@@ -245,6 +250,20 @@ def all_layers_with_CB_AND_Tr(self, U: GaugeConfiguration, omega_CB: torch.Tenso
         trace = suN.get_Tr(self.all_layers_with_CB(U=U, omega_CB=omega_CB,beta=beta))
         return trace
     #---
+    def all_layers_Laf(self, U: GaugeConfiguration, omega_CB: torch.Tensor, beta: torch.Tensor):
+        """
+        Network producing an output of shape (B, Ng, L1, ..., Ld, d)
+
+        Interpretation: for a given configuration `b` and fixed `a`, the outputs are the `tau_a L_a` applied to the function `f`.
+
+        """
+        U_PT = get_ParallelTransporters(U=U, K=self.K)
+        W = self.get_W(U=U) # set of locally transforming variables
+        Wprime = self.get_Wprime(U=U, W=W, U_PT=U_PT) # W' as in eq. III.11 of https://arxiv.org/pdf/2012.12901
+        W_CB = self.L_CB(W=W, Wprime=Wprime, omega_CB=omega_CB) # W after eq. 18 of https://arxiv.org/pdf/2401.06481
+        W_act = self.L_act(U=U, W=W_CB, act_fun=self.act_fun) # W after eq. 7 of https://arxiv.org/pdf/2012.12901
+        EU = self.L_exp(U=U, W=W_act, beta=beta)
+        W_res = self.get_W(U=EU) 
 #---
 
 

test/test-canonical_momenta.py

@@ -128,7 +128,8 @@ def perf(fun, info: str):
 CM = CanonicalMomenta(U=U)
 momenta_exp = perf(lambda: CM.LaRa_with_exp(f=f, U=U, f_is_real=f_is_real), "L_a & R_a arr from exp()")
 momenta_cr = perf(lambda: CM.LaRa_chain_rule(f=f, U=U), "L_a & R_a arr from chain rule")
-
+print(momenta_cr.shape)
+   
 
 # La_chain_rule_compiled = torch.compile(lambda U_conf: CM.La_chain_rule(f=f, U=U_conf.as_subclass(torch.Tensor)))
 # La_chain_rule_compiled(U)