Advanced Computing Platform for Theoretical Physics

Commit 39dd69ac authored by Lei Wang's avatar Lei Wang
Browse files

refactor hotrg_sweep; tensordot does not seem to provide speedup (although it...

refactor hotrg_sweep; tensordot does not seem to provide speedup (although it shows speed up in test_tensordot.py)
parent 512a2e45
......@@ -2,9 +2,19 @@ import time
import numpy as np
import torch
torch.manual_seed(42)
torch.set_num_threads(1)
from utils import gauge, trace
def renormalize(T, U):
T = torch.einsum('axob,amz,moyn,bnw->xwzy',(T, U, T, U))
#TU1 = torch.tensordot(T, U, dims=([0], [0]))
#TU2 = torch.tensordot(T, U, dims=([3], [1]))
#T = torch.tensordot(TU1, TU2, dims=([1,2,3], [1,3,0]))
#T = T.permute(0,3,1,2).contiguous()
return T
class Ising(torch.nn.Module):
def __init__(self, chi, niter, dtype=torch.float64, device='cpu', use_checkpoint=False):
super(Ising, self).__init__()
......@@ -17,7 +27,7 @@ class Ising(torch.nn.Module):
self.params = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(self.chi, self.chi, self.chi, dtype=dtype, device=device)) for _ in range(niter)])
def isometry(self, step):
def get_isometry(self, step):
chi = min(self.D**(2**(step//2)), self.chi)
chi_new = min(self.chi, chi**2)
return self.params[step][:chi, :chi, :chi_new]
......@@ -43,21 +53,19 @@ class Ising(torch.nn.Module):
U = vl if (el < er) else vr
self.set_isometry(step, U)
T = torch.einsum('axob,amz,moyn,bnw->xwzy',(T, U, T, U))
T = renormalize(T, U)
lnZ += torch.log(trace(T))/2**args.niter
return lnZ
def forward(self, T):
lnZ = 0.0
for n in range(self.niter):
f = torch.norm(T)
f = T.abs().max()
lnZ += 2**(-n)*torch.log(f)
T = T / f
U = self.isometry(n) # isometry for this step
#print (T.shape , U.shape)
T = torch.einsum('axob,amz,moyn,bnw->xwzy',(T, U, T, U))
U = self.get_isometry(n) # isometry for this step
T = renormalize(T, U)
lnZ += torch.log(trace(T))/2**args.niter
return lnZ
......@@ -66,12 +74,12 @@ if __name__=='__main__':
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument("-beta", type=float, default=0.44, help="beta")
parser.add_argument("-chi", type=int, default=20, help="chi")
parser.add_argument("-niter", type=int, default=20, help="niter")
parser.add_argument("-nsweeps", type=int, default=100, help="nsweeps")
parser.add_argument("-beta", type=float, default=0.44068679350977147, help="beta")
parser.add_argument("-chi", type=int, default=24, help="chi")
parser.add_argument("-niter", type=int, default=50, help="niter")
parser.add_argument("-nsweeps", type=int, default=20, help="nsweeps")
parser.add_argument("-init", default='random', choices=['random', 'hotrg'], help="initial isometry")
parser.add_argument("-exact", type=float, default=None, help="exact free energy")
parser.add_argument("-exact", type=float, default=-0.9296953983414709, help="exact free energy")
parser.add_argument("-float32", action='store_true', help="use float32")
parser.add_argument("-cuda", type=int, default=-1, help="GPU #")
args = parser.parse_args()
......@@ -129,7 +137,7 @@ if __name__=='__main__':
svd = time.time()
model.params[step].data = (U@V.t()).view(chi, chi, chi_new)
##print(forward-start, backward - forward, svd-backward)
#print(forward-start, backward - forward, svd-backward)
step += direction
#when reaching the end reverse sweep direction and report results
......@@ -141,4 +149,4 @@ if __name__=='__main__':
if args.exact is not None:
print (sweeps, loss.item(), (loss.item()-args.exact)/np.abs(args.exact) )
else:
print (sweeps, loss.item() )
print (sweeps, loss.item(), end-start)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment