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)