Add HF1e

gpu4pyscf/scf/__init__.py

@@ -20,13 +20,19 @@
 from . import dispersion
 
 def HF(mol, *args):
-    if mol.nelectron == 1 or mol.spin == 0:
-        return RHF(mol, *args)
+    if mol.spin == 0:
+        return hf.RHF(mol, *args)
+    elif mol.nelectron == 1:
+        from .rohf import HF1e
+        return HF1e(mol, *args)
     else:
         return UHF(mol, *args)
 
 def RHF(mol, *args):
     if mol.spin == 0:
         return hf.RHF(mol, *args)
+    elif mol.nelectron == 1:
+        from .rohf import HF1e
+        return HF1e(mol, *args)
     else:
         return ROHF(mol, *args)

gpu4pyscf/scf/rohf.py

@@ -18,6 +18,7 @@
 from pyscf.scf import rohf as rohf_cpu
 from gpu4pyscf.scf import hf, uhf
 from gpu4pyscf.lib.cupy_helper import tag_array, contract
+from gpu4pyscf.lib import logger
 
 
 def get_roothaan_fock(focka_fockb, dma_dmb, s):
@@ -215,3 +216,25 @@ def get_grad(self, mo_coeff, mo_occ, fock):
     def newton(self):
         from gpu4pyscf.scf.soscf import newton
         return newton(self)
+
+
+class HF1e(ROHF):
+    def kernel(self, *args):
+        logger.info(self, '\n')
+        logger.info(self, '******** 1 electron system ********')
+        h = self.get_hcore()
+        s = self.get_ovlp()
+        self.mo_energy, self.mo_coeff = self.eig(h, s)
+        self.mo_occ = self.get_occ(self.mo_energy, self.mo_coeff)
+        self.e_tot = self.mo_energy[0].real.get() + self.mol.energy_nuc()
+        if self.chkfile:
+            self.dump_chk({
+                'e_tot': self.e_tot,
+                'mo_energy': self.mo_energy,
+                'mo_coeff': self.mo_coeff,
+                'mo_occ': self.mo_occ
+            })
+        self.converged = True
+        self._finalize()
+        return self.e_tot
+    scf = kernel

gpu4pyscf/scf/tests/test_rhf.py

@@ -349,6 +349,11 @@ def test_initial_guess_tag(self):
         assert hasattr(dm, 'mo_coeff') and dm.mo_coeff.ndim == 2
         assert abs(cupy.einsum('ij,ji->', dm, s).get() - 24) < 1e-6
 
+    def test_1e(self):
+        mol = pyscf.M(atom='H', basis='ccpvdz', spin=1)
+        mf = mol.RHF().to_gpu().run()
+        self.assertAlmostEqual(mf.e_tot, -0.499278403419583, 9)
+
     # TODO:
     #test analyze
     #test mulliken_pop