Head and wings corrections for k-dTDA

momentGW/pbc/gw.py

@@ -3,8 +3,11 @@
 periodic systems.
 """
 
+from functools import reduce
+
 import numpy as np
 from dyson import MBLSE, Lehmann, MixedMBLSE
+from pyscf.pbc import tools
 
 from momentGW import energy, logging, util
 from momentGW.gw import GW
@@ -78,16 +81,57 @@ def build_se_static(self, integrals):
 
         Parameters
         ----------
-        integrals : KIntegrals
+        integrals : Integrals
             Integrals object.
 
         Returns
         -------
         se_static : numpy.ndarray
-            Static part of the self-energy at each k-point. If
-            `self.diagonal_se`, non-diagonal elements are set to zero.
+            Static part of the self-energy. If `self.diagonal_se`,
+            non-diagonal elements are set to zero.
         """
-        return super().build_se_static(integrals)
+
+        if getattr(self._scf, "xc", "hf") == "hf":
+            se_static = np.zeros_like(self._scf.make_rdm1(mo_coeff=self.mo_coeff))
+            if self.fc:
+                with util.SilentSCF(self._scf):
+                    vmf = self._scf.get_j() - self._scf.get_veff()
+                    dm = self._scf.make_rdm1(mo_coeff=self.mo_coeff)
+                    vk = integrals.get_k(dm, basis="ao")
+
+                s = self.cell.pbc_intor("int1e_ovlp", hermi=1, kpts=self.kpts)
+                madelung = tools.pbc.madelung(self.cell, self.kpts)
+                for k in range(len(self.kpts)):
+                    vk[k] += madelung * reduce(np.dot, (s[k], dm[k], s[k]))
+
+                se_static = vmf - vk * 0.5
+                se_static = util.einsum(
+                    "...pq,...pi,...qj->...ij", se_static, np.conj(self.mo_coeff), self.mo_coeff
+                )
+
+        else:
+            with util.SilentSCF(self._scf):
+                vmf = self._scf.get_j() - self._scf.get_veff()
+                dm = self._scf.make_rdm1(mo_coeff=self.mo_coeff)
+                vk = integrals.get_k(dm, basis="ao")
+
+            if self.fc:
+                s = self.cell.pbc_intor("int1e_ovlp", hermi=1, kpts=self.kpts)
+                madelung = tools.pbc.madelung(self.cell, self.kpts)
+                for k in range(len(self.kpts)):
+                    vk[k] += madelung * reduce(np.dot, (s[k], dm[k], s[k]))
+            se_static = vmf - vk * 0.5
+
+            se_static = util.einsum(
+                "...pq,...pi,...qj->...ij", se_static, np.conj(self.mo_coeff), self.mo_coeff
+            )
+
+        if self.diagonal_se:
+            se_static = util.einsum("...pq,pq->...pq", se_static, np.eye(se_static.shape[-1]))
+
+        se_static += util.einsum("...p,...pq->...pq", self.mo_energy, np.eye(se_static.shape[-1]))
+
+        return se_static
 
     def build_se_moments(self, nmom_max, integrals, **kwargs):
         """Build the moments of the self-energy.
@@ -110,9 +154,13 @@ def build_se_moments(self, nmom_max, integrals, **kwargs):
             Moments of the particle self-energy at each k-point. If
             `self.diagonal_se`, non-diagonal elements are set to zero.
         """
+        if self.fc:
+            fc = True
+        else:
+            fc = False
 
         if self.polarizability.lower() == "dtda":
-            tda = dTDA(self, nmom_max, integrals, **kwargs)
+            tda = dTDA(self, nmom_max, integrals, fc, **kwargs)
             return tda.kernel()
         if self.polarizability.lower() == "drpa":
             rpa = dRPA(self, nmom_max, integrals, **kwargs)
@@ -334,7 +382,7 @@ def make_rdm1(self, gf=None):
 
     @logging.with_timer("Energy")
     @logging.with_status("Calculating energy")
-    def energy_hf(self, gf=None, integrals=None):
+    def energy_hf(self, gf=None, integrals=None, **kwargs):
         """Calculate the one-body (Hartree--Fock) energy.
 
         Parameters
@@ -367,7 +415,7 @@ def energy_hf(self, gf=None, integrals=None):
                 "kpq,kpi,kqj->kij", self._scf.get_hcore(), self.mo_coeff.conj(), self.mo_coeff
             )
         rdm1 = self.make_rdm1()
-        fock = integrals.get_fock(rdm1, h1e)
+        fock = integrals.get_fock(rdm1, h1e, **kwargs)
 
         # Calculate the Hartree--Fock energy at each k-point
         e_1b = sum(energy.hartree_fock(rdm1[k], fock[k], h1e[k]) for k in self.kpts.loop(1))

momentGW/pbc/ints.py

@@ -398,7 +398,7 @@ def get_cderi_from_thc(self):
                 Lpx[ki, kj] = Lpx_k
                 Lia[ki, kj] = Lia_k
 
-                q = self.kpts.member(self.kpts.wrap_around(-self.kpts[q]))
+                invq = self.kpts.member(self.kpts.wrap_around(-self.kpts[q]))
 
                 block_switch = util.einsum("Pp,Pq,PQ->Qpq", coll_kj.conj(), coll_ki, cholesky_cou)
 
@@ -408,7 +408,7 @@ def get_cderi_from_thc(self):
                 )
                 tmp = util.einsum("Lpq,pi,qj->Lij", block_switch, coeffs[0].conj(), coeffs[1])
                 tmp = tmp.swapaxes(1, 2)
-                tmp = tmp.reshape(self.naux[q], -1)
+                tmp = tmp.reshape(self.naux[invq], -1)
                 Lai_k += tmp
 
                 Lai[ki, kj] = Lai_k
@@ -727,6 +727,12 @@ def get_fock(self, dm, h1e, **kwargs):
         """
         return super().get_fock(dm, h1e, **kwargs)
 
