Add analyze method for PBC-DFT methods (#559)

* Add analyze method for PBC-DFT methods

* improve logging msg

Author: sunqm

gpu4pyscf/pbc/scf/hf.py

@@ -336,6 +336,31 @@ def to_cpu(self):
         utils.to_cpu(self, out=mf)
         return mf
 
+    def analyze(self, verbose=logger.DEBUG, with_meta_lowdin=True, **kwargs):
+        '''Analyze the given SCF object:  print orbital energies, occupancies;
+        print orbital coefficients; Mulliken population analysis; Diople moment.
+        '''
+        from pyscf.scf.hf import mulliken_meta, mulliken_pop, MO_BASE
+        log = logger.new_logger(self, verbose)
+        cell = self.cell
+        mo_energy = self.mo_energy.get()
+        mo_occ = self.mo_occ.get()
+
+        if log.verbose >= logger.NOTE:
+            self.dump_scf_summary(log)
+            log.note('**** MO energy ****')
+            for i, c in enumerate(mo_occ):
+                log.note('MO #%-3d energy= %-18.15g occ= %g', i+MO_BASE, mo_energy[i], c)
+
+        s = self.get_ovlp().get()
+        dm = self.make_rdm1().get()
+        if with_meta_lowdin:
+            pop = mulliken_meta(cell, dm, s=s, verbose=log)
+        else:
+            pop = mulliken_pop(cell, dm, s=s, verbose=log)
+        dip = None
+        return pop, dip
+
 def normalize_dm_(mf, dm, s1e=None):
     '''
     Force density matrices integrated to the correct number of electrons.

gpu4pyscf/pbc/scf/khf.py

@@ -411,10 +411,11 @@ class KRHF(KSCF):
     check_sanity = pbchf.SCF.check_sanity
 
     def get_init_guess(self, cell=None, key='minao', s1e=None):
+        kpts = self.kpts
         if s1e is None:
-            s1e = self.get_ovlp(cell)
+            s1e = self.get_ovlp(cell, kpts)
         dm = mol_hf.SCF.get_init_guess(self, cell, key)
-        nkpts = len(self.kpts)
+        nkpts = len(kpts)
         if dm.ndim == 2:
             # dm[nao,nao] at gamma point -> dm_kpts[nkpts,nao,nao]
             dm = cp.repeat(dm[None,:,:], nkpts, axis=0)
@@ -446,3 +447,33 @@ def to_cpu(self):
         mf = khf_cpu.KRHF(self.cell)
         utils.to_cpu(self, out=mf)
         return mf
+
+    def analyze(self, verbose=None, **kwargs):
+        '''Analyze the given SCF object:  print orbital energies, occupancies;
+        print orbital coefficients; Mulliken population analysis; Dipole moment
+        '''
+        from pyscf.pbc.scf.khf import mulliken_meta
+        if verbose is None:
+            verbose = self.verbose
+        log = logger.new_logger(self, verbose)
+        mo_energy = self.mo_energy.get()
+        mo_occ = self.mo_occ.get()
+        cell = self.cell
+        kpts = self.kpts
+        if log.verbose >= logger.NOTE:
+            self.dump_scf_summary(log)
+            log.note('**** MO energy ****')
+            log.note('k-point                    nocc    HOMO/AU         LUMO/AU')
+            for k, kpt in enumerate(cell.get_scaled_kpts(kpts)):
+                nocc = np.count_nonzero(mo_occ[k])
+                homo = mo_energy[k,nocc-1]
+                lumo = mo_energy[k,nocc  ]
+                log.note('%2d (%6.3f %6.3f %6.3f) %2d   %15.9f %15.9f',
+                         k, kpt[0], kpt[1], kpt[2], nocc, homo, lumo)
+
+        log.note('**** Population analysis for atoms in the reference cell ****')
+        s = self.get_ovlp(kpts=kpts).get()
+        dm = self.make_rdm1().get()
+        pop, chg = mulliken_meta(cell, dm, kpts=kpts, s=s, verbose=verbose)
+        dip = None
+        return (pop, chg), dip

gpu4pyscf/pbc/scf/kuhf.py

@@ -223,7 +223,6 @@ def __init__(self, cell, kpts=None, exxdiv='ewald'):
     get_occ = get_occ
     energy_elec = energy_elec
     get_rho = khf.get_rho
-    analyze = NotImplemented
     canonicalize = canonicalize
 
     def get_init_guess(self, cell=None, key='minao', s1e=None):
@@ -323,3 +322,37 @@ def to_cpu(self):
         mf = kuhf_cpu.KUHF(self.cell)
         utils.to_cpu(self, out=mf)
         return mf
+
+    def analyze(self, verbose=None, **kwargs):
+        '''Analyze the given SCF object:  print orbital energies, occupancies;
+        print orbital coefficients; Mulliken population analysis; Dipole moment
+        '''
+        from pyscf.pbc.scf.kuhf import mulliken_meta
+        if verbose is None:
+            verbose = self.verbose
+        log = logger.new_logger(self, verbose)
+        mo_energy = self.mo_energy.get()
+        mo_occ = self.mo_occ.get()
+        cell = self.cell
+        kpts = self.kpts
+        if log.verbose >= logger.NOTE:
+            self.dump_scf_summary(log)
+            log.note('**** MO energy ****')
+            log.note('                           alpha                               | beta')
+            log.note('k-point                    nocc    HOMO/AU         LUMO/AU     | nocc    HOMO/AU         LUMO/AU')
+            for k, kpt in enumerate(cell.get_scaled_kpts(kpts)):
+                nocca = np.count_nonzero(mo_occ[0,k])
+                noccb = np.count_nonzero(mo_occ[1,k])
+                homoa = mo_energy[0,k,nocca-1]
+                homob = mo_energy[1,k,noccb-1]
+                lumoa = mo_energy[0,k,nocca  ]
+                lumob = mo_energy[1,k,noccb  ]
+                log.note('%2d (%6.3f %6.3f %6.3f) %2d   %15.9f %15.9f |%2d   %15.9f %15.9f',
+                         k, kpt[0], kpt[1], kpt[2], nocca, homoa, lumoa, noccb, homob, lumob)
+
+        log.note('**** Population analysis for atoms in the reference cell ****')
+        s = self.get_ovlp(kpts=kpts).get()
+        dm = self.make_rdm1().get()
+        pop, chg = mulliken_meta(cell, dm, kpts=kpts, s=s, verbose=verbose)
+        dip = None
+        return (pop, chg), dip

gpu4pyscf/pbc/scf/tests/test_pbc_scf_hf.py

@@ -27,17 +27,17 @@ class KnownValues(unittest.TestCase):
     def setUpClass(cls):
         L = 4
         n = 21
-        cell = pbcgto.Cell()
-        cell.build(unit = 'B',
-                   verbose = 7,
-                   output = '/dev/null',
-                   a = ((L,0,0),(0,L,0),(0,0,L)),
-                   mesh = [n,n,n],
-                   atom = [['He', (L/2.-.5,L/2.,L/2.-.5)],
-                           ['He', (L/2.   ,L/2.,L/2.+.5)]],
-                   basis = { 'He': [[0, (0.8, 1.0)],
-                                    [0, (1.0, 1.0)],
-                                    [0, (1.2, 1.0)]]})
+        cell = pbcgto.M(
+            unit = 'B',
+            verbose = 7,
+            output = '/dev/null',
+            a = ((L,0,0),(0,L,0),(0,0,L)),
+            mesh = [n,n,n],
+            atom = [['He', (L/2.-.5,L/2.,L/2.-.5)],
+                    ['He', (L/2.   ,L/2.,L/2.+.5)]],
+            basis = { 'He': [[0, (0.8, 1.0)],
+                             [0, (1.0, 1.0)],
+                             [0, (1.2, 1.0)]]})
         cls.cell = cell
 
     @classmethod
@@ -49,8 +49,13 @@ def test_rhf_exx_ewald(self):
         mf = scf.RHF(cell, exxdiv='ewald').run()
         self.assertAlmostEqual(mf.e_tot, -4.3511582284698633, 7)
         self.assertTrue(mf.mo_coeff.dtype == np.double)
+        pop = mf.analyze()[0][0]
+        self.assertAlmostEqual(lib.fp(pop), 0.011047586674983092, 5)
+
         kmf = scf.KRHF(cell, [[0,0,0]], exxdiv='ewald').run()
         self.assertAlmostEqual(mf.e_tot, kmf.e_tot, 8)
+        pop = kmf.analyze()[0][0]
+        self.assertAlmostEqual(lib.fp(pop), 0.011047586674983092, 5)
 
         # test bands
         np.random.seed(1)

gpu4pyscf/pbc/scf/tests/test_pbc_scf_uhf.py

@@ -52,6 +52,8 @@ def test_kuhf_bands(self):
         kmf_cpu = kmf.to_cpu().run()
         self.assertAlmostEqual(kmf.e_tot, kmf_cpu.e_tot, 8)
         self.assertAlmostEqual(kmf.e_tot, -4.021029656152094, 8)
+        pop = kmf.analyze()[0][0]
+        self.assertAlmostEqual(lib.fp(pop), 0.02897067698093582, 5)
 
         np.random.seed(1)
         kpts_bands = np.random.random((1,3))
@@ -64,6 +66,8 @@ def test_uhf_bands(self):
         mf_cpu = mf.to_cpu().run()
         self.assertAlmostEqual(mf.e_tot, mf_cpu.e_tot, 8)
         self.assertAlmostEqual(mf.e_tot, -3.9546467710639632, 7)
+        pop = mf.analyze()[0][0]
+        print(lib.fp(pop), -0.04691820429296646)
 
         np.random.seed(1)
         kpts_bands = np.random.random((1,3))

gpu4pyscf/pbc/scf/uhf.py

@@ -113,7 +113,6 @@ def get_init_guess(self, cell=None, key='minao', s1e=None):
     energy_elec = mol_uhf.UHF.energy_elec
     _finalize = mol_uhf.UHF._finalize
     get_rho = pbchf.get_rho
-    analyze = NotImplemented
     mulliken_pop = NotImplemented
     mulliken_meta = NotImplemented
     mulliken_meta_spin = NotImplemented
@@ -135,3 +134,33 @@ def to_cpu(self):
         mf = uhf_cpu.UHF(self.cell)
         utils.to_cpu(self, out=mf)
         return mf
+
+    def analyze(self, verbose=logger.DEBUG, with_meta_lowdin=True, **kwargs):
+        '''Analyze the given SCF object:  print orbital energies, occupancies;
+        print orbital coefficients; Mulliken population analysis; Diople moment.
+        '''
+        from pyscf.scf.hf import mulliken_meta, mulliken_pop, MO_BASE
+        log = logger.new_logger(self, verbose)
+        cell = self.cell
+        mo_energy = self.mo_energy.get()
+        mo_occ = self.mo_occ.get()
+
+        if log.verbose >= logger.NOTE:
+            self.dump_scf_summary(log)
+            mo_e_a, mo_e_b = mo_energy
+            mo_occ_a, mo_occ_b = mo_occ
+            nmo = len(mo_occ_a)
+            log.note('**** MO energy ****')
+            log.note('                             alpha | beta                alpha | beta')
+            for i in range(nmo):
+                log.note('MO #%-3d energy= %-18.15g | %-18.15g occ= %g | %g',
+                         i+MO_BASE, mo_e_a[i], mo_e_b[i], mo_occ_a[i], mo_occ_b[i])
+
+        s = self.get_ovlp().get()
+        dm = self.make_rdm1().get()
+        if with_meta_lowdin:
+            pop = mulliken_meta(cell, dm, s=s, verbose=log)
+        else:
+            pop = mulliken_pop(cell, dm, s=s, verbose=log)
+        dip = None
+        return pop, dip

gpu4pyscf/scf/hf.py

@@ -699,8 +699,8 @@ def eig(self, fock, s):
     init_direct_scf          = NotImplemented
     get_jk                   = _get_jk
     get_veff                 = NotImplemented
-    mulliken_meta            = hf_cpu.SCF.mulliken_meta
-    pop                      = hf_cpu.SCF.pop
+    mulliken_meta = pop      = NotImplemented
+    mulliken_pop             = NotImplemented
     _is_mem_enough           = NotImplemented
     density_fit              = NotImplemented
     newton                   = NotImplemented
@@ -717,8 +717,7 @@ def eig(self, fock, s):
     to_gks                   = NotImplemented
     to_ks                    = NotImplemented
     canonicalize             = NotImplemented
-    mulliken_pop             = NotImplemented
-    mulliken_meta            = NotImplemented
+    dump_scf_summary         = hf_cpu.dump_scf_summary
 
     smearing = smearing
 
@@ -744,13 +743,13 @@ def dip_moment(self, mol=None, dm=None, unit='Debye', origin=None,
                    verbose=logger.NOTE):
         if mol is None: mol = self.mol
         if dm is None: dm = self.make_rdm1()
-        return hf_cpu.dip_moment(mol, dm.get(), unit, origin, verbose)
+        return hf_cpu.dip_moment(mol, cupy.asnumpy(dm), unit, origin, verbose)
 
     def quad_moment(self, mol=None, dm=None, unit='DebyeAngstrom', origin=None,
                     verbose=logger.NOTE):
         if mol is None: mol = self.mol
         if dm is None: dm = self.make_rdm1()
-        return hf_cpu.quad_moment(mol, dm.get(), unit, origin, verbose)
+        return hf_cpu.quad_moment(mol, cupy.asnumpy(dm), unit, origin, verbose)
 
     def remove_soscf(self):
         lib.logger.warn('remove_soscf has no effect in current version')

gpu4pyscf/scf/rohf.py

@@ -97,7 +97,6 @@ class ROHF(hf.RHF):
     to_uks = NotImplemented
     to_gks = NotImplemented
     to_ks = NotImplemented
-    analyze = NotImplemented
     stability = NotImplemented
     mulliken_pop = NotImplemented
     mulliken_meta = NotImplemented

gpu4pyscf/scf/tests/test_rhf.py

@@ -20,7 +20,9 @@
 from pyscf import lib
 from gpu4pyscf import scf
 
-mol = pyscf.M(
+def setUpModule():
+    global mol, mol1
+    mol = pyscf.M(
     atom='''
 C  -0.65830719,  0.61123287, -0.00800148
 C   0.73685281,  0.61123287, -0.00800148
@@ -32,7 +34,7 @@
     output = '/dev/null'
 )
 
