Population analysis for KSCF (pyscf#2216)

* add mulliken_meta for kscf

* bug fix for j-build in get_jk

* bug fix for k-build in get_jk

* clean up mulliken_meta

* update k2gamma to_supercell_ao_integrals

* update pop test for k2gamma

* fix flake8

* update ref pop/chg value in kscf tests

* update comment

* Update the docstring of mulliken_meta

* Update docstring of mulliken_meta

* Update kuhf.py

---------

Co-authored-by: hongzhouye <>
Co-authored-by: Qiming Sun <[email protected]>

pyscf/pbc/scf/ghf.py

@@ -117,7 +117,7 @@ def __init__(self, cell, kpt=np.zeros(3),
     _finalize = mol_ghf.GHF._finalize
     analyze = lib.invalid_method('analyze')
     mulliken_pop = lib.invalid_method('mulliken_pop')
-    mulliken_meta = lib.invalid_method('mulliken_meta')
+    mulliken_meta = mol_ghf.GHF.mulliken_meta
     spin_square = mol_ghf.GHF.spin_square
     stability = mol_ghf.GHF.stability
 

pyscf/pbc/scf/kghf.py

@@ -136,38 +136,41 @@ def get_occ(mf, mo_energy_kpts=None, mo_coeff_kpts=None):
 
     return mo_occ_kpts
 
-def mulliken_meta(cell, dm_ao_kpts, verbose=logger.DEBUG,
+def _make_rdm1_meta(cell, dm_ao_kpts, kpts, pre_orth_method, s):
+    from pyscf.lo import orth
+    from pyscf.pbc.tools import k2gamma
+
+    kmesh = k2gamma.kpts_to_kmesh(cell, kpts-kpts[0])
+    nkpts, nso = dm_ao_kpts.shape[:2]
+    nao = nso // 2
+    scell, phase = k2gamma.get_phase(cell, kpts, kmesh)
+    s_sc = k2gamma.to_supercell_ao_integrals(cell, kpts, s, kmesh=kmesh, force_real=False)
+    orth_coeff = orth.orth_ao(scell, 'meta_lowdin', pre_orth_method, s=s_sc)[:,:nao] # cell 0 only
+    c_inv = np.dot(orth_coeff.T.conj(), s_sc)
+    c_inv = lib.einsum('aRp,Rk->kap', c_inv.reshape(nao,nkpts,nao), phase)
+    dm_aa = lib.einsum('kap,kpq,kbq->ab', c_inv, dm_ao_kpts[:,:nao,:nao], c_inv.conj())
+    dm_bb = lib.einsum('kap,kpq,kbq->ab', c_inv, dm_ao_kpts[:,nao:,nao:], c_inv.conj())
+
+    return (dm_aa, dm_bb)
+
+def mulliken_meta(cell, dm_ao_kpts, kpts, verbose=logger.DEBUG,
                   pre_orth_method=PRE_ORTH_METHOD, s=None):
     '''A modified Mulliken population analysis, based on meta-Lowdin AOs.
-
-    Note this function only computes the Mulliken population for the gamma
-    point density matrix.
+    The results are equivalent to the corresponding supercell calculation.
     '''
-    from pyscf.lo import orth
-    if s is None:
-        s = khf.get_ovlp(cell)
     log = logger.new_logger(cell, verbose)
-    log.note('Analyze output for *gamma point*.')
-    log.info('    To include the contributions from k-points, transform to a '
-             'supercell then run the population analysis on the supercell\n'
-             '        from pyscf.pbc.tools import k2gamma\n'
-             '        k2gamma.k2gamma(mf).mulliken_meta()')
-    log.note("KGHF mulliken_meta")
-    dm_ao_gamma = dm_ao_kpts[0,:,:]
-    nso = dm_ao_gamma.shape[-1]
-    nao = nso // 2
 
-    dm_ao_gamma_aa = dm_ao_gamma[:nao,:nao]
-    dm_ao_gamma_bb = dm_ao_gamma[nao:,nao:]
-    s_gamma = s[0,:,:].real
-    s_gamma_aa = s_gamma[:nao,:nao]
-    orth_coeff = orth.orth_ao(cell, 'meta_lowdin', pre_orth_method, s=s_gamma_aa)
-    c_inv = np.dot(orth_coeff.T, s_gamma_aa)
-    dm_aa = reduce(np.dot, (c_inv, dm_ao_gamma_aa, c_inv.T.conj()))
-    dm_bb = reduce(np.dot, (c_inv, dm_ao_gamma_bb, c_inv.T.conj()))
+    if s is None:
+        s = khf.get_ovlp(None, cell=cell, kpts=kpts)
+    if s is not None:
+        if s[0].shape == dm_ao_kpts[0].shape:   # s in SO
+            nao = dm_ao_kpts[0].shape[0]//2
+            s = lib.asarray(s[:,:nao,:nao], order='C') # keep only one spin sector
+
+    dm_aa, dm_bb = _make_rdm1_meta(cell, dm_ao_kpts, kpts, pre_orth_method, s)
 
     log.note(' ** Mulliken pop alpha/beta on meta-lowdin orthogonal AOs **')
-    return mol_uhf.mulliken_pop(cell, (dm_aa,dm_bb), np.eye(orth_coeff.shape[0]), log)
+    return mol_uhf.mulliken_pop(cell, (dm_aa,dm_bb), np.eye(dm_aa.shape[0]), log)
 
 def _cast_mol_init_guess(fn):
     def fn_init_guess(mf, cell=None, kpts=None):
@@ -259,12 +262,13 @@ def get_bands(self, kpts_band, cell=None, dm_kpts=None, kpts=None):
         raise NotImplementedError
 
     @lib.with_doc(mulliken_meta.__doc__)
-    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
+    def mulliken_meta(self, cell=None, dm=None, kpts=None, verbose=logger.DEBUG,
                       pre_orth_method=PRE_ORTH_METHOD, s=None):
         if cell is None: cell = self.cell
         if dm is None: dm = self.make_rdm1()
-        if s is None: s = self.get_ovlp(cell)
-        return mulliken_meta(cell, dm, s=s, verbose=verbose,
+        if kpts is None: kpts = self.kpts
+        if s is None: s = khf.get_ovlp(self, cell, kpts)
+        return mulliken_meta(cell, dm, kpts, s=s, verbose=verbose,
                              pre_orth_method=pre_orth_method)
 
     def mulliken_pop(self):

pyscf/pbc/scf/khf.py

@@ -278,31 +278,36 @@ def analyze(mf, verbose=None, with_meta_lowdin=WITH_META_LOWDIN,
         #return mf.mulliken_pop(mf.cell, dm, s=ovlp_ao, verbose=verbose)
 
 
-def mulliken_meta(cell, dm_ao_kpts, verbose=logger.DEBUG,
+def _make_rdm1_meta(cell, dm_ao_kpts, kpts, pre_orth_method, s):
+    from pyscf.lo import orth
+    from pyscf.pbc.tools import k2gamma
+
+    kmesh = k2gamma.kpts_to_kmesh(cell, kpts-kpts[0])
+    nkpts, nao = dm_ao_kpts.shape[:2]
+    scell, phase = k2gamma.get_phase(cell, kpts, kmesh)
+    s_sc = k2gamma.to_supercell_ao_integrals(cell, kpts, s, kmesh=kmesh, force_real=False)
+    orth_coeff = orth.orth_ao(scell, 'meta_lowdin', pre_orth_method, s=s_sc)[:,:nao] # cell 0 only
+    c_inv = np.dot(orth_coeff.T.conj(), s_sc)
+    c_inv = lib.einsum('aRp,Rk->kap', c_inv.reshape(nao,nkpts,nao), phase)
+    dm = lib.einsum('kap,kpq,kbq->ab', c_inv, dm_ao_kpts, c_inv.conj())
+
+    return dm
+
+
+def mulliken_meta(cell, dm_ao_kpts, kpts, verbose=logger.DEBUG,
                   pre_orth_method=PRE_ORTH_METHOD, s=None):
     '''A modified Mulliken population analysis, based on meta-Lowdin AOs.
-
-    Note this function only computes the Mulliken population for the gamma
-    point density matrix.
+    The results are equivalent to the corresponding supercell calculation.
     '''
-    from pyscf.lo import orth
-    if s is None:
-        s = get_ovlp(cell)
     log = logger.new_logger(cell, verbose)
-    log.note('Analyze output for *gamma point*')
-    log.info('    To include the contributions from k-points, transform to a '
-             'supercell then run the population analysis on the supercell\n'
-             '        from pyscf.pbc.tools import k2gamma\n'
-             '        k2gamma.k2gamma(mf).mulliken_meta()')
-    log.note("KRHF mulliken_meta")
-    dm_ao_gamma = dm_ao_kpts[0,:,:].real
-    s_gamma = s[0,:,:].real
-    orth_coeff = orth.orth_ao(cell, 'meta_lowdin', pre_orth_method, s=s_gamma)
-    c_inv = np.dot(orth_coeff.T, s_gamma)
-    dm = reduce(np.dot, (c_inv, dm_ao_gamma, c_inv.T.conj()))
+
+    if s is None:
+        s = get_ovlp(None, cell=cell, kpts=kpts)
+
+    dm = _make_rdm1_meta(cell, dm_ao_kpts, kpts, pre_orth_method, s)
 
     log.note(' ** Mulliken pop on meta-lowdin orthogonal AOs **')
-    return mol_hf.mulliken_pop(cell, dm, np.eye(orth_coeff.shape[0]), log)
+    return mol_hf.mulliken_pop(cell, dm, np.eye(dm.shape[0]), log)
 
 
 def canonicalize(mf, mo_coeff_kpts, mo_occ_kpts, fock=None):
@@ -633,7 +638,7 @@ def dump_chk(self, envs_or_file):
     def analyze(mf, verbose=None, with_meta_lowdin=WITH_META_LOWDIN, **kwargs):
         raise NotImplementedError
 
-    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
+    def mulliken_meta(self, cell=None, dm=None, kpts=None, verbose=logger.DEBUG,
                       pre_orth_method=PRE_ORTH_METHOD, s=None):
         raise NotImplementedError
 
@@ -738,12 +743,13 @@ def get_init_guess(self, cell=None, key='minao', s1e=None):
         return dm_kpts
 
     @lib.with_doc(mulliken_meta.__doc__)
-    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
+    def mulliken_meta(self, cell=None, dm=None, kpts=None, verbose=logger.DEBUG,
                       pre_orth_method=PRE_ORTH_METHOD, s=None):
         if cell is None: cell = self.cell
         if dm is None: dm = self.make_rdm1()
+        if kpts is None: kpts = self.kpts
         if s is None: s = self.get_ovlp(cell)
-        return mulliken_meta(cell, dm, s=s, verbose=verbose,
+        return mulliken_meta(cell, dm, kpts, s=s, verbose=verbose,
                              pre_orth_method=pre_orth_method)
 
     def nuc_grad_method(self):

pyscf/pbc/scf/krohf.py

@@ -211,15 +211,12 @@ def get_occ(mf, mo_energy_kpts=None, mo_coeff_kpts=None):
 
 
 @lib.with_doc(khf.mulliken_meta.__doc__)
-def mulliken_meta(cell, dm_ao_kpts, verbose=logger.DEBUG,
+def mulliken_meta(cell, dm_ao_kpts, kpts, verbose=logger.DEBUG,
                   pre_orth_method=PRE_ORTH_METHOD, s=None):
     '''Mulliken population analysis, based on meta-Lowdin AOs.
-
-    Note this function only computes the Mulliken population for the gamma
-    point density matrix.
     '''
     dm = dm_ao_kpts[0] + dm_ao_kpts[1]
-    return khf.mulliken_meta(cell, dm, verbose, pre_orth_method, s)
+    return khf.mulliken_meta(cell, dm, kpts, verbose, pre_orth_method, s)
 
 
 def canonicalize(mf, mo_coeff_kpts, mo_occ_kpts, fock=None):
@@ -362,12 +359,13 @@ def init_guess_by_chkfile(self, chk=None, project=True, kpts=None):
         if kpts is None: kpts = self.kpts
         return init_guess_by_chkfile(self.cell, chk, project, kpts)
 
-    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
+    def mulliken_meta(self, cell=None, dm=None, kpts=None, verbose=logger.DEBUG,
                       pre_orth_method=PRE_ORTH_METHOD, s=None):
         if cell is None: cell = self.cell
         if dm is None: dm = self.make_rdm1()
+        if kpts is None: kpts = self.kpts
         if s is None: s = self.get_ovlp(cell)
-        return mulliken_meta(cell, dm, s=s, verbose=verbose,
+        return mulliken_meta(cell, dm, kpts, s=s, verbose=verbose,
                              pre_orth_method=pre_orth_method)
 
     def stability(self,

pyscf/pbc/scf/kuhf.py

@@ -216,32 +216,52 @@ def energy_elec(mf, dm_kpts=None, h1e_kpts=None, vhf_kpts=None):
     return (e1+e_coul).real, e_coul.real
 
 
-def mulliken_meta(cell, dm_ao_kpts, verbose=logger.DEBUG,
+def _make_rdm1_meta(cell, dm_ao_kpts, kpts, pre_orth_method, s):
+    from pyscf.lo import orth
+    from pyscf.pbc.tools import k2gamma
+
+    kmesh = k2gamma.kpts_to_kmesh(cell, kpts-kpts[0])
+    nkpts, nao = dm_ao_kpts[0].shape[:2]
+    scell, phase = k2gamma.get_phase(cell, kpts, kmesh)
+    s_sc = k2gamma.to_supercell_ao_integrals(cell, kpts, s, kmesh=kmesh, force_real=False)
+    orth_coeff = orth.orth_ao(scell, 'meta_lowdin', pre_orth_method, s=s_sc)[:,:nao] # cell 0 only
+    c_inv = np.dot(orth_coeff.T.conj(), s_sc)
+    c_inv = lib.einsum('aRp,Rk->kap', c_inv.reshape(nao,nkpts,nao), phase)
+    dm_a = lib.einsum('kap,kpq,kbq->ab', c_inv, dm_ao_kpts[0], c_inv.conj())
+    dm_b = lib.einsum('kap,kpq,kbq->ab', c_inv, dm_ao_kpts[1], c_inv.conj())
+
+    return (dm_a, dm_b)
+
+
+def mulliken_meta(cell, dm_ao_kpts, kpts, verbose=logger.DEBUG,
                   pre_orth_method=PRE_ORTH_METHOD, s=None):
     '''A modified Mulliken population analysis, based on meta-Lowdin AOs.
-
-    Note this function only computes the Mulliken population for the gamma
-    point density matrix.
+    The results are equivalent to the corresponding supercell calculation.
     '''
-    from pyscf.lo import orth
+    log = logger.new_logger(cell, verbose)
+
     if s is None:
-        s = khf.get_ovlp(cell)
+        s = khf.get_ovlp(None, cell=cell, kpts=kpts)
+
+    dm_a, dm_b = _make_rdm1_meta(cell, dm_ao_kpts, kpts, pre_orth_method, s)
+
+    log.note(' ** Mulliken pop alpha/beta on meta-lowdin orthogonal AOs **')
+    return mol_uhf.mulliken_pop(cell, (dm_a,dm_b), np.eye(dm_a.shape[0]), log)
+
+
+def mulliken_meta_spin(cell, dm_ao_kpts, kpts, verbose=logger.DEBUG,
+                       pre_orth_method=PRE_ORTH_METHOD, s=None):
+    '''A modified Mulliken population analysis, based on meta-Lowdin AOs.
+    '''
     log = logger.new_logger(cell, verbose)
-    log.note('Analyze output for *gamma point*.')
-    log.info('    To include the contributions from k-points, transform to a '
-             'supercell then run the population analysis on the supercell\n'
-             '        from pyscf.pbc.tools import k2gamma\n'
-             '        k2gamma.k2gamma(mf).mulliken_meta()')
-    log.note("KUHF mulliken_meta")
-    dm_ao_gamma = dm_ao_kpts[:,0,:,:].real
-    s_gamma = s[0,:,:].real
-    orth_coeff = orth.orth_ao(cell, 'meta_lowdin', pre_orth_method, s=s_gamma)
-    c_inv = np.dot(orth_coeff.T, s_gamma)
-    dm_a = reduce(np.dot, (c_inv, dm_ao_gamma[0], c_inv.T.conj()))
-    dm_b = reduce(np.dot, (c_inv, dm_ao_gamma[1], c_inv.T.conj()))
+
+    if s is None:
+        s = khf.get_ovlp(None, cell=cell, kpts=kpts)
+
+    dm_a, dm_b = _make_rdm1_meta(cell, dm_ao_kpts, kpts, pre_orth_method, s)
 
     log.note(' ** Mulliken pop alpha/beta on meta-lowdin orthogonal AOs **')
-    return mol_uhf.mulliken_pop(cell, (dm_a,dm_b), np.eye(orth_coeff.shape[0]), log)
+    return mol_uhf.mulliken_spin_pop(cell, (dm_a,dm_b), np.eye(dm_a.shape[0]), log)
 
 
 def canonicalize(mf, mo_coeff_kpts, mo_occ_kpts, fock=None):
@@ -509,14 +529,25 @@ def init_guess_by_chkfile(self, chk=None, project=True, kpts=None):
         return init_guess_by_chkfile(self.cell, chk, project, kpts)
 
     @lib.with_doc(mulliken_meta.__doc__)
-    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
+    def mulliken_meta(self, cell=None, dm=None, kpts=None, verbose=logger.DEBUG,
                       pre_orth_method=PRE_ORTH_METHOD, s=None):
         if cell is None: cell = self.cell
         if dm is None: dm = self.make_rdm1()
+        if kpts is None: kpts = self.kpts
         if s is None: s = self.get_ovlp(cell)
-        return mulliken_meta(cell, dm, s=s, verbose=verbose,
+        return mulliken_meta(cell, dm, kpts, s=s, verbose=verbose,
                              pre_orth_method=pre_orth_method)
 
+    @lib.with_doc(mulliken_meta_spin.__doc__)
+    def mulliken_meta_spin(self, cell=None, dm=None, kpts=None, verbose=logger.DEBUG,
+                           pre_orth_method=PRE_ORTH_METHOD, s=None):
+        if cell is None: cell = self.cell
+        if dm is None: dm = self.make_rdm1()
+        if kpts is None: kpts = self.kpts
+        if s is None: s = self.get_ovlp(cell)
+        return mulliken_meta_spin(cell, dm, kpts, s=s, verbose=verbose,
+                                  pre_orth_method=pre_orth_method)
+
     def mulliken_pop(self):
         raise NotImplementedError
 

pyscf/pbc/scf/test/test_khf.py

@@ -62,13 +62,13 @@ def tearDownModule():
 
 class KnownValues(unittest.TestCase):
     def test_analyze(self):
-        rpop, rchg = kmf.analyze()[0] # pop at gamma point
+        rpop, rchg = kmf.analyze()[0]
         upop, uchg = kumf.analyze()[0]
         gpop, gchg = kgmf.analyze()[0]
         self.assertTrue(isinstance(rpop, np.ndarray) and rpop.ndim == 1)
         self.assertAlmostEqual(abs(upop[0]+upop[1]-rpop).max(), 0, 7)
         self.assertAlmostEqual(abs(gpop[0]+gpop[1]-rpop).max(), 0, 5)
-        self.assertAlmostEqual(lib.fp(rpop), 1.697446, 5)
+        self.assertAlmostEqual(lib.fp(rpop), 1.638430, 5)
 
     def test_kpt_vs_supercell_high_cost(self):
         # For large n, agreement is always achieved