def expectations(self, mpos) -> np.ndarray:
if len(mpos) < 3:
return np.array([self.expectation(mpo) for mpo in mpos])
assert 2 < len(mpos)
# id can be used as efficient hash because of `Matrix` implementation
mpo_ids = np.array([[id(m) for m in mpo] for mpo in mpos])
common_mpo_ids = mpo_ids[0].copy()
mpo0_unique_idx = np.where(np.sum(mpo_ids == common_mpo_ids, axis=0) == 1)[0][0]
common_mpo_ids[mpo0_unique_idx] = mpo_ids[1][mpo0_unique_idx]
x, unique_idx = np.where(mpo_ids != common_mpo_ids)
# should find one at each line
assert np.allclose(x, np.arange(len(mpos)))
common_mpo = list(mpos[0])
common_mpo[mpo0_unique_idx] = mpos[1][mpo0_unique_idx]
self_conj = self._expectation_conj()
environ = Environ()
environ.construct(self, self_conj, common_mpo, "l")
environ.construct(self, self_conj, common_mpo, "r")
res_list = []
for idx, mpo in zip(unique_idx, mpos):
l = environ.read("l", idx - 1)
r = environ.read("r", idx + 1)
path = self._expectation_path()
res = multi_tensor_contract(path, l, self[idx], mpo[idx], self_conj[idx], r)
res_list.append(float(res.real))
return np.array(res_list)
def hop(x):
# H*X
nonlocal count
count += 1
assert len(x) == xsize
tda_coeff = reshape_x(x)
res = [np.zeros_like(coeff) if coeff is not None else None for coeff in tda_coeff]
# fix ket and sweep bra and accumulate into res
for ims in range(site_num):
if tda_coeff[ims] is None:
assert tangent_u[ims] is None
continue
# mix-canonical mps
mps_tangent = merge(mps_l_cano, mps_r_cano, ims+1)
mps_tangent[ims] = tensordot(tangent_u[ims], tda_coeff[ims], (-1, 0))
mps_tangent_conj = mps_r_cano.copy()
environ = Environ(mps_tangent, mpo, "R", mps_conj=mps_tangent_conj)
for ims_conj in range(site_num):
ltensor = environ.GetLR(
"L", ims_conj-1, mps_tangent, mpo, itensor=None,
mps_conj=mps_tangent_conj,
method="System"
)
rtensor = environ.GetLR(
"R", ims_conj+1, mps_tangent, mpo, itensor=None,
mps_conj=mps_tangent_conj,
method="Enviro"
)
if tda_coeff[ims_conj] is not None:
# S-a l-S
# d
# O-b-O-f-O
# e
# S-c k-S
path = [
([0, 1], "abc, cek -> abek"),
([2, 0], "abek, bdef -> akdf"),
([1, 0], "akdf, lfk -> adl"),
]
out = multi_tensor_contract(
path, ltensor, asxp(mps_tangent[ims_conj]),
asxp(mpo[ims_conj]), rtensor
)
res[ims_conj] += asnumpy(tensordot(tangent_u[ims_conj], out,
([0,1], [0,1])))
# mps_conj combine
mps_tangent_conj[ims_conj] = mps_l_cano[ims_conj]
res = [mat for mat in res if mat is not None]
return np.concatenate(res, axis=None)
def expectation(self, mpo, self_conj=None) -> float:
if self_conj is None:
self_conj = self._expectation_conj()
environ = Environ()
environ.construct(self, self_conj, mpo, "r")
l = ones((1, 1, 1))
r = environ.read("r", 1)
path = self._expectation_path()
return float(multi_tensor_contract(path, l, self[0], mpo[0], self_conj[0], r).real)
def test_environ():
mps = Mps.random(holstein_model, 1, 10)
mpo = Mpo(holstein_model)
mps = mps.evolve(mpo, 10)
environ = Environ(mps, mpo)
for i in range(len(mps) - 1):
l = environ.read("L", i)
r = environ.read("R", i + 1)
e = complex(tensordot(l, r, axes=((0, 1, 2), (0, 1, 2)))).real
assert pytest.approx(e) == mps.expectation(mpo)
def test_environ_multi_mpo(mpdm):
mps = Mps.random(holstein_model, 1, 10)
if mpdm:
mps = MpDm.from_mps(mps)
mpo = Mpo(holstein_model)
mps = mps.evolve(mpo, 10)
environ = Environ(mps, mpo)
environ_multi_mpo = Environ(mps, [mpo])
for i in range(len(mps) - 1):
l = environ.read("L", i)
r = environ.read("R", i + 1)
l2 = environ_multi_mpo.read("L", i)
r2 = environ_multi_mpo.read("R", i + 1)
assert np.allclose(asnumpy(l), asnumpy(l2))
assert np.allclose(asnumpy(r), asnumpy(r2))
def optimize_mps_dmrg(mps, mpo):
"""
1 or 2 site optimization procedure
"""
method = mps.optimize_config.method
procedure = mps.optimize_config.procedure
inverse = mps.optimize_config.inverse
nroots = mps.optimize_config.nroots
assert method in ["2site", "1site"]
# print("optimization method", method)
nexciton = mps.nexciton
# construct the environment matrix
environ = Environ()
environ.construct(mps, mps, mpo, "L")
nMPS = len(mps)
# construct each sweep cycle scheme
if method == "1site":
loop = [["R", i]
for i in range(nMPS - 1, -1, -1)] + [["L", i]
for i in range(0, nMPS)]
else:
loop = [["R", i]
for i in range(nMPS - 1, 0, -1)] + [["L", i]
for i in range(1, nMPS)]
# initial matrix
ltensor = ones((1, 1, 1))
rtensor = ones((1, 1, 1))
energies = []
for isweep, (mmax, percent) in enumerate(procedure):
logger.debug(f"mmax, percent: {mmax}, {percent}")
logger.debug(f"energy: {mps.expectation(mpo)}")
logger.debug(f"{mps}")
for system, imps in loop:
if system == "R":
lmethod, rmethod = "Enviro", "System"
else:
lmethod, rmethod = "System", "Enviro"
if method == "1site":
lsite = imps - 1
addlist = [imps]
else:
lsite = imps - 2
addlist = [imps - 1, imps]
ltensor = environ.GetLR("L",
lsite,
mps,
mps,
mpo,
itensor=ltensor,
method=lmethod)
rtensor = environ.GetLR("R",
imps + 1,
mps,
mps,
mpo,
itensor=rtensor,
method=rmethod)
# get the quantum number pattern
qnmat, qnbigl, qnbigr = construct_qnmat(mps, mpo.ephtable,
mpo.pbond_list, addlist,
method, system)
cshape = qnmat.shape
# hdiag
tmp_ltensor = einsum("aba -> ba", ltensor)
tmp_MPOimps = einsum("abbc -> abc", mpo[imps])
tmp_rtensor = einsum("aba -> ba", rtensor)
if method == "1site":
# S-a c f-S
# O-b-O-g-O
# S-a c f-S
path = [([0, 1], "ba, bcg -> acg"), ([1, 0], "acg, gf -> acf")]
hdiag = multi_tensor_contract(path, tmp_ltensor, tmp_MPOimps,
tmp_rtensor)[(qnmat == nexciton)]
# initial guess b-S-c
# a
cguess = mps[imps][qnmat == nexciton]
else:
# S-a c d f-S
# O-b-O-e-O-g-O
# S-a c d f-S
tmp_MPOimpsm1 = einsum("abbc -> abc", mpo[imps - 1])
path = [
([0, 1], "ba, bce -> ace"),
([0, 1], "edg, gf -> edf"),
([0, 1], "ace, edf -> acdf"),
]
hdiag = multi_tensor_contract(path, tmp_ltensor, tmp_MPOimpsm1,
tmp_MPOimps,
tmp_rtensor)[(qnmat == nexciton)]
# initial guess b-S-c-S-e
# a d
cguess = tensordot(mps[imps - 1], mps[imps],
axes=1)[qnmat == nexciton]
cguess = cguess.asnumpy()
hdiag *= inverse
nonzeros = np.sum(qnmat == nexciton)
# print("Hmat dim", nonzeros)
def hop(c):
# convert c to initial structure according to qn pattern
cstruct = cvec2cmat(cshape, c, qnmat, nexciton)
if method == "1site":
# S-a l-S
# d
# O-b-O-f-O
# e
# S-c k-S
path = [
([0, 1], "abc, adl -> bcdl"),
([2, 0], "bcdl, bdef -> clef"),
([1, 0], "clef, lfk -> cek"),
]
cout = multi_tensor_contract(path, ltensor,
Matrix(cstruct), mpo[imps],
rtensor)
# for small matrices, check hermite:
# a=tensordot(ltensor, mpo[imps], ((1), (0)))
# b=tensordot(a, rtensor, ((4), (1)))
# c=b.transpose((0, 2, 4, 1, 3, 5))
# d=c.reshape(16, 16)
else:
# S-a l-S
# d g
# O-b-O-f-O-j-O
# e h
# S-c k-S
path = [
([0, 1], "abc, adgl -> bcdgl"),
([3, 0], "bcdgl, bdef -> cglef"),
([2, 0], "cglef, fghj -> clehj"),
([1, 0], "clehj, ljk -> cehk"),
]
cout = multi_tensor_contract(
path,
ltensor,
Matrix(cstruct),
mpo[imps - 1],
mpo[imps],
rtensor,
)
# convert structure c to 1d according to qn
return inverse * cout.asnumpy()[qnmat == nexciton]
if nroots != 1:
cguess = [cguess]
for iroot in range(nroots - 1):
cguess.append(np.random.random([nonzeros]) - 0.5)
precond = lambda x, e, *args: x / (hdiag.asnumpy() - e + 1e-4)
e, c = davidson(hop,
cguess,
precond,
max_cycle=100,
nroots=nroots,
max_memory=64000)
# scipy arpack solver : much slower than davidson
# A = spslinalg.LinearOperator((nonzeros,nonzeros), matvec=hop)
# e, c = spslinalg.eigsh(A,k=1, which="SA",v0=cguess)
# print("HC loops:", count[0])
# logger.debug(f"isweep: {isweep}, e: {e}")
energies.append(e)
cstruct = cvec2cmat(cshape, c, qnmat, nexciton, nroots=nroots)
if nroots == 1:
# direct svd the coefficient matrix
mt, mpsdim, mpsqn, compmps = renormalization_svd(
cstruct,
qnbigl,
qnbigr,
system,
nexciton,
Mmax=mmax,
percent=percent,
)
else:
# diagonalize the reduced density matrix
mt, mpsdim, mpsqn, compmps = renormalization_ddm(
cstruct,
qnbigl,
qnbigr,
system,
nexciton,
Mmax=mmax,
percent=percent,
)
if method == "1site":
mps[imps] = mt
if system == "L":
if imps != len(mps) - 1:
mps[imps + 1] = tensordot(compmps,
mps[imps + 1],
axes=1)
mps.qn[imps + 1] = mpsqn
else:
mps[imps] = tensordot(mps[imps], compmps, axes=1)
mps.qn[imps + 1] = [0]
else:
if imps != 0:
mps[imps - 1] = tensordot(mps[imps - 1],
compmps,
axes=1)
mps.qn[imps] = mpsqn
else:
mps[imps] = tensordot(compmps, mps[imps], axes=1)
mps.qn[imps] = [0]
else:
if system == "L":
mps[imps - 1] = mt
mps[imps] = compmps
else:
mps[imps] = mt
mps[imps - 1] = compmps
# mps.dim_list[imps] = mpsdim
mps.qn[imps] = mpsqn
energies = np.array(energies)
if nroots == 1:
logger.debug("Optimization complete, lowest energy = %g",
energies.min())
return energies
def _evolve_dmrg_tdvp_ps(self, mpo, evolve_dt) -> "Mps":
# PhysRevB.94.165116
# TDVP projector splitting
imag_time = np.iscomplex(evolve_dt)
if imag_time:
mps = self.copy()
mps_conj = mps
else:
mps = self.to_complex()
mps_conj = mps.conj() # another copy, so 3x memory is used.
# construct the environment matrix
environ = Environ()
# almost half is not used. Not a big deal.
environ.construct(mps, mps_conj, mpo, "L")
environ.construct(mps, mps_conj, mpo, "R")
# a workaround for https://github.com/scipy/scipy/issues/10164
if imag_time:
evolve_dt = -evolve_dt.imag
# used in calculating derivatives
coef = -1
else:
coef = 1j
# statistics for debug output
cmf_rk_steps = []
USE_RK = self.evolve_config.tdvp_ps_rk4
# sweep for 2 rounds
for i in range(2):
for imps in mps.iter_idx_list(full=True):
system = "L" if mps.left else "R"
ltensor = environ.read("L", imps - 1)
rtensor = environ.read("R", imps + 1)
shape = list(mps[imps].shape)
l_array = ltensor.array
r_array = rtensor.array
hop = hop_factory(l_array, r_array, mpo[imps].array, len(shape))
def hop_svt(ms):
# S-a l-S
#
# O-b - b-O
#
# S-c k-S
path = [([0, 1], "abc, ck -> abk"), ([1, 0], "abk, lbk -> al")]
HC = multi_tensor_contract(path, l_array, ms, r_array)
return HC
if USE_RK:
def func(t, y):
return hop(y.reshape(shape)).ravel() / coef
sol = solve_ivp(
func, (0, evolve_dt / 2.0), mps[imps].ravel().array, method="RK45"
)
cmf_rk_steps.append(len(sol.t))
mps_t = sol.y[:, -1]
else:
# Can't use the same func because here H should be Hermitian
def func(y):
return hop(y.reshape(shape)).ravel()
mps_t = expm_krylov(func, (evolve_dt / 2) / coef, mps[imps].ravel().array)
mps_t = mps_t.reshape(shape)
qnbigl, qnbigr = mps._get_big_qn(imps)
u, qnlset, v, qnrset = svd_qn.Csvd(
asnumpy(mps_t),
qnbigl,
qnbigr,
mps.qntot,
QR=True,
system=system,
full_matrices=False,
)
vt = v.T
if mps.is_left_canon and imps != 0:
mps[imps] = vt.reshape([-1] + shape[1:])
mps_conj[imps] = mps[imps].conj()
mps.qn[imps] = qnrset
rtensor = environ.GetLR(
"R", imps, mps, mps_conj, mpo, itensor=rtensor, method="System"
)
r_array = rtensor.array
# reverse update u site
shape_u = u.shape
if USE_RK:
def func_u(t, y):
return hop_svt(y.reshape(shape_u)).ravel() / coef
sol_u = solve_ivp(
func_u, (0, -evolve_dt / 2), u.ravel(), method="RK45"
)
cmf_rk_steps.append(len(sol_u.t))
mps_t = sol_u.y[:, -1]
else:
def func_u(y):
return hop_svt(y.reshape(shape_u)).ravel()
mps_t = expm_krylov(func_u, (-evolve_dt / 2) / coef, u.ravel())
mps_t = mps_t.reshape(shape_u)
mps[imps - 1] = tensordot(
mps[imps - 1].array,
mps_t,
axes=(-1, 0),
)
mps_conj[imps - 1] = mps[imps - 1].conj()
elif mps.is_right_canon and imps != len(mps) - 1:
mps[imps] = u.reshape(shape[:-1] + [-1])
mps_conj[imps] = mps[imps].conj()
mps.qn[imps + 1] = qnlset
ltensor = environ.GetLR(
"L", imps, mps, mps_conj, mpo, itensor=ltensor, method="System"
)
l_array = ltensor.array
# reverse update svt site
shape_svt = vt.shape
if USE_RK:
def func_svt(t, y):
return hop_svt(y.reshape(shape_svt)).ravel() / coef
sol_svt = solve_ivp(
func_svt, (0, -evolve_dt / 2), vt.ravel(), method="RK45"
)
cmf_rk_steps.append(len(sol_svt.t))
mps_t = sol_svt.y[:, -1]
else:
def func_svt(y):
return hop_svt(y.reshape(shape_svt)).ravel()
mps_t = expm_krylov(func_svt, (-evolve_dt / 2) / coef, vt.ravel())
mps_t = mps_t.reshape(shape_svt)
mps[imps + 1] = tensordot(
mps_t,
mps[imps + 1].array,
axes=(1, 0),
)
mps_conj[imps + 1] = mps[imps + 1].conj()
else:
mps[imps] = mps_t
mps_conj[imps] = mps[imps].conj()
mps._switch_direction()
if USE_RK:
steps_stat = stats.describe(cmf_rk_steps)
logger.debug(f"TDVP-PS CMF steps: {steps_stat}")
mps.evolve_config.stat = steps_stat
return mps
def _evolve_dmrg_tdvp_mctdhnew(self, mpo, evolve_dt) -> "Mps":
# new regularization scheme
# JCP 148, 124105 (2018)
# JCP 149, 044119 (2018)
# a workaround for https://github.com/scipy/scipy/issues/10164
imag_time = np.iscomplex(evolve_dt)
if imag_time:
evolve_dt = -evolve_dt.imag
# used in calculating derivatives
coef = -1
else:
coef = 1j
if self.is_left_canon:
assert self.check_left_canonical()
self.canonicalise()
mps = self.to_complex(inplace=True)
# construct the environment matrix
environ = Environ()
environ.construct(mps, mps.conj(), mpo, "R")
# initial matrix
ltensor = ones((1, 1, 1))
rtensor = ones((1, 1, 1))
new_mps = mps.metacopy()
# statistics for debug output
cmf_rk_steps = []
for imps in range(len(mps)):
shape = list(mps[imps].shape)
system = "L" if mps.left else "R"
qnbigl, qnbigr = mps._get_big_qn(imps)
u, s, qnlset, v, s, qnrset = svd_qn.Csvd(
mps[imps].asnumpy(),
qnbigl,
qnbigr,
mps.qntot,
system=system,
full_matrices=False,
)
vt = v.T
mps[imps] = u.reshape(shape[:-1] + [-1])
ltensor = environ.GetLR(
"L", imps - 1, mps, mps.conj(), mpo, itensor=ltensor, method="System"
)
rtensor = environ.GetLR(
"R", imps + 1, mps, mps.conj(), mpo, itensor=rtensor, method="Enviro"
)
epsilon = 1e-10
epsilon = np.sqrt(epsilon)
s = s + epsilon * np.exp(-s / epsilon)
svt = Matrix(np.diag(s).dot(vt))
rtensor = tensordot(rtensor, svt, axes=(2, 1))
rtensor = tensordot(Matrix(vt).conj(), rtensor, axes=(1, 0))
if imps != len(mps) - 1:
mps[imps + 1] = tensordot(svt, mps[imps + 1], axes=(-1, 0))
mps.qn[imps + 1] = qnlset
new_mps.qn[imps + 1] = qnlset.copy()
S_inv = xp.diag(1.0 / s)
hop = hop_factory(ltensor, rtensor, mpo[imps], len(shape))
func = integrand_func_factory(shape, hop, imps == len(mps) - 1, S_inv, coef)
sol = solve_ivp(
func, (0, evolve_dt), mps[imps].ravel().array, method="RK45"
)
cmf_rk_steps.append(len(sol.t))
ms = sol.y[:, -1].reshape(shape)
if imps == len(mps) - 1:
new_mps[imps] = ms * s[0]
else:
new_mps[imps] = ms
mps._switch_direction()
new_mps._switch_direction()
new_mps.canonicalise()
steps_stat = stats.describe(cmf_rk_steps)
logger.debug(f"TDVP-MCTDH CMF steps: {steps_stat}")
# new_mps.evolve_config.stat = steps_stat
return new_mps
def _evolve_dmrg_tdvp_mctdh(self, mpo, evolve_dt) -> "Mps":
# TDVP for original MCTDH
if self.is_right_canon:
assert self.check_right_canonical()
self.canonicalise()
# a workaround for https://github.com/scipy/scipy/issues/10164
imag_time = np.iscomplex(evolve_dt)
if imag_time:
evolve_dt = -evolve_dt.imag
# used in calculating derivatives
coef = -1
else:
coef = 1j
# qn for this method has not been implemented
self.use_dummy_qn = True
self.clear_qn()
mps = self.to_complex(inplace=True)
mps_conj = mps.conj()
environ = Environ()
environ.construct(mps, mps_conj, mpo, "R")
# initial matrix
ltensor = np.ones((1, 1, 1))
rtensor = np.ones((1, 1, 1))
new_mps = self.metacopy()
cmf_rk_steps = []
for imps in range(len(mps)):
ltensor = environ.GetLR(
"L", imps - 1, mps, mps_conj, mpo, itensor=ltensor, method="System"
)
rtensor = environ.GetLR(
"R", imps + 1, mps, mps_conj, mpo, itensor=rtensor, method="Enviro"
)
# density matrix
S = transferMat(mps, mps_conj, "R", imps + 1).asnumpy()
epsilon = 1e-8
w, u = scipy.linalg.eigh(S)
try:
w = w + epsilon * np.exp(-w / epsilon)
except FloatingPointError:
logger.warning(f"eigenvalue of density matrix contains negative value")
w -= 2 * w.min()
w = w + epsilon * np.exp(-w / epsilon)
# print
# "sum w=", np.sum(w)
# S = u.dot(np.diag(w)).dot(np.conj(u.T))
S_inv = xp.asarray(u.dot(np.diag(1.0 / w)).dot(np.conj(u.T)))
# pseudo inverse
# S_inv = scipy.linalg.pinvh(S,rcond=1e-2)
shape = mps[imps].shape
hop = hop_factory(ltensor, rtensor, mpo[imps], len(shape))
func = integrand_func_factory(shape, hop, imps == len(mps) - 1, S_inv, coef)
sol = solve_ivp(
func, (0, evolve_dt), mps[imps].ravel().array, method="RK45"
)
# print
# "CMF steps:", len(sol.t)
cmf_rk_steps.append(len(sol.t))
new_mps[imps] = sol.y[:, -1].reshape(shape)
new_mps[imps].check_lortho()
# print
# "orthogonal1", np.allclose(np.tensordot(MPSnew[imps],
# np.conj(MPSnew[imps]), axes=([0, 1], [0, 1])),
# np.diag(np.ones(MPSnew[imps].shape[2])))
steps_stat = stats.describe(cmf_rk_steps)
logger.debug(f"TDVP-MCTDH CMF steps: {steps_stat}")
return new_mps
def kernel(self, restart=False, include_psi0=False):
r"""calculate the roots
Parameters
----------
restart: bool, optional
if restart from the former converged root. Default is ``False``.
If ``restart = True``, ``include_psi0`` must be the same as the
former calculation.
include_psi0: bool, optional
if the basis of Hamiltonian includes the ground state
:math:`\Psi_0`. Default is ``False``.
Returns
-------
e: np.ndarray
the energy of the states, if ``include_psi0 = True``, the first
element is the ground state energy, otherwise, it is the energy of
the first excited state.
"""
# right canonical mps
mpo = self.hmpo
nroots = self.nroots
algo = self.algo
site_num = mpo.site_num
if not restart:
# make sure that M is not redundant near the edge
mps = self.mps.ensure_right_canon().canonicalise().normalize().canonicalise()
logger.debug(f"reference mps shape, {mps}")
mps_r_cano = mps.copy()
assert mps.to_right
tangent_u = []
for ims, ms in enumerate(mps):
shape = list(ms.shape)
u, s, vt = scipy.linalg.svd(ms.l_combine(), full_matrices=True)
rank = len(s)
if include_psi0 and ims == site_num-1:
tangent_u.append(u.reshape(shape[:-1]+[-1]))
else:
if rank < u.shape[1]:
tangent_u.append(u[:,rank:].reshape(shape[:-1]+[-1]))
else:
tangent_u.append(None) # the tangent space is None
mps[ims] = u[:,:rank].reshape(shape[:-1]+[-1])
vt = xp.einsum("i, ij -> ij", asxp(s), asxp(vt))
if ims == site_num-1:
assert vt.size == 1 and xp.allclose(vt, 1)
else:
mps[ims+1] = asnumpy(tensordot(vt, mps[ims+1], ([-1],[0])))
mps_l_cano = mps.copy()
mps_l_cano.to_right = False
mps_l_cano.qnidx = site_num-1
else:
mps_l_cano, mps_r_cano, tangent_u, tda_coeff_list = self.wfn
cguess = []
for iroot in range(len(tda_coeff_list)):
tda_coeff = tda_coeff_list[iroot]
x = [c.flatten() for c in tda_coeff if c is not None]
x = np.concatenate(x,axis=None)
cguess.append(x)
cguess = np.stack(cguess, axis=1)
xshape = []
xsize = 0
for ims in range(site_num):
if tangent_u[ims] is None:
xshape.append((0,0))
else:
if ims == site_num-1:
xshape.append((tangent_u[ims].shape[-1], 1))
else:
xshape.append((tangent_u[ims].shape[-1], mps_r_cano[ims+1].shape[0]))
xsize += np.prod(xshape[-1])
logger.debug(f"DMRG-TDA H dimension: {xsize}")
if USE_GPU:
oe_backend = "cupy"
else:
oe_backend = "numpy"
mps_tangent = mps_r_cano.copy()
environ = Environ(mps_tangent, mpo, "R")
hdiag = []
for ims in range(site_num):
ltensor = environ.GetLR(
"L", ims-1, mps_tangent, mpo, itensor=None,
method="System"
)
rtensor = environ.GetLR(
"R", ims+1, mps_tangent, mpo, itensor=None,
method="Enviro"
)
if tangent_u[ims] is not None:
u = asxp(tangent_u[ims])
tmp = oe.contract("abc, ded, bghe, agl, chl -> ld", ltensor, rtensor,
asxp(mpo[ims]), u, u, backend=oe_backend)
hdiag.append(asnumpy(tmp))
mps_tangent[ims] = mps_l_cano[ims]
hdiag = np.concatenate(hdiag, axis=None)
count = 0
# recover the vector-like x back to the ndarray tda_coeff
def reshape_x(x):
tda_coeff = []
offset = 0
for shape in xshape:
if shape == (0,0):
tda_coeff.append(None)
else:
size = np.prod(shape)
tda_coeff.append(x[offset:size+offset].reshape(shape))
offset += size
assert offset == xsize
return tda_coeff
def hop(x):
# H*X
nonlocal count
count += 1
assert len(x) == xsize
tda_coeff = reshape_x(x)
res = [np.zeros_like(coeff) if coeff is not None else None for coeff in tda_coeff]
# fix ket and sweep bra and accumulate into res
for ims in range(site_num):
if tda_coeff[ims] is None:
assert tangent_u[ims] is None
continue
# mix-canonical mps
mps_tangent = merge(mps_l_cano, mps_r_cano, ims+1)
mps_tangent[ims] = tensordot(tangent_u[ims], tda_coeff[ims], (-1, 0))
mps_tangent_conj = mps_r_cano.copy()
environ = Environ(mps_tangent, mpo, "R", mps_conj=mps_tangent_conj)
for ims_conj in range(site_num):
ltensor = environ.GetLR(
"L", ims_conj-1, mps_tangent, mpo, itensor=None,
mps_conj=mps_tangent_conj,
method="System"
)
rtensor = environ.GetLR(
"R", ims_conj+1, mps_tangent, mpo, itensor=None,
mps_conj=mps_tangent_conj,
method="Enviro"
)
if tda_coeff[ims_conj] is not None:
# S-a l-S
# d
# O-b-O-f-O
# e
# S-c k-S
path = [
([0, 1], "abc, cek -> abek"),
([2, 0], "abek, bdef -> akdf"),
([1, 0], "akdf, lfk -> adl"),
]
out = multi_tensor_contract(
path, ltensor, asxp(mps_tangent[ims_conj]),
asxp(mpo[ims_conj]), rtensor
)
res[ims_conj] += asnumpy(tensordot(tangent_u[ims_conj], out,
([0,1], [0,1])))
# mps_conj combine
mps_tangent_conj[ims_conj] = mps_l_cano[ims_conj]
res = [mat for mat in res if mat is not None]
return np.concatenate(res, axis=None)
if algo == "davidson":
if restart:
cguess = [cguess[:,i] for i in range(cguess.shape[1])]
else:
cguess = [np.random.random(xsize) - 0.5]
precond = lambda x, e, *args: x / (hdiag - e + 1e-4)
e, c = davidson(
hop, cguess, precond, max_cycle=100,
nroots=nroots, max_memory=64000
)
if nroots == 1:
c = [c]
c = np.stack(c, axis=1)
elif algo == "primme":
if not restart:
cguess = None
def multi_hop(x):
if x.ndim == 1:
return hop(x)
elif x.ndim == 2:
return np.stack([hop(x[:,i]) for i in range(x.shape[1])],axis=1)
else:
assert False
def precond(x):
if x.ndim == 1:
return np.einsum("i, i -> i", 1/(hdiag+1e-4), x)
elif x.ndim ==2:
return np.einsum("i, ij -> ij", 1/(hdiag+1e-4), x)
else:
assert False
A = scipy.sparse.linalg.LinearOperator((xsize,xsize),
matvec=multi_hop, matmat=multi_hop)
M = scipy.sparse.linalg.LinearOperator((xsize,xsize),
matvec=precond, matmat=precond)
e, c = primme.eigsh(A, k=min(nroots,xsize), which="SA",
v0=cguess,
OPinv=M,
method="PRIMME_DYNAMIC",
tol=1e-6)
else:
assert False
logger.debug(f"H*C times: {count}")
tda_coeff_list = []
for iroot in range(nroots):
tda_coeff_list.append(reshape_x(c[:,iroot]))
self.e = np.array(e)
self.wfn = [mps_l_cano, mps_r_cano, tangent_u, tda_coeff_list]
return self.e
def variational_compress(self, mpo=None, guess=None):
r"""Variational compress an mps/mpdm/mpo
Parameters
----------
mpo : renormalizer.mps.Mpo, optional
Default is ``None``. if mpo is not ``None``, the returned mps is
an approximation of ``mpo @ self``
guess : renormalizer.mps.MatrixProduct, optional
Initial guess of compressed mps/mpdm/mpo. Default is ``None``.
Note
----
the variational compress related configurations is defined in
``self`` if ``guess=None``, otherwise is defined in ``guess``
Returns
-------
mp : renormalizer.mps.MatrixProduct
a new compressed mps/mpdm/mpo, ``self`` is not overwritten.
``guess`` is overwritten.
"""
if mpo is None:
logger.info(
"Recommend to use svd to compress a single mps/mpo/mpdm.")
raise NotImplementedError
if guess is None:
# a minimal representation of self and mpo
compressed_mpo = mpo.copy().canonicalise().compress(
temp_m_trunc=self.compress_config.vguess_m[0])
compressed_mps = self.copy().canonicalise().compress(
temp_m_trunc=self.compress_config.vguess_m[1])
# the attributes of guess would be the same as self
guess = compressed_mpo.apply(compressed_mps)
mps = guess
mps.ensure_left_canon()
logger.info(f"initial guess bond dims: {mps.bond_dims}")
procedure = mps.compress_config.vprocedure
method = mps.compress_config.vmethod
environ = Environ(self, mpo, "L", mps_conj=mps.conj())
converged = False
for isweep, (mmax, percent) in enumerate(procedure):
logger.debug(f"isweep: {isweep}")
logger.debug(f"mmax, percent: {mmax}, {percent}")
logger.debug(f"mps bond dims: {mps.bond_dims}")
for imps in mps.iter_idx_list(full=True):
if method == "2site" and \
((mps.to_right and imps == mps.site_num-1)
or ((not mps.to_right) and imps == 0)):
break
if mps.to_right:
lmethod, rmethod = "System", "Enviro"
else:
lmethod, rmethod = "Enviro", "System"
if method == "1site":
lidx = imps - 1
cidx = [imps]
ridx = imps + 1
elif method == "2site":
if mps.to_right:
lidx = imps - 1
cidx = [imps, imps + 1]
ridx = imps + 2
else:
lidx = imps - 2
cidx = [imps - 1, imps] # center site
ridx = imps + 1
else:
assert False
logger.debug(f"optimize site: {cidx}")
ltensor = environ.GetLR("L",
lidx,
self,
mpo,
itensor=None,
method=lmethod,
mps_conj=mps.conj())
rtensor = environ.GetLR("R",
ridx,
self,
mpo,
itensor=None,
method=rmethod,
mps_conj=mps.conj())
# get the quantum number pattern
qnbigl, qnbigr, qnmat = mps._get_big_qn(cidx)
# center mo
cmo = [asxp(mpo[idx]) for idx in cidx]
cms = [asxp(self[idx]) for idx in cidx]
if method == "1site":
if cms[0].ndim == 3:
# S-a l-S
# d
# O-b-O-f-O
# e
# S-c k-S
path = [
([0, 1], "abc, cek -> abek"),
([2, 0], "abek, bdef -> akdf"),
([1, 0], "akdf, lfk -> adl"),
]
elif cms[0].ndim == 4:
# S-a l-S
# d
# O-b-O-f-O
# e
# S-c k-S
# g
path = [
([0, 2], "abc, bdef -> acdef"),
([2, 0], "acdef, cegk -> adfgk"),
([1, 0], "adfgk, lfk -> adgl"),
]
cout = multi_tensor_contract(path, ltensor, cms[0], cmo[0],
rtensor)
else:
if USE_GPU:
oe_backend = "cupy"
else:
oe_backend = "numpy"
if cms[0].ndim == 3:
# S-a l-S
# d g
# O-b-O-f-O-j-O
# e h
# S-c m k-S
cout = oe.contract(
"abc, bdef, fghj, ljk, cem, mhk -> adgl",
ltensor,
cmo[0],
cmo[1],
rtensor,
cms[0],
cms[1],
backend=oe_backend)
elif cms[0].ndim == 4:
# S-a l-S
# d g
# O-b-O-f-O-j-O
# e h
# S-c m k-S
# n p
cout = oe.contract(
"abc, bdef, fghj, ljk, cenm, mhpk -> adngpl",
ltensor,
cmo[0],
cmo[1],
rtensor,
cms[0],
cms[1],
backend=oe_backend)
# clean up the elements which do not meet the qn requirements
cout[qnmat != mps.qntot] = 0
mps._update_mps(cout, cidx, qnbigl, qnbigr, mmax, percent)
mps._switch_direction()
# check if convergence
if isweep > 0 and percent == 0 and \
mps.distance(mps_old) / np.sqrt(mps.dot(mps.conj()).real) < mps.compress_config.vrtol:
converged = True
break
mps_old = mps.copy()
if converged:
logger.info("Variational compress is converged!")
else:
logger.warning(
"Variational compress is not converged! Please increase the procedure!"
)
# remove the redundant bond dimension near the boundary of the MPS
mps.canonicalise()
logger.info(f"{mps}")
return mps
def optimize_mps(mps: Mps, mpo: Mpo, omega: float = None) -> Tuple[List, Mps]:
r""" DMRG ground state algorithm and state-averaged excited states algorithm
Parameters
----------
mps : renormalizer.mps.Mps
initial guess of mps
mpo : renormalizer.mps.Mpo
mpo of Hamiltonian
omega: float, optional
target the eigenpair near omega with special variational function
:math:(\hat{H}-\omega)^2. Default is `None`.
Returns
-------
energy : list
list of energy of each marco sweep.
:math:`[e_0, e_0, \cdots, e_0]` if ``nroots=1``.
:math:`[[e_0, \cdots, e_n], \dots, [e_0, \cdots, e_n]]` if ``nroots=n``.
mps : renormalizer.mps.Mps
optimized ground state mps. The input mps is overwritten and could not
be used anymore.
See Also
--------
renormalizer.utils.configs.OptimizeConfig : The optimization configuration.
"""
algo = mps.optimize_config.algo
method = mps.optimize_config.method
procedure = mps.optimize_config.procedure
inverse = mps.optimize_config.inverse
nroots = mps.optimize_config.nroots
assert method in ["2site", "1site"]
logger.info(f"optimization method: {method}")
logger.info(f"e_rtol: {mps.optimize_config.e_rtol}")
logger.info(f"e_atol: {mps.optimize_config.e_atol}")
if USE_GPU:
oe_backend = "cupy"
else:
oe_backend = "numpy"
# ensure that mps is left or right-canonical
# TODO: start from a mix-canonical MPS
if mps.is_left_canon:
mps.ensure_right_canon()
env = "R"
else:
mps.ensure_left_canon()
env = "L"
# in state-averged calculation, contains C of each state for better initial
# guess
averaged_ms = None
# the index of active site of the returned mps
res_mps_idx = None
# target eigenstate close to omega with (H-omega)^2
# construct the environment matrix
if omega is not None:
identity = Mpo.identity(mpo.model)
mpo = mpo.add(identity.scale(-omega))
environ = Environ(mps, [mpo, mpo], env)
else:
environ = Environ(mps, mpo, env)
macro_iteration_result = []
converged = False
for isweep, (mmax, percent) in enumerate(procedure):
logger.debug(f"isweep: {isweep}")
logger.debug(f"mmax, percent: {mmax}, {percent}")
logger.debug(f"{mps}")
micro_iteration_result = []
for imps in mps.iter_idx_list(full=True):
if method == "2site" and \
((mps.to_right and imps == mps.site_num-1)
or ((not mps.to_right) and imps == 0)):
break
if mps.to_right:
lmethod, rmethod = "System", "Enviro"
else:
lmethod, rmethod = "Enviro", "System"
if method == "1site":
lidx = imps - 1
cidx = [imps]
ridx = imps + 1
elif method == "2site":
if mps.to_right:
lidx = imps - 1
cidx = [imps, imps + 1]
ridx = imps + 2
else:
lidx = imps - 2
cidx = [imps - 1, imps] # center site
ridx = imps + 1
else:
assert False
logger.debug(f"optimize site: {cidx}")
if omega is None:
operator = mpo
else:
operator = [mpo, mpo]
ltensor = environ.GetLR("L",
lidx,
mps,
operator,
itensor=None,
method=lmethod)
rtensor = environ.GetLR("R",
ridx,
mps,
operator,
itensor=None,
method=rmethod)
# get the quantum number pattern
qnbigl, qnbigr, qnmat = mps._get_big_qn(cidx)
cshape = qnmat.shape
nonzeros = np.sum(qnmat == mps.qntot)
logger.debug(f"Hmat dim: {nonzeros}")
# center mo
cmo = [asxp(mpo[idx]) for idx in cidx]
if qnmat.size > 1000 and algo != "direct":
# iterative algorithm
# diagonal elements of H
if omega is None:
tmp_ltensor = xp.einsum("aba -> ba", ltensor)
tmp_cmo0 = xp.einsum("abbc -> abc", cmo[0])
tmp_rtensor = xp.einsum("aba -> ba", rtensor)
if method == "1site":
# S-a c f-S
# O-b-O-g-O
# S-a c f-S
path = [([0, 1], "ba, bcg -> acg"),
([1, 0], "acg, gf -> acf")]
hdiag = multi_tensor_contract(
path, tmp_ltensor, tmp_cmo0,
tmp_rtensor)[(qnmat == mps.qntot)]
else:
# S-a c d f-S
# O-b-O-e-O-g-O
# S-a c d f-S
tmp_cmo1 = xp.einsum("abbc -> abc", cmo[1])
path = [
([0, 1], "ba, bce -> ace"),
([0, 1], "edg, gf -> edf"),
([0, 1], "ace, edf -> acdf"),
]
hdiag = multi_tensor_contract(
path, tmp_ltensor, tmp_cmo0, tmp_cmo1,
tmp_rtensor)[(qnmat == mps.qntot)]
else:
if method == "1site":
# S-a d h-S
# O-b-O-f-O
# | e |
# O-c-O-g-O
# S-a d h-S
hdiag = oe.contract(
"abca, bdef, cedg, hfgh -> adh",
ltensor,
cmo[0],
cmo[0],
rtensor,
backend=oe_backend)[(qnmat == mps.qntot)]
else:
# S-a d h l-S
# O-b-O-f-O-j-O
# | e i |
# O-c-O-g-O-k-O
# S-a d h l-S
hdiag = oe.contract(
"abca, bdef, cedg, fhij, gihk, ljkl -> adhl",
ltensor,
cmo[0],
cmo[0],
cmo[1],
cmo[1],
rtensor,
backend=oe_backend)[(qnmat == mps.qntot)]
hdiag = asnumpy(hdiag * inverse)
# initial guess
if method == "1site":
# initial guess b-S-c
# a
if nroots == 1:
cguess = [asnumpy(mps[cidx[0]])[qnmat == mps.qntot]]
else:
cguess = []
if averaged_ms is not None:
for ms in averaged_ms:
cguess.append(asnumpy(ms)[qnmat == mps.qntot])
else:
# initial guess b-S-c-S-e
# a d
if nroots == 1:
cguess = [
asnumpy(
tensordot(mps[cidx[0]], mps[cidx[1]],
axes=1)[qnmat == mps.qntot])
]
else:
cguess = []
if averaged_ms is not None:
for ms in averaged_ms:
if mps.to_right:
cguess.append(
asnumpy(
tensordot(
ms, mps[cidx[1]],
axes=1)[qnmat == mps.qntot]))
else:
cguess.append(
asnumpy(
tensordot(
mps[cidx[0]], ms,
axes=1)[qnmat == mps.qntot]))
if omega is not None:
if method == "1site":
# S-a e j-S
# O-b-O-g-O
# | f |
# O-c-O-i-O
# S-d h k-S
expr = oe.contract_expression(
"abcd, befg, cfhi, jgik, aej -> dhk",
ltensor,
cmo[0],
cmo[0],
rtensor,
cshape,
constants=[0, 1, 2, 3])
else:
# S-a e j o-S
# O-b-O-g-O-l-O
# | f k |
# O-c-O-i-O-n-O
# S-d h m p-S
expr = oe.contract_expression(
"abcd, befg, cfhi, gjkl, ikmn, olnp, aejo -> dhmp",
ltensor,
cmo[0],
cmo[0],
cmo[1],
cmo[1],
rtensor,
cshape,
constants=[0, 1, 2, 3, 4, 5])
count = 0
def hop(x):
nonlocal count
count += 1
clist = []
if x.ndim == 1:
clist.append(x)
else:
for icol in range(x.shape[1]):
clist.append(x[:, icol])
res = []
for c in clist:
# convert c to initial structure according to qn pattern
cstruct = asxp(cvec2cmat(cshape, c, qnmat, mps.qntot))
if omega is None:
if method == "1site":
# S-a l-S
# d
# O-b-O-f-O
# e
# S-c k-S
path = [
([0, 1], "abc, adl -> bcdl"),
([2, 0], "bcdl, bdef -> clef"),
([1, 0], "clef, lfk -> cek"),
]
cout = multi_tensor_contract(
path, ltensor, cstruct, cmo[0], rtensor)
else:
# S-a l-S
# d g
# O-b-O-f-O-j-O
# e h
# S-c k-S
path = [
([0, 1], "abc, adgl -> bcdgl"),
([3, 0], "bcdgl, bdef -> cglef"),
([2, 0], "cglef, fghj -> clehj"),
([1, 0], "clehj, ljk -> cehk"),
]
cout = multi_tensor_contract(
path,
ltensor,
cstruct,
cmo[0],
cmo[1],
rtensor,
)
else:
cout = expr(cstruct, backend=oe_backend)
# convert structure c to 1d according to qn
res.append(asnumpy(cout)[qnmat == mps.qntot])
if len(res) == 1:
return inverse * res[0]
else:
return inverse * np.stack(res, axis=1)
if len(cguess) < nroots:
cguess.extend([
np.random.random([nonzeros]) - 0.5
for i in range(len(cguess), nroots)
])
if algo == "davidson":
precond = lambda x, e, *args: x / (hdiag - e + 1e-4)
e, c = davidson(hop,
cguess,
precond,
max_cycle=100,
nroots=nroots,
max_memory=64000)
# if one root, davidson return e as np.float
#elif algo == "arpack":
# # scipy arpack solver : much slower than pyscf/davidson
# A = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros), matvec=hop)
# e, c = scipy.sparse.linalg.eigsh(A, k=nroots, which="SA", v0=cguess)
# # scipy return numpy.array
# if nroots == 1:
# e = e[0]
#elif algo == "lobpcg":
# precond = lambda x: scipy.sparse.diags(1/(hdiag+1e-4)) @ x
# A = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
# matvec=hop, matmat=hop)
# M = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
# matvec=precond, matmat=hop)
# e, c = scipy.sparse.linalg.lobpcg(A, np.array(cguess).T,
# M=M, largest=False)
elif algo == "primme":
precond = lambda x: scipy.sparse.diags(1 /
(hdiag + 1e-4)) @ x
A = scipy.sparse.linalg.LinearOperator(
(nonzeros, nonzeros), matvec=hop, matmat=hop)
M = scipy.sparse.linalg.LinearOperator(
(nonzeros, nonzeros), matvec=precond, matmat=hop)
e, c = primme.eigsh(A,
k=min(nroots, nonzeros),
which="SA",
v0=np.array(cguess).T,
OPinv=M,
method="PRIMME_DYNAMIC",
tol=1e-6)
else:
assert False
logger.debug(f"use {algo}, HC hops: {count}")
else:
logger.debug(f"use direct eigensolver")
# direct algorithm
if omega is None:
if method == "1site":
# S-a l-S
# d
# O-b-O-f-O
# e
# S-c k-S
ham = oe.contract("abc,bdef,lfk->adlcek",
ltensor,
cmo[0],
rtensor,
backend=oe_backend)
ham = ham[:, :, :, qnmat == mps.qntot][
qnmat == mps.qntot, :] * inverse
else:
# S-a l-S
# d g
# O-b-O-f-O-j-O
# e h
# S-c k-S
ham = oe.contract("abc,bdef,fghj,ljk->adglcehk",
ltensor, cmo[0], cmo[1], rtensor)
ham = ham[:, :, :, :, qnmat == mps.qntot][
qnmat == mps.qntot, :] * inverse
else:
if method == "1site":
# S-a e j-S
# O-b-O-g-O
# | f |
# O-c-O-i-O
# S-d h k-S
ham = oe.contract("abcd, befg, cfhi, jgik -> aejdhk",
ltensor, cmo[0], cmo[0], rtensor)
ham = ham[:, :, :, qnmat == mps.qntot][
qnmat == mps.qntot, :] * inverse
else:
# S-a e j o-S
# O-b-O-g-O-l-O
# | f k |
# O-c-O-i-O-n-O
# S-d h m p-S
ham = oe.contract(
"abcd, befg, cfhi, gjkl, ikmn, olnp -> aejodhmp",
ltensor, cmo[0], cmo[0], cmo[1], cmo[1], rtensor)
ham = ham[:, :, :, :, qnmat == mps.qntot][
qnmat == mps.qntot, :] * inverse
w, v = scipy.linalg.eigh(asnumpy(ham))
if nroots == 1:
e = w[0]
c = v[:, 0]
else:
e = w[:nroots]
c = [
v[:, iroot] for iroot in range(min(nroots, v.shape[1]))
]
# if multi roots, both davidson and primme return np.ndarray
if nroots > 1:
e = e.tolist()
logger.debug(f"energy: {e}")
micro_iteration_result.append(e)
cstruct = cvec2cmat(cshape, c, qnmat, mps.qntot, nroots=nroots)
# store the "optimal" mps (usually in the middle of each sweep)
if res_mps_idx is not None and res_mps_idx == imps:
if nroots == 1:
res_mps = mps.copy()
res_mps._update_mps(cstruct, cidx, qnbigl, qnbigr, mmax,
percent)
else:
res_mps = [mps.copy() for i in range(len(cstruct))]
for iroot in range(len(cstruct)):
res_mps[iroot]._update_mps(cstruct[iroot], cidx,
qnbigl, qnbigr, mmax,
percent)
averaged_ms = mps._update_mps(cstruct, cidx, qnbigl, qnbigr, mmax,
percent)
mps._switch_direction()
res_mps_idx = micro_iteration_result.index(min(micro_iteration_result))
macro_iteration_result.append(micro_iteration_result[res_mps_idx])
# check if convergence
if isweep > 0 and percent == 0:
v1, v2 = sorted(macro_iteration_result)[:2]
if np.allclose(v1,
v2,
rtol=mps.optimize_config.e_rtol,
atol=mps.optimize_config.e_atol):
converged = True
break
logger.debug(
f"{isweep+1} sweeps are finished, lowest energy = {sorted(macro_iteration_result)[0]}"
)
if converged:
logger.info("DMRG is converged!")
else:
logger.warning("DMRG is not converged! Please increase the procedure!")
logger.info(
f"The lowest two energies: {sorted(macro_iteration_result)[:2]}.")
# remove the redundant basis near the edge
if nroots == 1:
res_mps = res_mps.normalize().ensure_left_canon().canonicalise()
logger.info(f"{res_mps}")
else:
res_mps = [
mp.normalize().ensure_left_canon().canonicalise() for mp in res_mps
]
logger.info(f"{res_mps[0]}")
return macro_iteration_result, res_mps