-mol1 = pyscf.M(
+    mol1 = pyscf.M(
     atom='''
 C  -1.20806619, -0.34108413, -0.00755148
 C   1.28636081, -0.34128013, -0.00668648
@@ -249,8 +251,17 @@ def test_rhf_d3(self):
         mf = scf.RHF(mol)
         mf.disp = 'd3bj'
         e_tot = mf.kernel()
+
+        #mf_ref = mol.RHF()
+        #mf_ref.disp = 'd3bj'
+        #e_ref = mf_ref.kernel()
+        #chg_ref = mf_ref.analyze()[0][1]
         e_ref = -151.1150439066
-        assert np.abs(e_tot - e_ref) < 1e-5
+        assert abs(e_tot - e_ref) < 1e-8
+
+        chg = mf.analyze()[0][1]
+        #assert abs(chg - chg_ref).max() < 1e-5
+        self.assertAlmostEqual(lib.fp(chg), -0.003225958206417059, 5)
 
     def test_rhf_d4(self):
         mf = scf.RHF(mol)
@@ -299,12 +310,15 @@ def test_rohf(self):
             O 0.00000000 0.00000000 0.60539399
             H 0.00000000 0.93467313 -1.18217476
             H 0.00000000 -0.93467313 -1.18217476''',
-            charge=1, spin=1, unit='B')
+            charge=1, spin=1, unit='B', verbose=5, output='/dev/null')
         mf = mol.ROHF().to_gpu().run()
         self.assertAlmostEqual(mf.e_tot, -107.61304925181142, 8)
         ref = mf.to_cpu().run()
         self.assertAlmostEqual(mf.e_tot, ref.e_tot, 8)
 
