Advanced Computing Platform for Theoretical Physics

commit大文件会使得服务器变得不稳定,请大家尽量只commit代码,不要commit大的文件。

Commit fc796aba authored by Lei Wang's avatar Lei Wang
Browse files

better control of gradient flow

parent d9aea05f
import torch
torch.set_num_threads(4)
def vmps(T, d, D, no_iter, Nepochs=5):
A = torch.nn.Parameter(0.01*torch.randn(D, d, D, dtype=T.dtype, device=T.device))
def mpsrg(A, T):
Asymm = (A + A.permute(2, 1, 0))*0.5
D, d = Asymm.shape[0], Asymm.shape[1]
#t0 = time.time()
B = torch.einsum('ldr,adcb,icj->lairbj', (Asymm, T, Asymm)).contiguous().view(D**2*d, D**2*d)
C = torch.einsum('ldr,idj->lirj', (Asymm, Asymm)).contiguous().view(D**2, D**2)
w, _ = torch.symeig(B, eigenvectors=True)
lnZ1 = torch.log(w.abs().max())
w, _ = torch.symeig(C, eigenvectors=True)
lnZ2 = torch.log(w.abs().max())
def mpsrg(B, C):
w, _ = torch.symeig(B, eigenvectors=True)
lnZ1 = torch.log(w.abs().max())
w, _ = torch.symeig(C, eigenvectors=True)
lnZ2 = torch.log(w.abs().max())
#lnZ1 = 0.0
#lnZ2 = 0.0
#for i in range(no_iter):
# s = B.norm(1)
# lnZ1 = lnZ1 + torch.log(s)/2**i
# B = B/s
# B = torch.mm(B, B)
#lnZ1 = 0.0
#lnZ2 = 0.0
#for i in range(no_iter):
# s = B.norm(1)
# lnZ1 = lnZ1 + torch.log(s)/2**i
# B = B/s
# B = torch.mm(B, B)
# s = C.norm(1)
# lnZ2 = lnZ2 + torch.log(s)/2**i
# C = C/s
# C = torch.mm(C, C)
return -lnZ1 + lnZ2
# s = C.norm(1)
# lnZ2 = lnZ2 + torch.log(s)/2**i
# C = C/s
# C = torch.mm(C, C)
return lnZ1 , lnZ2
def vmps(T, d, D, no_iter, Nepochs=5):
A = torch.nn.Parameter(0.01*torch.randn(D, d, D, dtype=T.dtype, device=T.device))
optimizer = torch.optim.LBFGS([A], max_iter=10)
def closure():
optimizer.zero_grad()
Asymm = (A + A.permute(2, 1, 0))*0.5
#t0 = time.time()
B = torch.einsum('ldr,adcb,icj->lairbj', (Asymm, T, Asymm)).contiguous().view(D**2*d, D**2*d)
C = torch.einsum('ldr,idj->lirj', (Asymm, Asymm)).contiguous().view(D**2, D**2)
#print ('einsum', time.time()- t0)
#print ((B-B.t()).abs().sum(), (C-C.t()).abs().sum())
#t0 = time.time()
lnZ1, lnZ2= mpsrg(B, C)
loss = mpsrg(A, T) # -lnZ
#print ('mpsrg', time.time()- t0)
loss = -lnZ1 + lnZ2
print (' loss', loss.item(), lnZ1.item(), lnZ2.item())
print (' loss', loss.item())
#t0 = time.time()
loss.backward(retain_graph=False)
#print ('backward', time.time()- t0)
......@@ -51,7 +50,7 @@ def vmps(T, d, D, no_iter, Nepochs=5):
loss = optimizer.step(closure)
print (' epoch, free energy', epoch, loss.item())
return -loss, A
return loss, A
if __name__=='__main__':
import time
......@@ -67,27 +66,18 @@ if __name__=='__main__':
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
def build_tensor(K):
c = torch.sqrt(torch.cosh(K))
s = torch.sqrt(torch.sinh(K))
M = torch.stack([torch.cat([c, s]), torch.cat([c, -s])])
T = torch.einsum('ai,aj,ak,al->ijkl', (M, M, M, M))
return T
#optimization
K = torch.tensor([args.beta], dtype=torch.float64, device=device)
T = build_tensor(K)
loss, A = vmps(T, 2, args.Dcut, args.Niter, args.Nepochs)
#for computing physical quantity
#build tensor
K = torch.tensor([args.beta], dtype=torch.float64, device=device, requires_grad=True)
T = build_tensor(K)
Asymm = (A + A.permute(2, 1, 0))*0.5
D, d = Asymm.shape[0], Asymm.shape[1]
B = torch.einsum('ldr,adcb,icj->lairbj', (Asymm, T, Asymm)).contiguous().view(D**2*d, D**2*d)
w, _ = torch.symeig(B, eigenvectors=True)
lnZ = torch.log(w.abs().max())
c = torch.sqrt(torch.cosh(K))
s = torch.sqrt(torch.sinh(K))
M = torch.stack([torch.cat([c, s]), torch.cat([c, -s])])
T = torch.einsum('ai,aj,ak,al->ijkl', (M, M, M, M))
#optimization
_, A = vmps(T.detach(), 2, args.Dcut, args.Niter, args.Nepochs)
#recompute lnZ using optimized A
lnZ = -mpsrg(A.detach(), T)
dlnZ = torch.autograd.grad(lnZ, K, create_graph=True)[0] # En = -d lnZ / d beta
print (-dlnZ.item())
#second order derivative evaluated in this way seems not correct, no effect of envioment
......
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