implementing Lie derivatives. TODO: batching

Author: simone-romiti

links/configuration.py

@@ -14,14 +14,14 @@ class ColorMatrix(torch.Tensor):
     Tensor representing a set of color matrices:
     (..., Nc, Nc)
     """
-    @staticmethod
-    def __new__(cls, x, *args, **kwargs):
-        instance = torch.Tensor._make_subclass(cls, x)
-        instance._Nc = instance.shape[-1]
-        return instance
+    # @staticmethod
+    # def __new__(cls, x, *args, **kwargs):
+    #     instance = torch.Tensor._make_subclass(cls, x)
+    #     return instance
 
     def __init__(self, x, *args, **kwargs):
-        pass
+        self._data = x
+        self._Nc = x.shape[-1]  # read from x, not instance
 
     def validate(self):
         """Check that it is actually a set of color matrice of size Nc \\times Nc"""
@@ -104,19 +104,18 @@ class GaugeConfiguration(ColorMatrix):
     Tensor representing a gauge configuration of links (in the fundamental representation) with shape:
     (B, L1, ..., Ld, d, Nc, Nc)
     """
-    @staticmethod
-    def __new__(cls, x, *args, **kwargs):
-        instance = super().__new__(cls, x, *args, **kwargs)
-        instance._batch_size = instance.shape[0]
-        instance._lattice_shape = instance.shape[1:-3]
-        instance._d = instance.shape[1 + len(instance._lattice_shape)]
-        instance._Nc = instance.shape[-1]
-        instance.validate()
-        return instance
+    # @staticmethod
+    # def __new__(cls, x, *args, **kwargs):
+    #     instance = torch.Tensor._make_subclass(cls, x)
+    #     return instance
 
     def __init__(self, x, *args, **kwargs):
-        pass
-
+        super().__init__(x, *args, **kwargs)
+        self._batch_size    = x.shape[0]
+        self._lattice_shape = x.shape[1:-3]
+        self._d             = x.shape[-3]
+        self.validate()
+    
     def validate(self):
         """Check that shape == (B, L1, ..., Ld, d, Nc, Nc)"""
         expected = (self.batch_size, *self.lattice_shape, self.n_dims, self.Nc, self.Nc)