+    def reciprocal_lattice(self):
+        """
+        Return the reciprocal lattice vectors.
+        """
+        return self.with_df.cell.reciprocal_vectors()
+
     @property
     def madelung(self):
         """

momentGW/pbc/tda.py

@@ -4,6 +4,8 @@
 
 import numpy as np
 import scipy.special
+from pyscf import lib
+from pyscf.pbc import dft
 
 from momentGW import logging, mpi_helper, util
 from momentGW.tda import dTDA as MoldTDA
@@ -34,6 +36,22 @@ class dTDA(MoldTDA):
         Default value is `None`.
     """
 
+    def __init__(
+        self,
+        gw,
+        nmom_max,
+        integrals,
+        fc=False,
+        mo_energy=None,
+        mo_occ=None,
+    ):
+        super().__init__(gw, nmom_max, integrals, mo_energy=mo_energy, mo_occ=mo_occ)
+        self.fc = fc
+        if self.fc:
+            q = np.array([1e-3, 0, 0]).reshape(1, 3)
+            self.q_abs = self.kpts.cell.get_abs_kpts(q)
+            self.qij = self.build_pert_term(self.q_abs[0])
+
     @logging.with_timer("Density-density moments")
     @logging.with_status("Constructing density-density moments")
     def build_dd_moments(self):
@@ -49,12 +67,20 @@ def build_dd_moments(self):
         kpts = self.kpts
         moments = np.zeros((self.nkpts, self.nkpts, self.nmom_max + 1), dtype=object)
 
+        if self.fc:
+            head = np.zeros((self.nkpts, self.nmom_max + 1), dtype=object)
+
         # Get the zeroth order moment
         for q in kpts.loop(1):
             for kj in kpts.loop(1, mpi=True):
                 kb = kpts.member(kpts.wrap_around(kpts[q] + kpts[kj]))
                 moments[q, kb, 0] += self.integrals.Lia[kj, kb] / self.nkpts
 
+                if q == 0 and self.fc:
+                    head[kj, 0] += (
+                        np.sqrt(4.0 * np.pi) / np.linalg.norm(self.q_abs[0])
+                    ) * self.qij[kj].conj()
+
         # Get the higher order moments
         for i in range(1, self.nmom_max + 1):
             for q in kpts.loop(1):
@@ -66,23 +92,42 @@ def build_dd_moments(self):
                         (self.mo_occ_w[kj], self.mo_occ_w[kb]),
                     )
                     moments[q, kb, i] += moments[q, kb, i - 1] * d.ravel()[None]
+                    if self.fc and q == 0:
+                        head[kj, i] += head[kj, i - 1] * d.ravel()
 
                 tmp = np.zeros((self.naux[q], self.naux[q]), dtype=complex)
+                tmp_head = np.zeros((self.naux[q]), dtype=complex)
+
                 for ki in kpts.loop(1, mpi=True):
                     ka = kpts.member(kpts.wrap_around(kpts[q] + kpts[ki]))
 
                     tmp += np.dot(moments[q, ka, i - 1], self.integrals.Lia[ki, ka].T.conj())
 
+                    if q == 0 and self.fc:
+                        tmp_head += util.einsum(
+                            "a,aP->P", head[ki, i - 1], self.integrals.Lia[ki, ki].T.conj()
+                        )
+
                 tmp = mpi_helper.allreduce(tmp)
                 tmp *= 2.0
                 tmp /= self.nkpts
 
+                tmp_head = mpi_helper.allreduce(tmp_head)
+                tmp_head *= 2.0
+                tmp_head /= self.nkpts
+
                 for kj in kpts.loop(1, mpi=True):
                     kb = kpts.member(kpts.wrap_around(kpts[q] + kpts[kj]))
 
                     moments[q, kb, i] += np.dot(tmp, self.integrals.Lai[kj, kb].conj())
 
