use solvent model for unrestricted

bytedance · Feb 6, 2024 · f64fcf3 · f64fcf3
1 parent eaf83a4
commit f64fcf3
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 3 deletions.
diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py
@@ -42,7 +42,7 @@ def pcm_for_scf(mf, solvent_obj=None, dm=None):
 # Inject PCM to SCF, TODO: add it to other methods later
 from gpu4pyscf import scf
 scf.hf.RHF.PCM = pcm_for_scf
-
+scf.uhf.UHF.PCM = pcm_for_scf
 # TABLE II,  J. Chem. Phys. 122, 194110 (2005)
 XI = {
     6: 4.84566077868,
@@ -313,7 +313,8 @@ def _get_vind(self, dms):
 
         nao = dms.shape[-1]
         dms = dms.reshape(-1,nao,nao)
-
+        if dms.shape[0] == 2:
+            dms = (dms[0] + dms[1]).reshape(-1,nao,nao)
         K = self._intermediates['K']
         R = self._intermediates['R']
         v_grids_e = self._get_v(dms)

diff --git a/gpu4pyscf/solvent/smd.py b/gpu4pyscf/solvent/smd.py
@@ -36,7 +36,7 @@ def smd_for_scf(mf, solvent_obj=None, dm=None):
 # Inject PCM to SCF, TODO: add it to other methods later
 from gpu4pyscf import scf
 scf.hf.RHF.SMD = smd_for_scf
-
+scf.uhf.UHF.SMD = smd_for_scf
 hartree2kcal = 627.5
 # see https://pubs.acs.org/doi/epdf/10.1021/jp810292n
 

diff --git a/gpu4pyscf/solvent/tests/test_pcm.py b/gpu4pyscf/solvent/tests/test_pcm.py
@@ -83,6 +83,16 @@ def test_SSVPE(self):
         print(f"Energy error in SS(V)PE: {numpy.abs(e_tot - -74.9689577454)}")
         assert numpy.abs(e_tot - -74.9689577454) < 1e-9
 
+    def test_uhf(self):
+        cm = pcm.PCM(mol)
+        cm.eps = epsilon
+        cm.verbose = 0
+        cm.lebedev_order = 29
+        cm.method = 'IEF-PCM'
+        mf = scf.UHF(mol).PCM(cm)
+        e_tot = mf.kernel()
+        assert numpy.abs(e_tot - -74.96901113434953) < 1e-9
+
 if __name__ == "__main__":
     print("Full Tests for PCMs")
     unittest.main()
diff --git a/gpu4pyscf/solvent/tests/test_smd.py b/gpu4pyscf/solvent/tests/test_smd.py
@@ -68,6 +68,13 @@ def test_smd_water(self):
         e_tot = mf.kernel()
         assert numpy.abs(e_tot - -76.0756052903) < 2e-4
 
+    def test_uhf(self):
+        mf = scf.UHF(mol)
+        mf = mf.SMD()
+        mf.with_solvent.solvent = 'water'
+        mf.with_solvent.sasa_ng = 590
+        e_tot = mf.kernel()
+        assert numpy.abs(e_tot - -76.07550951172617) < 2e-4
     # TODO: add more test for other molecules
 
 if __name__ == "__main__":