@@ -167,7 +166,10 @@ def from_hotstart(batchsize: int, L_mu: typing.List[int], Nc: int, seed: int,  d
         """
         d = len(L_mu)
         shape = (batchsize, *L_mu, d, Nc, Nc)
-        return GaugeConfiguration(suN.get_hotstart(shape=shape, seed=seed, dtype=dtype, device=device, requires_grad=requires_grad))
+        U_tensor = suN.get_hotstart(shape=shape, seed=seed, dtype=dtype, device=device, requires_grad=requires_grad)
+        U =  GaugeConfiguration(U_tensor)
+        return U
+
 
     def hotstart(self, seed: int) -> None:
         suN.apply_hotstart(U=self, seed=seed)
@@ -243,17 +245,21 @@ class LocallyGaugeCovariant(ColorMatrix):
 
     shape: (B, L1, ..., Ld, N_obs, Nc, Nc)
     """
-    @staticmethod
-    def __new__(cls, x, *args, **kwargs):
-        instance = super().__new__(cls, x, *args, **kwargs)
-        instance._batch_size = instance.shape[0]
-        instance._lattice_shape = instance.shape[1:-3]
-        instance._N_obs = instance.shape[-3]
-        instance.validate()
-        return instance
+    # @staticmethod
+    # def __new__(cls, x, *args, **kwargs):
+    #     instance = super().__new__(cls, x, *args, **kwargs)
+    #     instance._batch_size = instance.shape[0]
+    #     instance._lattice_shape = instance.shape[1:-3]
+    #     instance._N_obs = instance.shape[-3]
+    #     instance.validate()
+    #     return instance
 
     def __init__(self, x, *args, **kwargs):
-        pass
+        super().__init__(x, *args, **kwargs)
+        self._batch_size = self.shape[0]
+        self._lattice_shape = self.shape[1:-3]
+        self._N_obs = self.shape[-3]
+        self.validate()
 
     def validate(self):
         """Check that shape == (B, L1, ..., Ld, d, Nc, Nc)"""

links/lie_derivatives.py

@@ -0,0 +1,114 @@
+"""
+Implementation of Lie derivatives on functions of the gauge links,
+using autodifferentiation:
+
+$${ L_a f(U) = -i\\frac{d}{d \\omega} f( e^{-i \\omega \\tau_a } U ) |_{\\omega = 0} }$$
+$${ R_a f(U) = +i\\frac{d}{d \\omega} f( U e^{-i \\omega \\tau_a } ) |_{\\omega = 0} }$$
+
+The $L_a$ and $R_a$ are the canonical momenta associated to the link $U$,
+and satisfy:
+
+$${ [L_a , U] = - \\tau_a U }$$
+$${ [R_a , U] = + U \\tau_a }$$
+
+"""
+
+import typing
+import torch
+from torch.func import jvp
+
+import lattice_data_tools.links.suN as suN
+from lattice_data_tools.links.configuration import GaugeConfiguration, ColorMatrix
+
+
+
+class LieDerivatives:
+    """
+    Class for the calculation of the Lie derivatives (canonical momenta)
+    """
+    def __init__(self, U: GaugeConfiguration):
+        """
+        Initialization of tensors needed
+        
+        Main idea: this object is initialized before the training,
+        and allows to compute efficiently the derivatives with respect to $\\omega_a(x,\\mu)$ in $0$.
+        For instance:
+
+        $${ L_a(x,\\mu) f(..., U(x,\\mu), ...) = \\frac{d}{d\\omega} f(..., V_a U(x,\\mu), ...)|_{\\omega=0} }$$
+
+        where ${ V_a = exp(-i \\omega_a(x,\\mu) \\tau_a) }$
+
+        Input:
+            U: GaugeConfiguration object, used to infer the shape of `omega`, `dtype` and `device`.
+
+        where `Ng` is the number of generators of `SU(Nc)`.
+        """
+        self.Nc = U.Nc # number of colors of the group
+        self.Ng = U.Ng # number of generators in the algebra
+        self.omega_glob = torch.zeros(size=(1,), dtype=U.dtype, device=U.device, requires_grad=True)
+        omega_shape = (*U.shape[0:-2],) #, self.Ng) # (batch, x,\\mu) indices
+        # self.omega = torch.zeros(omega_shape, dtype=U.dtype, device=U.device, requires_grad=True)
+        # self.omega = self.omega_glob.expand(size=omega_shape) # 
+        self.omega = torch.zeros(size=omega_shape, dtype=U.dtype, device=U.device, requires_grad=True)
+        self.tau = suN.get_generators(Nc=self.Nc, device=U.device, dtype=U.dtype)
+        # omega_tau = torch.einsum("...a,aij->...aij", self.omega, tau)
+        #self.V = suN.get_exp_iA(A = - omega_tau) # exp(-i*omega*tau_a)
+        self.V = [suN.get_exp_iA(A = - torch.einsum("...,ij->...ij", self.omega, self.tau[a,:,:])) for a in range(self.Ng)] # exp(-i*omega*tau_a)
+    #---
+
+    def L_a(self, a: int, f: typing.Callable, U: GaugeConfiguration):
+        d, M = torch.linalg.eigh(self.tau[a,:,:].expand(*self.omega.shape, self.Nc, self.Nc))
+        phase = self.omega.unsqueeze(-1) * d
+        exp_iD = torch.diag_embed(torch.exp(1j*phase)) #.type(M.type())  # exp(d_k) for each eigenvalue d_k
+        Va = M @ exp_iD @ M.adjoint()   # U = M exp(iD) M^\\dagger
+        Va_U = GaugeConfiguration(Va @ U)
+        f_VaU = torch.Tensor(f(Va_U))
+        for i in range(200):
+            df_domega_Re = torch.autograd.grad(
+                    f_VaU.view(-1)[i].real,
+                    self.omega,
+                    retain_graph=True
+                )[0]
+            df_domega_Im = torch.autograd.grad(
+                    f_VaU.view(-1)[i].imag,
+                    self.omega,
+                    retain_graph=True
+                )[0]
+
+            print(i, df_domega_Re.shape, df_domega_Im.shape)
+
+        print(d.shape, self.omega.unsqueeze(-1).shape)
+        
+
+        V_a = suN.get_exp_iA(A = -torch.einsum("...,ij->...ij", omega_expanded, self.tau[a]))
+        print(torch.autograd.grad(V_a.view(-1)[0].real, omega)[0])
+        # Now omega -> V_a -> f_VaU is one connected graph
+        print("ciao", f_VaU[0,0].shape, f_VaU[0,0].real.requires_grad)
+        print(torch.autograd.grad(f_VaU[0,0].real, omega)[0])
+        f_real = torch.view_as_real(f_VaU[0,:])[..., 0]  # extracts real part, grad-safe
+        #grad = torch.autograd.grad(f_real, omega)[0]
+        #grad = torch.autograd.grad(f_VaU[0,:], omega)[0]
+        #print(grad)
+
+        # print(self.V[a].shape)
+        # f_VaU = f(self.V[a] @ U)
+        # print(f_VaU[0,:].requires_grad, self.omega_glob.shape)
+        # # f_VaU has shape (batchsize, 1)
+        # # self.omega has shape (batchsize,L1,...,Ld,d)
+        # print(f_VaU[0,:].shape, self.omega_glob.shape)
+        # print(torch.autograd.grad(f_VaU[0,:].real, self.omega_glob)[0])
+        # quit()
+        quit()
+        
+        return f_a
+
+
+        
+    def R_a(self, f: typing.Callable[[GaugeConfiguration], ColorMatrix], U: GaugeConfiguration):
+        pass
+        # f_prime_values = lambda V:  f(U @ V)
+        # df_domega = componentwise_autodiff(y=f_prime_values, x=self.omega)
+        # return +1j*df_domega
+
+
+

links/loops.py

@@ -3,28 +3,28 @@
 from lattice_data_tools.links.configuration import GaugeConfiguration, LocallyGaugeCovariant
 
 class WilsonLoopsGenerator:
-    def __init__(self, U: GaugeConfiguration):
-        """
-          U: batch of gauge configurations U^{(b)}(x, mu) as Nc \\times Nc matrices, with x=(x_1,...x_d).
-          It is passed as a multi-dimensional array U[b, L1,L2,...,Ld, d, Nc, Nc],
-          where N_c is the number of colors.
-        """
-        self.U = U
-    #---
-    def plaquettes(self):
+    """
+    Generation of Wilson loops from the gauge configuration.
+
+      U: batch of gauge configurations U^{(b)}(x, mu) as Nc \\times Nc matrices, with x=(x_1,...x_d).
+      It is passed as a multi-dimensional array U[b, L1,L2,...,Ld, d, Nc, Nc],
+      where N_c is the number of colors.
+    """
+    @staticmethod
+    def plaquettes(U: GaugeConfiguration):
         """
         Plaquettes (as loops, no trace or real part), assuming periodic boundary conditions on a L^d lattice.
 
         Returns: tensor of shape (batch_size, L1,...,Ld, N_plaq, Nc, Nc)
         """
-        assert(self.U.is_complex()) # check that the gauge configuration is complex-valued
-        d = self.U.shape[-3] # number of dimensions of the lattice
+        assert(U.is_complex()) # check that the gauge configuration is complex-valued
+        d = U.shape[-3] # number of dimensions of the lattice
         plaqs = []
         # loop over the positive-oriented plaquettes
         for mu in range(d):
             for nu in range(mu + 1, d):
-                U_mu = self.U[..., mu, :, :] # U_\\mu(\cdot) : all links along the direction \\mu
-                U_nu = self.U[..., nu, :, :] # U_\\nu(\cdot) : all links along the direction \\nu
+                U_mu = U[..., mu, :, :] # U_\\mu(\cdot) : all links along the direction \\mu
+                U_nu = U[..., nu, :, :] # U_\\nu(\cdot) : all links along the direction \\nu
                 # Spatial dim \\mu is dim \\mu+1 in the full tensor (batch at 0)
                 U_mu_fwd = torch.roll(U_mu, -1, dims=1+nu) # U_\\mu(x + \\nu): shifting **backwards**
                 U_nu_fwd = torch.roll(U_nu, -1, dims=1+mu) # U_\\nu(x + \\mu): shifting **backwards**
@@ -34,23 +34,28 @@ def plaquettes(self):
         plaqs = LocallyGaugeCovariant(torch.stack(plaqs, dim=-3))
         return plaqs
     #---
-    def Polyakov_loops(self):
+    @staticmethod
+    def Polyakov_loops(U: GaugeConfiguration):
         """
         Polyakov loops for each direction mu.
 
         Returns: A tensor of shape (batch, L1, ..., Ld, d, Nc, Nc)
         """
-        assert(self.U.is_complex())
-        d = self.U.shape[-3] # number of dimensions
-        lattice_shape = self.U.shape[1:-3] # shape of the lattice points grid
-        Poly = self.U.clone() # copy of the links. In the loop, each is extended to its Polyakov loop
+        assert(U.is_complex())
+        d = U.shape[-3] # number of dimensions
+        lattice_shape = U.shape[1:-3] # shape of the lattice points grid
+        poly_list = []
         for mu in range(d):
             L_mu = lattice_shape[mu] # extension of the lattice over the \\mu-th direction
+            P = U[..., mu, :, :]
             for k in range(1, L_mu):
-                Poly[..., mu, :, :] @= torch.roll(self.U, -k, dims=1+mu)[..., mu, :,:] # P_\\mu(x) --> P_\\mu(x)*U_\\mu(x+mu)
-        #-------
+                P @= torch.roll(U, -k, dims=1+mu)[..., mu, :,:] # P_\\mu(x) --> P_\\mu(x)*U_\\mu(x+mu)
+            #---
+            poly_list.append(P)
+        #---
+        Poly = torch.stack(poly_list, dim=-3)
         return LocallyGaugeCovariant(Poly)
-    #---
+