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

pyscf/pbc/tdscf/kuhf.py

@@ -17,41 +17,250 @@
 #
 
 from functools import reduce
-import numpy
+import numpy as np
 from pyscf import lib
 from pyscf.lib import logger
 from pyscf.tdscf import uhf
+from pyscf.tdscf._lr_eig import eigh as lr_eigh, eig as lr_eig
 from pyscf.pbc import scf
-from pyscf.pbc.tdscf.krhf import KTDBase, _get_e_ia, purify_krlyov_heff
-from pyscf.pbc.lib.kpts_helper import gamma_point
+from pyscf.pbc.tdscf.krhf import KTDBase, _get_e_ia
+from pyscf.pbc.lib.kpts_helper import is_gamma_point, get_kconserv_ria, conj_mapping
 from pyscf.pbc.scf import _response_functions  # noqa
 from pyscf import __config__
 
 REAL_EIG_THRESHOLD = getattr(__config__, 'pbc_tdscf_uhf_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) + (ai||jb)
+    B[i,a,j,b] = (ai||bj)
+
+    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_a, mo_b = mo = np.asarray(mf.mo_coeff)
+    mo_occ = np.asarray(mf.mo_occ)
+    kpts = mf.kpts
+    nkpts, nao, nmo = mo_a.shape
+    noccs = np.count_nonzero(mo_occ!=0, axis=2)
+    nocc = noccs[0,0]
+    nvir = nmo - nocc
+    assert np.all(noccs == nocc)
+    nocc_a = nocc_b = nocc
+    nvir_a = nvir_b = nvir
+    orbo_a = mo_a[:,:,:nocc]
+    orbo_b = mo_b[:,:,:nocc]
+    orbv_a = mo_a[:,:,nocc:]
+    orbv_b = mo_b[:,:,nocc:]
+
+    kconserv = get_kconserv_ria(cell, kpts)[kshift]
+    e_ia_a = np.asarray(_get_e_ia(mo_energy[0], mo_occ[0], kconserv)).astype(mo.dtype)
+    e_ia_b = np.asarray(_get_e_ia(mo_energy[1], mo_occ[1], kconserv)).astype(mo.dtype)
+    a_aa = np.diag(e_ia_a.ravel()).reshape(nkpts,nocc_a,nvir_a,nkpts,nocc_a,nvir_a)
+    a_bb = np.diag(e_ia_b.ravel()).reshape(nkpts,nocc_b,nvir_b,nkpts,nocc_b,nvir_b)
+    a_ab = np.zeros((nkpts,nocc_a,nvir_a,nkpts,nocc_b,nvir_b), dtype=a_aa.dtype)
+    b_aa = np.zeros_like(a_aa)
+    b_bb = np.zeros_like(a_bb)
+    b_ab = np.zeros_like(a_ab)
+    a = (a_aa, a_ab, a_bb)
+    b = (b_aa, b_ab, b_bb)
+    weight = 1./nkpts
+
+    def add_hf_(a, b, hyb=1):
+        eri_aa = mf.with_df.ao2mo_7d([mo_a,orbo_a,mo_a,mo_a], kpts)
+        eri_ab = mf.with_df.ao2mo_7d([mo_a,orbo_a,mo_b,mo_b], kpts)
+        eri_bb = mf.with_df.ao2mo_7d([mo_b,orbo_b,mo_b,mo_b], kpts)
+        eri_aa *= weight
+        eri_ab *= weight
+        eri_bb *= weight
+        eri_aa.reshape(nkpts,nkpts,nkpts,nmo,nocc_a,nmo,nmo)
+        eri_ab.reshape(nkpts,nkpts,nkpts,nmo,nocc_a,nmo,nmo)
+        eri_bb.reshape(nkpts,nkpts,nkpts,nmo,nocc_b,nmo,nmo)
+        a_aa, a_ab, a_bb = a
+        b_aa, b_ab, b_bb = b
+
+        for ki, ka in enumerate(kconserv):
+            for kj, kb in enumerate(kconserv):
+                a_aa[ki,:,:,kj] += np.einsum('aijb->iajb', eri_aa[ka,ki,kj,nocc_a:,:,:nocc_a,nocc_a:])
+                a_aa[ki,:,:,kj] -= np.einsum('jiab->iajb', eri_aa[kj,ki,ka,:nocc_a,:,nocc_a:,nocc_a:]) * hyb
+                a_bb[ki,:,:,kj] += np.einsum('aijb->iajb', eri_bb[ka,ki,kj,nocc_b:,:,:nocc_b,nocc_b:])
+                a_bb[ki,:,:,kj] -= np.einsum('jiab->iajb', eri_bb[kj,ki,ka,:nocc_b,:,nocc_b:,nocc_b:]) * hyb
+                a_ab[ki,:,:,kj] += np.einsum('aijb->iajb', eri_ab[ka,ki,kj,nocc_a:,:,:nocc_b,nocc_b:])
+
+            for kb, kj in enumerate(kconserv):
+                b_aa[ki,:,:,kj] += np.einsum('aibj->iajb', eri_aa[ka,ki,kb,nocc_a:,:,nocc_a:,:nocc_a])
+                b_aa[ki,:,:,kj] -= np.einsum('ajbi->iajb', eri_aa[ka,kj,kb,nocc_a:,:,nocc_a:,:nocc_a]) * hyb
+                b_bb[ki,:,:,kj] += np.einsum('aibj->iajb', eri_bb[ka,ki,kb,nocc_b:,:,nocc_b:,:nocc_b])
+                b_bb[ki,:,:,kj] -= np.einsum('ajbi->iajb', eri_bb[ka,kj,kb,nocc_b:,:,nocc_b:,:nocc_b]) * hyb
+                b_ab[ki,:,:,kj] += np.einsum('aibj->iajb', eri_ab[ka,ki,kb,nocc_a:,:,nocc_b:,:nocc_b])
+
+    if isinstance(mf, scf.hf.KohnShamDFT):
+        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):
+                rho0a = make_rho(0, ao, mask, xctype)
+                rho0b = make_rho(1, ao, mask, xctype)
+                rho = (rho0a, rho0b)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc[0,0] * weight
+
+                rho_o_a = lib.einsum('krp,kpi->kri', ao, orbo_a)
+                rho_v_a = lib.einsum('krp,kpi->kri', ao, orbv_a)
+                rho_o_b = lib.einsum('krp,kpi->kri', ao, orbo_b)
+                rho_v_b = lib.einsum('krp,kpi->kri', ao, orbv_b)
+                rho_ov_a = np.einsum('kri,kra->kria', rho_o_a, rho_v_a)
+                rho_ov_b = np.einsum('kri,kra->kria', rho_o_b, rho_v_b)
+                rho_vo_a = rho_ov_a.conj()[cmap]
+                rho_vo_b = rho_ov_b.conj()[cmap]
+                w_vo_aa = np.einsum('kria,r->kria', rho_vo_a, wfxc[0,0]) * (1/nkpts)
+                w_vo_ab = np.einsum('kria,r->kria', rho_vo_a, wfxc[0,1]) * (1/nkpts)
+                w_vo_bb = np.einsum('kria,r->kria', rho_vo_b, wfxc[1,1]) * (1/nkpts)
+
+                a_aa += lib.einsum('kria,lrjb->kialjb', w_vo_aa, rho_ov_a)
+                b_aa += lib.einsum('kria,lrjb->kialjb', w_vo_aa, rho_vo_a)
+
+                a_ab += lib.einsum('kria,lrjb->kialjb', w_vo_ab, rho_ov_b)
+                b_ab += lib.einsum('kria,lrjb->kialjb', w_vo_ab, rho_vo_b)
+
+                a_bb += lib.einsum('kria,lrjb->kialjb', w_vo_bb, rho_ov_b)
+                b_bb += lib.einsum('kria,lrjb->kialjb', w_vo_bb, rho_vo_b)
+
+        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):
+                rho0a = make_rho(0, ao, mask, xctype)
+                rho0b = make_rho(1, ao, mask, xctype)
+                rho = (rho0a, rho0b)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc * weight
+
+                rho_o_a = lib.einsum('kxrp,kpi->kxri', ao, orbo_a)
+                rho_v_a = lib.einsum('kxrp,kpi->kxri', ao, orbv_a)
+                rho_o_b = lib.einsum('kxrp,kpi->kxri', ao, orbo_b)
+                rho_v_b = lib.einsum('kxrp,kpi->kxri', ao, orbv_b)
+                rho_ov_a = np.einsum('kxri,kra->kxria', rho_o_a, rho_v_a[:,0])
+                rho_ov_b = np.einsum('kxri,kra->kxria', rho_o_b, rho_v_b[:,0])
+                rho_ov_a[:,1:4] += np.einsum('kri,kxra->kxria', rho_o_a[:,0], rho_v_a[:,1:4])
+                rho_ov_b[:,1:4] += np.einsum('kri,kxra->kxria', rho_o_b[:,0], rho_v_b[:,1:4])
+                rho_vo_a = rho_ov_a.conj()[cmap]
+                rho_vo_b = rho_ov_b.conj()[cmap]
+                w_vo_aa = np.einsum('xyr,kxria->kyria', wfxc[0,:,0], rho_vo_a) * (1/nkpts)
+                w_vo_ab = np.einsum('xyr,kxria->kyria', wfxc[0,:,1], rho_vo_a) * (1/nkpts)
+                w_vo_bb = np.einsum('xyr,kxria->kyria', wfxc[1,:,1], rho_vo_b) * (1/nkpts)
+
+                a_aa += lib.einsum('kxria,lxrjb->kialjb', w_vo_aa, rho_ov_a)
+                b_aa += lib.einsum('kxria,lxrjb->kialjb', w_vo_aa, rho_vo_a)
+
+                a_ab += lib.einsum('kxria,lxrjb->kialjb', w_vo_ab, rho_ov_b)
+                b_ab += lib.einsum('kxria,lxrjb->kialjb', w_vo_ab, rho_vo_b)
+
+                a_bb += lib.einsum('kxria,lxrjb->kialjb', w_vo_bb, rho_ov_b)
+                b_bb += lib.einsum('kxria,lxrjb->kialjb', w_vo_bb, rho_vo_b)
+
+        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):
+                rho0a = make_rho(0, ao, mask, xctype)
+                rho0b = make_rho(1, ao, mask, xctype)
+                rho = (rho0a, rho0b)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc * weight
+
+                rho_o_a = lib.einsum('kxrp,kpi->kxri', ao, orbo_a)
+                rho_o_b = lib.einsum('kxrp,kpi->kxri', ao, orbo_b)
+                rho_v_a = lib.einsum('kxrp,kpi->kxri', ao, orbv_a)
+                rho_v_b = lib.einsum('kxrp,kpi->kxri', ao, orbv_b)
+                rho_ov_a = np.einsum('kxri,kra->kxria', rho_o_a, rho_v_a[:,0])
+                rho_ov_b = np.einsum('kxri,kra->kxria', rho_o_b, rho_v_b[:,0])
+                rho_ov_a[:,1:4] += np.einsum('ri,xra->xria', rho_o_a[:,0], rho_v_a[:,1:4])
+                rho_ov_b[:,1:4] += np.einsum('ri,xra->xria', rho_o_b[:,0], rho_v_b[:,1:4])
+                tau_ov_a = np.einsum('kxri,kxra->kria', rho_o_a[:,1:4], rho_v_a[:,1:4]) * .5
+                tau_ov_b = np.einsum('kxri,kxra->kria', rho_o_b[:,1:4], rho_v_b[:,1:4]) * .5
+                rho_ov_a = np.vstack([rho_ov_a, tau_ov_a[:,np.newaxis]])
+                rho_ov_b = np.vstack([rho_ov_b, tau_ov_b[:,np.newaxis]])
+                rho_vo_a = rho_ov_a.conj()[cmap]
+                rho_vo_b = rho_ov_b.conj()[cmap]
+                w_vo_aa = np.einsum('xyr,kxria->kyria', wfxc[0,:,0], rho_vo_a) * (1/nkpts)
+                w_vo_ab = np.einsum('xyr,kxria->kyria', wfxc[0,:,1], rho_vo_a) * (1/nkpts)
+                w_vo_bb = np.einsum('xyr,kxria->kyria', wfxc[1,:,1], rho_vo_b) * (1/nkpts)
+
+                a_aa += lib.einsum('kxria,lxrjb->kilajb', w_vo_aa, rho_ov_a)
+                b_aa += lib.einsum('kxria,lxrjb->kilajb', w_vo_aa, rho_vo_a)
+
+                a_ab += lib.einsum('kxria,lxrjb->kilajb', w_vo_ab, rho_ov_b)
+                b_ab += lib.einsum('kxria,lxrjb->kilajb', w_vo_ab, rho_vo_b)
+
+                a_bb += lib.einsum('kxria,lxrjb->kilajb', w_vo_bb, rho_ov_b)
+                b_bb += lib.einsum('kxria,lxrjb->kilajb', w_vo_bb, rho_vo_b)
+    else:
+        add_hf_(a, b)
+
+    return a, b
+
 class TDA(KTDBase):