-        return moments
+                    if q == 0 and self.fc:
+                        head[kj, i] += util.einsum("P,Pa->a", tmp_head, self.integrals.Lia[kj, kj])
+
+        if self.fc:
+            return {"moments": moments, "head": head}
+        else:
+            return moments
 
     def kernel(self, exact=False):
         """
@@ -220,13 +265,33 @@ def build_se_moments(self, moments_dd):
             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)
 
+        if self.fc:
+            cell_vol = self.kpts.cell.vol
+            total_vol = cell_vol * self.nkpts
+            q0 = (6 * np.pi**2 / total_vol) ** (1 / 3)
+            norm_q_abs = np.linalg.norm(self.q_abs[0])
+            eta_head = np.zeros((self.nkpts, self.nmom_max + 1), dtype=object)
+            eta_wings = np.zeros((self.nkpts, self.nmom_max + 1), 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]))
-                    eta_aux += np.dot(moments_dd[q, kb, n], self.integrals.Lia[kj, kb].T.conj())
+                    if self.fc:
+                        eta_aux += np.dot(
+                            moments_dd["moments"][q, kb, n], self.integrals.Lia[kj, kb].T.conj()
+                        )
+                        if q == 0:
+                            eta_head[kb, n] += -(np.sqrt(4.0 * np.pi) / norm_q_abs) * np.sum(
+                                moments_dd["head"][kb, n] * self.qij[kb]
+                            )
+                            eta_wings[kb, n] += (np.sqrt(4.0 * np.pi) / norm_q_abs) * util.einsum(
+                                "Pa,a->P", moments_dd["moments"][q, kb, n], self.qij[kb]
+                            )
+                    else:
+                        eta_aux += np.dot(moments_dd[q, kb, n], self.integrals.Lia[kj, kb].T.conj())
 
                 eta_aux = mpi_helper.allreduce(eta_aux)
                 eta_aux *= 2.0
@@ -242,12 +307,61 @@ def build_se_moments(self, moments_dd):
                         Lp = self.integrals.Lpx[kp, kx][:, :, x]
                         subscript = f"P{pchar},Q{qchar},PQ->{pqchar}"
                         eta[kp, q][x, n] += util.einsum(subscript, Lp, Lp.conj(), eta_aux)
+                        if self.fc and q == 0:
+                            original = eta[kp, q][x, n]
+
+                            eta[kp, q][x, n] += (2 / np.pi) * (q0) * eta_head[kp, n] * original
+
+                            wing_tmp = util.einsum("Pp,P->p", Lp, eta_wings[kp, n])
+                            wing_tmp = wing_tmp.real * 2
+                            wing_tmp *= -(np.sqrt(cell_vol / (4 * (np.pi**3))) * q0**2)
+
+                            eta[kp, q][x, n] += util.einsum("p,pq->pq", wing_tmp, original)
 
         # Construct the self-energy moments
         moments_occ, moments_vir = self.convolve(eta)
 
         return moments_occ, moments_vir
 
+    def build_pert_term(self, qpt):
+        """
+        Compute 1/sqrt(Ω) * ⟨Ψ_{ik}|e^{iqr}|Ψ_{ak-q}⟩ at q-point index
+        q using perturbation theory.
+        """
+
+        coords, weights = dft.gen_grid.get_becke_grids(self.gw.cell, level=5)
+
+        qij = np.zeros((self.nkpts,), dtype=object)
+        for k in self.kpts.loop(1):
+            ao_p = dft.numint.eval_ao(self.gw.cell, coords, kpt=self.kpts[k], deriv=1)
+            ao, ao_grad = ao_p[0], ao_p[1:4]
+
+            ao_ao_grad = util.einsum("g,gm,xgn->xmn", weights, ao.conj(), ao_grad)
+            q_ao_ao_grad = util.einsum("x,xmn->mn", qpt, ao_ao_grad) * -1.0j
+            q_mo_mo_grad = util.einsum(
+                "mn,mi,na->ia",
+                q_ao_ao_grad,
+                self.integrals.mo_coeff_w[k][:, self.mo_occ_w[k] > 0].conj(),
+                self.integrals.mo_coeff_w[k][:, self.mo_occ_w[k] == 0],
+            )
+
+            d = lib.direct_sum(
+                "a-i->ia",
+                self.mo_energy_w[k][self.mo_occ_w[k] == 0],
+                self.mo_energy_w[k][self.mo_occ_w[k] > 0],
+            )
+            dens = q_mo_mo_grad / d
+            qij[k] = dens.flatten() / np.sqrt(self.gw.cell.vol)
+
+        return qij
+
+    build_dd_moments_exact = build_dd_moments
+
+    @property
+    def naux(self):
+        """Number of auxiliaries."""
+        return self.integrals.naux
+
     @property
     def nov(self):
         """Get the number of ov states in W."""

momentGW/util.py

@@ -473,7 +473,12 @@ def _fallback():
 
     # Reshape and transpose the output
     ct = ct.reshape(shape_ct, order="A")
-    c = ct.transpose(order_ct)
+    if ct.flags.f_contiguous:
+        c = np.asfortranarray(ct.transpose(order_ct))
+    elif ct.flags.c_contiguous:
+        c = np.ascontiguousarray(ct.transpose(order_ct))
+    else:
+        c = ct.transpose(order_ct)
 
     return c