Add chkfile support for pysisyphus (pyscf#203)

* Add chkfile support for pysisyphus

* Turn off chkfile by default

* Accidentally pushed default settings

* Fix small errors according to lint

* Add comment about tight thresholds

gpu4pyscf/grad/rhf.py

@@ -696,4 +696,7 @@ def extra_force(self, atom_id, envs):
         '''
         return 0
 
-Grad = Gradients
+Grad = Gradients
+
+from gpu4pyscf import scf
+scf.hf.RHF.Gradients = lib.class_as_method(Gradients)

gpu4pyscf/scf/hf.py

@@ -15,17 +15,18 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-import time
 import copy
 import ctypes
-import contextlib
+import sys
+import h5py
 import numpy as np
 import cupy
 import scipy.linalg
 from functools import reduce
 from pyscf import gto
 from pyscf import lib as pyscf_lib
-from pyscf.scf import hf, jk, _vhf
+from pyscf.scf import hf, _vhf
+from pyscf.scf import chkfile
 from gpu4pyscf import lib
 from gpu4pyscf.lib.cupy_helper import (eigh, load_library, tag_array,
                                        return_cupy_array, cond)
@@ -402,7 +403,7 @@ def _kernel(mf, conv_tol=1e-10, conv_tol_grad=None,
 
     h1e = cupy.asarray(mf.get_hcore(mol))
     s1e = cupy.asarray(mf.get_ovlp(mol))
-    
+
     vhf = mf.get_veff(mol, dm)
     e_tot = mf.energy_tot(dm, h1e, vhf)
     logger.info(mf, 'init E= %.15g', e_tot)
@@ -419,7 +420,12 @@ def _kernel(mf, conv_tol=1e-10, conv_tol_grad=None,
         _, mf_diis.Corth = mf.eig(fock, s1e)
     else:
         mf_diis = None
-
+
+    if dump_chk and mf.chkfile:
+        # Explicit overwrite the mol object in chkfile
+        # Note in pbc.scf, mf.mol == mf.cell, cell is saved under key "mol"
+        chkfile.save_mol(mol, mf.chkfile)
+
     for cycle in range(mf.max_cycle):
         t0 = log.init_timer()
         dm_last = dm
@@ -441,6 +447,15 @@ def _kernel(mf, conv_tol=1e-10, conv_tol_grad=None,
         t1 = log.timer_debug1('total', *t0)
         logger.info(mf, 'cycle= %d E= %.15g  delta_E= %4.3g  |ddm|= %4.3g',
                     cycle+1, e_tot, e_tot-last_hf_e, norm_ddm)
+
+        if dump_chk:
+            local_variables = locals()
+            for key in local_variables:
+                value = local_variables[key]
+                if (type(value) is cupy.ndarray):
+                    local_variables[key] = cupy.asnumpy(value)
+            mf.dump_chk(local_variables)
+
         e_diff = abs(e_tot-last_hf_e)
         norm_gorb = cupy.linalg.norm(mf.get_grad(mo_coeff, mo_occ, f))
         if(e_diff < conv_tol and norm_gorb < conv_tol_grad):
@@ -549,8 +564,8 @@ def __call__(self, mol_or_geom, **kwargs):
             dm0 = None
             if cupy.array_equal(self._last_mol_fp, mol.ao_loc):
                 dm0 = self.make_rdm1()
-            else:
-                raise NotImplementedError
+            elif self.chkfile and h5py.is_hdf5(self.chkfile):
+                dm0 = self.from_chk(self.chkfile)
         self.mo_coeff = None  # To avoid last mo_coeff being used by SOSCF
         e_tot = self.kernel(dm0=dm0, **kwargs)
         self._last_mol_fp = mol.ao_loc
@@ -581,7 +596,30 @@ class SCF(pyscf_lib.StreamObject):
     _keys               = hf.SCF._keys
 
     # methods
-    __init__                 = hf.SCF.__init__
+    def __init__(self, mol):
+        if not mol._built:
+            sys.stderr.write('Warning: %s must be initialized before calling SCF.\n'
+                             'Initialize %s in %s\n' % (mol, mol, self))
+            mol.build()
+        self.mol = mol
+        self.verbose = mol.verbose
+        self.max_memory = mol.max_memory
+        self.stdout = mol.stdout
+
+        # The chkfile part is different from pyscf, we turn off chkfile by default.
+        self.chkfile = None
+
+##################################################
+# don't modify the following attributes, they are not input options
+        self.mo_energy = None
+        self.mo_coeff = None
+        self.mo_occ = None
+        self.e_tot = 0
+        self.converged = False
+        self.scf_summary = {}
+
+        self._opt = {None: None}
+        self._eri = None # Note: self._eri requires large amount of memory
 
     def check_sanity(self):
         s1e = self.get_ovlp()
@@ -601,14 +639,14 @@ def check_sanity(self):
     get_fock                 = hf.SCF.get_fock
     get_occ                  = hf.SCF.get_occ
     get_grad                 = hf.SCF.get_grad
-    dump_chk                 = NotImplemented
+    dump_chk                 = hf.SCF.dump_chk
     init_guess_by_minao      = hf.SCF.init_guess_by_minao
     init_guess_by_atom       = hf.SCF.init_guess_by_atom
     init_guess_by_huckel     = hf.SCF.init_guess_by_huckel
     init_guess_by_mod_huckel = hf.SCF.init_guess_by_mod_huckel
     init_guess_by_1e         = hf.SCF.init_guess_by_1e
-    init_guess_by_chkfile    = NotImplemented
-    from_chk                 = NotImplemented
+    init_guess_by_chkfile    = hf.SCF.init_guess_by_chkfile
+    from_chk                 = hf.SCF.from_chk
     get_init_guess           = hf.SCF.get_init_guess
     make_rdm1                = hf.SCF.make_rdm1
     make_rdm2                = hf.SCF.make_rdm2
@@ -682,7 +720,6 @@ class RHF(SCF):
     get_init_guess = return_cupy_array(hf.RHF.get_init_guess)
     init_direct_scf = NotImplemented
     make_rdm2 = NotImplemented
-    dump_chk = NotImplemented
     newton = NotImplemented
     x2c = x2c1e = sfx2c1e = NotImplemented
     to_rhf = NotImplemented

gpu4pyscf/scf/tests/test_rhf.py

@@ -16,6 +16,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import unittest
+import tempfile
 import numpy as np
 import cupy
 import pyscf
@@ -264,11 +265,23 @@ def test_rhf_d4(self):
         e_ref = -151.09634038447925
         assert np.abs(e_tot - e_ref) < 1e-5
 
+    def test_chkfile(self):
+        ftmp = tempfile.NamedTemporaryFile(dir = pyscf.lib.param.TMPDIR)
+        mf = scf.RHF(mol)
+        mf.chkfile = ftmp.name
+        mf.kernel()
+        dm_stored = mf.make_rdm1(mf.mo_coeff, mf.mo_occ)
+        dm_stored = cupy.asnumpy(dm_stored)
+
+        mf_copy = scf.RHF(mol)
+        mf_copy.chkfile = ftmp.name
+        dm_loaded = mf_copy.init_guess_by_chkfile()
+        assert np.allclose(dm_stored, dm_loaded, atol = 1e-14) # Since we reload the MO coefficients, the density matrix should be identical up to numerical noise.
+
     # TODO:
     #test analyze
     #test mulliken_pop
     #test mulliken_meta
-    #test chkfile
     #test stability
     #test newton
     #test x2c

gpu4pyscf/scf/tests/test_uhf.py

@@ -16,6 +16,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import unittest
+import tempfile
 import numpy as np
 import cupy
 import pyscf
@@ -272,13 +273,28 @@ def test_uhf_d4(self):
         print('pyscf - qchem ', e_tot - e_ref)
         assert np.abs(e_tot - e_ref) < 1e-5
     '''
+
+    def test_chkfile(self):
+        ftmp = tempfile.NamedTemporaryFile(dir = pyscf.lib.param.TMPDIR)
+        mf = scf.UHF(mol)
+        mf.chkfile = ftmp.name
+        mf.kernel()
+        dma_stored, dmb_stored = mf.make_rdm1(mf.mo_coeff, mf.mo_occ)
+        dma_stored, dmb_stored = cupy.asnumpy(dma_stored), cupy.asnumpy(dmb_stored)
+
+        mf_copy = scf.UHF(mol)
+        mf_copy.chkfile = ftmp.name
+        dma_loaded, dmb_loaded = mf_copy.init_guess_by_chkfile()
+        assert np.allclose(dma_stored, dma_loaded, atol = 1e-14) # Since we reload the MO coefficients, the density matrix should be identical up to numerical noise.
+        assert np.allclose(dmb_stored, dmb_loaded, atol = 1e-14)
+        assert not np.allclose(dma_stored, dmb_loaded, atol = 1e-1) # Just to make sure alpha and beta electron are different in the test system
+
     # TODO:
     #test analyze
     #test mulliken_pop
     #test mulliken_spin_pop
     #test mulliken_meta
     #test mulliken_meta_spin
-    #test chkfile
     #test stability
     #test newton
     #test x2c

gpu4pyscf/scf/uhf.py

@@ -20,11 +20,10 @@
 import cupy
 from pyscf.scf import uhf
 from pyscf import lib as pyscf_lib
-from gpu4pyscf.scf.hf import _get_jk, eigh, damping, level_shift, _kernel
+from gpu4pyscf.scf.hf import eigh, damping, level_shift
 from gpu4pyscf.scf import hf
 from gpu4pyscf.lib import logger
 from gpu4pyscf.lib.cupy_helper import tag_array
-from gpu4pyscf import lib
 from gpu4pyscf.scf import diis
 
 def make_rdm1(mo_coeff, mo_occ, **kwargs):
@@ -199,7 +198,7 @@ def get_grad(self, mo_coeff, mo_occ, fock=None):
     init_guess_by_huckel     = uhf.UHF.init_guess_by_huckel
     init_guess_by_mod_huckel = uhf.UHF.init_guess_by_mod_huckel
     init_guess_by_1e         = uhf.UHF.init_guess_by_1e
-    init_guess_by_chkfile    = NotImplemented
+    init_guess_by_chkfile    = uhf.UHF.init_guess_by_chkfile
 
     analyze            = NotImplemented
     mulliken_pop       = NotImplemented
@@ -224,7 +223,6 @@ def get_grad(self, mo_coeff, mo_occ, fock=None):
     energy_elec = energy_elec
 
     make_rdm2 = NotImplemented
-    dump_chk = NotImplemented
     newton = NotImplemented
     x2c = x2c1e = sfx2c1e = NotImplemented
     to_rhf = NotImplemented