-    conv_tol = getattr(__config__, 'pbc_tdscf_rhf_TDA_conv_tol', 1e-6)
 
-    def gen_vind(self, mf, kshift):
-        '''Compute Ax'''
-        kconserv = self.kconserv[kshift]
+    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.
+        '''
+        kconserv = get_kconserv_ria(mf.cell, mf.kpts)[kshift]
+
         mo_coeff = mf.mo_coeff
         mo_occ = mf.mo_occ
         nkpts = len(mo_occ[0])
         nao, nmo = mo_coeff[0][0].shape
-        occidxa = [numpy.where(mo_occ[0][k]> 0)[0] for k in range(nkpts)]
-        occidxb = [numpy.where(mo_occ[1][k]> 0)[0] for k in range(nkpts)]
-        viridxa = [numpy.where(mo_occ[0][k]==0)[0] for k in range(nkpts)]
-        viridxb = [numpy.where(mo_occ[1][k]==0)[0] for k in range(nkpts)]
+        occidxa = [mo_occ[0][k]> 0 for k in range(nkpts)]
+        occidxb = [mo_occ[1][k]> 0 for k in range(nkpts)]
+        viridxa = [mo_occ[0][k]==0 for k in range(nkpts)]
+        viridxb = [mo_occ[1][k]==0 for k in range(nkpts)]
         orboa = [mo_coeff[0][k][:,occidxa[k]] for k in range(nkpts)]
         orbob = [mo_coeff[1][k][:,occidxb[k]] for k in range(nkpts)]
-        orbva = [mo_coeff[0][kconserv[k]][:,viridxa[kconserv[k]]] for k in range(nkpts)]
-        orbvb = [mo_coeff[1][kconserv[k]][:,viridxb[kconserv[k]]] for k in range(nkpts)]
+        orbva = [mo_coeff[0][k][:,viridxa[k]] for k in range(nkpts)]
+        orbvb = [mo_coeff[1][k][:,viridxb[k]] for k in range(nkpts)]
+        dtype = np.result_type(*mo_coeff[0])
 
         moe = scf.addons.mo_energy_with_exxdiv_none(mf)
         e_ia_a = _get_e_ia(moe[0], mo_occ[0], kconserv)
         e_ia_b = _get_e_ia(moe[1], mo_occ[1], kconserv)
-        hdiag = numpy.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])
+        hdiag = np.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])
 
         mem_now = lib.current_memory()[0]
         max_memory = max(2000, self.max_memory*.8-mem_now)
@@ -60,32 +269,31 @@ class TDA(KTDBase):
         def vind(zs):
             nz = len(zs)
             zs = [_unpack(z, mo_occ, kconserv) for z in zs]
-            dmov = numpy.empty((2,nz,nkpts,nao,nao), dtype=numpy.complex128)
+            dms = np.empty((2,nz,nkpts,nao,nao), dtype=dtype)
             for i in range(nz):
                 dm1a, dm1b = zs[i]
-                for k in range(nkpts):
-                    dmov[0,i,k] = reduce(numpy.dot, (orboa[k], dm1a[k], orbva[k].conj().T))
-                    dmov[1,i,k] = reduce(numpy.dot, (orbob[k], dm1b[k], orbvb[k].conj().T))
+                for k, kp in enumerate(kconserv):
+                    dms[0,i,kp] = lib.einsum('ov,pv,qo->pq', dm1a[k], orbva[kp], orboa[k].conj())
+                    dms[1,i,kp] = lib.einsum('ov,pv,qo->pq', dm1b[k], orbvb[kp], orbob[k].conj())
 
             with lib.temporary_env(mf, exxdiv=None):
-                dmov = dmov.reshape(2,nz,nkpts,nao,nao)
-                v1ao = vresp(dmov, kshift)
+                v1ao = vresp(dms, kshift)
                 v1ao = v1ao.reshape(2,nz,nkpts,nao,nao)
 
             v1s = []
             for i in range(nz):
                 dm1a, dm1b = zs[i]
-                v1as = []
-                v1bs = []
-                for k in range(nkpts):
-                    v1a = reduce(numpy.dot, (orboa[k].conj().T, v1ao[0,i,k], orbva[k]))
-                    v1b = reduce(numpy.dot, (orbob[k].conj().T, v1ao[1,i,k], orbvb[k]))
+                v1as = [None] * nkpts
+                v1bs = [None] * nkpts
+                for k, kp in enumerate(kconserv):
+                    v1a = lib.einsum('pq,qo,pv->ov', v1ao[0,i,kp], orboa[k], orbva[kp].conj())
+                    v1b = lib.einsum('pq,qo,pv->ov', v1ao[1,i,kp], orbob[k], orbvb[kp].conj())
                     v1a += e_ia_a[k] * dm1a[k]
                     v1b += e_ia_b[k] * dm1b[k]
-                    v1as.append(v1a.ravel())
-                    v1bs.append(v1b.ravel())
-                v1s.append( numpy.concatenate(v1as + v1bs) )
-            return lib.asarray(v1s).reshape(nz,-1)
+                    v1as[k] = v1a.ravel()
+                    v1bs[k] = v1b.ravel()
+                v1s.append( np.concatenate(v1as + v1bs) )
+            return np.stack(v1s)
 
         return vind, hdiag
 
@@ -94,18 +302,18 @@ 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_a = _get_e_ia(mo_energy[0], mo_occ[0], kconserv)
         e_ia_b = _get_e_ia(mo_energy[1], mo_occ[1], kconserv)
-        e_ia = numpy.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])
+        e_ia = np.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])
 
         nov = e_ia.size
         nstates = min(nstates, nov)
-        e_threshold = numpy.sort(e_ia)[nstates-1]
+        e_threshold = np.partition(e_ia, nstates-1)[nstates-1]
         e_threshold += self.deg_eia_thresh
 
-        idx = numpy.where(e_ia <= e_threshold)[0]
-        x0 = numpy.zeros((idx.size, nov))
+        idx = np.where(e_ia <= e_threshold)[0]
+        x0 = np.zeros((idx.size, nov))
         for i, j in enumerate(idx):
             x0[i, j] = 1  # Koopmans' excitations
         return x0
@@ -123,18 +331,16 @@ class TDA(KTDBase):
         mo_occ = mf.mo_occ
 
         def pickeig(w, v, nroots, envs):
-            idx = numpy.where(w > self.positive_eig_threshold)[0]
+            idx = np.where(w > self.positive_eig_threshold)[0]
             return w[idx], v[:,idx], idx
 
         log = logger.Logger(self.stdout, self.verbose)
-        precision = self.cell.precision * 1e-2
-        hermi = 1
 
         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)
@@ -144,15 +350,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, hermi, 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 )
             self.xy.append( [(_unpack(xi, mo_occ, kconserv),  # (X_alpha, X_beta)
@@ -165,30 +366,34 @@ class TDA(KTDBase):
 CIS = KTDA = TDA
 
 
-class TDHF(TDA):
-    def gen_vind(self, mf, kshift):
-        kconserv = self.kconserv[kshift]
+class TDHF(KTDBase):
+
+    get_ab = TDA.get_ab
+
+    def gen_vind(self, mf, kshift=0):
+        assert kshift == 0
+
         mo_coeff = mf.mo_coeff
         mo_occ = mf.mo_occ
         nkpts = len(mo_occ[0])
         nao, nmo = mo_coeff[0][0].shape
-        occidxa = [numpy.where(mo_occ[0][k]> 0)[0] for k in range(nkpts)]
-        occidxb = [numpy.where(mo_occ[1][k]> 0)[0] for k in range(nkpts)]
-        viridxa = [numpy.where(mo_occ[0][k]==0)[0] for k in range(nkpts)]
-        viridxb = [numpy.where(mo_occ[1][k]==0)[0] for k in range(nkpts)]
+        occidxa = [mo_occ[0][k]> 0 for k in range(nkpts)]
+        occidxb = [mo_occ[1][k]> 0 for k in range(nkpts)]
+        viridxa = [mo_occ[0][k]==0 for k in range(nkpts)]
+        viridxb = [mo_occ[1][k]==0 for k in range(nkpts)]
         orboa = [mo_coeff[0][k][:,occidxa[k]] for k in range(nkpts)]
         orbob = [mo_coeff[1][k][:,occidxb[k]] for k in range(nkpts)]
-        orbva = [mo_coeff[0][kconserv[k]][:,viridxa[kconserv[k]]] for k in range(nkpts)]
-        orbvb = [mo_coeff[1][kconserv[k]][:,viridxb[kconserv[k]]] for k in range(nkpts)]
+        orbva = [mo_coeff[0][k][:,viridxa[k]] for k in range(nkpts)]
+        orbvb = [mo_coeff[1][k][:,viridxb[k]] for k in range(nkpts)]
+        dtype = np.result_type(*mo_coeff[0])
 
+        kconserv = np.arange(nkpts)
         moe = scf.addons.mo_energy_with_exxdiv_none(mf)
         e_ia_a = _get_e_ia(moe[0], mo_occ[0], kconserv)
         e_ia_b = _get_e_ia(moe[1], mo_occ[1], kconserv)
-        hdiag = numpy.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])
-        hdiag = numpy.hstack((hdiag, -hdiag))
-        tot_x_a = sum(x.size for x in e_ia_a)
-        tot_x_b = sum(x.size for x in e_ia_b)
-        tot_x = tot_x_a + tot_x_b
+        hdiag = np.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])
+        tot_x = hdiag.size
+        hdiag = np.hstack((hdiag, -hdiag))
 
         mem_now = lib.current_memory()[0]
         max_memory = max(2000, self.max_memory*.8-mem_now)
@@ -198,53 +403,52 @@ class TDHF(TDA):
             nz = len(xys)
             x1s = [_unpack(x[:tot_x], mo_occ, kconserv) for x in xys]
             y1s = [_unpack(x[tot_x:], mo_occ, kconserv) for x in xys]
-            dmov = numpy.empty((2,nz,nkpts,nao,nao), dtype=numpy.complex128)
+            dms = np.empty((2,nz,nkpts,nao,nao), dtype=dtype)
             for i in range(nz):
                 xa, xb = x1s[i]
                 ya, yb = y1s[i]
                 for k in range(nkpts):
-                    dmx = reduce(numpy.dot, (orboa[k], xa[k]  , orbva[k].conj().T))
-                    dmy = reduce(numpy.dot, (orbva[k], ya[k].T, orboa[k].conj().T))
-                    dmov[0,i,k] = dmx + dmy  # AX + BY
-                    dmx = reduce(numpy.dot, (orbob[k], xb[k]  , orbvb[k].conj().T))
-                    dmy = reduce(numpy.dot, (orbvb[k], yb[k].T, orbob[k].conj().T))
-                    dmov[1,i,k] = dmx + dmy  # AX + BY
+                    dms[0,i,k]  = lib.einsum('ov,pv,qo->pq', xa[k], orbva[k], orboa[k].conj())
+                    dms[1,i,k]  = lib.einsum('ov,pv,qo->pq', xb[k], orbvb[k], orbob[k].conj())
+                    dms[0,i,k] += lib.einsum('ov,qv,po->pq', ya[k], orbva[k].conj(), orboa[k])
+                    dms[1,i,k] += lib.einsum('ov,qv,po->pq', yb[k], orbvb[k].conj(), orbob[k])
 
             with lib.temporary_env(mf, exxdiv=None):
-                dmov = dmov.reshape(2,nz,nkpts,nao,nao)
-                v1ao = vresp(dmov, kshift)
+                v1ao = vresp(dms, kshift)
                 v1ao = v1ao.reshape(2,nz,nkpts,nao,nao)
 
             v1s = []
             for i in range(nz):
                 xa, xb = x1s[i]
                 ya, yb = y1s[i]
-                v1xsa = []
-                v1xsb = []
-                v1ysa = []
-                v1ysb = []
+                v1xsa = [0] * nkpts
+                v1xsb = [0] * nkpts
+                v1ysa = [0] * nkpts
+                v1ysb = [0] * nkpts
                 for k in range(nkpts):
-                    v1xa = reduce(numpy.dot, (orboa[k].conj().T, v1ao[0,i,k], orbva[k]))
-                    v1xb = reduce(numpy.dot, (orbob[k].conj().T, v1ao[1,i,k], orbvb[k]))
-                    v1ya = reduce(numpy.dot, (orbva[k].conj().T, v1ao[0,i,k], orboa[k])).T
-                    v1yb = reduce(numpy.dot, (orbvb[k].conj().T, v1ao[1,i,k], orbob[k])).T
-                    v1xa+= e_ia_a[k] * xa[k]
-                    v1xb+= e_ia_b[k] * xb[k]
-                    v1ya+= e_ia_a[k] * ya[k]
-                    v1yb+= e_ia_b[k] * yb[k]
-                    v1xsa.append(v1xa.ravel())
-                    v1xsb.append(v1xb.ravel())
-                    v1ysa.append(-v1ya.ravel())
-                    v1ysb.append(-v1yb.ravel())
-                v1s.append( numpy.concatenate(v1xsa + v1xsb + v1ysa + v1ysb) )
-            return numpy.hstack(v1s).reshape(nz,-1)
+                    v1xa = lib.einsum('pq,qo,pv->ov', v1ao[0,i,k], orboa[k], orbva[k].conj())
+                    v1xb = lib.einsum('pq,qo,pv->ov', v1ao[1,i,k], orbob[k], orbvb[k].conj())
+                    v1ya = lib.einsum('pq,po,qv->ov', v1ao[0,i,k], orboa[k].conj(), orbva[k])
+                    v1yb = lib.einsum('pq,po,qv->ov', v1ao[1,i,k], orbob[k].conj(), orbvb[k])
+                    v1xa += e_ia_a[k] * xa[k]
+                    v1xb += e_ia_b[k] * xb[k]
+                    v1ya += e_ia_a[k] * ya[k]
+                    v1yb += e_ia_b[k] * yb[k]
+                    v1xsa[k] += v1xa.ravel()
+                    v1xsb[k] += v1xb.ravel()
+                    v1ysa[k] -= v1ya.ravel()
+                    v1ysb[k] -= v1yb.ravel()
+                v1s.append( np.concatenate(v1xsa + v1xsb + v1ysa + v1ysb) )
+            return np.stack(v1s)
 
         return vind, hdiag
 
     def init_guess(self, mf, kshift, nstates=None, wfnsym=None):
         x0 = TDA.init_guess(self, mf, kshift, nstates)
-        y0 = numpy.zeros_like(x0)
-        return numpy.asarray(numpy.block([[x0, y0], [y0, x0.conj()]]))
+        y0 = np.zeros_like(x0)
+        return np.hstack([x0, y0])
+
+    get_precond = uhf.TDHF.get_precond
 
     def kernel(self, x0=None):
         '''TDHF diagonalization with non-Hermitian eigenvalue solver
