def canonicalize(forward, wfn0, M = 0): """ Canonicalize the wavefunction. Parameters ---------- forward : int 0 for left and 1 for right. wfn0 : ndarray current MPS. M : int bond dimension Returns ------- mps0 : ndarray canonicalized mps. gaug : ndarray gauge, i.e. sigma * v """ if forward: mps0, s, wfn1 = svd("ij,k", wfn0, M) gaug = einsum("ij, jk -> ik", diag(s), wfn1) else: wfn1, s, mps0 = svd("i,jk", wfn0, M) gaug = einsum("ij, jk -> ik", wfn1, diag(s)) return mps0, gaug
def compute_sigmavector_twosite(lmpo, rmpo, lopr, ropr, wfn0): """ Compute the sigma vector, i.e. sigma = H * c used for Davidson algorithm, in the twosite algorithm _L N M R_ |___|_|___| |___|_|___| """ scr1 = einsum("laL,lnmr->aLnmr", lopr, wfn0) scr2 = einsum("aLnmr,anNb->LmNbr", scr1, lmpo) scr3 = einsum("LmNbr,bmMc->LNMrc", scr2, rmpo) sgv0 = einsum("LNMcr,rcR->LNMR", scr3, ropr) return sgv0
def compute_diagonal_elements_twosite(lmpo, rmpo, lopr, ropr): lmpo_diag = einsum('annb -> anb', lmpo) rmpo_diag = einsum('bmmc -> bmc', rmpo) lopr = einsum('lal->la', lopr) ropr = einsum('rcr->rc', ropr) scr1 = einsum('la, anb->lnb', lopr, lmpo) scr2 = einsum('lnb, bmc->lnmc', scr1, rmpo) scr3 = einsum('lnmc, rc -> lnmr', scr2, ropr) return scr3
def renormalize(forward, mpo0, opr0, bra0, ket0): """ Renormalized the block opr. Parameters ---------- forward : int 0 for left and 1 for right. mpo0 : ndarray MPO. opr0 : ndarray block opr. bra0 : ndarray upper MPS. ket0 : ndarray down MPS Returns ------- opr1 : ndarray renormalized block opr. """ if forward: scr1 = einsum('laL, LNR -> laNR', opr0, bra0.conj()) scr2 = einsum('laNR, anNb-> lnbR ', scr1, mpo0) opr1 = einsum('lnbR, lnr-> rbR ', scr2, ket0) else: scr1 = einsum('LNR, rbR -> rbNL', bra0, opr0) scr2 = einsum('rbNL, anNb -> rnaL', scr1, mpo0) opr1 = einsum('rnaL, lnr-> laL ', scr2, ket0) return opr1
def compute_diagonal_elements(mpo0, lopr, ropr): """ Compute the diagonal elements of sandwich <L|MPO|R>, used as preconditioner of Davidson algorithm. Math ---------- _l n r_ |___|___| |_l | r_| n Parameters ---------- mpo0 : ndarray The MPO. lopr : ndarray left block operators ropr : ndarray right block operators Returns ------- diag : ndarray The diagonal element, stored as 'lnr'. """ mpo0_diag = einsum('annb -> anb', mpo0) lopr_diag = einsum('lal -> la', lopr) mpo0_diag = einsum('rbr -> rb', ropr) scr1 = einsum('la, anb -> lnb', lopr_diag, mpo0_diag) scr2 = einsum('lnb, rb -> lnr', scr1, ropr_diag) diag = scr2 # ZHC NOTE check the SDcopy of upcast options return diag
def compute_sigmavector(mpo0, lopr, ropr, wfn0): """ Compute the sigma vector, i.e. sigma = H * c used for Davidson algorithm. Math ---------- _L N R_ |___|___| |___|___| Parameters ---------- mpo0 : ndarray The MPO. lopr : ndarray left block operators ropr : ndarray right block operators wfn0 : ndarray The current MPS. (wavefunction for desired roots) Returns ------- sgv0 : ndarray The sigma vector, stored as LNR. """ # ZHC NOTE the contraction order and stored structure may be optimized. scr1 = einsum('laL, lnr -> rnaL', lopr, wfn0) scr2 = einsum('rnaL, anNb -> rbNL', scr1, mpo0) sgv0 = einsum('rbNL, rbR -> LNR', scr2, ropr) return sgv0
def optimize_onesite(forward, mpo0, lopr, ropr, wfn0, wfn1, M = 0): """ Optimization for onesite algorithm. Parameters ---------- forward : int 0 for left and 1 for right. mpo0 : ndarray MPO. lopr : ndarray left block opr. ropr : ndarray right block opr. wfn0 : ndarray MPS for canoicalization. wfn1 : ndarray MPS. M : int bond dimension Returns ------- energy : float or list of floats The energy of desired root(s). """ # def davidson(x0, diag_flat): # """ # Davidson algorithm. # Parameters # ---------- # x0 : ndarray # initial state. # diag_flat : ndarray # precomputed diagonal elements, 1D array. # Returns # ------- # energy : float or list of floats # The energy of desired root(s). # coeff : ndarray or list of ndarray # The wavefunction. # """ diag_flat = compute_diagonal_elements(mpo0, lopr, ropr).ravel() mps_shape = wfn0.shape def compute_sigma_flat(x): return compute_sigmavector(mpo0, lopr, ropr, x.reshape(mps_shape)).ravel() def compute_precond_flat(dx, e, x0): return dx / (diag_flat - e) energy, wfn0 = lib.linalg_helper.davidson(compute_sigma_flat, wfn0.ravel(), compute_precond_flat) wfn0 = wfn.reshape(mps_shape) if forward: wfn0, gaug = canonicalize(1, wfn0, M) # wfn0 R => lmps gaug wfn1 = einsum("ij,jkl->ikl", gaug, wfn1) lopr = renormalize(1, mpo0, lopr, wfn0, wfn0) else: wfn0, gaug = canonicalize(0, wfn0, M) # wfn0 R => lmps gaug wfn1 = einsum("ijk,kl->ijl", wfn1, gaug) ropr = renormalize(0, mpo0, ropr, wfn0, wfn0) return energy, wfn0, wfn1, lopr, ropr
def optimize_twosite(forward, lmpo, rmpo, lopr, ropr, lwfn, rwfn, M=0) """ Optimization for twosite algorithm. Parameters ---------- M : int bond dimension Returns ------- energy : float or list of floats The energy of desired root(s). """ wfn2 = einsum("lnr, rms -> lnms", lwfn, rwfn) diag = compute_diagonal_elements(lmpo, rmpo, lopr, ropr) mps_shape = wfn2.shape def compute_sigma_flat(x): return compute_sigmavector(mpo0, lopr, ropr, x.reshape(mps_shape)).ravel() def compute_precond_flat(dx, e, x0): return dx / (diag_flat - e) energy, wfn0 = lib.linalg_helper.davidson(compute_sigma_flat, wfn2.ravel(), compute_precond_flat) wfn0 = wfn0.reshape(mps_shape) if forward: wfn0, gaug = canonicalize(1, wfn0, M) # wfn0 R => lmps gaug wfn1 = einsum("ij,jkl->ikl", gaug, wfn1)