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, dwt = svd("ij, k", wfn0, M)
gaug = einsum("ij, jk -> ik", s, wfn1)
else:
wfn1, s, mps0, dwt = svd("i, jk", wfn0, M)
gaug = einsum("ij, jk -> ik", wfn1, s)
return mps0, gaug
def dot_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 -> Lanmr", lopr, wfn0)
scr2 = einsum("Lanmr, aNnb -> LNbmr", scr1, lmpo)
scr3 = einsum("LNbmr, bMmc -> LNMcr", scr2, rmpo)
sgv0 = einsum("LNMcr, Rcr -> LNMR", scr3, ropr)
return sgv0
def initialize_heisenberg(N, h, J, M):
"""
Initialize the MPS, MPO, lopr and ropr.
"""
# MPS
mpss = sMPX.rand([2] * N, D=M, bc='obc', seed=0)
normalize_factor = 1.0 / gMPX.norm(mpss)
mpss = gMPX.mul(normalize_factor, mpss)
# make MPS right canonical
for i in xrange(N - 1, 0, -1):
mpss[i], gaug = canonicalize(0, mpss[i], M=M)
mpss[i - 1] = einsum("ijk, kl -> ijl", mpss[i - 1], gaug)
# MPO
mpos = np.asarray(heisenberg_mpo(N, h, J))
# lopr
loprs = [COO(np.array([[[1.0]]]))]
# ropr
roprs = [COO(np.array([[[1.0]]]))]
for i in xrange(N - 1, 0, -1):
roprs.append(
renormalize(0, mpos[i], roprs[-1], mpss[i].conj(), mpss[i]))
# NOTE the loprs and roprs should be list currently to support pop()!
return mpss, mpos, loprs, roprs
def diag_onesite(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 = COO(einsum('annb -> anb', mpo0))
#lopr_diag = COO(einsum('lal -> la', lopr))
#ropr_diag = COO(einsum('rbr -> br', ropr))
mpo0_diag = sh.diagonal(mpo0, axes=[1, 2])
lopr_diag = sh.diagonal(lopr, axes=[0, 2])
ropr_diag = sh.diagonal(ropr, axes=[0, 2])
scr1 = einsum('la, anb -> lnb', lopr_diag, mpo0_diag)
diag = einsum('lnb, rb -> lnr', scr1, ropr_diag)
#diag = scr2
# ZHC NOTE check the SDcopy of upcast options
return diag
def optimize_twosite(forward, lmpo, rmpo, lopr, ropr, lwfn, rwfn, M, tol):
"""
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 = diag_twosite(lmpo, rmpo, lopr, ropr)
mps_shape = wfn2.shape
def dot_flat(x):
return dot_twosite(lmpo, rmpo, lopr, ropr,
x.reshape(mps_shape)).ravel()
def compute_precond_flat(dx, e, x0):
return dx / (diag_flat - e)
energy, wfn0 = linalg_helper.davidson(dot_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)
lopr = renormalize(1, mpo0, lopr, wfn0.conj(), wfn0)
return energy, wfn0, wfn1, lopr
else:
wfn0, gaug = canonicalize(0, wfn0, M) # wfn0 R => lmps gaug
wfn1 = einsum("ijk, kl -> ijl", wfn1, gaug)
ropr = renormalize(0, mpo0, ropr, wfn0.conj(), wfn0)
return energy, wfn0, wfn1, ropr
def diag_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-> cr', ropr)
scr1 = einsum('la, anb -> lnb', lopr, lmpo)
scr2 = einsum('lnb, bmc -> lnmc', scr1, rmpo)
diag = einsum('lnmc, cr -> lnmr', scr2, ropr)
return diag
def test_einsum_outer_prod(shape_x, shape_y, descr):
density = 0.4
#np.random.seed(2)
np.set_printoptions(3, linewidth=1000, suppress=True)
x = sparse.random(shape_x, density, format='coo')
x_d = x.todense()
y = sparse.random(shape_y, density, format='coo')
y_d = y.todense()
sout = sh.einsum(descr, x, y)
out = np.einsum(descr, x_d, y_d)
assert_eq(sout, out)
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. should already be conjugated.
ket0 : ndarray
down MPS
Returns
-------
opr1 : ndarray
renormalized block opr.
"""
if forward:
scr = einsum('Lal, lnr -> Lanr', opr0, ket0)
scr = einsum('Lanr, aNnb -> LNbr', scr, mpo0)
opr1 = einsum('LNR, LNbr -> Rbr', bra0, scr)
else:
scr = einsum('LNR, Rbr -> LNbr', bra0, opr0)
scr = einsum('LNbr, aNnb -> Lanr', scr, mpo0)
opr1 = einsum('Lanr, lnr-> Lal ', scr, ket0)
return opr1
def dot_onesite(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 -> Lanr', lopr, wfn0)
scr2 = einsum('Lanr, aNnb -> LNbr', scr1, mpo0)
sgv0 = einsum('LNbr, Rbr -> LNR', scr2, ropr)
return sgv0
def optimize_onesite(forward, mpo0, lopr, ropr, wfn0, wfn1, M, tol):
"""
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 canonicalization.
wfn1 : ndarray
MPS.
M : int
bond dimension
Returns
-------
energy : float or list of floats
The energy of desired root(s).
"""
#diag_flat = diag_onesite(mpo0, lopr, ropr).ravel()
diag_flat = diag_onesite(mpo0, lopr, ropr).ravel().todense()
mps_shape = wfn0.shape
def dot_flat(x):
#return dot_onesite(mpo0, lopr, ropr, x.reshape(mps_shape)).ravel()
return dot_onesite(mpo0, lopr, ropr,
COO(x.reshape(mps_shape))).ravel().todense()
def compute_precond_flat(dx, e, x0):
#return COO(dx.todense() / (diag_flat.todense() - e))
return dx / (diag_flat - e)
#dot_func = sparse.coo.common.dot
dot_func = np.dot
#energy, wfn0 = linalg_helper.davidson(dot_flat, wfn0.ravel(), compute_precond_flat, tol = tol, dot = dot_func)
energy, wfn0 = linalg_helper.davidson(dot_flat,
wfn0.ravel().todense(),
compute_precond_flat,
tol=tol,
dot=dot_func)
#wfn0 = wfn0.reshape(mps_shape)
wfn0 = COO(wfn0.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.conj(), wfn0)
return energy, wfn0, wfn1, lopr
else:
wfn0, gaug = canonicalize(0, wfn0, M) # wfn0 R => lmps gaug
wfn1 = einsum("ijk,kl->ijl", wfn1, gaug)
ropr = renormalize(0, mpo0, ropr, wfn0.conj(), wfn0)
return energy, wfn0, wfn1, ropr
def einsum(idx, *tensors, **kwargs):
if any(isinstance(a, sparse.coo.COO) for a in tensors):
return sparse_helper.einsum(idx, *tensors)
else:
return np.einsum(idx, *tensors, **kwargs)