@@ -258,23 +462,23 @@ class TDHF(TDA):
         mf = self._scf
         mo_occ = mf.mo_occ
 
-        real_system = (gamma_point(self._scf.kpts) and
-                       self._scf.mo_coeff[0][0].dtype == numpy.double)
+        real_system = (is_gamma_point(self._scf.kpts) and
+                       self._scf.mo_coeff[0][0].dtype == np.double)
+
+        if any(k != 0 for k in self.kshift_lst):
+            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) &
+            realidx = np.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
                                   (w.real > self.positive_eig_threshold))[0]
             return lib.linalg_helper._eigs_cmplx2real(w, v, realidx, real_system)
 
-        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)
@@ -284,15 +488,10 @@ class TDHF(TDA):
             else:
                 x0k = x0[i]
 
-            converged, w, 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, 0, log),
-                                        verbose=self.verbose)
+            converged, w, 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 )
 
             e = []
@@ -300,13 +499,14 @@ class TDHF(TDA):
             for i, z in enumerate(x1):
                 xs, ys = z.reshape(2,-1)
                 norm = lib.norm(xs)**2 - lib.norm(ys)**2
-                if norm > 0:
-                    norm = 1/numpy.sqrt(norm)
-                    xs *= norm
-                    ys *= norm
-                    e.append(w[i])
-                    xy.append((_unpack(xs, mo_occ, kconserv), _unpack(ys, mo_occ, kconserv)))
-            self.e.append( numpy.array(e) )
+                if norm < 0:
+                    log.warn('TDDFT amplitudes |X| smaller than |Y|')
+                norm = abs(norm)**-.5
+                xs *= norm
+                ys *= norm
+                e.append(w[i])
+                xy.append((_unpack(xs, mo_occ, kconserv), _unpack(ys, mo_occ, kconserv)))
+            self.e.append( np.array(e) )
             self.xy.append( xy )
 
         log.timer(self.__class__.__name__, *cpu0)
