test_grad.py

'''
Different ways of computing gradient of infinite matrix product trace 
Z = Tr(T*T...*T)
'''

import torch 
torch.manual_seed(42)

N = 10
Niter = 30
A = torch.randn(N, N, dtype=torch.float64)
T = torch.nn.Parameter(A@A.t()) # a symmetric positive definite transfer matrix 

lnZ = 0.0 
M = T.clone() # M is the converged envioment tensor
for i in range(Niter):
    s = M.norm() 
    lnZ = lnZ + torch.log(s)/2**i
    M = M/s
    M = M@M 
lnZ = lnZ + torch.log(torch.trace(M))/(2**Niter)

#(1) impurity grad 
impurity_grad = (M@M).t()/torch.trace(M@T@M)

#(2) autograd on lnZ
lnZ.backward()
lnZ_grad = T.grad.clone()

#(3) direct compute dominant eigenvalue 
w, v = torch.symeig(T, eigenvectors=True)
T.grad.zero_()
loss = torch.log(w[-1])
loss.backward()
exact_grad = ((T.grad + T.grad.t())/2) # need to symmetrize since it is an upper triangular matrix

# (4) use leading eigenvector of the transfer matrix $T$
#c.f. sec.3.1  https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
#or, HF theorem 
eigenvector_grad= torch.ger(v[:,-1], v[:, -1])/w[-1] # outer product of the leading eigenvector and its transpose

#(5) lnZ using envioment 
w, v = torch.symeig(T, eigenvectors=True)
T.grad.zero_()

#try this 
e = v[:, -1] #environment
#and this 
#e = v[:, -1].detach() #environment

loss = torch.log(e@T@e) 
loss.backward()
environment_grad = T.grad.clone()

print ((impurity_grad-exact_grad).abs().max().item())
print ((lnZ_grad-exact_grad).abs().max().item())
print ((eigenvector_grad-exact_grad).abs().max().item())
print ((environment_grad-exact_grad).abs().max().item())