Recent Changes - Search:

HomePage

PmWiki

pmwiki.org

Tutorials /

MPS-DMRG

This is a DMRG code for finite systems, using MPS and MPO representations for the wavefunction and Hamiltonian. As such it is very flexible - it is easy to change the Hamiltonian to do a different model. You can get the source code here

This Python code follows closely how a 'serious' MPS code works, although it is not as fast as the C++ code in the Matrix Product Toolkit. The initial version of the code was very slow, and this was due to the inefficient implementation of the numpy.einsum() function that doesn’t optimize the tensor contractions. Breaking up the expression into pair-wise contractions speeds up the code by a factor 100x. Some other aspects that differ from a 'professional' MPS code:

  1. The MPO is stored in a dense structure. This is inappropriate for an MPO because it only has very few non-zero matrix elements (eg for the Heisenberg model, the MPO is 5x5 on a 2x2 local Hilbert space, for a total of 100 matrix elements. But all except for 12 of the matrix elements are zero.)
  2. The np.einsum function is not particularly efficient, it doesn't use an optimized matrix multiply to do the tensor contractions (but this will probably be fixed in a later version of NumPy).
  3. It would be better to control the number of iterations used by the eigensolver, either by restricting it to a small number (eg 4 - 10 iterations) or selecting an appropriate tolerance. This can be done adaptively, eg using the truncation error or wavefunction fidelity.
  4. No symmetries are implemented. The Heisenberg model has {$SU(2)$} symmetry, which improves the efficiency by around a factor 100x. Using just {$U(1)$} improves the efficiency by about a factor 10x.
  1. #####################################################
  2. # Simple DMRG code using MPS/MPO representations    #
  3. # Ian McCulloch August 2017                         #
  4. #####################################################
  5.  
  6. import numpy as np
  7. import scipy
  8. import scipy.sparse.linalg
  9. import scipy.sparse as sparse
  10. import math
  11.  
  12. ## initial E and F matrices for the left and right vacuum states
  13. def initial_E(W):
  14.     E = np.zeros((W.shape[0],1,1))
  15.     E[0] = 1
  16.     return E
  17.  
  18. def initial_F(W):
  19.     F = np.zeros((W.shape[1],1,1))
  20.     F[-1] = 1
  21.     return F
  22.  
  23. def contract_from_right(W, A, F, B):
  24.     # the einsum function doesn't appear to optimize the contractions properly,
  25.     # so we split it into individual summations in the optimal order
  26.     #return np.einsum("abst,sij,bjl,tkl->aik",W,A,F,B, optimize=True)
  27.     Temp = np.einsum("sij,bjl->sbil", A, F)
  28.     Temp = np.einsum("sbil,abst->tail", Temp, W)
  29.     return np.einsum("tail,tkl->aik", Temp, B)
  30.  
  31. def contract_from_left(W, A, E, B):
  32.     # the einsum function doesn't appear to optimize the contractions properly,
  33.     # so we split it into individual summations in the optimal order
  34.     # return np.einsum("abst,sij,aik,tkl->bjl",W,A,E,B, optimize=True)
  35.     Temp = np.einsum("sij,aik->sajk", A, E)
  36.     Temp = np.einsum("sajk,abst->tbjk", Temp, W)
  37.     return np.einsum("tbjk,tkl->bjl", Temp, B)
  38.  
  39. # construct the F-matrices for all sites except the first
  40. def construct_F(Alist, MPO, Blist):
  41.     F = [initial_F(MPO[-1])]
  42.  
  43.     for i in range(len(MPO)-1, 0, -1):
  44.         F.append(contract_from_right(MPO[i], Alist[i], F[-1], Blist[i]))
  45.     return F
  46.  
  47. def construct_E(Alist, MPO, Blist):
  48.     return [initial_E(MPO[0])]
  49.  
  50.  
  51.  
  52. # 2-1 coarse-graining of two site MPO into one site
  53. def coarse_grain_MPO(W, X):
  54.     return np.reshape(np.einsum("abst,bcuv->acsutv",W,X),
  55.                       [W.shape[0], X.shape[1],
  56.                        W.shape[2]*X.shape[2],
  57.                        W.shape[3]*X.shape[3]])
  58.  
  59. def product_W(W1, W2):
  60.     return np.reshape(np.einsum("abst,cdtu->acbdsu", W1, W2), [W1.shape[0]*W2.shape[0],
  61.                                                                W1.shape[1]*W2.shape[1],
  62.                                                                W1.shape[2],W2.shape[3]])
  63.  
  64. def product_MPO(M1, M2):
  65.     assert len(M1) == len(M2)
  66.     Result = []
  67.     for i in range(0, len(M1)):
  68.         Result.append(product_W(M1[i], M2[i]))
  69.     return Result
  70.  
  71.  
  72. # 2-1 coarse-graining of two-site MPS into one site
  73. def coarse_grain_MPS(A,B):
  74.     return np.reshape(np.einsum("sij,tjk->stik",A,B),
  75.                       [A.shape[0]*B.shape[0], A.shape[1], B.shape[2]])
  76.  
  77. def fine_grain_MPS(A, dims):
  78.     assert A.shape[0] == dims[0] * dims[1]
  79.     Theta = np.transpose(np.reshape(A, dims + [A.shape[1], A.shape[2]]),
  80.                          (0,2,1,3))
  81.     M = np.reshape(Theta, (dims[0]*A.shape[1], dims[1]*A.shape[2]))
  82.     U, S, V = np.linalg.svd(M, full_matrices=0)
  83.     U = np.reshape(U, (dims[0], A.shape[1], -1))
  84.     V = np.transpose(np.reshape(V, (-1, dims[1], A.shape[2])), (1,0,2))
  85.     # assert U is left-orthogonal
  86.     # assert V is right-orthogonal
  87.     #print(np.dot(V[0],np.transpose(V[0])) + np.dot(V[1],np.transpose(V[1])))
  88.     return U, S, V
  89.  
  90. def truncate_MPS(U, S, V, m):
  91.     m = min(len(S), m)
  92.     trunc = np.sum(S[m:])
  93.     S = S[0:m]
  94.     U = U[:,:,0:m]
  95.     V = V[:,0:m,:]
  96.     return U,S,V,trunc,m
  97.  
  98. def Expectation(AList, MPO, BList):
  99.     E = [[[1]]]
  100.     for i in range(0,len(MPO)):
  101.         E = contract_from_left(MPO[i], AList[i], E, BList[i])
  102.     return E[0][0][0]
  103.  
  104. class HamiltonianMultiply(sparse.linalg.LinearOperator):
  105.     def __init__(self, E, W, F):
  106.         self.E = E
  107.         self.W = W
  108.         self.F = F
  109.         self.dtype = np.dtype('d')
  110.         self.req_shape = [W.shape[2], E.shape[1], F.shape[2]]
  111.         self.size = self.req_shape[0]*self.req_shape[1]*self.req_shape[2]
  112.         self.shape = [self.size, self.size]
  113.  
  114.     def _matvec(self, A):
  115.         # the einsum function doesn't appear to optimize the contractions properly,
  116.         # so we split it into individual summations in the optimal order
  117.         #R = np.einsum("aij,sik,abst,bkl->tjl",self.E,np.reshape(A, self.req_shape),
  118.         #              self.W,self.F, optimize=True)
  119.         R = np.einsum("aij,sik->ajsk", self.E, np.reshape(A, self.req_shape))
  120.         R = np.einsum("ajsk,abst->bjtk", R, self.W)
  121.         R = np.einsum("bjtk,bkl->tjl", R, self.F)
  122.         return np.reshape(R, -1)
  123.  
  124. ## optimize a single site given the MPO matrix W, and tensors E,F
  125. def optimize_site(A, W, E, F):
  126.     H = HamiltonianMultiply(E,W,F)
  127.     # we choose tol=1E-8 here, which is OK for small calculations.
  128.     # to bemore robust, we should take the tol -> 0 towards the end
  129.     # of the calculation.
  130.     E,V = sparse.linalg.eigsh(H,1,v0=A,which='SA', tol=1E-8)
  131.     return (E[0],np.reshape(V[:,0], H.req_shape))
  132.  
  133. def optimize_two_sites(A, B, W1, W2, E, F, m, dir):
  134.     W = coarse_grain_MPO(W1,W2)
  135.     AA = coarse_grain_MPS(A,B)
  136.     H = HamiltonianMultiply(E,W,F)
  137.     E,V = sparse.linalg.eigsh(H,1,v0=AA,which='SA')
  138.     AA = np.reshape(V[:,0], H.req_shape)
  139.     A,S,B = fine_grain_MPS(AA, [A.shape[0], B.shape[0]])
  140.     A,S,B,trunc,m = truncate_MPS(A,S,B,m)
  141.     if (dir == 'right'):
  142.         B = np.einsum("ij,sjk->sik", np.diag(S), B)
  143.     else:
  144.         assert dir == 'left'
  145.         A = np.einsum("sij,jk->sik", A, np.diag(S))
  146.     return E[0], A, B, trunc, m
  147.  
  148. def two_site_dmrg(MPS, MPO, m, sweeps):
  149.     E = construct_E(MPS, MPO, MPS)
  150.     F = construct_F(MPS, MPO, MPS)
  151.     F.pop()
  152.     for sweep in range(0,int(sweeps/2)):
  153.         for i in range(0, len(MPS)-2):
  154.             Energy,MPS[i],MPS[i+1],trunc,states = optimize_two_sites(MPS[i],MPS[i+1],
  155.                                                                      MPO[i],MPO[i+1],
  156.                                                                      E[-1], F[-1], m, 'right')
  157.             print("Sweep {:} Sites {:},{:}    Energy {:16.12f}    States {:4} Truncation {:16.12f}"
  158.                      .format(sweep*2,i,i+1, Energy, states, trunc))
  159.             E.append(contract_from_left(MPO[i], MPS[i], E[-1], MPS[i]))
  160.             F.pop();
  161.         for i in range(len(MPS)-2, 0, -1):
  162.             Energy,MPS[i],MPS[i+1],trunc,states = optimize_two_sites(MPS[i],MPS[i+1],
  163.                                                                      MPO[i],MPO[i+1],
  164.                                                                      E[-1], F[-1], m, 'left')
  165.             print("Sweep {} Sites {},{}    Energy {:16.12f}    States {:4} Truncation {:16.12f}"
  166.                      .format(sweep*2+1,i,i+1, Energy, states, trunc))
  167.             F.append(contract_from_right(MPO[i+1], MPS[i+1], F[-1], MPS[i+1]))
  168.             E.pop();
  169.     return MPS
  170.            
  171.  
  172. d=2   # local bond dimension
  173. N=100 # number of sites
  174.  
  175. InitialA1 = np.zeros((d,1,1))
  176. InitialA1[0,0,0] = 1
  177. InitialA2 = np.zeros((d,1,1))
  178. InitialA2[1,0,0] = 1
  179.  
  180. ## initial state |01010101>
  181. MPS = [InitialA1, InitialA2] * int(N/2)
  182.  
  183. ## Local operators
  184. I = np.identity(2)
  185. Z = np.zeros((2,2))
  186. Sz = np.array([[0.5,  0  ],
  187.              [0  , -0.5]])
  188. Sp = np.array([[0, 0],
  189.              [1, 0]])
  190. Sm = np.array([[0, 1],
  191.              [0, 0]])
  192.  
  193. ## Hamiltonian MPO
  194. W = np.array([[I, Sz, 0.5*Sp,  0.5*Sm,  Z],
  195.               [Z,  Z,      Z,       Z, Sz],
  196.               [Z,  Z,      Z,       Z, Sm],
  197.               [Z,  Z,      Z,       Z, Sp],
  198.               [Z,  Z,      Z,       Z,  I]])
  199.  
  200. Wfirst = np.array([[I, Sz, 0.5*Sp, 0.5*Sm,   Z]])
  201.  
  202. Wlast = np.array([[Z], [Sz], [Sm], [Sp], [I]])
  203.  
  204. # the complete MPO
  205. MPO = [Wfirst] + ([W] * (N-2)) + [Wlast]
  206.  
  207. HamSquared = product_MPO(MPO, MPO)
  208.  
  209. # 8 sweeps with m=10 states
  210. MPS = two_site_dmrg(MPS, MPO, 10, 8)
  211.  
  212. Energy = Expectation(MPS, MPO, MPS)
  213. print("Final energy expectation value {}".format(Energy))
  214.  
  215. H2 = Expectation(MPS, HamSquared, MPS)
  216. print("variance = {:16.12f}".format(H2 - Energy*Energy))
Get Code
Edit - History - Print - Recent Changes - Search
Page last modified on September 14, 2021, at 04:35 AM