@@ -318,17 +518,17 @@ def _unpack(vo, mo_occ, kconserv):
     za = []
     zb = []
     p1 = 0
+    no_a_kpts = [np.count_nonzero(occ) for occ in mo_occ[0]]
+    no_b_kpts = [np.count_nonzero(occ) for occ in mo_occ[1]]
     for k, occ in enumerate(mo_occ[0]):
-        no = numpy.count_nonzero(occ > 0)
-        no1 = numpy.count_nonzero(mo_occ[0][kconserv[k]] > 0)
-        nv = occ.size - no1
+        no = no_a_kpts[k]
+        nv = occ.size - no_a_kpts[kconserv[k]]
         p0, p1 = p1, p1 + no * nv
         za.append(vo[p0:p1].reshape(no,nv))
 
     for k, occ in enumerate(mo_occ[1]):
-        no = numpy.count_nonzero(occ > 0)
-        no1 = numpy.count_nonzero(mo_occ[1][kconserv[k]] > 0)
-        nv = occ.size - no1
+        no = no_b_kpts[k]
+        nv = occ.size - no_b_kpts[kconserv[k]]
         p0, p1 = p1, p1 + no * nv
         zb.append(vo[p0:p1].reshape(no,nv))
     return za, zb
@@ -336,55 +536,3 @@ def _unpack(vo, mo_occ, kconserv):
 
 scf.kuhf.KUHF.TDA  = lib.class_as_method(KTDA)
 scf.kuhf.KUHF.TDHF = lib.class_as_method(KTDHF)
-
-
-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 = [37]*3
-    cell.build()
-    mf = scf.KUHF(cell, cell.make_kpts([2,1,1])).set(exxdiv=None)
-#    mf.with_df = df.DF(cell, cell.make_kpts([2,1,1]))
-#    mf.with_df.auxbasis = 'weigend'
-#    mf.with_df._cderi = 'eri3d-df.h5'
-#    mf.with_df.build(with_j3c=False)
-    mf.run()
-
-    td = TDA(mf)
-    td.verbose = 5
-    td.nstates = 5
-    print(td.kernel()[0][0] * 27.2114)
-
-    td = TDHF(mf)
-    td.verbose = 5
-    td.nstates = 5
-    print(td.kernel()[0][0] * 27.2114)
-
-    cell.spin = 2
-    mf = scf.KUHF(cell, cell.make_kpts([2,1,1])).set(exxdiv=None)
-    mf.run()
-
-    td = TDA(mf)
-    td.verbose = 5
-    td.nstates = 5
-    print(td.kernel()[0][0] * 27.2114)
-
-    td = TDHF(mf)
-    td.verbose = 5
-    td.nstates = 5
-    print(td.kernel()[0][0] * 27.2114)