+        chg = mf.analyze()[0][1]
+        self.assertAlmostEqual(lib.fp(chg), -0.0705568646397904, 5)
+
     # TODO:
     #test analyze
     #test mulliken_pop

gpu4pyscf/scf/tests/test_uhf.py

@@ -22,7 +22,9 @@
 from gpu4pyscf.lib.multi_gpu import num_devices
 import pytest
 
-mol = pyscf.M(
+def setUpModule():
+    global mol, mol1
+    mol = pyscf.M(
     atom='''
 C  -0.65830719,  0.61123287, -0.00800148
 C   0.73685281,  0.61123287, -0.00800148
@@ -35,7 +37,7 @@
     output = '/dev/null'
 )
 
-mol1 = pyscf.M(
+    mol1 = pyscf.M(
     atom='''
 C  -1.20806619, -0.34108413, -0.00755148
 C   1.28636081, -0.34128013, -0.00668648
@@ -258,9 +260,12 @@ def test_uhf_scf_fast(self):
         mol1 = mol.copy()
         mol1.basis = 'sto3g'
         mol1.build(False, False)
-        e_tot = mol1.UHF().to_gpu().kernel()
+        mf = mol1.UHF().to_gpu()
+        e_tot = mf.kernel()
         e_ref = -148.8650361770461
         assert np.abs(e_tot - e_ref) < 1e-5
+        chg = mf.analyze()[0][1]
+        self.assertAlmostEqual(lib.fp(chg), 0.022191785654920748, 5)
 
     @pytest.mark.slow
     def test_uhf_d3bj(self):