From ebc20bae3851601eac6e22c78ca0a80e25265590 Mon Sep 17 00:00:00 2001 From: jjren Date: Sat, 5 Aug 2023 14:07:18 +0800 Subject: [PATCH 1/5] add functionality in mp/mps/mpo --- renormalizer/model/basis.py | 5 +- renormalizer/mps/mp.py | 112 +++++++++++++++++------------------- renormalizer/mps/mpdm.py | 18 +----- renormalizer/mps/mpo.py | 25 ++++++++ renormalizer/mps/mps.py | 43 +++++++++++--- 5 files changed, 119 insertions(+), 84 deletions(-) diff --git a/renormalizer/model/basis.py b/renormalizer/model/basis.py index 9a5f1f83..57083798 100644 --- a/renormalizer/model/basis.py +++ b/renormalizer/model/basis.py @@ -588,7 +588,7 @@ def op_mat(self, op: Union[Op, str]): # operators currently not having analytical matrix elements else: - logger.warning("Note that the quadrature part is not fully tested!") + logger.debug("Note that the quadrature part is not fully tested!") op_symbol = "*".join(op_symbol.split()) # potential operators @@ -643,7 +643,7 @@ def quad(self, expr): if s != "": expr = sp.sympify(s)*expr expr = expr.subs({sL:self.L, sxi:self.xi}) - print(expr) + logger.debug(f"operator expr: {expr}") expr = sp.lambdify([x, sibas, sjbas], expr, "numpy") mat = np.zeros((self.nbas, self.nbas)) @@ -651,6 +651,7 @@ def quad(self, expr): for jbas in range(self.nbas): val, error = scipy.integrate.quad(lambda x: expr(x, ibas, jbas), self.xi, self.xf) + logger.debug(f"quadrature value and error: {val}, {error}") mat[ibas, jbas] = val return mat diff --git a/renormalizer/mps/mp.py b/renormalizer/mps/mp.py index cd0bfd64..502cb37f 100644 --- a/renormalizer/mps/mp.py +++ b/renormalizer/mps/mp.py @@ -47,11 +47,7 @@ def load(cls, model: Model, fname: str): mp.dtype = backend.real_dtype mp.append(mt) - mp.qn = [] - for i in range(nsites+1): - subqn = npload[f"subqn_{i}"].astype(int).tolist() - mp.qn.append(subqn) - + mp.qn = npload["qn"] mp.qnidx = int(npload["qnidx"]) mp.qntot = npload["qntot"].astype(int) mp.to_right = bool(npload["to_right"]) @@ -353,65 +349,65 @@ def mp_norm(self) -> float: res = np.sqrt(res) return float(res) - - def add(self, other: "MatrixProduct"): - assert np.all(self.qntot == other.qntot) - assert self.site_num == other.site_num - + + def add(self, others: List["MatrixProduct"]): + + logger.info(f"new_mps:{type(others)}") + + if not isinstance(others, list): + others = [others] + new_mps = self.metacopy() - if other.dtype == backend.complex_dtype: - new_mps.dtype = backend.complex_dtype - if self.is_complex: + for other in others: + assert np.all(self.qntot == other.qntot) + assert self.site_num == other.site_num + if other.dtype == backend.complex_dtype: + new_mps.dtype = backend.complex_dtype + + if new_mps.is_complex: new_mps.to_complex(inplace=True) new_mps.compress_config.update(self.compress_config) - if self.is_mps: # MPS - new_mps[0] = dstack([self[0], other[0]]) - for i in range(1, self.site_num - 1): - mta = self[i] - mtb = other[i] - pdim = mta.shape[1] - assert pdim == mtb.shape[1] - new_ms = zeros( - [mta.shape[0] + mtb.shape[0], pdim, mta.shape[2] + mtb.shape[2]], - dtype=new_mps.dtype, - ) - new_ms[: mta.shape[0], :, : mta.shape[2]] = mta - new_ms[mta.shape[0] :, :, mta.shape[2] :] = mtb - new_mps[i] = new_ms - - new_mps[-1] = vstack([self[-1], other[-1]]) - elif self.is_mpo or self.is_mpdm: # MPO - new_mps[0] = concatenate((self[0], other[0]), axis=3) - for i in range(1, self.site_num - 1): - mta = self[i] - mtb = other[i] - pdimu = mta.shape[1] - pdimd = mta.shape[2] - assert pdimu == mtb.shape[1] - assert pdimd == mtb.shape[2] - - new_ms = zeros( - [ - mta.shape[0] + mtb.shape[0], - pdimu, - pdimd, - mta.shape[3] + mtb.shape[3], - ], - dtype=new_mps.dtype, - ) - new_ms[: mta.shape[0], :, :, : mta.shape[3]] = mta[:, :, :, :] - new_ms[mta.shape[0] :, :, :, mta.shape[3] :] = mtb[:, :, :, :] - new_mps[i] = new_ms + new_mps[0] = concatenate([self[0]] + [other[0] for other in others], axis=-1) + + for i in range(1, self.site_num - 1): + mts = [self[i]] + [other[i] for other in others] + pdim = self[i].shape[1:-1] + for mt in mts: + assert pdim == mt.shape[1:-1] + new_ms = zeros( + (sum([mt.shape[0] for mt in mts]),) + pdim + + (sum([mt.shape[-1] for mt in mts]),), dtype=new_mps.dtype, + ) + start_first = 0 + start_last = 0 + for mt in mts: + if len(pdim) == 1: + new_ms[start_first:start_first+mt.shape[0], :, start_last:start_last+mt.shape[-1]] = mt + elif len(pdim) == 2: + new_ms[start_first:start_first+mt.shape[0], :, :, start_last:start_last+mt.shape[-1]] = mt + else: + assert False - new_mps[-1] = concatenate((self[-1], other[-1]), axis=0) - else: - assert False + start_first += mt.shape[0] + start_last += mt.shape[-1] + new_mps[i] = new_ms + + new_mps[-1] = concatenate([self[-1]] + [other[-1] for other in others], axis=0) #assert self.qnidx == other.qnidx - new_mps.move_qnidx(other.qnidx) - new_mps.to_right = other.to_right - new_mps.qn = [np.concatenate([qn1, qn2]) for qn1, qn2 in zip(self.qn, other.qn)] + original_qnidx = [] + for other in others: + original_qnidx.append(other.qnidx) + other.move_qnidx(self.qnidx) + + new_mps.to_right = self.to_right + qn_all = [self.qn] + [other.qn for other in others] + new_mps.qn = [np.concatenate(qns) for qns in zip(*qn_all)] + + for i, other in enumerate(others): + other.move_qnidx(original_qnidx[i]) + # qn at the boundary should have dimension 1 new_mps.qn[0] = np.zeros((1, new_mps.qn[0].shape[1]), dtype=int) new_mps.qn[-1] = np.zeros((1, new_mps.qn[0].shape[1]), dtype=int) @@ -419,7 +415,7 @@ def add(self, other: "MatrixProduct"): new_mps.canonicalise() new_mps.compress() return new_mps - + def compress(self, temp_m_trunc=None, ret_s=False): """ inp: canonicalise MPS (or MPO) diff --git a/renormalizer/mps/mpdm.py b/renormalizer/mps/mpdm.py index 14427814..7aff5955 100644 --- a/renormalizer/mps/mpdm.py +++ b/renormalizer/mps/mpdm.py @@ -26,25 +26,9 @@ def ground_state(cls, model, max_entangled): @classmethod def from_mps(cls, mps: Mps): - mpo = cls() - mpo.model = mps.model - for ms in mps: - mo = np.zeros(tuple([ms.shape[0]] + [ms.shape[1]] * 2 + [ms.shape[2]])) - for iaxis in range(ms.shape[1]): - mo[:, iaxis, iaxis, :] = ms[:, iaxis, :].array - mpo.append(mo) - + mpo = super().from_mps(mps) mpo.coeff = mps.coeff - - mpo.optimize_config = mps.optimize_config mpo.evolve_config = mps.evolve_config - mpo.compress_add = mps.compress_add - - mpo.qn = [qn.copy() for qn in mps.qn] - mpo.qntot = mps.qntot - mpo.qnidx = mps.qnidx - mpo.to_right = mps.to_right - mpo.compress_config = mps.compress_config.copy() return mpo @classmethod diff --git a/renormalizer/mps/mpo.py b/renormalizer/mps/mpo.py index 9a22c3fa..510e6330 100644 --- a/renormalizer/mps/mpo.py +++ b/renormalizer/mps/mpo.py @@ -247,6 +247,31 @@ def identity(cls, model: Model): mpo.build_empty_qn() return mpo + @classmethod + def from_mps(cls, mps): + mpo = cls() + mpo.model = mps.model + for ms in mps: + mo = np.zeros(tuple([ms.shape[0]] + [ms.shape[1]] * 2 + [ms.shape[2]])) + for iaxis in range(ms.shape[1]): + mo[:, iaxis, iaxis, :] = ms[:, iaxis, :].array + mpo.append(mo) + + mpo.optimize_config = mps.optimize_config + mpo.compress_add = mps.compress_add + + if mpo.is_mpo: + assert np.allclose(mps.coeff, 1) + # currently, only used when qn is zeros + for qn in mps.qn: + assert np.allclose(qn, np.zeros_like(qn)) + mpo.qn = [qn.copy() for qn in mps.qn] + mpo.qntot = mps.qntot + mpo.qnidx = mps.qnidx + mpo.to_right = mps.to_right + mpo.compress_config = mps.compress_config.copy() + return mpo + def __init__(self, model: Model = None, terms: Union[Op, List[Op]] = None, offset: Quantity = Quantity(0), ): """ diff --git a/renormalizer/mps/mps.py b/renormalizer/mps/mps.py index 53c27890..905b21d1 100644 --- a/renormalizer/mps/mps.py +++ b/renormalizer/mps/mps.py @@ -390,6 +390,28 @@ def from_dense(cls, model, wfn: np.ndarray): mp.append(residual_wfn) mp.build_empty_qn() return mp + + @classmethod + def from_mpo(cls, mpo: Mpo): + # convert diagonal mpo to mps, which usually happens in + # potential energy surface + mps = cls() + mps.model = mpo.model + for mo in mpo: + ms = np.zeros(tuple([mo.shape[0]] + [mo.shape[1]] + [mo.shape[3]])) + for iaxis in range(mo.shape[1]): + ms[:, iaxis, :] = mo[:, iaxis, iaxis, :].array + mps.append(ms) + + mps.coeff = 1 + if mpo.is_mpo: + logger.warning("Note that the qn part is directly inherited from mpo, make sure it is what you want!") + mps.qn = [qn.copy() for qn in mpo.qn] + mps.qntot = mpo.qntot + mps.qnidx = mpo.qnidx + mps.to_right = None + mps.compress_config = mpo.compress_config.copy() + return mps def __init__(self): super().__init__() @@ -1785,13 +1807,20 @@ def __setitem__(self, key, value): return super().__setitem__(key, value) - def add(self, other): - if not np.allclose(self.coeff, other.coeff): - self.scale(self.coeff, inplace=True) - other.scale(other.coeff, inplace=True) - self.coeff = 1 - other.coeff = 1 - return super().add(other) + def add(self, others): + """ + support add many mpss together in a batch way + """ + if not isinstance(others, list): + others = [others] + + for other in others: + if not np.allclose(self.coeff, other.coeff): + self.scale(self.coeff, inplace=True) + other.scale(other.coeff, inplace=True) + self.coeff = 1 + other.coeff = 1 + return super().add(others) def distance(self, other) -> float: if not np.allclose(self.coeff, other.coeff): From 0a69285afb8ec780ce0030da2427e5268fc1aff5 Mon Sep 17 00:00:00 2001 From: jjren Date: Sat, 5 Aug 2023 14:43:03 +0800 Subject: [PATCH 2/5] update --- renormalizer/mps/mps.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/renormalizer/mps/mps.py b/renormalizer/mps/mps.py index 905b21d1..76547bc5 100644 --- a/renormalizer/mps/mps.py +++ b/renormalizer/mps/mps.py @@ -405,11 +405,11 @@ def from_mpo(cls, mpo: Mpo): mps.coeff = 1 if mpo.is_mpo: - logger.warning("Note that the qn part is directly inherited from mpo, make sure it is what you want!") + logger.debug("Note that the qn part is directly inherited from mpo, make sure it is what you want!") mps.qn = [qn.copy() for qn in mpo.qn] mps.qntot = mpo.qntot mps.qnidx = mpo.qnidx - mps.to_right = None + mps.to_right = mpo.to_right mps.compress_config = mpo.compress_config.copy() return mps From d33f09da765d44c2d1bc00bdc1dcc331ecd692c2 Mon Sep 17 00:00:00 2001 From: jjren Date: Thu, 10 Aug 2023 13:24:52 +0800 Subject: [PATCH 3/5] add ExpDVR basis --- renormalizer/model/basis.py | 178 +++++++++++++++++++++++++++++------- renormalizer/mps/mp.py | 5 +- renormalizer/mps/mpo.py | 14 ++- renormalizer/mps/mps.py | 4 +- 4 files changed, 160 insertions(+), 41 deletions(-) diff --git a/renormalizer/model/basis.py b/renormalizer/model/basis.py index 57083798..8dcacafe 100644 --- a/renormalizer/model/basis.py +++ b/renormalizer/model/basis.py @@ -123,12 +123,16 @@ class BasisSHO(BasisSet): dof: The name of the DoF contained in the basis set. The type could be anything that can be hashed. omega (float): the frequency of the oscillator. nbas (int): number of dimension of the basis set (highest occupation number of the harmonic oscillator) - x0 (float): the origin of the harmonic oscillator. Default = 0. + x0 (float): the origin of the harmonic oscillator. Default = 0. dvr (bool): whether to use discrete variable representation. Default = False. general_xp_power (bool): whether calculate :math:`x` and :math:`x^2` (or :math:`p` and :math:`p^2`) through general expression for :math:`x`power or :math:`p` power. This is not efficient because :math:`x` and :math:`x^2` (or :math:`p` and :math:`p^2`) have been hard-coded already. The option is only used for testing. + + Note + ----- + Currently, only the mass reduced coordinates are supported. """ is_phonon = True @@ -333,7 +337,19 @@ def op_mat(self, op: Union[Op, str]): # n is designed for occupation number of the SHO basis mat = np.diag(np.arange(self.nbas)) else: - raise ValueError(f"op_symbol:{op_symbol} is not supported. ") + op_symbol = "*".join(op_symbol.split()) + + if "dx" not in op_symbol: + op_symbol = op_symbol.replace("^", "**") + x = sp.symbols("x") + expr = sp.lambdify(x, op_symbol, "numpy") + mat = np.diag(expr(self.dvr_x)) + if not self.dvr: + mat = self.dvr_v @ mat @ self.dvr_v.T + else: + #mat = self.quad(op_symbol, complex_func=True) + #mat = self.dvr_v @ mat @ self.dvr_v.conj().T + raise ValueError(f"op_symbol:{op_symbol} is not supported. Quadrature is not implemented yet!") self._recursion_flag -= 1 return mat * op_factor @@ -388,6 +404,129 @@ def copy(self, new_dof): return self.__class__(new_dof, self.nbas) +def quad(bas, expr, complex_func=False): + x,sL,sxi,sibas,sjbas = sp.symbols("x sL sxi sibas sjbas") + bra = bas.eigenfunc + if complex_func: + bra = bra.replace("I","-I") + ket = bas.eigenfunc.replace("ibas","jbas") + expr = "*".join((bra,expr,ket)) + + expr = expr.replace("dx^2", "dx*dx") + expr = expr.replace("dx**2", "dx*dx") + + expr_s = expr.split("dx") + expr_s = [s.rstrip('*') for s in expr_s] + expr_s = [s.lstrip('*') for s in expr_s] + expr_s = [s.replace("^", "**") for s in expr_s] + if len(expr_s) == 1: + expr = sp.sympify(expr_s[0]) + else: + expr = sp.sympify(expr_s[-1]) + for s in expr_s[::-1][1:]: + expr = sp.diff(expr, x) + if s != "": + expr = sp.sympify(s)*expr + expr = expr.subs({sL:bas.L, sxi:bas.xi}) + logger.debug(f"operator expr: {expr}") + expr = sp.lambdify([x, sibas, sjbas], expr, "numpy") + + if complex_func: + mat = np.zeros((bas.nbas, bas.nbas), dtype=np.complex128) + else: + mat = np.zeros((bas.nbas, bas.nbas), dtype=np.float64) + for ibas in range(bas.nbas): + for jbas in range(bas.nbas): + val, error = scipy.integrate.quad(lambda x: expr(x, ibas, jbas), + bas.xi, bas.xf, complex_func=complex_func) + logger.debug(f"ibas, jbas, quadrature value and error: {ibas},{jbas}, {val}, {error}") + mat[ibas, jbas] = val + return mat + + +class BasisExpDVR(BasisSet): + r""" + Exponential DVR + See Phys. Rep. 324, 1–105 (2000). B.4.5 + J. Chem. Phys. 96, 1 (1992). + usually used in angular DoF with periodic boundary condition + """ + is_phonon = True + quad = quad + + def __init__(self, dof, nbas, xi, xf, force_quadrature=False): + if nbas % 2 == 0: + raise ValueError("In EXP-DVR, the number of basis should be odd!") + + self.dvr = True + self.L = xf - xi + self.xi = xi + self.xf = xf + self.dvr_x = xi + np.arange(1,nbas+1) * self.L / nbas + self.force_quadrature = force_quadrature # for debug + + self.dvr_v = np.zeros((nbas, nbas), dtype=np.complex128) + for ik, k in enumerate(range(-(nbas // 2), nbas//2+1)): + vec = 1/np.sqrt(self.L)*np.exp(2*np.pi*1j*k*(self.dvr_x-xi)/self.L)*np.sqrt((self.L/nbas)) + self.dvr_v[:,ik] = vec #(grid,basis) + assert np.allclose(np.eye(nbas), self.dvr_v.T.conj() @ self.dvr_v) + assert np.allclose(np.eye(nbas), self.dvr_v @ self.dvr_v.T.conj()) + super().__init__(dof, nbas, [0] * nbas) + + + def __str__(self): + return f"BasisExpDVR(xi: {self.xi}, xf: {self.xf}, nbas: {self.nbas})" + + def op_mat(self, op: Union[Op, str]): + + if not isinstance(op, Op): + op = Op(op, None) + op_symbol, op_factor = op.symbol, op.factor + + # partialx is deprecated + op_symbol = op_symbol.replace("partialx", "dx") + + if op_symbol == "I": + mat = np.eye(self.nbas) + + elif op_symbol == "dx": # and not self.force_quadrature: + mat = np.zeros((self.nbas, self.nbas)) + for i in range(1, self.nbas+1): + for j in range(1, i): + res = np.pi / self.L * (-1)**(i-j) / np.sin(np.pi*(i-j)/self.nbas) + mat[i-1, j-1] = res + mat[j-1, i-1] = -res + + elif op_symbol in ["dx^2", "dx dx"]: # and not self.force_quadrature: + mat = np.zeros((self.nbas, self.nbas)) + for i in range(1, self.nbas+1): + for j in range(1, i+1): + if i == j: + res = -np.pi**2/3/self.L**2*(self.nbas**2-1) + mat[i-1, i-1] = res + else: + res = -2*np.pi**2 / self.L**2 * (-1)**(i-j) * \ + np.cos(np.pi*(i-j)/self.nbas) / np.sin(np.pi*(i-j)/self.nbas)**2 + mat[i-1, j-1] = res + mat[j-1, i-1] = res + else: + op_symbol = "*".join(op_symbol.split()) + + if "dx" not in op_symbol and not self.force_quadrature: + op_symbol = op_symbol.replace("^", "**") + x = sp.symbols("x") + expr = sp.lambdify(x, op_symbol, "numpy") + mat = np.diag(expr(self.dvr_x)) + else: + mat = self.quad(op_symbol, complex_func=True) + mat = self.dvr_v @ mat @ self.dvr_v.conj().T + + return mat * op_factor + + @property + def eigenfunc(self): + return f"sqrt(1/sL) * exp(2*I*pi*(sibas-{self.nbas//2})*(x-sxi)/sL)" + class BasisSineDVR(BasisSet): r""" Sine DVR basis (particle-in-a-box) for vibrational, angular, and @@ -424,6 +563,7 @@ class BasisSineDVR(BasisSet): """ is_phonon = True + quad = quad def __init__(self, dof, nbas, xi, xf, endpoint=False, quadrature=False, dvr=False): @@ -447,6 +587,8 @@ def __init__(self, dof, nbas, xi, xf, endpoint=False, quadrature=False, self.dvr_x = xi + tmp * self.L / (nbas+1) self.dvr_v = np.sqrt(2/(nbas+1)) * \ np.sin(np.tensordot(tmp, tmp, axes=0)*np.pi/(nbas+1)) + assert np.allclose(self.dvr_v, self.dvr_v.T) + assert np.allclose(self.dvr_v @ self.dvr_v, np.eye(nbas)) self.quadrature = quadrature self.dvr = dvr @@ -625,37 +767,6 @@ def op_mat(self, op: Union[Op, str]): def eigenfunc(self): return "sqrt(2/sL) * sin((sibas+1)*pi*(x-sxi)/sL)" - def quad(self, expr): - x,sL,sxi,sibas,sjbas = sp.symbols("x sL sxi sibas sjbas") - bra = self.eigenfunc - ket = self.eigenfunc.replace("ibas","jbas") - expr = "*".join((bra,expr,ket)) - expr_s = expr.split("dx") - expr_s = [s.rstrip('*') for s in expr_s] - expr_s = [s.lstrip('*') for s in expr_s] - expr_s = [s.replace("^", "**") for s in expr_s] - if len(expr_s) == 1: - expr = sp.sympify(expr_s[0]) - else: - expr = sp.sympify(expr_s[-1]) - for s in expr_s[::-1][1:]: - expr = sp.diff(expr, x) - if s != "": - expr = sp.sympify(s)*expr - expr = expr.subs({sL:self.L, sxi:self.xi}) - logger.debug(f"operator expr: {expr}") - expr = sp.lambdify([x, sibas, sjbas], expr, "numpy") - - mat = np.zeros((self.nbas, self.nbas)) - for ibas in range(self.nbas): - for jbas in range(self.nbas): - val, error = scipy.integrate.quad(lambda x: expr(x, ibas, jbas), - self.xi, self.xf) - logger.debug(f"quadrature value and error: {val}, {error}") - mat[ibas, jbas] = val - return mat - - def _du(self): # int_0^L u=x-xi du=\frac{\partial}{\partial u} mat = np.zeros((self.nbas, self.nbas)) @@ -1028,3 +1139,4 @@ def x_power_k(k, m, n): def p_power_k(k,m,n): # return x_power_k(k,m,n) * (1j)**(m-n) + diff --git a/renormalizer/mps/mp.py b/renormalizer/mps/mp.py index 502cb37f..f7115b6b 100644 --- a/renormalizer/mps/mp.py +++ b/renormalizer/mps/mp.py @@ -344,6 +344,7 @@ def mp_norm(self) -> float: """ res = self.conj().dot(self).real if res < 0: + logger.info(f"mp_norm^2 < 0: {res}") assert np.abs(res) < 1e-8 res = 0 res = np.sqrt(res) @@ -352,8 +353,6 @@ def mp_norm(self) -> float: def add(self, others: List["MatrixProduct"]): - logger.info(f"new_mps:{type(others)}") - if not isinstance(others, list): others = [others] @@ -602,7 +601,7 @@ def variational_compress(self, mpo=None, guess=None): # check convergence if isweep > 0 and percent == 0: error = mps.distance(mps_old) / np.sqrt(mps.dot(mps.conj()).real) - logger.info(f"Variation compress relative error: {error}") + logger.info(f"Variational compress relative error: {error}") if error < mps.compress_config.vrtol: logger.info("Variational compress is converged!") break diff --git a/renormalizer/mps/mpo.py b/renormalizer/mps/mpo.py index 510e6330..9777863f 100644 --- a/renormalizer/mps/mpo.py +++ b/renormalizer/mps/mpo.py @@ -251,8 +251,10 @@ def identity(cls, model: Model): def from_mps(cls, mps): mpo = cls() mpo.model = mps.model + mpo.dtype = mps.dtype for ms in mps: - mo = np.zeros(tuple([ms.shape[0]] + [ms.shape[1]] * 2 + [ms.shape[2]])) + mo = np.zeros(tuple([ms.shape[0]] + [ms.shape[1]] * 2 + + [ms.shape[2]]), dtype=mps.dtype) for iaxis in range(ms.shape[1]): mo[:, iaxis, iaxis, :] = ms[:, iaxis, :].array mpo.append(mo) @@ -272,7 +274,8 @@ def from_mps(cls, mps): mpo.compress_config = mps.compress_config.copy() return mpo - def __init__(self, model: Model = None, terms: Union[Op, List[Op]] = None, offset: Quantity = Quantity(0), ): + def __init__(self, model: Model = None, terms: Union[Op, List[Op]] = None, + offset: Quantity = Quantity(0), dtype=None): """ todo: document @@ -297,8 +300,11 @@ def __init__(self, model: Model = None, terms: Union[Op, List[Op]] = None, offse raise ValueError("Terms all have factor 0.") table, factor = _terms_to_table(model, terms, -self.offset) - - self.dtype = factor.dtype + + if dtype is None: + self.dtype = factor.dtype + else: + self.dtype = dtype mpo_symbol, self.qn, self.qntot, self.qnidx, self.symbolic_out_ops_list, self.primary_ops = construct_symbolic_mpo(table, factor) # print(_format_symbolic_mpo(mpo_symbol)) diff --git a/renormalizer/mps/mps.py b/renormalizer/mps/mps.py index 76547bc5..89c628e7 100644 --- a/renormalizer/mps/mps.py +++ b/renormalizer/mps/mps.py @@ -397,8 +397,10 @@ def from_mpo(cls, mpo: Mpo): # potential energy surface mps = cls() mps.model = mpo.model + mps.dtype = mpo.dtype for mo in mpo: - ms = np.zeros(tuple([mo.shape[0]] + [mo.shape[1]] + [mo.shape[3]])) + ms = np.zeros(tuple([mo.shape[0]] + [mo.shape[1]] + [mo.shape[3]]), + dtype=mps.dtype) for iaxis in range(mo.shape[1]): ms[:, iaxis, :] = mo[:, iaxis, iaxis, :].array mps.append(ms) From 2236e712ad610463937b63e43a187eab410758b2 Mon Sep 17 00:00:00 2001 From: jjren Date: Thu, 31 Aug 2023 15:12:02 +0800 Subject: [PATCH 4/5] add shift algorithm for excited states --- renormalizer/mps/gs.py | 104 +++++++++++++++++++++++++++--- renormalizer/mps/tests/test_gs.py | 44 +++++++++++++ 2 files changed, 138 insertions(+), 10 deletions(-) diff --git a/renormalizer/mps/gs.py b/renormalizer/mps/gs.py index 4e4828f1..5333d8c7 100644 --- a/renormalizer/mps/gs.py +++ b/renormalizer/mps/gs.py @@ -45,7 +45,8 @@ def construct_mps_mpo(model, mmax, nexciton, offset=Quantity(0)): return mps, mpo -def optimize_mps(mps: Mps, mpo: Union[Mpo, StackedMpo], omega: float = None) -> Tuple[List, Mps]: +def optimize_mps(mps: Mps, mpo: Union[Mpo, StackedMpo], omega: float = None, + proj_mpss: List = None) -> Tuple[List, Mps]: r"""DMRG ground state algorithm and state-averaged excited states algorithm Parameters @@ -57,6 +58,11 @@ def optimize_mps(mps: Mps, mpo: Union[Mpo, StackedMpo], omega: float = None) -> omega: float, optional target the eigenpair near omega with special variational function :math:(\hat{H}-\omega)^2. Default is `None`. + proj_mpss: list + the mps and coefficient ([(mps_1,alpha_1),...]) that is used to shift the original Hamiltonian + :math:`H' = H_0 + \sum_j \alpha_j |mps_j> environ = [Environ(mps, item, env) for item in mpo.mpos] else: environ = Environ(mps, mpo, env) + + if proj_mpss is not None: + proj_mpss = [proj_mps.scale(np.sqrt(alpha)) for proj_mps, alpha in proj_mpss] + + if proj_mpss is not None: + identity = Mpo.identity(mps.model) + proj_environ = [Environ(proj_mps, identity, env, + mps_conj=mps.conj()) for proj_mps in proj_mpss] + else: + proj_environ = None macro_iteration_result = [] # Idx of the active site with lowest energy for each sweep @@ -128,7 +144,8 @@ def optimize_mps(mps: Mps, mpo: Union[Mpo, StackedMpo], omega: float = None) -> logger.debug(f"{mps}") - micro_iteration_result, res_mps, mpo = single_sweep(mps, mpo, environ, omega, percent, opt_e_idx) + micro_iteration_result, res_mps, mpo = single_sweep(mps, mpo, environ, + omega, percent, opt_e_idx, proj_mpss, proj_environ) opt_e = min(micro_iteration_result) macro_iteration_result.append(opt_e[0]) @@ -171,7 +188,9 @@ def single_sweep( environ: Environ, omega: float, percent: float, - last_opt_e_idx: int + last_opt_e_idx: int, + proj_mpss: List = None, + proj_environ: List = None, ): method = mps.optimize_config.method @@ -235,10 +254,52 @@ def single_sweep( cmo = [[asxp(mpo_item[idx]) for idx in cidx] for mpo_item in mpo.mpos] else: cmo = [asxp(mpo[idx]) for idx in cidx] + + if proj_environ is not None: + identity = Mpo.identity(mps.model) + projectors = [] + for idx in range(len(proj_environ)): + proj_ltensor = proj_environ[idx].GetLR("L", lidx, + proj_mpss[idx], identity, + itensor=None, method=lmethod, mps_conj=mps.conj()) + proj_rtensor = proj_environ[idx].GetLR("R", ridx, + proj_mpss[idx], identity, + itensor=None, method=rmethod, mps_conj=mps.conj()) + # remove the dummy axis as a result of identity MPO + proj_ltensor = proj_ltensor.reshape(proj_ltensor.shape[0], + proj_ltensor.shape[-1]) + proj_rtensor = proj_rtensor.reshape(proj_rtensor.shape[0], + proj_rtensor.shape[-1]) + if method == "1site": + # a c + # e + # b d + tmp = oe.contract( + "ab,cd,bed->aec", + asxp(proj_ltensor), asxp(proj_rtensor), + asxp(proj_mpss[idx][cidx[0]]), + backend=OE_BACKEND + ) + else: + # a c + # e h + # b g d + tmp = oe.contract( + "ab,cd,beg,ghd->aehc", + asxp(proj_ltensor), asxp(proj_rtensor), + asxp(proj_mpss[idx][cidx[0]]), + asxp(proj_mpss[idx][cidx[1]]), + backend=OE_BACKEND + ) + projectors.append(tmp[qn_mask]) + + else: + projectors = None use_direct_eigh = np.prod(cshape) < 1000 or mps.optimize_config.algo == "direct" if use_direct_eigh: - e, c = eigh_direct(mps, qn_mask, ltensor, rtensor, cmo, omega) + e, c = eigh_direct(mps, qn_mask, ltensor, rtensor, cmo, omega, + projectors) else: # the iterative approach # generate initial guess @@ -268,7 +329,8 @@ def single_sweep( cguess.extend( [np.random.rand(guess_dim) - 0.5 for i in range(len(cguess), nroots)] ) - e, c = eigh_iterative(mps, qn_mask, ltensor, rtensor, cmo, omega, cguess) + e, c = eigh_iterative(mps, qn_mask, ltensor, rtensor, cmo, omega, + cguess, projectors) # if multi roots, both davidson and primme return np.ndarray if nroots > 1: @@ -370,6 +432,7 @@ def eigh_direct( rtensor: Union[xp.ndarray, List[xp.ndarray]], cmo: List[xp.ndarray], omega: float, + projectors: List[xp.ndarray] = None, ): if isinstance(ltensor, list): assert isinstance(rtensor, list) @@ -377,8 +440,15 @@ def eigh_direct( ham = sum([get_ham_direct(mps, qn_mask, ltensor_item, rtensor_item, cmo_item, omega) for ltensor_item, rtensor_item, cmo_item in zip(ltensor, rtensor, cmo)]) else: ham = get_ham_direct(mps, qn_mask, ltensor, rtensor, cmo, omega) - inverse = mps.optimize_config.inverse - w, v = scipy.linalg.eigh(asnumpy(ham) * inverse) + + ham *= mps.optimize_config.inverse + + if projectors is not None: + for proj in projectors: + proj_ham = tensordot(proj, proj.conj(), 0) + ham += proj_ham + + w, v = scipy.linalg.eigh(asnumpy(ham)) nroots = mps.optimize_config.nroots if nroots == 1: @@ -389,7 +459,6 @@ def eigh_direct( c = [v[:, iroot] for iroot in range(min(nroots, v.shape[1]))] return e, c - def get_ham_iterative( mps: Mps, qn_mask: np.ndarray, @@ -407,12 +476,14 @@ def get_ham_iterative( 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) + else: # S-a c d f-S # O-b-O-e-O-g-O @@ -474,6 +545,7 @@ def eigh_iterative( cmo: List[xp.ndarray], omega: float, cguess: List[np.ndarray], + projectors: List[xp.ndarray] = None, ): # iterative algorithm inverse = mps.optimize_config.inverse @@ -485,7 +557,14 @@ def eigh_iterative( expr = func_sum([expr_item for hdiag_item, expr_item in ham]) else: hdiag, expr = get_ham_iterative(mps, qn_mask, ltensor, rtensor, cmo, omega) - + + if projectors is not None: + proj_hdiag = xp.zeros(hdiag.shape) + for proj in projectors: + proj_hdiag += proj * proj.conj() + proj_hdiag = asnumpy(proj_hdiag) + hdiag += proj_hdiag + count = 0 def hop(x): @@ -502,8 +581,13 @@ def hop(x): # convert c to initial structure according to qn pattern cstruct = asxp(cvec2cmat(c, qn_mask)) cout = expr(cstruct) * inverse + cout = cout[qn_mask] + c = asxp(c) + if projectors is not None: + for proj in projectors: + cout += proj.conj().dot(c) * proj # convert structure c to 1d according to qn - res.append(asnumpy(cout)[qn_mask]) + res.append(asnumpy(cout)) if len(res) == 1: return res[0] diff --git a/renormalizer/mps/tests/test_gs.py b/renormalizer/mps/tests/test_gs.py index 4e129564..d18b7a25 100644 --- a/renormalizer/mps/tests/test_gs.py +++ b/renormalizer/mps/tests/test_gs.py @@ -13,8 +13,11 @@ from renormalizer.tests.parameter import holstein_model from renormalizer.utils.configs import OFS from renormalizer.mps.tests import cur_dir +import logging +logger = logging.getLogger(__name__) + nexciton = 1 procedure = [[10, 0.4], [20, 0.2], [30, 0.1], [40, 0], [40, 0]] @@ -59,6 +62,47 @@ def test_multistate(method, algo): assert np.allclose(energy[-1], energy_std) assert np.allclose(expectation, energy_std) +@pytest.mark.parametrize("method", ( + "1site", + "2site", +)) +@pytest.mark.parametrize("algo", ( + "davidson", + pytest.param("primme", marks=pytest.mark.skipif(primme is None, reason="primme not installed")) +)) +def test_multistate_projector(method, algo): + model = holstein_model.switch_scheme(4) + mps, mpo = construct_mps_mpo(model, procedure[0][0], nexciton) + proj_mpss = [] + nstates = 4 + for istate in range(nstates): + mps = Mps.random(model, nexciton, procedure[0][0], percent=1.0) + mps.optimize_config.procedure = procedure + mps.optimize_config.nroots = 1 + mps.optimize_config.method = method + mps.optimize_config.algo = algo + mps.optimize_config.e_atol = 1e-6 + mps.optimize_config.e_rtol = 1e-6 + energy, mps = optimize_mps(mps, mpo, proj_mpss=proj_mpss) + proj_mpss.append((mps, 0.005)) + + proj_mpss = [x[0] for x in proj_mpss] + # for debug + #ham_mat = np.zeros([nstates, nstates]) + #for i in range(nstates): + # for j in range(i+1): + # logger.info(f"{i}, {j}, orth:{proj_mpss[i].conj().dot(proj_mpss[j])}") + # ham_mat[i,j] = ham_mat[j,i] = proj_mpss[j].expectation(mpo, proj_mpss[i].conj()) + #e, v = scipy.linalg.eigh(ham_mat) + #logger.info(f"{ham_mat}") + #logger.info(f"e={e}") + expectation = [mp.expectation(mpo) for mp in proj_mpss] + energy_std = np.array([0.08401412, 0.08449771, 0.08449801, 0.08449945]) + holstein_model.gs_zpe + logger.info(f"{expectation}") + logger.info(f"{energy_std}") + + assert np.allclose(expectation, energy_std) + @pytest.mark.parametrize("method", ( "1site", From d49fdae8e8adb27973424ce67063b1ce4c1d1343 Mon Sep 17 00:00:00 2001 From: jjren Date: Sun, 17 Sep 2023 21:28:10 +0800 Subject: [PATCH 5/5] support time correlation function calculation of nuclear Hamiltonian --- renormalizer/model/basis.py | 10 + renormalizer/mps/gs.py | 2 + renormalizer/mps/mpdm.py | 42 +- renormalizer/mps/mpo.py | 12 + renormalizer/mps/mps.py | 4 +- renormalizer/mps/tda.py | 54 +- renormalizer/mps/thermalprop.py | 40 +- renormalizer/tcf/__init__.py | 2 + renormalizer/tcf/base.py | 922 ++++++++++++++++++++++++++++++++ 9 files changed, 1061 insertions(+), 27 deletions(-) create mode 100644 renormalizer/tcf/__init__.py create mode 100644 renormalizer/tcf/base.py diff --git a/renormalizer/model/basis.py b/renormalizer/model/basis.py index 8dcacafe..d1fff885 100644 --- a/renormalizer/model/basis.py +++ b/renormalizer/model/basis.py @@ -352,6 +352,10 @@ def op_mat(self, op: Union[Op, str]): raise ValueError(f"op_symbol:{op_symbol} is not supported. Quadrature is not implemented yet!") self._recursion_flag -= 1 + + if np.allclose(mat, mat.conj()): + mat = mat.real + return mat * op_factor def copy(self, new_dof): @@ -520,6 +524,9 @@ def op_mat(self, op: Union[Op, str]): else: mat = self.quad(op_symbol, complex_func=True) mat = self.dvr_v @ mat @ self.dvr_v.conj().T + + if np.allclose(mat, mat.conj()): + mat = mat.real return mat * op_factor @@ -761,6 +768,9 @@ def op_mat(self, op: Union[Op, str]): if self.dvr and self._recursion_flag == 0: mat = self.dvr_v.T @ mat @ self.dvr_v + if np.allclose(mat, mat.conj()): + mat = mat.real + return mat * op_factor @property diff --git a/renormalizer/mps/gs.py b/renormalizer/mps/gs.py index 5333d8c7..7510d7b8 100644 --- a/renormalizer/mps/gs.py +++ b/renormalizer/mps/gs.py @@ -637,6 +637,8 @@ def hop(x): method="PRIMME_DYNAMIC", tol=1e-6, ) + if nroots == 1 and isinstance(e, np.ndarray): + e = e[0] else: assert False logger.debug(f"use {algo}, HC hops: {count}") diff --git a/renormalizer/mps/mpdm.py b/renormalizer/mps/mpdm.py index 7aff5955..e0cd5dc0 100644 --- a/renormalizer/mps/mpdm.py +++ b/renormalizer/mps/mpdm.py @@ -66,6 +66,46 @@ def evolve_exact(self, h_mpo, evolve_dt, space): new_mpdm = self.apply(MPOprop, canonicalise=True) new_mpdm.coeff *= np.exp(-1.0j * h_mpo.offset * evolve_dt) return new_mpdm + + @classmethod + def max_entangled_mpdm(cls, model, qnblock, normalize=True): + r''' only for single multi_dof state or spin state + ''' + if isinstance(qnblock, int): + qnblock = np.array([qnblock]) + + # todo add assert + mpdm = cls() + mpdm.model = model + qn_size = model.qn_size + qntot = np.zeros(qn_size, dtype=int) + mpdmqn = [np.zeros((1, qn_size), dtype=int)] + for iba, ba in enumerate(model.basis): + pdim = ba.nbas + if ba.is_phonon: + ms = np.eye(pdim).reshape(1,pdim,pdim,1) + if normalize: + ms /= np.sqrt(pdim) + elif ba.multi_dof or ba.is_spin: + ms = np.zeros(pdim) + ms[(np.array(ba.sigmaqn)==qnblock).all(axis=1)] = 1 + ms = np.diag(ms).reshape(1,pdim,pdim,1) + if normalize: + nonzero = np.count_nonzero((np.array(ba.sigmaqn)==qnblock).all(axis=1)) + ms /= np.sqrt(nonzero) + qntot += qnblock + else: + assert False + + mpdm.append(ms) + mpdmqn.append(qntot.reshape(1,-1)) + mpdm.qn = mpdmqn + mpdm.qn[-1] = np.zeros((1, qn_size), dtype=int) + mpdm.qntot = qntot + mpdm.qnidx = mpdm.site_num-1 + mpdm.to_right = False + + return mpdm def todense(self): # explicitly call to MPO because MPS is firstly inherited @@ -106,7 +146,7 @@ def _expectation_path(self): return path def conj_trans(self): - raise NotImplementedError + #raise NotImplementedError logger.warning("using conj_trans on mpdm leads to dummy qn") new_mpdm: "MpDmBase" = super().conj_trans() new_mpdm.coeff = new_mpdm.coeff.conjugate() diff --git a/renormalizer/mps/mpo.py b/renormalizer/mps/mpo.py index 9777863f..ceef88df 100644 --- a/renormalizer/mps/mpo.py +++ b/renormalizer/mps/mpo.py @@ -506,6 +506,18 @@ def is_hermitian(self): def __matmul__(self, other): return self.apply(other) + + def __add__(self, other: Union["Mpo", float, complex]): + if isinstance(other, Mpo): + return self.add(other) + iden = self.identity(self.model).scale(other) + return self.add(iden) + + def __sub__(self, other: Union["Mpo", float, complex]): + if isinstance(other, Mpo): + return self.add(other.scale(-1)) + iden = self.identity(self.model).scale(-other) + return self.add(iden) class StackedMpo: diff --git a/renormalizer/mps/mps.py b/renormalizer/mps/mps.py index 89c628e7..d36b2930 100644 --- a/renormalizer/mps/mps.py +++ b/renormalizer/mps/mps.py @@ -613,6 +613,8 @@ def expand_bond_dimension(self, hint_mpo=None, coef=1e-10, include_ex=True): """ # expander m target m_target = self.compress_config.bond_dim_max_value - self.bond_dims_mean + if m_target <= 0: + return self # will be restored at exit self.compress_config.bond_dim_max_value = m_target if self.compress_config.criteria is not CompressCriteria.fixed: @@ -645,7 +647,7 @@ def expand_bond_dimension(self, hint_mpo=None, coef=1e-10, include_ex=True): lastone = self + ex_state else: - lastone = self + lastone = self.copy().normalize("mps_and_coeff") expander_list: List["MatrixProduct"] = [] cumulated_m = 0 while True: diff --git a/renormalizer/mps/tda.py b/renormalizer/mps/tda.py index a6afec76..6b7580d7 100644 --- a/renormalizer/mps/tda.py +++ b/renormalizer/mps/tda.py @@ -44,7 +44,8 @@ def __init__(self, model, hmpo, mps, nroots=1, algo=None): self.model = model self.hmpo = hmpo - self.mps = mps # mps will be overwritten inplace + self.psi0 = mps + self.mps = mps.copy() # mps will be overwritten inplace self.nroots = nroots if algo is None: if primme is not None: @@ -276,23 +277,30 @@ def multi_hop(x): return np.stack([hop(x[:,i]) for i in range(x.shape[1])],axis=1) else: assert False - + + e0 = self.psi0.expectation(self.hmpo) + tmp = hdiag-e0 + tmp[np.where(np.abs(tmp)<1e-6)] = 1e-6 + def precond(x): if x.ndim == 1: - return np.einsum("i, i -> i", 1/(hdiag+1e-4), x) + return np.einsum("i, i -> i", 1/tmp, x) elif x.ndim ==2: - return np.einsum("i, ij -> ij", 1/(hdiag+1e-4), x) + return np.einsum("i, ij -> ij", 1/tmp, 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) + e, c, stats = primme.eigsh(A, k=min(nroots,xsize), which="SA", + v0=cguess, + OPinv=M, + method="PRIMME_DYNAMIC", + tol=1e-6, + return_stats=True, + return_history=False) + logger.debug(f"primme statistics: {stats}") else: assert False @@ -351,7 +359,33 @@ def load_wfn(self, model): tda_coeff_list.append(tda_coeff) self.wfn = [mps_l_cano, mps_r_cano, tangent_u, tda_coeff_list] - + + def trans_gs_ex(self, mpo, mps=None): + """ + calculate transition matrix element between gs and tda_root + + """ + if mps is None: + mps = self.psi0 + mps_conj = mps.conj() + + mps_l_cano, mps_r_cano, tangent_u, tda_coeff_list = self.wfn + + trans_tot = [] + for iroot in range(self.nroots): + tda_coeff = tda_coeff_list[iroot] + + trans = 0 + for ims in range(mps_l_cano.site_num): + if tangent_u[ims] is None: + assert tda_coeff[ims] is None + continue + mps_tangent = merge(mps_l_cano, mps_r_cano, ims+1) + mps_tangent[ims] = tensordot(tangent_u[ims], tda_coeff[ims],[-1,0]) + trans += mps_tangent.expectation(mpo, mps_conj) + trans_tot.append(trans) + return trans_tot + def analysis_1ordm(self): r""" calculate one-orbital reduced density matrix of each tda root diff --git a/renormalizer/mps/thermalprop.py b/renormalizer/mps/thermalprop.py index 71edfc56..78e177fc 100644 --- a/renormalizer/mps/thermalprop.py +++ b/renormalizer/mps/thermalprop.py @@ -45,11 +45,18 @@ def __init__( job_name: str = None, properties: Property = None, auto_expand: bool = True, + include_ex: bool = True, ): self.init_mpdm: MpDm = init_mpdm.canonicalise() if h_mpo_model is None: h_mpo_model = self.init_mpdm.model - self.h_mpo = Mpo(h_mpo_model) + if isinstance(h_mpo_model, Mpo): + self.h_mpo = h_mpo_model + else: + if "h_mpo" in h_mpo_model.mpos.keys(): + self.h_mpo = h_mpo_model.mpos["h_mpo"] + else: + self.h_mpo = Mpo(h_mpo_model) logger.info(f"Bond dim of h_mpo: {self.h_mpo.bond_dims}") self.exact = exact assert space in ["GS", "EX"] @@ -60,6 +67,7 @@ def __init__( self._vn_entropy_array = [] self.properties = properties self.auto_expand = auto_expand + self.include_ex = include_ex super().__init__(evolve_config=evolve_config, dump_mps=dump_mps, dump_dir=dump_dir, job_name=job_name) @@ -67,7 +75,8 @@ def __init__( def init_mps(self): self.init_mpdm.evolve_config = self.evolve_config if self.evolve_config.is_tdvp and self.auto_expand: - self.init_mpdm = self.init_mpdm.expand_bond_dimension(self.h_mpo) + self.init_mpdm = self.init_mpdm.expand_bond_dimension(self.h_mpo, + include_ex=self.include_ex) return self.init_mpdm def process_mps(self, mps): @@ -76,17 +85,18 @@ def process_mps(self, mps): if self.exact: # skip the fuss for efficiency return - for attr_str in ["e_occupations", "ph_occupations"]: - attr = getattr(mps, attr_str) - logger.info(f"{attr_str}: {attr}") - self_array = getattr(self, f"_{attr_str}_array") - self_array.append(attr) - vn_entropy = mps.calc_bond_entropy() - self._vn_entropy_array.append(vn_entropy) - logger.info(f"vn entropy: {vn_entropy}") - logger.info( - f"Energy: {new_energy}, total electron: {self._e_occupations_array[-1].sum()}" - ) + # this only works for e-ph problem + #for attr_str in ["e_occupations", "ph_occupations"]: + # attr = getattr(mps, attr_str) + # logger.info(f"{attr_str}: {attr}") + # self_array = getattr(self, f"_{attr_str}_array") + # self_array.append(attr) + #vn_entropy = mps.calc_bond_entropy() + #self._vn_entropy_array.append(vn_entropy) + #logger.info(f"vn entropy: {vn_entropy}") + #logger.info( + # f"Energy: {new_energy}, total electron: {self._e_occupations_array[-1].sum()}" + #) # calculate other properties defined in Property if self.properties is not None: @@ -103,8 +113,8 @@ def evolve_exact(self, old_mpdm: MpDm, evolve_dt): return new_mpdm def evolve_prop(self, old_mpdm, evolve_dt): - h_mpo = Mpo(self.h_mpo.model, offset=Quantity(self.energies[-1])) - return old_mpdm.evolve(h_mpo, evolve_dt) + #h_mpo = Mpo(self.h_mpo.model, offset=Quantity(self.energies[-1])) + return old_mpdm.evolve(self.h_mpo, evolve_dt) def evolve_single_step(self, evolve_dt): old_mpdm = self.latest_mps diff --git a/renormalizer/tcf/__init__.py b/renormalizer/tcf/__init__.py new file mode 100644 index 00000000..633f8661 --- /dev/null +++ b/renormalizer/tcf/__init__.py @@ -0,0 +1,2 @@ +# -*- coding: utf-8 -*- + diff --git a/renormalizer/tcf/base.py b/renormalizer/tcf/base.py new file mode 100644 index 00000000..bd676a4d --- /dev/null +++ b/renormalizer/tcf/base.py @@ -0,0 +1,922 @@ +# -*- coding: utf-8 -*- +''' +only works for multi_electron state +''' + +from renormalizer.utils import TdMpsJob +from renormalizer.mps.backend import xp +from renormalizer.mps.matrix import asxp +from renormalizer.mps import Mpo, Mps, MpDm, gs, ThermalProp +from renormalizer.mps.mps import BraKetPair +from renormalizer.utils import OptimizeConfig, EvolveConfig, CompressConfig, constant, CompressCriteria +from renormalizer.model import Op, Model +#from renormalizer.transport.kubo import current_op +from renormalizer.model import basis as ba + +import numpy as np +import logging +import itertools +import os + +logger = logging.getLogger(__name__) + + +def single_mol_model(fdusin, fnac, projector=0): + w0 = [] + d0 = [] # project on PES0 normal coordinates + w1 = [] + d1 = [] # project on PES1 normal coordinates + s021 = [] + s120 = [] + with open(fdusin, "r") as f: + lines = f.readlines() + read_start = False + for line in lines: + if line[:6] == "------": + if read_start: + break + else: + read_start = True + continue + if read_start: + split_line = line.split() + w0.append(float(split_line[3])) + d0.append(float(split_line[4])) + w1.append(float(split_line[9])) + d1.append(float(split_line[10])) + + nmodes = len(w0) + + start_s120 = start_s021 = False + for iline, line in enumerate(lines): + if line.rstrip().lstrip() == "BEGIN_DUSH_DATA_1": + start_s021 = True # s021: S matrix (x_0, x_1) + start_s120 = False # s120: S matrix (x_1, x_0) + if line.rstrip().lstrip() == "BEGIN_DUSH_DATA_2": + start_s120 = True + start_s021 = False + if start_s120 or start_s021: + if line.split()[0] == "MODE": + res = [] + for subline in lines[iline+1:iline+int(np.ceil(nmodes/10))+1]: + res.extend([float(i) for i in subline.split()]) + if start_s120: + s120.append(res) + elif start_s021: + s021.append(res) + nac = [] + with open(fnac, "r") as f: + lines = f.readlines() + for iline, line in enumerate(lines): + split_line = line.split() + if len(split_line) > 0: + if split_line[0] == "m" and split_line[1] == "n": + for subline in lines[iline+2:iline+2+nmodes]: + nac.append(float(subline.split()[4])) + + s021 = np.stack(s021, axis=0) + s120 = np.stack(s120, axis=0) + assert np.allclose(s021.dot(s021.T), np.eye(nmodes), atol=1e-6) + assert np.allclose(s120.dot(s120.T), np.eye(nmodes), atol=1e-6) + assert np.allclose(s120, s021.T) + nmodes -= projector + w0 = np.array(w0[projector:]) * constant.cm2au + w1 = np.array(w1[projector:]) * constant.cm2au + d0 = np.array(d0[projector:]) + d1 = np.array(d1[projector:]) + nac = np.array(nac[projector:]) + + s021 = s021[projector:, projector:] + s120 = s120[projector:, projector:] + + return w0, w1, d0, d1, nac, s021, s120 + + +def abs_dipole_op(model): + """ + excitation between 0 -> 1 + dipole contains dipole transition integral + """ + dipole_terms = [] + assert "dipole" in model.para.keys() + for key, value in model.para["dipole"].items(): + siteidx = model.dof_to_siteidx[key[0]] + ba = model.basis[siteidx] + idx0 = ba.dof.index(key[0]) + idx1 = ba.dof.index(key[1]) + if ba.sigmaqn[idx0] == 1 and ba.sigmaqn[idx1] == 0: + dof = list(key) + else: + dof = list(key[::-1]) + print("dof", dof, value) + dipole_terms.append(Op("a^\dagger a", dof, value, qn=[1,0])) + + return Mpo(model, terms=dipole_terms) + +def emi_dipole_op(model): + # the constant coupling operator (condon approximation) + return abs_dipole_op(model).conj_trans() + +def nac_op(model): + # the momentum operator coupling + # the quantum number is gs: 0 ex:1 + if "nac_mpo" in model.mpos.keys(): + logger.info("load nac_mpo form model.mpos") + return model.mpos["nac_mpo"] + else: + nac_terms = [] + for key, value in model.para["nac"].items(): + nac_terms.append(Op("a^\dagger a partialx", ["gs","ex", key], + factor=-value, qn=[0,-1,0])) + return Mpo(model, terms=nac_terms) + +def spin_op(model): + # spin operator + spin_terms = [] + for key, value in model.para["spin"].items(): + spin_terms.append(Op(value, key, factor=1, qn=0)) + return Mpo(model, terms=spin_terms) + +class CorrFuncBase(TdMpsJob): + r'''Abstract class + + ''' + def __init__( + self, + model, + imps_qntot, + temperature, + optimize_config=None, + compress_config=None, + evolve_config=None, + dump_mps=None, + dump_dir=None, + job_name=None, + ): + + self.model = model + if "h_mpo" in model.mpos.keys(): + logger.info("load h_mpo form model.mpos") + self.h_mpo = model.mpos["h_mpo"] + else: + self.h_mpo = Mpo(model) + self.temperature = temperature + self.compress_config = compress_config + if self.compress_config is None: + self.compress_config = CompressConfig() + self.optimize_config = optimize_config + if self.optimize_config is None: + self.optimize_config = OptimizeConfig() + self.imps_qntot = imps_qntot + self.dump_dir = dump_dir + self.job_name = job_name + self.imps = self.init_imps() + self.op_b = self.init_op_b() + self.op_a = self.init_op_a() + + self.e_imps = self.imps.expectation(self.h_mpo) + + self._autocorr = [] + self._autocorr_time = [] + + self.e_ket = self.e_bra = 0 + + super(CorrFuncBase, self).__init__(evolve_config=evolve_config, + dump_mps=dump_mps, dump_dir=dump_dir, job_name=job_name) + + self.latest_mps.ket_mps.evolve_config = self.latest_mps.bra_mps.evolve_config = self.evolve_config + self.latest_mps.ket_mps.compress_config = self.latest_mps.bra_mps.compress_config = self.compress_config + + self.e_ket = self.latest_mps.ket_mps.expectation(self.h_mpo) + self.e_bra = self.latest_mps.bra_mps.expectation(self.h_mpo) + logger.debug(f"e_ket:{self.e_ket}, e_bra:{self.e_bra}") + self.ket_mpo = self.h_mpo - self.e_ket + self.bra_mpo = self.h_mpo - self.e_bra + + def pruner(self, braket_pair): + bra_mps, ket_mps = braket_pair + bra_mps.compress_config = ket_mps.compress_config = self.compress_config + if self.evolve_config.is_tdvp: + logger.debug("expand ket mps.") + ket_mps = ket_mps.expand_bond_dimension(self.h_mpo, + include_ex=False) + logger.debug("expand bra mps.") + bra_mps = bra_mps.expand_bond_dimension(self.h_mpo, + include_ex=False) + logger.debug("compress ket mps.") + ket_mps.canonicalise().compress().normalize("mps_only") + logger.debug("compress bra mps.") + bra_mps.canonicalise().compress().normalize("mps_only") + + return BraKetPair(bra_mps, ket_mps, braket_pair.mpo) + + + def init_op_a(self): + raise NotImplementedError + + def init_op_b(self): + raise NotImplementedError + + def init_imps(self): + raise NotImplementedError + + @property + def autocorr(self): + return np.array(self._autocorr) + + @property + def autocorr_time(self): + return np.array(self._autocorr_time) + + + def get_dump_dict(self): + dump_dict = dict() + dump_dict["correlation function time"] = self.autocorr_time + dump_dict["ACF"] = self.autocorr + dump_dict["time series"] = self.evolve_times + #if self.properties is not None: + # for prop_str in self.properties.prop_res.keys(): + # dump_dict[prop_str] = self.properties.prop_res[prop_str] + + return dump_dict + + +class FTCorrFuncBase(CorrFuncBase): + r'''Abstract class + ''' + + def __init__( + self, + model, + imps_qntot, + temperature, + insteps, + ievolve_config=None, + icompress_config=None, + compress_config=None, + evolve_config=None, + dump_mps=None, + dump_dir=None, + job_name=None, + ): + + self.insteps = insteps + self.ievolve_config = ievolve_config + self.icompress_config = icompress_config + if self.icompress_config is None: + self.icompress_config = CompressConfig() + if self.ievolve_config is None: + self.ievolve_config = EvolveConfig() + + super(FTCorrFuncBase, self).__init__( + model, + imps_qntot, + temperature, + compress_config=compress_config, + evolve_config=evolve_config, + dump_mps=dump_mps, + dump_dir=dump_dir, + job_name=job_name, + ) + + + def init_mps(self): + ket_mps = self.op_b.apply(self.imps) + ket_mps.normalize("mps_norm_to_coeff") + bra_mps = self.imps + + return self.pruner(BraKetPair(bra_mps, ket_mps, mpo=self.op_a.conj_trans())) + + def init_imps(self): + mpdm = MpDm.max_entangled_mpdm(self.model, self.imps_qntot) + mpdm.compress_config = self.icompress_config + tp = ThermalProp( + mpdm, self.h_mpo, evolve_config=self.ievolve_config, + dump_dir=self.dump_dir, job_name=self.job_name, include_ex=False + ) + if tp._defined_output_path: + try: + logger.info( + f"load density matrix from {self._thermal_dump_path}" + ) + mpdm = MpDm.load(self.model, self._thermal_dump_path) + logger.info(f"density matrix loaded:{mpdm}") + return mpdm + except FileNotFoundError: + logger.debug(f"no file found in {self._thermal_dump_path}") + logger.info(f"calculate mpdm from scratch.") + + tp.evolve(None, self.insteps, self.temperature.to_beta() / 2j) + mpdm = tp.latest_mps + if tp._defined_output_path: + mpdm.dump(self._thermal_dump_path) + + return mpdm + + def evolve_single_step(self, evolve_dt): + bra_mps, ket_mps = self.latest_mps + bra_mps = bra_mps.evolve(self.bra_mpo, evolve_dt) + ket_mps = ket_mps.evolve(self.ket_mpo, evolve_dt) + return BraKetPair(bra_mps, ket_mps, mpo=self.latest_mps.mpo) + + def process_mps(self, braket_pair): + t = self.evolve_times[-1] + self._autocorr.append(braket_pair.ft * np.exp(1.0j*t*(self.e_bra - self.e_ket))) + self._autocorr_time.append(t) + + #bra_mps, ket_mps = braket_pair + #bra_rdm = bra_mps.calc_reduced_density_matrix() + #ket_rdm = ket_mps.calc_reduced_density_matrix() + #self.e_rdm.append([bra_rdm, ket_rdm]) + + @property + def _thermal_dump_path(self): + assert self._defined_output_path + return os.path.join(self.dump_dir, self.job_name + "_impo.npz") + +class ZTCorrFuncBase(CorrFuncBase): + r'''Abstract class + ''' + + def init_imps(self): + mmax = self.optimize_config.procedure[0][0] + imps = Mps.random(self.model, self.imps_qntot, mmax, percent=1) + imps.optimize_config = self.optimize_config + energies, imps = gs.optimize_mps(imps, self.h_mpo) + #e_rdm = imps.calc_reduced_density_matrix() + #logger.info(f"e_rdm: {e_rdm}") + return imps + + def process_mps(self, braket_pair): + if not self.evolve_times[-1] == 0: + last_bra_mps, last_ket_mps = self.latest_mps + bra_mps, ket_mps = braket_pair + t_ket = self.evolve_times[-1] + t_bra = self.evolve_times[-2] + self._autocorr.append(BraKetPair(last_bra_mps, ket_mps).ft * + np.exp(-1.0j*t_ket*self.e_ket) * + np.exp(-1.0j*t_bra*self.e_bra) * + np.exp(1.0j*(t_bra+t_ket)*self.e_imps)) + self._autocorr_time.append(t_bra + t_ket) + + t = self.evolve_times[-1] + self._autocorr.append(braket_pair.ft * + np.exp(-1.0j*t*(self.e_ket+self.e_bra)) * + np.exp(1.0j*2*t*self.e_imps)) + self._autocorr_time.append(2*t) + + bra_mps, ket_mps = braket_pair + #bra_rdm = bra_mps.calc_reduced_density_matrix() + #ket_rdm = ket_mps.calc_reduced_density_matrix() + #self.e_rdm.append([bra_rdm, ket_rdm]) + +class ZTCorrFuncBase_state_to_state(ZTCorrFuncBase): + def __init__(self, broad_func, acf_stos_compress_config, *args, **kw): + # broadening function + self.broad_func = broad_func + assert acf_stos_compress_config is not None + self.acf_stos_compress_config = acf_stos_compress_config + self.acf_stos = None + self.rate = [0,] + super().__init__(*args, **kw) + + def process_mps(self, braket_pair): + bra_mps, ket_mps = braket_pair + + mpdm_compress_config = self.acf_stos_compress_config + + if not self.evolve_times[-1] == 0: + last_bra_mps, last_ket_mps = self.latest_mps + t_ket = self.evolve_times[-1] + t_bra = self.evolve_times[-2] + + mpdm_odd = MpDm.from_bra_ket(ket_mps, last_bra_mps.conj()) + mpdm_odd = mpdm_odd.scale( + np.exp(-1.0j*t_ket*self.e_ket) * + np.exp(-1.0j*t_bra*self.e_bra) * + np.exp(1.0j*(t_bra+t_ket)*self.e_imps)*self.broad_func(t_bra+t_ket)) + self._autocorr_time.append(t_bra + t_ket) + mpdm_odd = mpdm_odd.diagonal() + self._autocorr.append(mpdm_odd.sum()) + mpdm_odd.compress_config = mpdm_compress_config + mpdm_odd.canonicalise().compress() + + t = self.evolve_times[-1] + mpdm_even = MpDm.from_bra_ket(ket_mps, bra_mps.conj()) + mpdm_even = mpdm_even.scale( + np.exp(-1.0j*t*self.e_ket) * + np.exp(-1.0j*t*self.e_bra) * + np.exp(1.0j*2*t*self.e_imps)*self.broad_func(2*t)) + self._autocorr_time.append(2*t) + mpdm_even = mpdm_even.diagonal() + self._autocorr.append(mpdm_even.sum()) + mpdm_even.compress_config = mpdm_compress_config + mpdm_even.canonicalise().compress() + + if self.autocorr_time[-1] == 0: + # the 0 evolve step + self.acf_stos = mpdm_even + else: + dt = self.autocorr_time[-1] - self.autocorr_time[-2] + if self.autocorr_time[-3] == 0: + # the first evolve step + new_acf_stos = self.acf_stos.scale(dt/2) + mpdm_odd.scale(dt) + mpdm_even.scale(dt) + else: + new_acf_stos = self.acf_stos + mpdm_odd.scale(dt) + mpdm_even.scale(dt) + new_acf_stos.compress_config = mpdm_compress_config + new_acf_stos.canonicalise().compress() + self.acf_stos = new_acf_stos + current_rate = self.acf_stos.sum() + self.rate.extend([current_rate,]*2) + logger.debug(f"acf_stos:{self.acf_stos}") + #bra_rdm = bra_mps.calc_reduced_density_matrix() + #ket_rdm = ket_mps.calc_reduced_density_matrix() + #self.e_rdm.append([bra_rdm, ket_rdm]) + + + def dump_dict(self): + super().dump_dict() + if self.acf_stos is not None: + mps_path = os.path.join(self.dump_dir, + self.job_name+"_acf_stos" + ".npz") + self.acf_stos.dump(mps_path) + + def get_dump_dict(self): + super_dump_dict = super().get_dump_dict() + super_dump_dict["rate"] = self.rate + + return super_dump_dict + +class FTCorrFuncBase_state_to_state(FTCorrFuncBase): + def __init__(self, broad_func, acf_stos_compress_config, *args, **kw): + self.broad_func = broad_func + assert acf_stos_compress_config is not None + self.acf_stos_compress_config = acf_stos_compress_config + self.acf_stos = None + self.rate = [0,] + super().__init__(*args, **kw) + + + def process_mps(self, braket_pair): + t = self.evolve_times[-1] + bra_mps, ket_mps = braket_pair + mpdm_compress_config = self.acf_stos_compress_config + + bra_mps = self.op_a.apply(bra_mps) + bra_mps.canonicalise() + bra_mps.compress() + + mpdm = ket_mps.apply(bra_mps.conj_trans()).scale( + np.exp(1.0j*t*(self.e_bra - self.e_ket)) * self.broad_func(t)) + mpdm = mpdm.diagonal() + self._autocorr_time.append(t) + self._autocorr.append(mpdm.sum()) + mpdm.compress_config = mpdm_compress_config + logger.debug(f"mpdm:{mpdm}") + mpdm.canonicalise() + mpdm.compress() + + if self.autocorr_time[-1] == 0: + self.acf_stos = mpdm + # the 0 evolve step + else: + dt = self.autocorr_time[-1] - self.autocorr_time[-2] + if self.autocorr_time[-2] == 0: + # the first evolve step + new_acf_stos = self.acf_stos.scale(dt/2) + mpdm.scale(dt) + else: + new_acf_stos = self.acf_stos + mpdm.scale(dt) + new_acf_stos.compress_config = mpdm_compress_config + new_acf_stos.canonicalise().compress() + self.acf_stos = new_acf_stos + self.rate.append(self.acf_stos.sum()) + logger.debug(f"acf_stos:{self.acf_stos}") + + #bra_rdm = bra_mps.calc_reduced_density_matrix() + #ket_rdm = ket_mps.calc_reduced_density_matrix() + #self.e_rdm.append([bra_rdm, ket_rdm]) + + def dump_dict(self): + super().dump_dict() + if self.acf_stos is not None: + mps_path = os.path.join(self.dump_dir, + self.job_name+"_acf_stos" + ".npz") + self.acf_stos.dump(mps_path) + + def get_dump_dict(self): + super_dump_dict = super().get_dump_dict() + super_dump_dict["rate"] = self.rate + + return super_dump_dict + +def analysis_dominant_config(mps, nconfigs=1): + # analysis the dominant configuration an mps + # mps should be real + + config_visited = [] + ci_coeff_list = [] + while len(config_visited) < nconfigs: + mps_rank1 = mps.canonicalise().compress(temp_m_trunc=1) + # get config with the largest coeff + config = [] + for ims, ms in enumerate(mps_rank1): + ms = ms.array.flatten()**2 + quanta = int(np.argmax(ms)) + config.append(quanta) + + while config in config_visited: + # random mutate + idx = np.random.randint(mps.site_num) + quant = np.random.randint(mps.model.pbond_list[idx]) + config[idx] = quant + + config_visited.append(config) + + sentinel = xp.ones(1) + for ims, ms in enumerate(mps): + sentinel = xp.tensordot(sentinel, asxp(ms[:,config[ims],:]), + axes=(0,0)) + + ci_coeff_list.append(sentinel.item()*mps.coeff) + condition = {} + for idx in range(len(config)): + dofname = mps.model.basis[idx].dofs + condition[dofname[0]] = config[idx] + mps = mps - Mps.hartree_product_state(mps.model, condition).scale(ci_coeff_list[-1]) + + return config_visited, ci_coeff_list + +def monte_carlo_sampling(mps, vec, prob=None, nconfigs=1, converge=None, bag=None): + + if converge is None: + converge = (mps.sum()/mps.coeff).real + + np.random.seed() + if bag is None: + bag = dict() + tot = 0 + else: + assert isinstance(bag, dict) + tot = np.sum(list(bag.values())).real + logger.info(f"current tot: {tot}") + + coeff = ci_coeff(mps, vec) + if vec not in bag: + bag[vec] = coeff + tot += coeff.real + + while len(bag) < nconfigs and tot < converge: + isite = np.random.randint(mps.site_num) + if prob is None: + quant = np.random.randint(mps.model.pbond_list[isite]) + else: + p = np.random.random() + for ival, val in enumerate(prob[isite]): + if p < val: + break + quant = ival + # if not mutated, continue + if quant == vec[isite]: + continue + vec_new = list(vec) + vec_new[isite] = quant + vec_new = tuple(vec_new) + if vec_new in bag: + coeff_new = bag[vec_new] + else: + coeff_new = ci_coeff(mps, vec_new) + bag[vec_new] = coeff_new + tot += coeff_new.real + + # if accept the mutation + p = min(1, coeff_new.real/coeff.real) + if np.random.random() < p: + vec = vec_new + coeff = coeff_new + return bag + +def independent_sample(mps, prob, nconfigs=1, thresh=0): + np.random.seed() + bag = dict() + while len(bag) < nconfigs: + p = np.random.random(size=mps.site_num) + config = [] + for isite in range(mps.site_num): + for ival, val in enumerate(prob[isite]): + if p[isite] < val: + break + config.append(ival) + config = tuple(config) + if config in bag: + continue + coeff = ci_coeff(mps, config) + if coeff.real > thresh: + bag[config] = coeff + return bag + +def ci_coeff(mps, config): + sentinel = xp.ones(1) + for ims, ms in enumerate(mps): + sentinel = xp.tensordot(sentinel, asxp(ms[:,config[ims],:]), + axes=(0,0)) + return sentinel.item() + + +class ZTAACorrFuncBase(ZTCorrFuncBase): + r''' Abstract class + Note: please make sure the operator A is real + ''' + def init_mps(self): + ket_mps = self.op_b.apply(self.imps) + ket_mps.normalize("mps_norm_to_coeff") + + bra_mps = ket_mps.copy() + + return self.pruner(BraKetPair(bra_mps, ket_mps)) + + def init_op_a(self): + return self.init_op_b() + + def evolve_single_step(self, evolve_dt): + _, ket_mps = self.latest_mps + ket_mps = ket_mps.evolve(self.ket_mpo, evolve_dt) + bra_mps = ket_mps.conj() + return BraKetPair(bra_mps, ket_mps) + +class ZTAACorrFuncBase_state_to_state(ZTAACorrFuncBase, + ZTCorrFuncBase_state_to_state): + pass + +class ZTABCorrFuncBase(ZTCorrFuncBase): + r'''Abstract class + ''' + + def init_mps(self): + ket_mps = self.op_b.apply(self.imps) + ket_mps.normalize("mps_norm_to_coeff") + bra_mps = self.op_a.apply(self.imps) + bra_mps.normalize("mps_norm_to_coeff") + return self.pruner(BraKetPair(bra_mps, ket_mps)) + + def evolve_single_step(self, evolve_dt): + bra_mps, ket_mps = self.latest_mps + ket_mps = ket_mps.evolve(self.ket_mpo, evolve_dt) + bra_mps = bra_mps.evolve(self.bra_mpo, -evolve_dt) + return BraKetPair(bra_mps, ket_mps) + + +class ZTabs(ZTAACorrFuncBase): + + def init_op_b(self): + return abs_dipole_op(self.model) + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + +class ZTabs_TFD(ZTAACorrFuncBase): + + def init_op_b(self): + return abs_dipole_op(self.model) + + def init_imps(self): + assert self.imps_qntot == 0 + e_dofs = self.model.e_dofs + for dof in e_dofs: + siteidx = self.model.dof_to_siteidx[dof] + ba = self.model.basis[siteidx] + idx = ba.dof_name_map[dof] + if ba.sigmaqn[idx] == 0: + init_condition = {dof:idx} + break + + imps = Mps.hartree_product_state(self.model, condition=init_condition) + return imps + + +class ZTemi(ZTAACorrFuncBase): + + def init_op_b(self): + return emi_dipole_op(self.model) + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + + +class FTabs(FTCorrFuncBase): + + def init_op_b(self): + return abs_dipole_op(self.model) + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + + +class FTemi(FTCorrFuncBase): + + def init_op_b(self): + return emi_dipole_op(self.model) + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + +class ZTnr(ZTAACorrFuncBase): + + def init_op_b(self): + return nac_op(self.model) + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + +class ZTnr_state_to_state(ZTAACorrFuncBase_state_to_state): + + def init_op_b(self): + return nac_op(self.model) + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + +class FTnr(FTCorrFuncBase): + + def init_op_b(self): + return nac_op(self.model) + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + +class FTnr_state_to_state(FTCorrFuncBase_state_to_state): + + def init_op_b(self): + return nac_op(self.model) + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + +class FTcurrent_current_Holstein(FTCorrFuncBase): + def init_op_b(self): + j_op, _ = current_op(self.model, None) + return j_op + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + + +def nac_op_agg(model): + nac_terms = [] + for key, value in model.para["nac"].items(): + nac_terms.append(Op("a^\dagger a partialx", list(key), + factor=-value, qn=[0,-1,0])) + return Mpo(model, terms=nac_terms) + +class ZTnr_agg(ZTAACorrFuncBase): + + def init_op_b(self): + return nac_op_agg(self.model) + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + +# isc rate +#################################################################3 +def isc_op(model): + isc_terms = [] + for key, value in model.para["isc"].items(): + # the singlet state has qn=1, triplet state has qn=0 + op = Op("a^\dagger a", list(key), factor=value, qn=[0,-1]) + isc_terms.append(op) + return Mpo(model, terms=isc_terms) + +class ZTisc(ZTAACorrFuncBase): + + def init_op_b(self): + return isc_op(self.model) + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() + +class FTisc(FTCorrFuncBase): + + def init_op_b(self): + return isc_op(self.model) + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 1 + return super().init_imps() +################################################################ + +def risc_op(model): + raise NotImplementedError + +class ZT_sbm_spin_spin(ZTAACorrFuncBase): + + def init_op_b(self): + return spin_op(self.model) + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + +class FT_sbm_spin_spin(FTCorrFuncBase): + + def init_op_b(self): + return spin_op(self.model) + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + +class sbm_spin(ZTABCorrFuncBase): + + def init_op_a(self): + return spin_op(self.model) + + def init_op_b(self): + return Mpo.identity(self.model) + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + +################################################################ +class ZTIRAA(ZTAACorrFuncBase): + + def init_op_b(self): + return IR_op(self.model,"B") + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + +class FTIRAA(FTCorrFuncBase): + + def init_op_b(self): + return IR_op(self.model,"B") + + def init_op_a(self): + return self.init_op_b() + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + +class ZTIRAB(ZTABCorrFuncBase): + + def init_op_b(self): + return IR_op(self.model, "B") + + def init_op_a(self): + return IR_op(self.model, "A") + + def init_imps(self): + assert self.imps_qntot == 0 + return super().init_imps() + +def IR_op(model, op_symbol): + return model.mpos["IR"+op_symbol] + +#class PhotoPhysics(CorFuncTdMpsJobBase): +# r""" Photophysics properties of molecular aggregates with or without DRE +# """ +# def __init__( +# self, +# model, +# job_type, +# temperature, +# optimize_config=None, +# evolve_config=None, +# dump_dir=None, +# job_name=None, +# ): +# +# if optimize_config is None: