pyscf 2.6.2__py3-none-macosx_11_0_arm64.whl → 2.7.0__py3-none-macosx_11_0_arm64.whl

pyscf/pbc/tdscf/kproxy_supercell.py

@@ -13,7 +13,7 @@ regular `numpy.linalg.eig` can be used to retrieve TD roots. Several variants of
  * (this module) `pyscf.pbc.tdscf.kproxy_supercell`: PBC implementation constructing supercells. Works with an arbitrary number of
    k-points but has an overhead due to ignoring the momentum conservation law. In addition, works only with
    time reversal invariant (TRI) models: i.e. the k-point grid has to be aligned and contain at least one TRI momentum.
- * `pyscf.pbc.tdscf.kproxy`: same as the above but respect the momentum conservation and, thus, diagonlizes smaller
+ * `pyscf.pbc.tdscf.kproxy`: same as the above but respect the momentum conservation and, thus, diagonalizes smaller
    matrices (the performance gain is the total number of k-points in the model).
 """
 
@@ -495,7 +495,7 @@ def orb2ov(space, nocc, nmo):
     """
     Converts orbital active space specification into ov-pairs space spec.
     Args:
-        space (ndarray): the obital space. Basis order: [k, orb=o+v];
+        space (ndarray): the orbital space. Basis order: [k, orb=o+v];
         nocc (Iterable): the numbers of occupied orbitals per k-point;
         nmo (Iterable): the total numbers of orbitals per k-point;
 

pyscf/pbc/tdscf/krhf.py

@@ -26,10 +26,11 @@ from pyscf import lib
 from pyscf.lib import linalg_helper
 from pyscf.lib import logger
 from pyscf.tdscf import rhf
+from pyscf.tdscf._lr_eig import eigh as lr_eigh, eig as lr_eig
 from pyscf.pbc import scf
 from pyscf.pbc.tdscf.rhf import TDBase
 from pyscf.pbc.scf import _response_functions  # noqa
-from pyscf.pbc.lib.kpts_helper import gamma_point, get_kconserv_ria
+from pyscf.pbc.lib.kpts_helper import is_gamma_point, get_kconserv_ria, conj_mapping
 from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
 from pyscf.pbc import df as pbcdf
 from pyscf.data import nist
@@ -37,8 +38,139 @@ from pyscf import __config__
 
 REAL_EIG_THRESHOLD = getattr(__config__, 'pbc_tdscf_rhf_TDDFT_pick_eig_threshold', 1e-3)
 
+def get_ab(mf, kshift=0):
+    r'''A and B matrices for TDDFT response function.
+
+    A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ia||bj)
+    B[i,a,j,b] = (ia||jb)
+
+    Ref: Chem Phys Lett, 256, 454
+
+    Kwargs:
+        kshift : integer
+            The index of the k-point that represents the transition between
+            k-points in the excitation coefficients.
+    '''
+    cell = mf.cell
+    mo_energy = scf.addons.mo_energy_with_exxdiv_none(mf)
+    mo = numpy.asarray(mf.mo_coeff)
+    mo_occ = numpy.asarray(mf.mo_occ)
+    kpts = mf.kpts
+    nkpts, nao, nmo = mo.shape
+    noccs = numpy.count_nonzero(mo_occ==2, axis=1)
+    nocc = noccs[0]
+    nvir = nmo - nocc
+    assert all(noccs == nocc)
+    orbo = mo[:,:,:nocc]
+    orbv = mo[:,:,nocc:]
+
+    kconserv = get_kconserv_ria(cell, kpts)[kshift]
+    e_ia = numpy.asarray(_get_e_ia(mo_energy, mo_occ, kconserv)).astype(mo.dtype)
+    a = numpy.diag(e_ia.ravel()).reshape(nkpts,nocc,nvir,nkpts,nocc,nvir)
+    b = numpy.zeros_like(a)
+    weight = 1./nkpts
+
+    def add_hf_(a, b, hyb=1):
+        eri = mf.with_df.ao2mo_7d([orbo,mo,mo,mo], kpts)
+        eri *= weight
+        eri = eri.reshape(nkpts,nkpts,nkpts,nocc,nmo,nmo,nmo)
+        for ki, ka in enumerate(kconserv):
+            for kj, kb in enumerate(kconserv):
+                a[ki,:,:,kj] += numpy.einsum('iabj->iajb', eri[ki,ka,kb,:nocc,nocc:,nocc:,:nocc]) * 2
+                a[ki,:,:,kj] -= numpy.einsum('ijba->iajb', eri[ki,kj,kb,:nocc,:nocc,nocc:,nocc:]) * hyb
+
+            for kb, kj in enumerate(kconserv):
+                b[ki,:,:,kj] += numpy.einsum('iajb->iajb', eri[ki,ka,kj,:nocc,nocc:,:nocc,nocc:]) * 2
+                b[ki,:,:,kj] -= numpy.einsum('ibja->iajb', eri[ki,kb,kj,:nocc,nocc:,:nocc,nocc:]) * hyb
+
+    if isinstance(mf, scf.hf.KohnShamDFT):
+        assert kshift == 0
+        ni = mf._numint
+        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, cell.spin)
+
+        add_hf_(a, b, hyb)
+        if omega != 0:  # For RSH
+            raise NotImplementedError
+
+        xctype = ni._xc_type(mf.xc)
+        dm0 = mf.make_rdm1(mo, mo_occ)
+        make_rho = ni._gen_rho_evaluator(cell, dm0, hermi=1, with_lapl=False)[0]
+        mem_now = lib.current_memory()[0]
+        max_memory = max(2000, mf.max_memory*.8-mem_now)
+        cmap = conj_mapping(cell, kpts)
+
+        if xctype == 'LDA':
+            ao_deriv = 0
+            for ao, _, mask, weight, coords \
+                    in ni.block_loop(cell, mf.grids, nao, ao_deriv, kpts, None, max_memory):
+                rho = make_rho(0, ao, mask, xctype)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc[0,0] * weight
+
+                rho_o = lib.einsum('krp,kpi->kri', ao, orbo)
+                rho_v = lib.einsum('krp,kpi->kri', ao, orbv)
+                rho_ov = numpy.einsum('kri,kra->kria', rho_o, rho_v)
+                w_ov = numpy.einsum('kria,r->kria', rho_ov, wfxc) * (2/nkpts)
+                rho_vo = rho_ov.conj()[cmap]
+                a += lib.einsum('kria,lrjb->kialjb', w_ov, rho_vo)
+                b += lib.einsum('kria,lrjb->kialjb', w_ov, rho_ov)
+
+        elif xctype == 'GGA':
+            ao_deriv = 1
+            for ao, _, mask, weight, coords \
+                    in ni.block_loop(cell, mf.grids, nao, ao_deriv, kpts, None, max_memory):
+                rho = make_rho(0, ao, mask, xctype)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc * weight
+                rho_o = lib.einsum('kxrp,kpi->kxri', ao, orbo)
+                rho_v = lib.einsum('kxrp,kpi->kxri', ao, orbv)
+                rho_ov = numpy.einsum('kxri,kra->kxria', rho_o, rho_v[:,0])
+                rho_ov[:,1:4] += numpy.einsum('kri,kxra->kxria', rho_o[:,0], rho_v[:,1:4])
+                w_ov = numpy.einsum('xyr,kxria->kyria', wfxc, rho_ov) * (2/nkpts)
+                rho_vo = rho_ov.conj()[cmap]
+                a += lib.einsum('kxria,lxrjb->kialjb', w_ov, rho_vo)
+                b += lib.einsum('kxria,lxrjb->kialjb', w_ov, rho_ov)
+
+        elif xctype == 'HF':
+            pass
+
+        elif xctype == 'NLC':
+            raise NotImplementedError('NLC')
+
+        elif xctype == 'MGGA':
+            ao_deriv = 1
+            for ao, _, mask, weight, coords \
+                    in ni.block_loop(cell, mf.grids, nao, ao_deriv, kpts, None, max_memory):
+                rho = make_rho(0, ao, mask, xctype)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc * weight
+                rho_o = lib.einsum('kxrp,kpi->kxri', ao, orbo)
+                rho_v = lib.einsum('kxrp,kpi->kxri', ao, orbv)
+                rho_ov = numpy.einsum('kxri,kra->kxria', rho_o, rho_v[:,0])
+                rho_ov[:,1:4] += numpy.einsum('kri,kxra->kxria', rho_o[:,0], rho_v[:,1:4])
+                tau_ov = numpy.einsum('kxri,kxra->kria', rho_o[:,1:4], rho_v[:,1:4]) * .5
+                rho_ov = numpy.concatenate([rho_ov, tau_ov[:,numpy.newaxis]], axis=1)
+                w_ov = numpy.einsum('xyr,kxria->kyria', wfxc, rho_ov) * (2/nkpts)
+                rho_vo = rho_ov.conj()[cmap]
+                a += lib.einsum('kxria,lxrjb->kialjb', w_ov, rho_vo)
+                b += lib.einsum('kxria,lxrjb->kialjb', w_ov, rho_ov)
+    else:
+        add_hf_(a, b)
+
+    return a, b
+
 class KTDBase(TDBase):
-    _keys = {'kconserv', 'kshift_lst'}
+    '''
+    Attributes:
+        kshift_lst : list of integers
+            Each element in the list is the index of the k-point that
+            represents the transition between k-points in the excitation
+            coefficients.
+    '''
+
+    conv_tol = getattr(__config__, 'pbc_tdscf_rhf_TDA_conv_tol', 1e-4)
+
+    _keys = {'kshift_lst'}
 
     def __init__(self, mf, kshift_lst=None):
         assert isinstance(mf, scf.khf.KSCF)
@@ -46,8 +178,6 @@ class KTDBase(TDBase):
         warn_pbc2d_eri(mf)
 
         if kshift_lst is None: kshift_lst = [0]
-
-        self.kconserv = get_kconserv_ria(mf.cell, mf.kpts)
         self.kshift_lst = kshift_lst
 
     def dump_flags(self, verbose=None):
@@ -67,7 +197,6 @@ class KTDBase(TDBase):
         log.info('conv_tol = %g', self.conv_tol)
         log.info('eigh lindep = %g', self.lindep)
         log.info('eigh level_shift = %g', self.level_shift)
-        log.info('eigh max_space = %d', self.max_space)
         log.info('eigh max_cycle = %d', self.max_cycle)
         log.info('chkfile = %s', self.chkfile)
         log.info('max_memory %d MB (current use %d MB)',
@@ -79,7 +208,7 @@ class KTDBase(TDBase):
     def check_sanity(self):
         TDBase.check_sanity(self)
         mf = self._scf
-        if any([k != 0 for k in self.kshift_lst]):
+        if any(k != 0 for k in self.kshift_lst):
             if mf.rsjk is not None or not isinstance(mf.with_df, pbcdf.df.DF):
                 logger.error(self, 'Solutions with non-zero kshift for %s are '
                              'only supported by GDF/RSDF')
@@ -98,14 +227,22 @@ class KTDBase(TDBase):
     get_nto = lib.invalid_method('get_nto')
 
 class TDA(KTDBase):
-    conv_tol = getattr(__config__, 'pbc_tdscf_rhf_TDA_conv_tol', 1e-6)
 
-    def gen_vind(self, mf, kshift):
-        # exxdiv corrections are kept in hdiag while excluding them when calling
-        # the contractions between two-electron integrals and X/Y amplitudes.
-        # See also the relevant comments in function pbc.tdscf.rhf.TDA.gen_vind
+    @lib.with_doc(get_ab.__doc__)
+    def get_ab(self, mf=None, kshift=0):
+        if mf is None: mf = self._scf
+        return get_ab(mf, kshift)
+
+    def gen_vind(self, mf, kshift=0):
+        '''Compute Ax
+
+        Kwargs:
+            kshift : integer
+                The index of the k-point that represents the transition between
+                k-points in the excitation coefficients.
+        '''
         singlet = self.singlet
-        kconserv = self.kconserv[kshift]
+        kconserv = get_kconserv_ria(mf.cell, mf.kpts)[kshift]
 
         mo_coeff = mf.mo_coeff
         mo_occ = mf.mo_occ
@@ -114,7 +251,7 @@ class TDA(KTDBase):
         occidx = [numpy.where(mo_occ[k]==2)[0] for k in range(nkpts)]
         viridx = [numpy.where(mo_occ[k]==0)[0] for k in range(nkpts)]
         orbo = [mo_coeff[k][:,occidx[k]] for k in range(nkpts)]
-        orbv = [mo_coeff[kconserv[k]][:,viridx[kconserv[k]]] for k in range(nkpts)]
+        orbv = [mo_coeff[k][:,viridx[k]] for k in range(nkpts)]
         e_ia = _get_e_ia(scf.addons.mo_energy_with_exxdiv_none(mf), mo_occ, kconserv)
         hdiag = numpy.hstack([x.ravel() for x in e_ia])
 
@@ -127,10 +264,10 @@ class TDA(KTDBase):
             z1s = [_unpack(z, mo_occ, kconserv) for z in zs]
             dmov = numpy.empty((nz,nkpts,nao,nao), dtype=numpy.complex128)
             for i in range(nz):
-                for k in range(nkpts):
+                for k, kp in enumerate(kconserv):
                     # *2 for double occupancy
                     dm1 = z1s[i][k] * 2
-                    dmov[i,k] = reduce(numpy.dot, (orbo[k], dm1, orbv[k].conj().T))
+                    dmov[i,k] = reduce(numpy.dot, (orbo[k], dm1, orbv[kp].conj().T))
 
             with lib.temporary_env(mf, exxdiv=None):
                 v1ao = vresp(dmov, kshift)
@@ -138,8 +275,8 @@ class TDA(KTDBase):
             for i in range(nz):
                 dm1 = z1s[i]
                 v1 = []
-                for k in range(nkpts):
-                    v1vo = reduce(numpy.dot, (orbo[k].conj().T, v1ao[i,k], orbv[k]))
+                for k, kp in enumerate(kconserv):
+                    v1vo = reduce(numpy.dot, (orbo[k].conj().T, v1ao[i,k], orbv[kp]))
                     v1vo += e_ia[k] * dm1[k]
                     v1.append(v1vo.ravel())
                 v1s.append( numpy.concatenate(v1) )
@@ -151,13 +288,13 @@ class TDA(KTDBase):
 
         mo_energy = mf.mo_energy
         mo_occ = mf.mo_occ
-        kconserv = self.kconserv[kshift]
+        kconserv = get_kconserv_ria(mf.cell, mf.kpts)[kshift]
         e_ia = numpy.concatenate( [x.reshape(-1) for x in
                                    _get_e_ia(mo_energy, mo_occ, kconserv)] )
 
         nov = e_ia.size
         nstates = min(nstates, nov)
-        e_threshold = numpy.sort(e_ia)[nstates-1]
+        e_threshold = numpy.partition(e_ia, nstates-1)[nstates-1]
         e_threshold += self.deg_eia_thresh
 
         idx = numpy.where(e_ia <= e_threshold)[0]
@@ -188,13 +325,12 @@ class TDA(KTDBase):
             return w[idx], v[:,idx], idx
 
         log = logger.Logger(self.stdout, self.verbose)
-        precision = self.cell.precision * 1e-2
 
         self.converged = []
         self.e = []
         self.xy = []
         for i,kshift in enumerate(self.kshift_lst):
-            kconserv = self.kconserv[kshift]
+            kconserv = get_kconserv_ria(mf.cell, mf.kpts)[kshift]
 
             vind, hdiag = self.gen_vind(self._scf, kshift)
             precond = self.get_precond(hdiag)
@@ -204,14 +340,10 @@ class TDA(KTDBase):
             else:
                 x0k = x0[i]
 
-            converged, e, x1 = \
-                    lib.davidson1(vind, x0k, precond,
-                                  tol=self.conv_tol,
-                                  max_cycle=self.max_cycle,
-                                  nroots=self.nstates, lindep=self.lindep,
-                                  max_space=self.max_space, pick=pickeig,
-                                  fill_heff=purify_krlyov_heff(precision, 0, log),
-                                  verbose=self.verbose)
+            converged, e, x1 = lr_eigh(
+                vind, x0k, precond, tol_residual=self.conv_tol, lindep=self.lindep,
+                nroots=self.nstates, pick=pickeig, max_cycle=self.max_cycle,
+                max_memory=self.max_memory, verbose=log)
             self.converged.append( converged )
             self.e.append( e )
             # 1/sqrt(2) because self.x is for alpha excitation amplitude and 2(X^+*X) = 1
@@ -223,14 +355,21 @@ class TDA(KTDBase):
 CIS = KTDA = TDA
 
 
-class TDHF(TDA):
-    def gen_vind(self, mf, kshift):
+class TDHF(KTDBase):
+
+    @lib.with_doc(get_ab.__doc__)
+    def get_ab(self, mf=None, kshift=0):
+        if mf is None: mf = self._scf
+        return get_ab(mf, kshift)
+
+    def gen_vind(self, mf, kshift=0):
         '''
         [ A   B ][X]
         [-B* -A*][Y]
         '''
+        assert kshift == 0
+
         singlet = self.singlet
-        kconserv = self.kconserv[kshift]
 
         mo_coeff = mf.mo_coeff
         mo_occ = mf.mo_occ
@@ -239,7 +378,9 @@ class TDHF(TDA):
         occidx = [numpy.where(mo_occ[k]==2)[0] for k in range(nkpts)]
         viridx = [numpy.where(mo_occ[k]==0)[0] for k in range(nkpts)]
         orbo = [mo_coeff[k][:,occidx[k]] for k in range(nkpts)]
-        orbv = [mo_coeff[kconserv[k]][:,viridx[kconserv[k]]] for k in range(nkpts)]
+        orbv = [mo_coeff[k][:,viridx[k]] for k in range(nkpts)]
+
+        kconserv = numpy.arange(nkpts)
         e_ia = _get_e_ia(scf.addons.mo_energy_with_exxdiv_none(mf), mo_occ, kconserv)
         hdiag = numpy.hstack([x.ravel() for x in e_ia])
         tot_x = hdiag.size
@@ -253,31 +394,36 @@ class TDHF(TDA):
             nz = len(xys)
             z1xs = [_unpack(xy[:tot_x], mo_occ, kconserv) for xy in xys]
             z1ys = [_unpack(xy[tot_x:], mo_occ, kconserv) for xy in xys]
-            dmov = numpy.empty((nz,nkpts,nao,nao), dtype=numpy.complex128)
+            dmov = numpy.zeros((nz,nkpts,nao,nao), dtype=numpy.complex128)
             for i in range(nz):
                 for k in range(nkpts):
                     # *2 for double occupancy
                     dmx = z1xs[i][k] * 2
                     dmy = z1ys[i][k] * 2
-                    dmov[i,k] = reduce(numpy.dot, (orbo[k], dmx, orbv[k].T.conj()))
-                    dmov[i,k]+= reduce(numpy.dot, (orbv[k], dmy.T, orbo[k].T.conj()))
+                    dmov[i,k] += reduce(numpy.dot, (orbo[k], dmx  , orbv[k].T.conj()))
+                    dmov[i,k] += reduce(numpy.dot, (orbv[k], dmy.T, orbo[k].T.conj()))
 
             with lib.temporary_env(mf, exxdiv=None):
-                v1ao = vresp(dmov, kshift)
+                v1ao = vresp(dmov, kshift) # = <mb||nj> Xjb + <mj||nb> Yjb
             v1s = []
             for i in range(nz):
                 dmx = z1xs[i]
                 dmy = z1ys[i]
-                v1xs = []
-                v1ys = []
+                v1xs = [0] * nkpts
+                v1ys = [0] * nkpts
                 for k in range(nkpts):
+                    # AX + BY
+                    # = <ib||aj> Xjb + <ij||ab> Yjb
+                    # = (<mb||nj> Xjb + <mj||nb> Yjb) Cmi* Cna
                     v1x = reduce(numpy.dot, (orbo[k].T.conj(), v1ao[i,k], orbv[k]))
+                    # (B*)X + (A*)Y
+                    # = <ab||ij> Xjb + <aj||ib> Yjb
+                    # = (<mb||nj> Xjb + <mj||nb> Yjb) Cma* Cni
                     v1y = reduce(numpy.dot, (orbv[k].T.conj(), v1ao[i,k], orbo[k])).T
-                    v1x+= e_ia[k] * dmx[k]
-                    v1y+= e_ia[k] * dmy[k]
-                    v1xs.append(v1x.ravel())
-                    v1ys.append(-v1y.ravel())
-                # v1s += v1xs + v1ys
+                    v1x += e_ia[k] * dmx[k]
+                    v1y += e_ia[k].conj() * dmy[k]
+                    v1xs[k] += v1x.ravel()
+                    v1ys[k] -= v1y.ravel()
                 v1s.append( numpy.concatenate(v1xs + v1ys) )
             return lib.asarray(v1s).reshape(nz,-1)
         return vind, hdiag
@@ -285,7 +431,9 @@ class TDHF(TDA):
     def init_guess(self, mf, kshift, nstates=None):
         x0 = TDA.init_guess(self, mf, kshift, nstates)
         y0 = numpy.zeros_like(x0)
-        return numpy.asarray(numpy.block([[x0, y0], [y0, x0.conj()]]))
+        return numpy.hstack([x0, y0])
+
+    get_precond = rhf.TDHF.get_precond
 
     def kernel(self, x0=None):
         '''TDHF diagonalization with non-Hermitian eigenvalue solver
@@ -299,9 +447,21 @@ class TDHF(TDA):
         mf = self._scf
         mo_occ = mf.mo_occ
 
-        real_system = (gamma_point(self._scf.kpts) and
+        real_system = (is_gamma_point(self._scf.kpts) and
                        self._scf.mo_coeff[0].dtype == numpy.double)
 
+        if any(k != 0 for k in self.kshift_lst):
+            # It's not clear how to define the Y matrix for kshift!=0 .
+            # When the A tensor is constructed against the X(kshift) matrix,
+            # the diagonal terms e_ia are calculated as e_i[k] - e_k[k+kshift].
+            # Given the k-conserve relation in the A tensor, the j-b indices in
+            # the A tensor should follow j[k'], b[k'+kshift]. This leads to the
+            # j-b indices in the B tensor being defined as (j[k'+shift], b[k']).
+            # To form the square A-B-B-A matrix, the diagonal terms for the
+            # -A* part need to be constructed as e_i[k+kshift] - e_a[k], which
+            # conflict to the diagonal terms of the A tensor.
+            raise RuntimeError('kshift != 0 for TDHF')
+
         # We only need positive eigenvalues
         def pickeig(w, v, nroots, envs):
             realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
@@ -309,8 +469,6 @@ class TDHF(TDA):
             return lib.linalg_helper._eigs_cmplx2real(w, v, realidx, real_system)
 
         log = logger.Logger(self.stdout, self.verbose)
-        precision = self.cell.precision * 1e-2
-        hermi = 0
 
         def norm_xy(z, kconserv):
             x, y = z.reshape(2,-1)
@@ -324,7 +482,7 @@ class TDHF(TDA):
         self.e = []
         self.xy = []
         for i,kshift in enumerate(self.kshift_lst):
-            kconserv = self.kconserv[kshift]
+            kconserv = get_kconserv_ria(mf.cell, mf.kpts)[kshift]
 
             vind, hdiag = self.gen_vind(self._scf, kshift)
             precond = self.get_precond(hdiag)
@@ -334,15 +492,10 @@ class TDHF(TDA):
             else:
                 x0k = x0[i]
 
-            converged, e, x1 = \
-                    lib.davidson_nosym1(vind, x0k, precond,
-                                        tol=self.conv_tol,
-                                        max_cycle=self.max_cycle,
-                                        nroots=self.nstates,
-                                        lindep=self.lindep,
-                                        max_space=self.max_space, pick=pickeig,
-                                        fill_heff=purify_krlyov_heff(precision, hermi, log),
-                                        verbose=self.verbose)
+            converged, e, x1 = lr_eig(
+                vind, x0k, precond, tol_residual=self.conv_tol, lindep=self.lindep,
+                nroots=self.nstates, pick=pickeig, max_cycle=self.max_cycle,
+                max_memory=self.max_memory, verbose=log)
             self.converged.append( converged )
             self.e.append( e )
             self.xy.append( [norm_xy(z, kconserv) for z in x1] )
@@ -366,87 +519,16 @@ def _get_e_ia(mo_energy, mo_occ, kconserv=None):
 def _unpack(vo, mo_occ, kconserv):
     z = []
     p1 = 0
-    for k, occ in enumerate(mo_occ):
-        no = numpy.count_nonzero(occ > 0)
-        no1 = numpy.count_nonzero(mo_occ[kconserv[k]] > 0)
-        nv = occ.size - no1
+    no_kpts = [numpy.count_nonzero(occ) for occ in mo_occ]
+    for k, no in enumerate(no_kpts):
+        kp = kconserv[k]
+        nv = mo_occ[kp].size - no_kpts[kp]
         p0, p1 = p1, p1 + no * nv
         z.append(vo[p0:p1].reshape(no,nv))
     return z
 
-def purify_krlyov_heff(precision, hermi, log):
-    def fill_heff(heff, xs, ax, xt, axt, dot):
-        if hermi == 1:
-            heff = linalg_helper._fill_heff_hermitian(heff, xs, ax, xt, axt, dot)
-        else:
-            heff = linalg_helper._fill_heff(heff, xs, ax, xt, axt, dot)
-        space = len(axt)
-        # TODO: PBC integrals has larger errors than molecule systems.
-        # purify the effective Hamiltonian with symmetry and other
-        # possible conditions.
-        if abs(heff[:space,:space].imag).max() < precision:
-            log.debug('Remove imaginary part of the Krylov space effective Hamiltonian')
-            heff[:space,:space].imag = 0
-        return heff
-    return fill_heff
-
 
 scf.khf.KRHF.TDA  = lib.class_as_method(KTDA)
 scf.khf.KRHF.TDHF = lib.class_as_method(KTDHF)
 scf.krohf.KROHF.TDA  = None
 scf.krohf.KROHF.TDHF = None
-
-
-if __name__ == '__main__':
-    from pyscf.pbc import gto
-    from pyscf.pbc import scf
-    from pyscf.pbc import df
-    cell = gto.Cell()
-    cell.unit = 'B'
-    cell.atom = '''
-    C  0.          0.          0.
-    C  1.68506879  1.68506879  1.68506879
-    '''
-    cell.a = '''
-    0.          3.37013758  3.37013758
-    3.37013758  0.          3.37013758
-    3.37013758  3.37013758  0.
-    '''
-
-    cell.basis = 'gth-szv'
-    cell.pseudo = 'gth-pade'
-    cell.mesh = [25]*3
-    cell.build()
-    mf = scf.KRHF(cell, cell.make_kpts([2,1,1])).set(exxdiv=None)
-    #mf.with_df = df.MDF(cell, cell.make_kpts([2,1,1]))
-    #mf.with_df.auxbasis = 'weigend'
-    #mf.with_df._cderi = 'eri3d-mdf.h5'
-    #mf.with_df.build(with_j3c=False)
-    mf.run()
-    #mesh=9  -8.65192427146353
-    #mesh=12 -8.65192352289817
-    #mesh=15 -8.6519235231529
-    #MDF mesh=5 -8.6519301815144
-
-    td = TDA(mf)
-    td.verbose = 5
-    print(td.kernel()[0][0] * 27.2114)
-    #mesh=9  [ 6.0073749   6.09315355  6.3479901 ]
-    #mesh=12 [ 6.00253282  6.09317929  6.34799109]
-    #mesh=15 [ 6.00253396  6.09317949  6.34799109]
-    #MDF mesh=5 [ 6.09317489  6.09318265  6.34798637]
-
-    #from pyscf.pbc import tools
-    #scell = tools.super_cell(cell, [2,1,1])
-    #mf = scf.RHF(scell).run()
-    #td = rhf.TDA(mf)
-    #td.verbose = 5
-    #print(td.kernel()[0] * 27.2114)
-
-    td = TDHF(mf)
-    td.verbose = 5
-    print(td.kernel()[0][0] * 27.2114)
-    #mesh=9  [ 6.03860914  6.21664545  8.20305225]
-    #mesh=12 [ 6.03868259  6.03860343  6.2167623 ]
-    #mesh=15 [ 6.03861321  6.03861324  6.21675868]
-    #MDF mesh=5 [ 6.03861693  6.03861775  6.21675694]

pyscf/pbc/tdscf/krhf_slow_supercell.py

@@ -94,7 +94,7 @@ class PhysERI(PeriodicMFMixin, TDERIMatrixBlocks):
         """
         Retrieves the merged diagonal block with specified or all possible k-index pairs.
         Args:
-            pairs (Iterable): pairs of k-points to assmble;
+            pairs (Iterable): pairs of k-points to assemble;
 
         Returns:
             The diagonal block.

pyscf/pbc/tdscf/krks.py

@@ -43,7 +43,7 @@ RPA = KTDDFT = TDDFT
 def _rebuild_df(td):
     log = lib.logger.new_logger(td)
     mf = td._scf
-    if any([k != 0 for k in td.kshift_lst]):
+    if any(k != 0 for k in td.kshift_lst):
         if isinstance(mf.with_df, df.df.DF):
             if mf.with_df._j_only:
                 log.warn(f'Non-zero kshift is requested for {td.__class__.__name__}, '
@@ -57,47 +57,3 @@ dft.krks.KRKS.TDDFT = lib.class_as_method(TDDFT)
 dft.kroks.KROKS.TDA   = None
 dft.kroks.KROKS.TDHF  = None
 dft.kroks.KROKS.TDDFT = None
-
-
-if __name__ == '__main__':
-    from pyscf.pbc import gto
-    from pyscf.pbc import dft
-    cell = gto.Cell()
-    cell.unit = 'B'
-    cell.atom = '''
-    C  0.          0.          0.
-    C  1.68506879  1.68506879  1.68506879
-    '''
-    cell.a = '''
-    0.          3.37013758  3.37013758
-    3.37013758  0.          3.37013758
-    3.37013758  3.37013758  0.
-    '''
-    cell.basis = 'gth-szv'
-    cell.pseudo = 'gth-pade'
-    cell.mesh = [25]*3
-    cell.build()
-
-    mf = dft.KRKS(cell, cell.make_kpts([2,1,1]))
-    #mf.with_df = df.MDF(cell, cell.make_kpts([2,1,1]))
-    #mf.with_df.auxbasis = 'weigend'
-    #mf.with_df._cderi = 'eri3d-mdf.h5'
-    #mf.with_df.build(with_j3c=False)
-    mf.xc = 'lda,'
-    mf.kernel()
-#mesh=12 -10.3077341607895
-#mesh=5  -10.3086623157515
-
-    td = TDDFT(mf)
-    td.nstates = 5
-    td.verbose = 5
-    print(td.kernel()[0][0] * 27.2114)
-#mesh=12 [ 6.08108297  6.10231481  6.10231478  6.38355803  6.38355804]
-#MDF mesh=5 [ 6.07919157  6.10251718  6.10253961  6.37202499  6.37565246]
-
-    td = TDA(mf)
-    td.singlet = False
-    td.verbose = 5
-    print(td.kernel()[0][0] * 27.2114)
-#mesh=12 [ 4.01539192  5.1750807   5.17508071]
-#MDF mesh=5 [ 4.01148649  5.18043397  5.18043459]