Finite size corrections

momentGW/pbc/tda.py

@@ -6,6 +6,8 @@
 import scipy.special
 from pyscf import lib
 from pyscf.agf2 import mpi_helper
+from pyscf.pbc import dft
+from pyscf.pbc.gw.krgw_ac import get_qij
 
 from momentGW.tda import TDA as MolTDA
 
@@ -214,11 +216,162 @@ def build_se_moments(self, moments_dd):
         moments_occ, moments_vir = self.convolve(eta)
         cput1 = lib.logger.timer(self.gw, "constructing SE moments", *cput1)
 
+        if self.gw.fc:
+            moments_dd_fc = self.build_dd_moments_fc()
+
+            moments_occ_fc, moments_vir_fc = self.build_se_moments_fc(*moments_dd_fc)
+
+            moments_occ += moments_occ_fc
+            moments_vir += moments_vir_fc
+
+            cput1 = lib.logger.timer(self.gw, "fc correction", *cput1)
+
         return moments_occ, moments_vir
 
     def build_dd_moments_exact(self):
         raise NotImplementedError
 
+    def build_dd_moments_fc(self):
+        """
+        Build the moments of the "head" (G=0, G'=0) and "wing"
+        (G=P, G'=0) density-density response.
+        """
+
+        kpts = self.kpts
+        integrals = self.integrals
+
+        # q-point for q->0 finite size correction
+        qpt = np.array([1e-3, 0.0, 0.0])
+        qpt = self.kpts.get_abs_kpts(qpt)
+
+        # 1/sqrt(Ω) * ⟨Ψ_{ik}|e^{iqr}|Ψ_{ak-q}⟩
+        qij = self.build_pert_term(qpt)
+
+        # Build the DD moments for the "head" (G=0, G'=0) correction
+        moments_head = np.zeros((self.nkpts, self.nmom_max + 1), dtype=complex)
+        for k in kpts.loop(1):
+            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],
+            )
+            dn = np.ones_like(d)
+            for n in range(self.nmom_max + 1):
+                moments_head[k, n] = lib.einsum("ia,ia,ia->", qij[k], qij[k].conj(), dn)
+                dn *= d
+
+        # Build the DD moments for the "wing" (G=P, G'=0) correction
+        moments_wing = np.zeros((self.nkpts, self.nmom_max + 1), dtype=object)
+        for k in kpts.loop(1):
+            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],
+            )
+            dn = np.ones_like(d)
+            for n in range(self.nmom_max + 1):
+                moments_wing[k, n] = lib.einsum(
+                    "Lx,x,x->L",
+                    integrals.Lia[k, k],
+                    qij[k].conj().ravel(),
+                    dn.ravel(),
+                )
+                dn *= d
+
+        moments_head *= -4.0 * np.pi
+        moments_wing *= -4.0 * np.pi
+
+        return moments_head, moments_wing
+
+    def build_se_moments_fc(self, moments_head, moments_wing):
+        """
+        Build the moments of the self-energy corresponding to the
+        "wing" (G=P, G'=0) and "head" (G=0, G'=0) density-density
+        response via convolution.
+        """
+
+        kpts = self.kpts
+        integrals = self.integrals
+        moms = np.arange(self.nmom_max + 1)
+
+        # Construct the self-energy moments for the "head" (G=0, G'=0)
+        moments_occ_h = np.zeros((self.nkpts, self.nmom_max + 1, self.nmo, self.nmo), dtype=complex)
+        moments_vir_h = np.zeros((self.nkpts, self.nmom_max + 1, self.nmo, self.nmo), dtype=complex)
+        for n in moms:
+            fp = scipy.special.binom(n, moms)
+            fh = fp * (-1) ** moms
+            for k in kpts.loop(1):
+                eo = np.power.outer(self.mo_energy_g[k][self.mo_occ_g[k] > 0], n - moms)
+                to = lib.einsum("t,kt,t->", fh, eo, moments_head[k])
+                moments_occ_h[k, n] = np.diag([to] * self.nmo)
+
+                ev = np.power.outer(self.mo_energy_g[k][self.mo_occ_g[k] == 0], n - moms)
+                tv = lib.einsum("t,kt,t->", fp, ev, moments_head[k])
+                moments_vir_h[k, n] = np.diag([tv] * self.nmo)
+
+        # Construct the self-energy moments for the "wing" (G=P, G'=0)
+        moments_occ_w = np.zeros((self.nkpts, self.nmom_max + 1, self.nmo, self.nmo), dtype=complex)
+        moments_vir_w = np.zeros((self.nkpts, self.nmom_max + 1, self.nmo, self.nmo), dtype=complex)
+        for n in moms:
+            fp = scipy.special.binom(n, moms)
+            fh = fp * (-1) ** moms
+            for k in kpts.loop(1):
+                eta = np.zeros(
+                    (self.integrals.nmo_g[k], self.nmom_max + 1, self.nmo), dtype=complex
+                )
+                for t in moms:
+                    eta[:, t] = lib.einsum("L,Lpx->xp", moments_wing[k, t], integrals.Lpx[k, k])
+
+                eo = np.power.outer(self.mo_energy_g[k][self.mo_occ_g[k] > 0], n - moms)
+                to = lib.einsum("t,kt,ktp->p", fh, eo, eta[self.mo_occ_g[k] > 0])
+                moments_occ_w[k, n] = np.diag(to)
+
+                ev = np.power.outer(self.mo_energy_g[k][self.mo_occ_g[k] == 0], n - moms)
+                tv = lib.einsum("t,kt,ktp->p", fp, ev, eta[self.mo_occ_g[k] == 0])
+                moments_vir_w[k, n] = np.diag(tv)
+
+        moments_occ = moments_occ_h + moments_occ_w
+        moments_vir = moments_vir_h + moments_vir_w
+
+        factor = -2.0 / np.pi * (6.0 * np.pi**2 / self.gw.cell.vol / self.nkpts) ** (1.0 / 3.0)
+        moments_occ *= factor
+        moments_vir *= factor
+
+        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)
+        ngrid = len(coords)
+
+        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 = lib.einsum("g,gm,xgn->xmn", weights, ao.conj(), ao_grad)
+            q_ao_ao_grad = lib.einsum("x,xmn->mn", qpt, ao_ao_grad) * -1.0j
+            q_mo_mo_grad = lib.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 / np.sqrt(self.gw.cell.vol)
+
+        return qij
+
     @property
     def naux(self):
         """Number of auxiliaries."""