Recent Changes - Search:

HomePage

PmWiki

pmwiki.org

Tutorials /

MPS-TEBD

This code extends the MPS-DMRG code to calculate the spectral function of the s=1/2 Heisenberg model using TEBD. You can download the code here

  1. #####################################################
  2. # Simple TEBD code using MPS/MPO representations    #
  3. # Ian McCulloch 25/5/2020                           #
  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. import cmath
  12. from mpl_toolkits import mplot3d
  13. import matplotlib.pyplot as plt
  14. import sys
  15. import copy
  16.  
  17. np.set_printoptions(threshold=sys.maxsize)
  18.  
  19. # pseudo-inverse of a diagonal matrix.
  20. # We want to cut off values smaller than the threshold. We use
  21. # Tikhonov regularization, s -> s / (s^2 + a^2), where a is the cutoff
  22. def pseudo_inverse(S):
  23.     a2 = 1e-7 ** 2
  24.     return np.array([x / (x**2 + a2) for x in S])
  25.  
  26. def norm_frob(A):
  27.     return np.sum(np.multiply(np.conj(A),A))
  28.  
  29. ## take an MPS in right orthogonal form and construct the canonical form
  30. ## Lambda_0 Gamma_0 Lambda_1 Gamma_1 ... Gamma_N Lambda_N+1
  31. def canonicalize(MPS):
  32.     Lambda = [np.array([1])]
  33.     Gamma = []
  34.     V = np.array([[1]])
  35.     lastS = np.array([1])
  36.     for i in range(len(MPS)):
  37.         # reshape into m x (dm)
  38.         Temp = np.einsum("ij,sjk->sik", V, MPS[i])
  39.         s = np.shape(Temp)
  40.         A = np.reshape(Temp, (s[0]*s[1], s[2]))
  41.         U, S, V = np.linalg.svd(A, full_matrices=0)
  42.         Gamma = Gamma + [np.einsum("i,sij->sij", pseudo_inverse(lastS), np.reshape(U, (s[0], s[1], np.shape(U)[1])))]
  43.         Lambda = Lambda + [S]
  44.         V = np.einsum("i,ij->ij", S, V)
  45.         lastS = S
  46.     # assert V is 1x1 identity matrix
  47.     return Lambda, Gamma
  48.  
  49. def orthonormalize(MPS):
  50.     # sweep left-to-right with SVD's
  51.     V = np.array([[1]])
  52.     for i in range(len(MPS)):
  53.         s = np.shape(MPS[i])
  54.         A = np.reshape(np.einsum("ij,sjk->sik",V,MPS[i]), (s[0]*s[1], s[2]))
  55.         U,S,V = np.linalg.svd(A, full_matrices=0)
  56.         MPS[i] = np.reshape(U, (s[0], s[1], np.shape(U)[1]))
  57.         V = np.einsum("i,ij->ij", S, V)
  58.     # sweep right to left
  59.     U = np.array([[1]])
  60.     for i in range(len(MPS)-1, -1, -1):
  61.         s = np.shape(MPS[i])
  62.         A = np.reshape(np.einsum("sij,jk->isk", MPS[i], U), (s[1], s[0]*np.shape(U)[1]))
  63.         U,S,V = np.linalg.svd(A, full_matrices=0)
  64.         MPS[i] = np.einsum("isk->sik", np.reshape(V, (np.shape(V)[0], s[0], s[2])))
  65.         U = np.einsum("ij,j->ij", U, S)
  66.     # the final U is the norm, which we discard
  67.     return MPS    
  68.  
  69. def TEBD_Gate(Lambda0, Gamma1, Lambda1, Gamma2, Lambda2, Gate, m):
  70.     # contract the outside lambda's into Gamma
  71.     Gamma1 = np.einsum("i,sij->sij", Lambda0, Gamma1)
  72.     Gamma2 = np.einsum("sij,j->sij", Gamma2, Lambda2)
  73.     # apply the centre lambda
  74.     Gamma2 = np.einsum("i,sij->sij", Lambda1, Gamma2)
  75.     # make 2-body tensor Gamma1 Lambda1 Gamma2
  76.     Temp = np.einsum("sij,tjk->stik", Gamma1, Gamma2)
  77.     # apply the 2-body gate and reorder to apply the svd
  78.     TShape = np.shape(Temp)
  79.     GShape = (TShape[0],TShape[1],TShape[0],TShape[1])
  80.     Temp = np.einsum("uvst,stik->uivk", np.reshape(Gate,GShape), Temp)
  81.     # SVD
  82.     s = np.shape(Temp)
  83.     A = np.reshape(Temp, (s[0]*s[1], s[2]*s[3]))
  84.     U, S, V = np.linalg.svd(A, full_matrices=0)
  85.     # truncate
  86.     m = min(len(S), m)
  87.     S = S[0:m]
  88.     U = U[:,0:m]
  89.     V = V[0:m,:]
  90.     # reshape back to Gamma
  91.     Gamma1 = np.reshape(U, (s[0], s[1], np.shape(U)[1]))
  92.     Gamma2 = np.einsum("isj->sij", np.reshape(V, (np.shape(V)[0], s[2], s[3])))
  93.     # invert the outer Lambda matrices
  94.     Gamma1 = np.einsum("i,sij->sij", pseudo_inverse(Lambda0), Gamma1)
  95.     Gamma2 = np.einsum("sij,j->sij", Gamma2, pseudo_inverse(Lambda2))
  96.     return Gamma1, S, Gamma2
  97.  
  98. def EvolveEvenBonds(Lambda, Gamma, Gate, m):
  99.     for i in range(0, len(Gamma)-1, 2):
  100.         Gamma[i], Lambda[i+1], Gamma[i+1] = TEBD_Gate(Lambda[i], Gamma[i], Lambda[i+1], Gamma[i+1], Lambda[i+2], Gate, m)
  101.     return Lambda, Gamma
  102.  
  103. def EvolveOddBonds(Lambda, Gamma, Gate, m):
  104.     for i in range(1, len(Gamma)-1, 2):
  105.         Gamma[i], Lambda[i+1], Gamma[i+1] = TEBD_Gate(Lambda[i], Gamma[i], Lambda[i+1], Gamma[i+1], Lambda[i+2], Gate, m)
  106.     return Lambda, Gamma
  107.  
  108. def MPS_from_canonical(Lambda, Gamma):
  109.     MPS = []
  110.     for i in range(len(Gamma)):
  111.         MPS = MPS + [np.einsum("i,sij->sij", Lambda[i], Gamma[i])]
  112.     return MPS
  113.  
  114. # evolve a timestep, using simple 2nd order Lie-Suzuki-Trotter decomposition
  115. def EvolveTimestep(Lambda, Gamma, H, m, dt):
  116.     Lambda,Gamma = EvolveEvenBonds(Lambda, Gamma, scipy.linalg.expm(complex(0,-dt/2)*H), m)
  117.     Lambda,Gamma = EvolveOddBonds(Lambda, Gamma, scipy.linalg.expm(complex(0,-dt)*H), m)
  118.     Lambda,Gamma = EvolveEvenBonds(Lambda, Gamma, scipy.linalg.expm(complex(0,-dt/2)*H), m)
  119.     return Lambda,Gamma
  120.    
  121. def LocalExpectationValue(Lambda, Gamma, Site, Op):
  122.     G = np.einsum("i,sij->sij", Lambda[Site], Gamma[Site])
  123.     G = np.einsum("sij,j->sij", G, Lambda[Site+1])
  124.     H = np.einsum("st,tij->sij", Op, G)
  125.     return np.einsum("sij,sij", np.conj(G), H)
  126.  
  127. # return an array of the local expectation value at every site
  128. def Density(Lambda, Gamma, Op):
  129.     return np.array([LocalExpectationValue(Lambda, Gamma, n, Op) for n in range(0,len(Gamma))])
  130.  
  131. ## initial E and F matrices for the left and right vacuum states
  132. def initial_E(W):
  133.     E = np.zeros((W.shape[0],1,1))
  134.     E[0] = 1
  135.     return E
  136.  
  137. def initial_F(W):
  138.     F = np.zeros((W.shape[1],1,1))
  139.     F[-1] = 1
  140.     return F
  141.  
  142. def contract_from_right(W, A, F, B):
  143.     # the einsum function doesn't appear to optimize the contractions properly,
  144.     # so we split it into individual summations in the optimal order
  145.     #return np.einsum("abst,sij,bjl,tkl->aik",W,A,F,B, optimize=True)
  146.     Temp = np.einsum("sij,bjl->sbil", np.conj(A), F)
  147.     Temp = np.einsum("sbil,abst->tail", Temp, W)
  148.     return np.einsum("tail,tkl->aik", Temp, B)
  149.  
  150. def contract_from_left(W, A, E, B):
  151.     # the einsum function doesn't appear to optimize the contractions properly,
  152.     # so we split it into individual summations in the optimal order
  153.     # return np.einsum("abst,sij,aik,tkl->bjl",W,A,E,B, optimize=True)
  154.     Temp = np.einsum("sij,aik->sajk", np.conj(A), E)
  155.     Temp = np.einsum("sajk,abst->tbjk", Temp, W)
  156.     return np.einsum("tbjk,tkl->bjl", Temp, B)
  157.  
  158. # construct the F-matrices for all sites except the first
  159. def construct_F(Alist, MPO, Blist):
  160.     F = [initial_F(MPO[-1])]
  161.  
  162.     for i in range(len(MPO)-1, 0, -1):
  163.         F.append(contract_from_right(MPO[i], Alist[i], F[-1], Blist[i]))
  164.     return F
  165.  
  166. def construct_E(Alist, MPO, Blist):
  167.     return [initial_E(MPO[0])]
  168.  
  169.  
  170.  
  171. # 2-1 coarse-graining of two site MPO into one site
  172. def coarse_grain_MPO(W, X):
  173.     return np.reshape(np.einsum("abst,bcuv->acsutv",W,X),
  174.                       [W.shape[0], X.shape[1],
  175.                        W.shape[2]*X.shape[2],
  176.                        W.shape[3]*X.shape[3]])
  177.  
  178. def product_W(W1, W2):
  179.     return np.reshape(np.einsum("abst,cdtu->acbdsu", W1, W2), [W1.shape[0]*W2.shape[0],
  180.                                                                W1.shape[1]*W2.shape[1],
  181.                                                                W1.shape[2],W2.shape[3]])
  182.  
  183. def product_MPO(M1, M2):
  184.     assert len(M1) == len(M2)
  185.     Result = []
  186.     for i in range(0, len(M1)):
  187.         Result.append(product_W(M1[i], M2[i]))
  188.     return Result
  189.  
  190.  
  191. # 2-1 coarse-graining of two-site MPS into one site
  192. def coarse_grain_MPS(A,B):
  193.     return np.reshape(np.einsum("sij,tjk->stik",A,B),
  194.                       [A.shape[0]*B.shape[0], A.shape[1], B.shape[2]])
  195.  
  196. def fine_grain_MPS(A, dims):
  197.     assert A.shape[0] == dims[0] * dims[1]
  198.     Theta = np.transpose(np.reshape(A, dims + [A.shape[1], A.shape[2]]),
  199.                          (0,2,1,3))
  200.     M = np.reshape(Theta, (dims[0]*A.shape[1], dims[1]*A.shape[2]))
  201.     U, S, V = np.linalg.svd(M, full_matrices=0)
  202.     U = np.reshape(U, (dims[0], A.shape[1], -1))
  203.     V = np.transpose(np.reshape(V, (-1, dims[1], A.shape[2])), (1,0,2))
  204.     # assert U is left-orthogonal
  205.     # assert V is right-orthogonal
  206.     #print(np.dot(V[0],np.transpose(V[0])) + np.dot(V[1],np.transpose(V[1])))
  207.     return U, S, V
  208.  
  209. def truncate_MPS(U, S, V, m):
  210.     m = min(len(S), m)
  211.     trunc = np.sum(S[m:])
  212.     S = S[0:m]
  213.     U = U[:,:,0:m]
  214.     V = V[:,0:m,:]
  215.     return U,S,V,trunc,m
  216.  
  217. def Expectation(AList, MPO, BList):
  218.     E = [[[1]]]
  219.     for i in range(0,len(MPO)):
  220.         E = contract_from_left(MPO[i], AList[i], E, BList[i])
  221.     return E[0][0][0]
  222.  
  223. class HamiltonianMultiply(sparse.linalg.LinearOperator):
  224.     def __init__(self, E, W, F):
  225.         self.E = E
  226.         self.W = W
  227.         self.F = F
  228.         self.dtype = np.dtype('d')
  229.         self.req_shape = [W.shape[2], E.shape[1], F.shape[2]]
  230.         self.size = self.req_shape[0]*self.req_shape[1]*self.req_shape[2]
  231.         self.shape = [self.size, self.size]
  232.  
  233.     def _matvec(self, A):
  234.         # the einsum function doesn't appear to optimize the contractions properly,
  235.         # so we split it into individual summations in the optimal order
  236.         #R = np.einsum("aij,sik,abst,bkl->tjl",self.E,np.reshape(A, self.req_shape),
  237.         #              self.W,self.F, optimize=True)
  238.         R = np.einsum("aij,sik->ajsk", self.E, np.reshape(A, self.req_shape))
  239.         R = np.einsum("ajsk,abst->bjtk", R, self.W)
  240.         R = np.einsum("bjtk,bkl->tjl", R, self.F)
  241.         return np.reshape(R, -1)
  242.  
  243. ## optimize a single site given the MPO matrix W, and tensors E,F
  244. def optimize_site(A, W, E, F):
  245.     H = HamiltonianMultiply(E,W,F)
  246.     # we choose tol=1E-8 here, which is OK for small calculations.
  247.     # to bemore robust, we should take the tol -> 0 towards the end
  248.     # of the calculation.
  249.     E,V = sparse.linalg.eigsh(H,1,v0=A,which='SA', tol=1E-8)
  250.     return (E[0],np.reshape(V[:,0], H.req_shape))
  251.  
  252. def optimize_two_sites(A, B, W1, W2, E, F, m, dir):
  253.     W = coarse_grain_MPO(W1,W2)
  254.     AA = coarse_grain_MPS(A,B)
  255.     H = HamiltonianMultiply(E,W,F)
  256.     E,V = sparse.linalg.eigsh(H,1,v0=AA,which='SA')
  257.     AA = np.reshape(V[:,0], H.req_shape)
  258.     A,S,B = fine_grain_MPS(AA, [A.shape[0], B.shape[0]])
  259.     A,S,B,trunc,m = truncate_MPS(A,S,B,m)
  260.     if (dir == 'right'):
  261.         B = np.einsum("ij,sjk->sik", np.diag(S), B)
  262.     else:
  263.         assert dir == 'left'
  264.         A = np.einsum("sij,jk->sik", A, np.diag(S))
  265.     return E[0], A, B, trunc, m
  266.  
  267. def two_site_dmrg(MPS, MPO, m, sweeps):
  268.     E = construct_E(MPS, MPO, MPS)
  269.     F = construct_F(MPS, MPO, MPS)
  270.     F.pop()
  271.     for sweep in range(0,int(sweeps/2)):
  272.         for i in range(0, len(MPS)-2):
  273.             Energy,MPS[i],MPS[i+1],trunc,states = optimize_two_sites(MPS[i],MPS[i+1],
  274.                                                                      MPO[i],MPO[i+1],
  275.                                                                      E[-1], F[-1], m, 'right')
  276.             print("Sweep {:} Sites {:},{:}    Energy {:16.12f}    States {:4} Truncation {:16.12f}"
  277.                      .format(sweep*2,i,i+1, Energy, states, trunc))
  278.             E.append(contract_from_left(MPO[i], MPS[i], E[-1], MPS[i]))
  279.             F.pop();
  280.         for i in range(len(MPS)-2, 0, -1):
  281.             Energy,MPS[i],MPS[i+1],trunc,states = optimize_two_sites(MPS[i],MPS[i+1],
  282.                                                                      MPO[i],MPO[i+1],
  283.                                                                      E[-1], F[-1], m, 'left')
  284.             print("Sweep {} Sites {},{}    Energy {:16.12f}    States {:4} Truncation {:16.12f}"
  285.                      .format(sweep*2+1,i,i+1, Energy, states, trunc))
  286.             F.append(contract_from_right(MPO[i+1], MPS[i+1], F[-1], MPS[i+1]))
  287.             E.pop();
  288.     return MPS
  289.  
  290.  
  291.            
  292.  
  293. d=2   # local bond dimension
  294. N=40 # number of sites
  295.  
  296. InitialA1 = np.zeros((d,1,1))
  297. InitialA1[0,0,0] = 1
  298. InitialA2 = np.zeros((d,1,1))
  299. InitialA2[1,0,0] = 1
  300.  
  301. ## initial state |01010101>
  302. MPS = [InitialA1, InitialA2] * int(N/2)
  303.  
  304. ## Local operators
  305. I = np.identity(2)
  306. Z = np.zeros((2,2))
  307. Sz = np.array([[0.5,  0  ],
  308.              [0  , -0.5]])
  309. Sp = np.array([[0, 1],
  310.              [0, 0]])
  311. Sm = np.array([[0, 0],
  312.              [1, 0]])
  313.  
  314. ## Hamiltonian MPO
  315. W = np.array([[I, Sz, 0.5*Sp,  0.5*Sm,  Z],
  316.               [Z,  Z,      Z,       Z, Sz],
  317.               [Z,  Z,      Z,       Z, Sm],
  318.               [Z,  Z,      Z,       Z, Sp],
  319.               [Z,  Z,      Z,       Z,  I]])
  320.  
  321. Wfirst = np.array([[I, Sz, 0.5*Sp, 0.5*Sm,   Z]])
  322.  
  323. Wlast = np.array([[Z], [Sz], [Sm], [Sp], [I]])
  324.  
  325. # the complete MPO
  326. MPO = [Wfirst] + ([W] * (N-2)) + [Wlast]
  327.  
  328. HamSquared = product_MPO(MPO, MPO)
  329.  
  330. m = 10
  331. sweeps = 8
  332. MPS = two_site_dmrg(MPS, MPO, m, sweeps)
  333.  
  334. Energy = Expectation(MPS, MPO, MPS)
  335. print("Final energy expectation value {}".format(Energy))
  336.  
  337. H2 = Expectation(MPS, HamSquared, MPS)
  338. print("variance = {:16.12f}".format(H2 - Energy*Energy))
  339.  
  340. # sage the groundstate, we need it later
  341. GS = copy.deepcopy(MPS)
  342.  
  343. # Act on the MPS with a local operator at some site
  344. MPS[N//2] = np.einsum("st,tij->sij", Sp, MPS[N//2])
  345.  
  346. # orthonormalize the MPS
  347. MPS = orthonormalize(MPS)
  348.  
  349. # canonicalize in preparation for TEND
  350. Lambda,Gamma = canonicalize(MPS)
  351.  
  352. # 2 sites of the Hamiltonian.  Full Hamiltonian is sum over translations of H2site
  353. H2site = 0.5 * (np.kron(Sp, Sm) + np.kron(Sm, Sp)) + np.kron(Sz, Sz)
  354.  
  355. # matrix of Sz values as a function of time and position.  Rows are timesteps, columns are position
  356. SzMatrix = [Density(Lambda, Gamma, Sz)]
  357.  
  358. # for the correlation function, we want f(x,t) = <0|S-(t) S+(0)|0> = <0|S- e^{iHt} S+(0)|0> (ignoring phase factor)
  359. IdentityMPO = [np.array([[I]]) for i in range(N)]
  360. Correlator = []
  361. print(I)
  362. for i in range(N):
  363.     c = copy.deepcopy(IdentityMPO)
  364.     c[i][0,0] = Sm
  365.     Correlator = Correlator + [c]
  366.  
  367. # initial correlation
  368. Correlation = [ [Expectation(GS, Correlator[n], MPS) for n in range(N)] ]
  369.  
  370. # timestep
  371. dt = 0.02
  372. NT = 600 # number of timesteps
  373. Skip = 1 # only calculate correlations after this many timesteps
  374. for i in range(NT):
  375.     print("timestep {} / {}".format(i,NT))
  376.     Lambda,Gamma = EvolveTimestep(Lambda, Gamma, H2site, 20, dt)
  377.     SzMatrix = SzMatrix + [Density(Lambda, Gamma, Sz)]
  378.     # correlation function
  379.     if (i % Skip == 0):
  380.         MPS = MPS_from_canonical(Lambda, Gamma)
  381.         Correlation = Correlation + [ [Expectation(GS, Correlator[n], MPS) for n in range(N)] ]
  382.    
  383.  
  384. X = np.outer(np.linspace(0,N-1,N), np.ones(NT+1))
  385. T = np.outer(np.ones(N), np.linspace(0,dt,NT+1))
  386.  
  387. SzMatrix = np.transpose(np.array(SzMatrix).real)
  388.  
  389. #np.set_printoptions(threshold=sys.maxsize)
  390. #print(SzMatrix)
  391.  
  392. fig = plt.figure()
  393. ax = plt.axes(projection='3d')
  394. ax.plot_surface(X, T, SzMatrix, cmap='viridis', edgecolor='none')
  395. ax.set_xlabel("X")
  396. ax.set_ylabel("Time")
  397. ax.set_zlabel("Sz")
  398. ax.set_title('Evolution of the magnetization')
  399. fig.show()
  400.  
  401. # correlation function
  402. Correlation = np.transpose(np.array(Correlation))
  403.  
  404. figc = plt.figure()
  405. axc = plt.axes(projection='3d')
  406. axc.plot_surface(X, T, np.absolute(Correlation), cmap='viridis', edgecolor='none')
  407. axc.set_xlabel("X")
  408. axc.set_ylabel("Time")
  409. axc.set_zlabel("Sz")
  410. axc.set_title('Time-dependent correlation')
  411. figc.show()
  412.  
  413. # Calculate the spectral function
  414. # time evolution is symmetric, so we can add the -ve time component
  415. Correlation = np.concatenate((np.flip(Correlation[:,1:], axis=1), Correlation), axis=1)
  416.  
  417. # FFT
  418. Spec = np.absolute(np.fft.fft2(Correlation))
  419.  
  420. # we didn't correct for the groundstate energy, so just plot by hand the relevant part
  421. Spec = Spec[:,68:52:-1]
  422. NumOmega=np.shape(Spec)[1]
  423.  
  424. K = np.outer(np.linspace(0,np.pi*2,N), np.ones(NumOmega))
  425. W = np.outer(np.ones(N), np.linspace(0,NumOmega*np.pi/(dt*NT), NumOmega))
  426.  
  427. print(np.shape(Spec))
  428. print(np.shape(K))
  429. print(np.shape(W))
  430.  
  431. fig2 = plt.figure()
  432. ax2 = plt.axes(projection='3d')
  433. ax2.plot_surface(K, W, np.absolute(Spec), cmap='viridis', edgecolor='none')
  434. ax2.set_xlabel("k")
  435. ax2.set_ylabel("Energy")
  436. ax2.set_zlabel("Spectral density")
  437. ax2.set_title('Spectral function')
  438. fig2.show()
  439. input()
Get Code
Edit - History - Print - Recent Changes - Search
Page last modified on May 25, 2020, at 02:31 AM