diff --git a/momentGW/bse.py b/momentGW/bse.py index a0ef8a6f..697d1849 100644 --- a/momentGW/bse.py +++ b/momentGW/bse.py @@ -91,11 +91,21 @@ def __init__(self, gw, **kwargs): # Attributes self.gf = None + @property + def polarizability_name(self): + """Get the polarizability name.""" + return { + "drpa": "dRPA", + "drpa-exact": "dRPA", + "dtda": "dTDA", + "thc-dtda": "THC-dTDA", + "tdax": "TDAx", + }[self.polarizability.lower()] + @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-BSE" + return f"{self.polarizability_name}-BSE" @logging.with_timer("Integral construction") @logging.with_status("Constructing integrals") @@ -526,8 +536,7 @@ def __init__(self, gw, **kwargs): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-cpBSE" + return f"{self.polarizability_name}-cpBSE" @logging.with_timer("Dynamic polarizability moments") @logging.with_status("Constructing dynamic polarizability moments") diff --git a/momentGW/evgw.py b/momentGW/evgw.py index 6cda563f..ae2efd65 100644 --- a/momentGW/evgw.py +++ b/momentGW/evgw.py @@ -206,8 +206,7 @@ class evGW(GW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-evG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" + return f"{self.polarizability_name}-evG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" def check_convergence(self, mo_energy, mo_energy_prev, th, th_prev, tp, tp_prev): """Check for convergence, and print a summary of changes. diff --git a/momentGW/fsgw.py b/momentGW/fsgw.py index fc6001dc..924568ef 100644 --- a/momentGW/fsgw.py +++ b/momentGW/fsgw.py @@ -219,5 +219,4 @@ class fsGW(GW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-fsGW" + return f"{self.polarizability_name}-fsGW" diff --git a/momentGW/gw.py b/momentGW/gw.py index e1d715fc..564f1373 100644 --- a/momentGW/gw.py +++ b/momentGW/gw.py @@ -11,7 +11,7 @@ from momentGW.fock import FockLoop, search_chempot from momentGW.ints import Integrals from momentGW.rpa import dRPA -from momentGW.tda import dTDA +from momentGW.tda import TDAx, dTDA def kernel( @@ -138,11 +138,21 @@ class GW(BaseGW): _kernel = kernel + @property + def polarizability_name(self): + """Get the polarizability name.""" + return { + "drpa": "dRPA", + "drpa-exact": "dRPA", + "dtda": "dTDA", + "thc-dtda": "THC-dTDA", + "tdax": "TDAx", + }[self.polarizability.lower()] + @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-G0W0" + return f"{self.polarizability_name}-G0W0" @logging.with_timer("Static self-energy") @logging.with_status("Building static self-energy") @@ -235,6 +245,10 @@ def build_se_moments(self, nmom_max, integrals, **kwargs): tda = thc.dTDA(self, nmom_max, integrals, **kwargs) return tda.kernel() + elif self.polarizability.lower() == "tdax": + tda = TDAx(self, nmom_max, integrals, **kwargs) + return tda.kernel() + else: raise NotImplementedError diff --git a/momentGW/pbc/evgw.py b/momentGW/pbc/evgw.py index 8b4f44ca..b2bae302 100644 --- a/momentGW/pbc/evgw.py +++ b/momentGW/pbc/evgw.py @@ -88,8 +88,7 @@ class evKGW(KGW, evGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-evKG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" + return f"{self.polarizability_name}-evKG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" def check_convergence(self, mo_energy, mo_energy_prev, th, th_prev, tp, tp_prev): """Check for convergence, and print a summary of changes. diff --git a/momentGW/pbc/fsgw.py b/momentGW/pbc/fsgw.py index 4908ce95..46471f2b 100644 --- a/momentGW/pbc/fsgw.py +++ b/momentGW/pbc/fsgw.py @@ -91,5 +91,4 @@ class fsKGW(KGW, fsGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-fsKGW" + return f"{self.polarizability_name}-fsKGW" diff --git a/momentGW/pbc/gw.py b/momentGW/pbc/gw.py index f0991ea6..4820d595 100644 --- a/momentGW/pbc/gw.py +++ b/momentGW/pbc/gw.py @@ -13,7 +13,7 @@ from momentGW.pbc.fock import FockLoop, search_chempot_unconstrained from momentGW.pbc.ints import KIntegrals from momentGW.pbc.rpa import dRPA -from momentGW.pbc.tda import dTDA +from momentGW.pbc.tda import TDAx, dTDA class KGW(BaseKGW, GW): @@ -66,8 +66,7 @@ class KGW(BaseKGW, GW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-KG0W0" + return f"{self.polarizability_name}-KG0W0" @logging.with_timer("Static self-energy") @logging.with_status("Building static self-energy") @@ -120,6 +119,9 @@ def build_se_moments(self, nmom_max, integrals, **kwargs): elif self.polarizability.lower() == "thc-dtda": tda = thc.dTDA(self, nmom_max, integrals, **kwargs) return tda.kernel() + elif self.polarizability.lower() == "tdax": + tda = TDAx(self, nmom_max, integrals, **kwargs) + return tda.kernel() else: raise NotImplementedError diff --git a/momentGW/pbc/qsgw.py b/momentGW/pbc/qsgw.py index 39acf88d..9877834d 100644 --- a/momentGW/pbc/qsgw.py +++ b/momentGW/pbc/qsgw.py @@ -106,8 +106,7 @@ class qsKGW(KGW, qsGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-qsKGW" + return f"{self.polarizability_name}-qsKGW" @staticmethod def project_basis(matrix, ovlp, mo1, mo2): diff --git a/momentGW/pbc/scgw.py b/momentGW/pbc/scgw.py index 862f3f26..fb1f282b 100644 --- a/momentGW/pbc/scgw.py +++ b/momentGW/pbc/scgw.py @@ -80,5 +80,4 @@ class scKGW(KGW, scGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-KG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" + return f"{self.polarizability_name}-KG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" diff --git a/momentGW/pbc/tda.py b/momentGW/pbc/tda.py index 0ade92d6..65a16eb4 100644 --- a/momentGW/pbc/tda.py +++ b/momentGW/pbc/tda.py @@ -265,3 +265,108 @@ def kpts(self): def nkpts(self): """Get the number of k-points.""" return self.gw.nkpts + + +class TDAx(dTDA): + """ + Compute the self-energy moments using TDA (with exchange) with + periodic boundary conditions. + + Parameters + ---------- + gw : BaseKGW + GW object. + nmom_max : int + Maximum moment number to calculate. + integrals : KIntegrals + Density-fitted integrals at each k-point. + mo_energy : dict, optional + Molecular orbital energies at each k-point. Keys are "g" and + "w" for the Green's function and screened Coulomb interaction, + respectively. If `None`, use `gw.mo_energy` for both. Default + value is `None`. + mo_occ : dict, optional + Molecular orbital occupancies at each k-point. Keys are "g" + and "w" for the Green's function and screened Coulomb + interaction, respectively. If `None`, use `gw.mo_occ` for both. + Default value is `None`. + """ + + @logging.with_timer("Self-energy moments") + @logging.with_status("Constructing self-energy moments") + def build_se_moments(self, moments_dd): + """Build the moments of the self-energy via convolution. + + Parameters + ---------- + moments_dd : numpy.ndarray + Moments of the density-density response at each k-point. + + Returns + ------- + moments_occ : numpy.ndarray + Moments of the occupied self-energy at each k-point. + moments_vir : numpy.ndarray + Moments of the virtual self-energy at each k-point. + """ + + # Get the sizes + nocc = self.integrals.nocc + nvir = self.integrals.nvir + kpts = self.kpts + + # Setup dependent on diagonal SE + if self.gw.diagonal_se: + pqchar = pchar = qchar = "p" + eta_shape = lambda k: (self.mo_energy_g[k].size, self.nmom_max + 1, self.nmo) + else: + pqchar, pchar, qchar = "pq", "p", "q" + eta_shape = lambda k: (self.mo_energy_g[k].size, self.nmom_max + 1, self.nmo, self.nmo) + eta = np.zeros((self.nkpts, self.nkpts), dtype=object) + + # Get the moments + for n in range(self.nmom_max + 1): + for kp, kx, ki in kpts.loop(3): + ka = kpts.conserve(kp, kx, ki) + q = kpts.member(kpts.wrap_around(kpts[ki] - kpts[ka])) + + if not isinstance(eta[kp, q], np.ndarray): + eta[kp, q] = np.zeros(eta_shape(kx), dtype=complex) + + Lia = self.integrals.Lia[ki, ka] + Lia = Lia.reshape(Lia.shape[0], nocc[ki], nvir[ka]) + + Lxa = self.integrals.Lia[kx, ka] + Lxa = Lxa.reshape(Lxa.shape[0], nocc[kx], nvir[ka]) + + Lix = self.integrals.Lia[ki, kx] + Lix = Lix.reshape(Lix.shape[0], nocc[ki], nvir[kx]) + + Lpi = self.integrals.Lpx[kp, ki][:, :, : nocc[ki]] + Lpa = self.integrals.Lpx[kp, ka][:, :, nocc[ka] :] + + moment = moments_dd[q, ka, n] + moment = moment.reshape(Lia.shape) + + for x in range(self.mo_energy_g[kx].size): + Lp = self.integrals.Lpx[kp, kx][:, :, x] + + v = util.einsum("Pia,Pq->iaq", Lia, Lp) * 2.0 + if self.mo_occ_g[kx][x] > 0: + La = Lxa[:, x] + v -= util.einsum("Pa,Pqi->iaq", La, Lpi) + else: + Li = Lix[:, :, x - nocc[kx]] + v -= util.einsum("Pi,Pqa->iaq", Li, Lpa) + + eta[kp, q][x, n] += util.einsum( + f"P{pchar},Pia,ia{qchar}->{pqchar}", + Lp, + moment, + v.conj(), + ) + + # Construct the self-energy moments + moments_occ, moments_vir = self.convolve(eta) + + return moments_occ, moments_vir diff --git a/momentGW/pbc/uhf/evgw.py b/momentGW/pbc/uhf/evgw.py index 52db4a32..e64adaf4 100644 --- a/momentGW/pbc/uhf/evgw.py +++ b/momentGW/pbc/uhf/evgw.py @@ -89,8 +89,7 @@ class evKUGW(KUGW, evKGW, evUGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-evKUGW" + return f"{self.polarizability_name}-evKUG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" def check_convergence(self, mo_energy, mo_energy_prev, th, th_prev, tp, tp_prev): """Check for convergence, and print a summary of changes. diff --git a/momentGW/pbc/uhf/fsgw.py b/momentGW/pbc/uhf/fsgw.py index f01c4501..6a5f31c0 100644 --- a/momentGW/pbc/uhf/fsgw.py +++ b/momentGW/pbc/uhf/fsgw.py @@ -92,5 +92,4 @@ class fsKUGW(KUGW, fsKGW, fsUGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-fsKUGW" + return f"{self.polarizability_name}-fsKUGW" diff --git a/momentGW/pbc/uhf/gw.py b/momentGW/pbc/uhf/gw.py index 646afde9..9084181b 100644 --- a/momentGW/pbc/uhf/gw.py +++ b/momentGW/pbc/uhf/gw.py @@ -66,8 +66,7 @@ class KUGW(BaseKUGW, KGW, UGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-KUG0W0" + return f"{self.polarizability_name}-KUG0W0" @logging.with_timer("Static self-energy") @logging.with_status("Building static self-energy") diff --git a/momentGW/pbc/uhf/qsgw.py b/momentGW/pbc/uhf/qsgw.py index edcb4e03..3baa65fa 100644 --- a/momentGW/pbc/uhf/qsgw.py +++ b/momentGW/pbc/uhf/qsgw.py @@ -108,8 +108,7 @@ class qsKUGW(KUGW, qsKGW, qsUGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-qsKUGW" + return f"{self.polarizability_name}-qsKUGW" @staticmethod def project_basis(matrix, ovlp, mo1, mo2): diff --git a/momentGW/pbc/uhf/scgw.py b/momentGW/pbc/uhf/scgw.py index 9703e577..e35ea768 100644 --- a/momentGW/pbc/uhf/scgw.py +++ b/momentGW/pbc/uhf/scgw.py @@ -81,5 +81,4 @@ class scKUGW(KUGW, scKGW, scUGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-scKUG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" + return f"{self.polarizability_name}-scKUG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" diff --git a/momentGW/pbc/uhf/tda.py b/momentGW/pbc/uhf/tda.py index 8899276f..f847cca1 100644 --- a/momentGW/pbc/uhf/tda.py +++ b/momentGW/pbc/uhf/tda.py @@ -12,9 +12,8 @@ class dTDA(KdTDA, MolUdTDA): """ - Compute the self-energy moments using dTDA and numerical - integration with periodic boundary conditions and unrestricted - references. + Compute the self-energy moments using dTDA with periodic boundary + conditions and unrestricted references. Parameters ---------- @@ -269,3 +268,110 @@ def nov(self): [np.sum(occ == 0) for occ in self.mo_occ_w[1]], ), ) + + +class TDAx(dTDA): + """ + Compute the self-energy moments using TDA (with exchange) with + periodic boundary conditions and unrestricted references. + + Parameters + ---------- + gw : BaseKUGW + GW object. + nmom_max : int + Maximum moment number to calculate. + integrals : KUIntegrals + Integrals object. + mo_energy : dict, optional + Molecular orbital energies at each k-point for each spin channel. + Keys are "g" and "w" for the Green's function and screened + Coulomb interaction, respectively. If `None`, use `gw.mo_energy` + for both. Default value is `None`. + mo_occ : dict, optional + Molecular orbital occupancies at each k-point for each spin + channel. Keys are "g" and "w" for the Green's function and + screened Coulomb interaction, respectively. If `None`, use + `gw.mo_occ` for both. Default value is `None`. + """ + + @logging.with_timer("Self-energy moments") + @logging.with_status("Constructing self-energy moments") + def build_se_moments(self, moments_dd): + """Build the moments of the self-energy via convolution. + + Parameters + ---------- + moments_dd : numpy.ndarray + Moments of the density-density response at each k-point. + + Returns + ------- + moments_occ : numpy.ndarray + Moments of the occupied self-energy at each k-point for each + spin channel. + moments_vir : numpy.ndarray + Moments of the virtual self-energy at each k-point for each + spin channel. + """ + + # Get the sizes + nocc = self.integrals.nocc # noqa # FIXME + nvir = self.integrals.nvir # noqa # FIXME + kpts = self.kpts + + # Setup dependent on diagonal SE + if self.gw.diagonal_se: + pqchar = pchar = qchar = "p" + eta_shape = lambda s, k: (self.mo_energy_g[s][k].size, self.nmom_max + 1, self.nmo[s]) + else: + pqchar, pchar, qchar = "pq", "p", "q" + eta_shape = lambda s, k: ( + self.mo_energy_g[s][k].size, + self.nmom_max + 1, + self.nmo[s], + self.nmo[s], + ) + eta = np.zeros((2, self.nkpts, self.nkpts), dtype=object) + + # Get the moments in (aux|aux) and rotate to (mo|mo) + for n in range(self.nmom_max + 1): + for q in kpts.loop(1): + eta_aux = 0 + for kj in kpts.loop(1, mpi=True): + kb = kpts.member(kpts.wrap_around(kpts[q] + kpts[kj])) + Lia = np.concatenate( + [ + self.integrals[0].Lia[kj, kb], + self.integrals[1].Lia[kj, kb], + ], + axis=1, + ) + eta_aux += np.dot(moments_dd[q, kb, n], Lia.T.conj()) + + eta_aux = mpi_helper.allreduce(eta_aux) + eta_aux /= self.nkpts + + for kp in kpts.loop(1, mpi=True): + kx = kpts.member(kpts.wrap_around(kpts[kp] - kpts[q])) + + for s in range(2): + if not isinstance(eta[s, kp, q], np.ndarray): + eta[s, kp, q] = np.zeros(eta_shape(s, kx), dtype=eta_aux.dtype) + + for x in range(self.mo_energy_g[s][kx].size): + Lp = self.integrals[s].Lpx[kp, kx][:, :, x] + subscript = f"P{pchar},Q{qchar},PQ->{pqchar}" + eta[s, kp, q][x, n] += util.einsum(subscript, Lp, Lp.conj(), eta_aux) + + # Construct the self-energy moments + moments_occ = [None, None] + moments_vir = [None, None] + moments_occ[0], moments_vir[0] = self.convolve( + eta[0], mo_energy_g=self.mo_energy_g[0], mo_occ_g=self.mo_occ_g[0] + ) + moments_occ[1], moments_vir[1] = self.convolve( + eta[1], mo_energy_g=self.mo_energy_g[1], mo_occ_g=self.mo_occ_g[1] + ) + + return tuple(moments_occ), tuple(moments_vir) diff --git a/momentGW/qsgw.py b/momentGW/qsgw.py index 16cb16b6..97b45340 100644 --- a/momentGW/qsgw.py +++ b/momentGW/qsgw.py @@ -280,8 +280,7 @@ class qsGW(GW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-qsGW" + return f"{self.polarizability_name}-qsGW" @staticmethod def project_basis(matrix, ovlp, mo1, mo2): diff --git a/momentGW/scgw.py b/momentGW/scgw.py index a03dbf24..20dfaeed 100644 --- a/momentGW/scgw.py +++ b/momentGW/scgw.py @@ -192,5 +192,4 @@ class scGW(evGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-G{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" + return f"{self.polarizability_name}-G{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" diff --git a/momentGW/tda.py b/momentGW/tda.py index 9e48e24f..2cd384fb 100644 --- a/momentGW/tda.py +++ b/momentGW/tda.py @@ -343,3 +343,121 @@ def build_dd_moment_inv(self): moment = np.dot(u, Liadinv) return moment + + +class TDAx(dTDA): + """ + Compute the self-energy moments using TDA (with exchange). + + Parameters + ---------- + gw : BaseGW + GW object. + nmom_max : int + Maximum moment number to calculate. + integrals : Integrals + Integrals object. + mo_energy : dict, optional + Molecular orbital energies. Keys are "g" and "w" for the Green's + function and screened Coulomb interaction, respectively. + If `None`, use `gw.mo_energy` for both. Default value is `None`. + mo_occ : dict, optional + Molecular orbital occupancies. Keys are "g" and "w" for the + Green's function and screened Coulomb interaction, respectively. + If `None`, use `gw.mo_occ` for both. Default value is `None`. + """ + + @logging.with_timer("Self-energy moments") + @logging.with_status("Constructing self-energy moments") + def build_se_moments(self, moments_dd): + """Build the moments of the self-energy via convolution. + + Parameters + ---------- + moments_dd : numpy.ndarray + Moments of the density-density response. + + Returns + ------- + moments_occ : numpy.ndarray + Moments of the occupied self-energy. + moments_vir : numpy.ndarray + Moments of the virtual self-energy. + """ + + # Setup dependent on diagonal SE + q0, q1 = self.mpi_slice(self.mo_energy_g.size) + if self.gw.diagonal_se: + eta = np.zeros((q1 - q0, self.nmo)) + pq = p = q = "p" + else: + eta = np.zeros((q1 - q0, self.nmo, self.nmo)) + pq, p, q = "pq", "p", "q" + + # Initialise output moments + moments_occ = np.zeros((self.nmom_max + 1, self.nmo, self.nmo)) + moments_vir = np.zeros((self.nmom_max + 1, self.nmo, self.nmo)) + + # Unpack the integrals + Lia = self.integrals.Lia.reshape( + self.integrals.naux, + self.integrals.nocc_w, + self.integrals.nvir_w, + ) + + # Get the moments + for n in range(self.nmom_max + 1): + moment = moments_dd[n].reshape(Lia.shape) + for x in range(q1 - q0): + Lp = self.integrals.Lpx[:, :, x] + v = util.einsum("Pia,Pq->iaq", Lia, Lp) * 2.0 + if self.mo_occ_g[x + q0] > 0: + v -= util.einsum( + "Pa,Pqi->iaq", + Lia[:, x + q0], + self.integrals.Lpx[:, :, self.mo_occ_g > 0], + ) + else: + v -= util.einsum( + "Pi,Pqa->iaq", + Lia[:, :, x + q0 - Lia.shape[1]], + self.integrals.Lpx[:, :, self.mo_occ_g == 0], + ) + eta[x] = util.einsum(f"P{p},Pia,ia{q}->{pq}", Lp, moment, v) + + # Construct the self-energy moments for this order only to + # save memory + moments_occ_n, moments_vir_n = self.convolve(eta[:, None], eta_orders=[n]) + moments_occ += moments_occ_n + moments_vir += moments_vir_n + + return moments_occ, moments_vir + + @logging.with_timer("Dynamic polarizability moments") + @logging.with_status("Constructing dynamic polarizability moments") + def build_dp_moments(self): + """ + Build the moments of the dynamic polarizability for optical + spectra calculations. + + Returns + ------- + moments : numpy.ndarray + Moments of the dynamic polarizability. + """ + raise NotImplementedError + + @logging.with_timer("Inverse density-density moment") + @logging.with_status("Constructing inverse density-density moment") + def build_dd_moment_inv(self): + r""" + Build the first inverse (`n=-1`) moment of the density-density + response. + + Returns + ------- + moment : numpy.ndarray + First inverse (`n=-1`) moment of the density-density + response. + """ + raise NotImplementedError diff --git a/momentGW/uhf/evgw.py b/momentGW/uhf/evgw.py index 0c41c4b7..6b76f96a 100644 --- a/momentGW/uhf/evgw.py +++ b/momentGW/uhf/evgw.py @@ -85,8 +85,7 @@ class evUGW(UGW, evGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-evUG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" + return f"{self.polarizability_name}-evUG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" def check_convergence(self, mo_energy, mo_energy_prev, th, th_prev, tp, tp_prev): """Check for convergence, and print a summary of changes. diff --git a/momentGW/uhf/fsgw.py b/momentGW/uhf/fsgw.py index a9cc5a33..23b371dc 100644 --- a/momentGW/uhf/fsgw.py +++ b/momentGW/uhf/fsgw.py @@ -88,5 +88,4 @@ class fsUGW(UGW, fsGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-fsUGW" + return f"{self.polarizability_name}-fsUGW" diff --git a/momentGW/uhf/gw.py b/momentGW/uhf/gw.py index 48d1d979..9dbe57a0 100644 --- a/momentGW/uhf/gw.py +++ b/momentGW/uhf/gw.py @@ -13,7 +13,7 @@ from momentGW.uhf.fock import FockLoop from momentGW.uhf.ints import UIntegrals from momentGW.uhf.rpa import dRPA -from momentGW.uhf.tda import dTDA +from momentGW.uhf.tda import TDAx, dTDA class UGW(BaseUGW, GW): @@ -70,8 +70,7 @@ class UGW(BaseUGW, GW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-UG0W0" + return f"{self.polarizability_name}-UG0W0" @logging.with_timer("Static self-energy") @logging.with_status("Building static self-energy") @@ -129,6 +128,10 @@ def build_se_moments(self, nmom_max, integrals, **kwargs): rpa = dRPA(self, nmom_max, integrals, **kwargs) return rpa.kernel() + elif self.polarizability.lower() == "tdax": + tda = TDAx(self, nmom_max, integrals, **kwargs) + return tda.kernel() + else: raise NotImplementedError diff --git a/momentGW/uhf/qsgw.py b/momentGW/uhf/qsgw.py index 922c0e97..325c46b0 100644 --- a/momentGW/uhf/qsgw.py +++ b/momentGW/uhf/qsgw.py @@ -103,8 +103,7 @@ class qsUGW(UGW, qsGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-qsUGW" + return f"{self.polarizability_name}-qsUGW" @staticmethod def project_basis(matrix, ovlp, mo1, mo2): diff --git a/momentGW/uhf/scgw.py b/momentGW/uhf/scgw.py index 11f7691b..d6e362d3 100644 --- a/momentGW/uhf/scgw.py +++ b/momentGW/uhf/scgw.py @@ -77,5 +77,4 @@ class scUGW(UGW, scGW): @property def name(self): """Get the method name.""" - polarizability = self.polarizability.upper().replace("DTDA", "dTDA").replace("DRPA", "dRPA") - return f"{polarizability}-UG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" + return f"{self.polarizability_name}-UG{'0' if self.g0 else ''}W{'0' if self.w0 else ''}" diff --git a/momentGW/uhf/tda.py b/momentGW/uhf/tda.py index eda8afd2..36966a8a 100644 --- a/momentGW/uhf/tda.py +++ b/momentGW/uhf/tda.py @@ -244,3 +244,135 @@ def nov(self): np.sum(self.mo_occ_w[0] > 0) * np.sum(self.mo_occ_w[0] == 0), np.sum(self.mo_occ_w[1] > 0) * np.sum(self.mo_occ_w[1] == 0), ) + + +class TDAx(dTDA): + """ + Compute the self-energy moments using TDA (with exchange). + + Parameters + ---------- + gw : BaseUGW + GW object. + nmom_max : int + Maximum moment number to calculate. + integrals : UIntegrals + Integrals object. + mo_energy : dict, optional + Molecular orbital energies for each spin. Keys are "g" and "w" + for the Green's function and screened Coulomb interaction, + respectively. If `None`, use `gw.mo_energy` for both. Default + value is `None`. + mo_occ : dict, optional + Molecular orbital occupancies for each spin. Keys are "g" and + "w" for the Green's function and screened Coulomb interaction, + respectively. If `None`, use `gw.mo_occ` for both. Default + value is `None`. + """ + + @logging.with_timer("Self-energy moments") + @logging.with_status("Constructing self-energy moments") + def build_se_moments(self, moments_dd): + """Build the moments of the self-energy via convolution. + + Parameters + ---------- + moments_dd : numpy.ndarray + Moments of the density-density response for each spin + channel. + + Returns + ------- + moments_occ : numpy.ndarray + Moments of the occupied self-energy for each spin channel. + moments_vir : numpy.ndarray + Moments of the virtual self-energy for each spin channel. + """ + + # Setup dependent on diagonal SE + a0, a1 = self.mpi_slice(self.mo_energy_g[0].size) + b0, b1 = self.mpi_slice(self.mo_energy_g[1].size) + if self.gw.diagonal_se: + eta = [ + np.zeros((a1 - a0, self.nmo[0])), + np.zeros((b1 - b0, self.nmo[1])), + ] + pq = p = q = "p" + else: + eta = [ + np.zeros((a1 - a0, self.nmo[0], self.nmo[0])), + np.zeros((b1 - b0, self.nmo[1], self.nmo[1])), + ] + pq, p, q = "pq", "p", "q" + + # Initialise output moments + moments_occ = [ + np.zeros((self.nmom_max + 1, self.nmo[0], self.nmo[0])), + np.zeros((self.nmom_max + 1, self.nmo[1], self.nmo[1])), + ] + moments_vir = [ + np.zeros((self.nmom_max + 1, self.nmo[0], self.nmo[0])), + np.zeros((self.nmom_max + 1, self.nmo[1], self.nmo[1])), + ] + + # Unpack the integrals + Lia = ( + self.integrals[0].Lia.reshape( + self.integrals[0].naux, + self.integrals[0].nocc_w, + self.integrals[0].nvir_w, + ), + self.integrals[1].Lia.reshape( + self.integrals[1].naux, + self.integrals[1].nocc_w, + self.integrals[1].nvir_w, + ), + ) + + # Get the moments + for n in range(self.nmom_max + 1): + moment = [ + moments_dd[n].ravel()[: Lia[0].size].reshape(Lia[0].shape), + moments_dd[n].ravel()[Lia[0].size :].reshape(Lia[1].shape), + ] + for ss, (q0, q1) in enumerate([(a0, a1), (b0, b1)]): + os = (ss + 1) % 2 + for x in range(q1 - q0): + Lp = self.integrals[ss].Lpx[:, :, x] + v_ss = util.einsum("Pia,Pq->iaq", Lia[ss], Lp) + v_os = util.einsum("Pia,Pq->iaq", Lia[os], Lp) + if self.mo_occ_g[ss][x + q0] > 0: + v_ss -= util.einsum( + "Pa,Pqi->iaq", + Lia[ss][:, x + q0], + self.integrals[ss].Lpx[:, :, self.mo_occ_g[ss] > 0], + ) + else: + v_ss -= util.einsum( + "Pi,Pqa->iaq", + Lia[ss][:, :, x + q0 - Lia[ss].shape[1]], + self.integrals[ss].Lpx[:, :, self.mo_occ_g[ss] == 0], + ) + eta[ss][x] += util.einsum(f"P{p},Pia,ia{q}->{pq}", Lp, moment[ss], v_ss) + eta[ss][x] += util.einsum(f"P{p},Pia,ia{q}->{pq}", Lp, moment[os], v_os) + + # Construct the self-energy moments for this order only + # to save memory + moments_occ_n, moments_vir_n = self.convolve( + eta[0][:, None], + eta_orders=[n], + mo_energy_g=self.mo_energy_g[0], + mo_occ_g=self.mo_occ_g[0], + ) + moments_occ[0] += moments_occ_n + moments_vir[0] += moments_vir_n + moments_occ_n, moments_vir_n = self.convolve( + eta[1][:, None], + eta_orders=[n], + mo_energy_g=self.mo_energy_g[1], + mo_occ_g=self.mo_occ_g[1], + ) + moments_occ[1] += moments_occ_n + moments_vir[1] += moments_vir_n + + return tuple(moments_occ), tuple(moments_vir)