dimer_covering.py

'''
This code solves counting problem [1] using differentiable TRG programming [2] 
In a nutshell, it computes the maximum eigenvalue of tranfer matrix via variational principle
[1] Vanderstraeten et al, PRE 98, 042145 (2018)
[2] Levin and Nave, PRL 99, 120601 (2007)
'''
import torch 
torch.set_num_threads(4)
torch.manual_seed(42)

from vmps import vmps as contraction

if __name__=='__main__':
    import time
    import argparse
    parser = argparse.ArgumentParser(description='')
    parser.add_argument("-D", type=int, default=2, help="D")
    parser.add_argument("-Dcut", type=int, default=24, help="Dcut")
    parser.add_argument("-Niter", type=int, default=50, help="Niter")

    parser.add_argument("-float32", action='store_true', help="use float32")
    parser.add_argument("-lanczos", action='store_true', help="lanczos")
    parser.add_argument("-cuda", type=int, default=-1, help="use GPU")
    args = parser.parse_args()
    device = torch.device("cpu" if args.cuda<0 else "cuda:"+str(args.cuda))
    dtype = torch.float32 if args.float32 else torch.float64

    if args.lanczos: print ('use lanczos')

    d = 2 # fixed

    D = args.D
    Dcut = args.Dcut
    Niter = args.Niter

    B = 0.01* torch.randn(d, D, D, D, D, dtype=dtype, device=device)
    #symmetrize initial boundary PEPS
    B = (B + B.permute(0, 4, 2, 3, 1))/2. 
    B = (B + B.permute(0, 1, 3, 2, 4))/2. 
    B = (B + B.permute(0, 3, 4, 1, 2))/2. 
    B = (B + B.permute(0, 2, 1, 4, 3))/2. 
    A = torch.nn.Parameter(B.view(d, D**4))
    
    #boundary MPS
    A1 = torch.nn.Parameter(0.01*torch.randn(Dcut, D**2*d, Dcut, dtype=dtype, device=device))
    A2 = torch.nn.Parameter(0.01*torch.randn(Dcut, D**2, Dcut, dtype=dtype, device=device))

    #dimer covering
    T = torch.zeros(d, d, d, d, d, d, dtype=dtype, device=device)
    T[0, 0, 0, 0, 0, 1] = 1.0 
    T[0, 0, 0, 0, 1, 0] = 1.0 
    T[0, 0, 0, 1, 0, 0] = 1.0 
    T[0, 0, 1, 0, 0, 0] = 1.0 
    T[0, 1, 0, 0, 0, 0] = 1.0 
    T[1, 0, 0, 0, 0, 0] = 1.0 
    T = T.view(d, d**4, d)

    optimizer = torch.optim.LBFGS([A], max_iter=20)
    #optimizer = torch.optim.Adam([A])
    
    def closure():
        optimizer.zero_grad()

        T1 = torch.einsum('xa,xby,yc' , (A,T,A)).view(D,D,D,D, d,d,d,d, D,D,D,D).permute(0,4,8, 1,5,9, 2,6,10, 3,7,11).contiguous().view(D**2*d, D**2*d, D**2*d, D**2*d)
        #double layer 
        T2 = (A.t()@A).view(D, D, D, D, D, D, D, D).permute(0,4, 1,5, 2,6, 3,7).contiguous().view(D**2, D**2, D**2, D**2)
        t0=time.time()
        lnT = contraction(T1, D**2*d, Dcut, Niter, A1, use_lanczos=args.lanczos)
        lnZ = contraction(T2, D**2, Dcut, Niter, A2, use_lanczos=args.lanczos)
        loss = (-lnT + lnZ)
        print ('    contraction done {:.3f}s'.format(time.time()-t0))
        print ('    total loss', loss.item())
        #print ('    loss, error', loss.item(), error1.item(), error2.item())

        t0=time.time()
        loss.backward()
        print ('    backward done {:.3f}s'.format(time.time()-t0))
        return loss
    
    for epoch in range(100):
        loss = optimizer.step(closure)
        print ('epoch, residual entropy', epoch, -loss.item())