initial commit

Author: RealBunbun

.gitignore

@@ -0,0 +1,5 @@
+.DS_Store
+
+__py_cache__/
+
+Result/

args.py

@@ -0,0 +1,27 @@
+import argparse
+
+parser = argparse.ArgumentParser(description="")
+
+parser.add_argument("-model", type=str, default="XY", help="interacting model [XY]")
+parser.add_argument("-tau", type=float, default=0.025, help="initial inverse temperature [0.025]")
+parser.add_argument("-NBeta", type=int, default=10, help="number of temperature point [10]")
+
+parser.add_argument("-D", type=int, default=20, help="bond dimensions [20]")
+parser.add_argument("-depth", type=int, default=3, help="sweep depth [3]")
+parser.add_argument("-Niter", type=int, default=10, help="max Iteration of single isometry optimization [10]")
+parser.add_argument("-Nsweep", type=int, default=3, help="number of sweeps [3]")
+parser.add_argument("-opti", type=int, default=1, help="optimization flag [1]")
+
+parser.add_argument("-use_float32", action="store_true", help="use float32")
+parser.add_argument("-cuda", type=int, default=-1, help="GPU ID [default:-1 (CPU)] ")
+parser.add_argument("-rdir", type=str, default="./Result/", help="result folder")
+
+args = parser.parse_args()
+
+print('---------------+---------')
+print('%15s|%8s'%('arguments','values'))
+print('---------------+---------')
+for arg in vars(args):
+    print('%15s|%8s'%(arg, getattr(args, arg)))
+print('---------------+---------')
+print('\n')

dTRG.py

@@ -0,0 +1,263 @@
+import os
+import sys
+import re
+import torch
+import time
+import numpy as np
+from utils import kronecker_product as kron
+from torch.utils.checkpoint import checkpoint
+from datetime import datetime
+from expm import expm
+
+class iXTRG(torch.nn.Module):
+    def __init__(self,args,NBeta,dtype,device):
+        super(iXTRG, self).__init__()
+
+        # model
+        self.NBeta = NBeta  # no. of isometries on different layers, starts from 1
+        self.tau = torch.tensor([args.tau], dtype=dtype, device=device)
+        self.model = args.model
+
+        # algorithm
+        self.D = args.D 
+        self.opti = args.opti
+        self.Niter = args.Niter
+        self.Nsweep = args.Nsweep
+        self.depth = args.depth
+
+        # io
+        self.dtype = dtype
+        self.device = device
+        self.rdir = args.rdir
+        
+        self.fName = self.rdir+ 'Fe_'+self.model+'_tau'+str(self.tau.item())+'_NBeta'+str(self.NBeta)+\
+                     '_Dc'+str(self.D)+'_opti'+str(self.opti)+\
+                     '_Niter'+str(self.Niter)+'_Nswp'+str(self.Nsweep)+\
+                     '_depth'+str(self.depth)+'_device'+str(args.cuda)+'.pt'
+       
+        # create dir if not existed
+        if os.path.isdir(self.rdir) is False: os.makedirs(self.rdir)
+
+        # Module Parameters
+        """
+        [Wl_1, Wl_2, ..., Wl_NBeta, Wr_1, Wr_2, ..., Wr_NBeta]
+        """
+        """
+        2nd-order Trotter-Suzuki Decomposition!!!
+        """
+        nb = self.NBeta-1
+        Dl = min(4**(nb+2), self.D)
+        Dr = min(4**(nb+2), self.D)
+        D2 = min(4**(nb+3), self.D)
+        Wl.append(torch.randn(Dl,Dr,D2,dtype=dtype,device='cpu'))
+            
+        Dl = min(4**(nb+1), self.D)
+        Dr = min(4**(nb+1), self.D)
+        D2 = min(4**(nb+2), self.D)
+        Wr.append(torch.randn(Dl,Dr,D2,dtype=dtype,device='cpu'))
+        
+        self.params = torch.nn.ParameterList(
+                        [torch.nn.Parameter(_.to(device)) for _ in Wl+Wr]
+                        )
+        self.params.append(torch.nn.Parameter(self.tau))
+         
+    def turn_on_grad(self,i):
+        for j, p in enumerate(self.parameters()):
+            if (j==i): p.requires_grad = True
+            else:      p.requires_grad = False
+
+    def getHamilton(self):
+        # spin operator 
+        sz = torch.tensor([[1, 0],[0,-1]], dtype=dtype, device=device)/2
+        sx = torch.tensor([[0, 1],[1, 0]], dtype=dtype, device=device)/2
+        sy = torch.tensor([[0,-1],[1, 0]], dtype=dtype, device=device)/2
+
+        if self.model == "Ising":
+            H = kron(sz,sz)
+        elif self.model == "XY":
+            H = kron(sx,sx) - kron(sy,sy)
+        return H
+
+    def getisometry(self, step, sizeT):
+        D = min(self.D, sizeT)
+        D_new = min(self.D, D**2)
+        return self.params[step][:D,:D,:D_new]
+
+    def getMaxEigBiLayer(self, Ta, Tb):
+        # initial boundary 'vector'
+        Vl = torch.einsum('ijkl,mjkl->im',Ta,Ta)
+        normFactor = torch.norm(Vl)
+        Vl = Vl/normFactor
+
+        # Power method
+        matchCnt = 0
+        for _ in range(400):
+            Vl = torch.einsum('im,ijkl,mnkl->jn',Vl,Ta,Ta)
+            Vl = torch.einsum('im,ijkl,mnkl->jn',Vl,Tb,Tb)
+            if (torch.norm(Vl)-normFactor)/normFactor<1e-8: matchCnt +=1 
+            if matchCnt==5: break
+            normFactor = torch.norm(Vl)
+            Vl = Vl/normFactor
+            if _==399: print('Eig not well converged!',end=' ')
+        return normFactor
+
+
+    def initRho(self, trotter_order=2):
+        tau,device,dtype,model = self.params[-1],self.device,self.dtype,self.model
+        #print("Generate initial rho({}) via Trotter decomp.\n".format(tau))
+        if trotter_order==1:
+            # get Hamiltonian
+            H = self.getHamilton()
+            # local trotter gate
+            rho = expm(-tau*H).view(2,2,2,2)
+            rho = torch.einsum('ijkl->ikjl',rho).contiguous().view(4,4)
+            # svd & truncate the 0 values
+            U,S,V = torch.svd(rho)
+            # trotter gate in form of two tensor contraction
+            hl = (U@torch.diag(torch.sqrt(S))).view(2,2,4)
+            hr = (V@torch.diag(torch.sqrt(S))).view(2,2,4)
+            # local tensor of initial mpo, index order: [l,r,d,u] 
+            Ta = torch.tensordot(hr,hl,([0],[1])).permute(1,3,2,0).contiguous()
+            Tb = torch.tensordot(hr,hl,([1],[0])).permute(1,3,0,2).contiguous()
+        elif trotter_order==2:
+            # get Hamiltonian
+            H = self.getHamilton()
+            # local trotter gate
+            rho = expm(-tau*H).view(2,2,2,2)
+            rho = torch.einsum('ijkl->ikjl',rho).contiguous().view(4,4)
+            # half local trotter gate
+            rho_half = expm(-tau*H/2).view(2,2,2,2)
+            rho_half = torch.einsum('ijkl->ikjl',rho_half).contiguous().view(4,4)
+            # svd
+            U,S,V = torch.svd(rho)
+            U2,S2,V2 = torch.svd(rho_half)
+            # trotter gate in form of two tensor contraction
+            hl = (U@torch.diag(torch.sqrt(S))).view(2,2,4)
+            hr = (V@torch.diag(torch.sqrt(S))).view(2,2,4)
+            hl2 = (U2@torch.diag(torch.sqrt(S2))).view(2,2,4)
+            hr2 = (V2@torch.diag(torch.sqrt(S2))).view(2,2,4)
+            # local tensor of initial mpo, index order: [l,r,d,u] 
+            Ta = torch.einsum('ijk,jlm,lno->okmin',hr2,hl,hr2).contiguous().view(16,4,2,2)
+            Tb = torch.einsum('ijk,jlm,lno->mokin',hl2,hr,hl2).contiguous().view(4,16,2,2)
+        else:
+            raise Exception('only 1st and 2nd trotter are available!')
+        return Ta,Tb 
+
+
+    def forward(self, nlayer):
+        if nlayer == 0:
+            [Ta, Tb] = self.initRho()
+            lnZ = 0.0
+        else:
+            Ta = Tas[nlayer-1].to(device)
+            Tb = Tbs[nlayer-1].to(device)
+            lnZ = lnZs[nlayer-1].to(device)
+
+        for nbeta in range(nlayer,self.NBeta):
+           Sizea = list(Ta.size())
+           Sizeb = list(Tb.size())
+
+           # obtain Isometry tensor from antisymmetric tensor
+           Wa = self.getisometry(nbeta,Sizea[0])
+           Wb = self.getisometry(nbeta+self.NBeta,Sizeb[0])
+
+           # evolution Ta & Tb
+           Ta = checkpoint(self.rgtens,Wa,Ta,Wb)
+           Tb = checkpoint(self.rgtens,Wb,Tb,Wa)
+           lnZ = 2*lnZ + torch.log(torch.norm(Ta))\
+                       + torch.log(torch.norm(Tb))
+           Ta = Ta/torch.norm(Ta)
+           Tb = Tb/torch.norm(Tb)
+
+           if len(Tas)<nbeta+1: Tas.append(Ta.detach().to('cpu'))
+           else:                Tas[nbeta]=Ta.detach().to('cpu')
+           if len(Tbs)<nbeta+1: Tbs.append(Tb.detach().to('cpu'))
+           else:                Tbs[nbeta]=Tb.detach().to('cpu')
+           if len(lnZs)<nbeta+1: lnZs.append(lnZ.detach().to('cpu'))
+           else:                 lnZs[nbeta]=lnZ.detach().to('cpu')
+
+        # free energy
+        ee = self.getMaxEigBiLayer(Ta,Tb)
+        lnZ = 1/2*(2*lnZ + torch.log(ee))
+        return lnZ
+
+    def rgtens(self, Wa, T, Wb):
+        T = torch.einsum('ijk,ilmn,jopm,loq->kqpn',Wa,T,T,Wb)
+        return T
+
+    def update_single_layer(self, nlayer):
+        loss_old = 0
+        for niter in range(self.Niter):
+            for ii in range(2):
+                print('\t\tW(%02d-%02d),'%(nlayer,ii),end=' ')
+                self.turn_on_grad(nlayer+self.NBeta*ii)
+        
+                self.zero_grad()
+                loss = self.forward(nlayer)
+                loss.backward()
+                
+                with torch.no_grad():
+                   E = self.params[nlayer+self.NBeta*ii].grad
+                   D, D, D_new = E.shape
+                   E = E.view(D**2, D_new)
+                   # perform MERA update
+                   U,S,V = torch.svd(E)
+                   self.params[nlayer+self.NBeta*ii].data = (U@V.t()).view(D,D,D_new)       
+                   print('%+.15f'%(-loss.item()/2**(self.NBeta+1)/self.tau), end=' ')
+            if abs((loss_old - loss.item())/loss.item())<1e-8: break
+            else: 
+                loss_old = loss.item()
+            print(' ')
+        return loss.item()
+        
+
+"""
+===================================================================================
+Main entry of the program
+===================================================================================
+"""
+if __name__=="__main__":
+    from args import args
+    device = torch.device("cpu" if args.cuda<0 else "cuda:"+str(args.cuda))
+    dtype = torch.float32 if args.use_float32 else torch.float64
+    FeS = []  # init thermal quantities
+
+    Wl = []; Wr = []
+    Tas = []; Tbs = []; lnZs = []
+    for nbeta in range(args.NBeta):  # main loop, cooling down the system
+        beta = 2**(nbeta+1)*args.tau 
+        model = iXTRG(args, nbeta+1, dtype, device)    
+        print('\n '+'-'*40)
+        print('Optimizing lnZ(%.5f)'%(2*beta))
+        print('\nInitialization:')
+        loss = model.update_single_layer(nbeta)
+        
+        opti_flag = 0   # no sweep by default
+        if(nbeta+1 >= model.opti) :
+           opti_flag = 1
+
+        depth = min(model.depth,nbeta)
+        # call sweep optimization
+        if opti_flag:
+            print('\noptimiztion:')
+            for nswp in range(args.Nsweep):
+                print('#sweep {}:'.format(nswp))
+                # nlayer is No. of layer to be optimized, from 1 to 
+                for nlayer in range(nbeta-depth,nbeta+1):
+                    loss = model.update_single_layer(nlayer) 
+                    print(' ')
+        
+        # saving isometric tensor for next iteration
+        Wl = [m.detach().to('cpu') for m in model.params[:nbeta+1]]
+        Wr = [m.detach().to('cpu') for m in model.params[nbeta+1:-1]]
+        
+        # remove redundant isometries (to save memory)
+        if (nbeta-depth >= 1):
+           Wl[nbeta-depth-1] = torch.tensor([0.])
+           Wr[nbeta-depth-1] = torch.tensor([0.])
+        
+        del model.params
+        torch.cuda.empty_cache()
+         
+    # save data    
+    torch.save(FeS, model.fName)

expm.py

@@ -0,0 +1,80 @@
+#copied from https://github.com/geoopt/geoopt/blob/master/geoopt/linalg/_expm.py
+
+# author: @nsde
+# maintainer: @ferrine
+
+
+import torch.jit
+
+
+@torch.jit.script
+def torch_pade13(A):  # pragma: no cover
+    # avoid torch select operation and unpack coefs
+    (b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13) = (
+        64764752532480000.0,
+        32382376266240000.0,
+        7771770303897600.0,
+        1187353796428800.0,
+        129060195264000.0,
+        10559470521600.0,
+        670442572800.0,
+        33522128640.0,
+        1323241920.0,
+        40840800.0,
+        960960.0,
+        16380.0,
+        182.0,
+        1.0,
+    )
+
+    ident = torch.eye(A.shape[1], dtype=A.dtype, device=A.device)
+    A2 = torch.matmul(A, A)
+    A4 = torch.matmul(A2, A2)
+    A6 = torch.matmul(A4, A2)
+    U = torch.matmul(
+        A,
+        torch.matmul(A6, b13 * A6 + b11 * A4 + b9 * A2)
+        + b7 * A6
+        + b5 * A4
+        + b3 * A2
+        + b1 * ident,
+    )
+    V = (
+        torch.matmul(A6, b12 * A6 + b10 * A4 + b8 * A2)
+        + b6 * A6
+        + b4 * A4
+        + b2 * A2
+        + b0 * ident
+    )
+    return U, V
+
+
+@torch.jit.script
+def matrix_2_power(x, p):  # pragma: no cover
+    for _ in range(int(p)):
+        x = x @ x
+    return x
+
+
+@torch.jit.script
+def expm(A):  # pragma: no cover
+    # no checks, this is private implementation
+    # but A should be a matrix
+    A_fro = torch.norm(A)
+
+    # Scaling step
+
+    n_squarings = torch.clamp(
+        torch.ceil(torch.log(A_fro / 5.371920351148152).div(0.6931471805599453)), min=0
+    )
+    scaling = 2.0 ** n_squarings
+    Ascaled = A / scaling
+
+    # Pade 13 approximation
+    U, V = torch_pade13(Ascaled)
+    P = U + V
+    Q = -U + V
+
+    R, _ = torch.solve(P, Q)  # solve P = Q*R
+    expmA = matrix_2_power(R, n_squarings)
+    return expmA