From b715ea9ce1a8a7437d10bd800154c97d0aacee2d Mon Sep 17 00:00:00 2001
From: henryw7 <60555480+henryw7@users.noreply.github.com>
Date: Tue, 3 Sep 2024 13:28:54 -0700
Subject: [PATCH] Add chkfile support for pysisyphus (#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 | 5 ++-
gpu4pyscf/scf/hf.py | 61 ++++++++++++++++++++++++++-------
gpu4pyscf/scf/tests/test_rhf.py | 15 +++++++-
gpu4pyscf/scf/tests/test_uhf.py | 18 +++++++++-
gpu4pyscf/scf/uhf.py | 6 ++--
5 files changed, 86 insertions(+), 19 deletions(-)
diff --git a/gpu4pyscf/grad/rhf.py b/gpu4pyscf/grad/rhf.py
index 44c81c62..f500b317 100644
--- a/gpu4pyscf/grad/rhf.py
+++ b/gpu4pyscf/grad/rhf.py
@@ -696,4 +696,7 @@ def extra_force(self, atom_id, envs):
'''
return 0
-Grad = Gradients
\ No newline at end of file
+Grad = Gradients
+
+from gpu4pyscf import scf
+scf.hf.RHF.Gradients = lib.class_as_method(Gradients)
diff --git a/gpu4pyscf/scf/hf.py b/gpu4pyscf/scf/hf.py
index b0cd7e24..15fe67af 100644
--- a/gpu4pyscf/scf/hf.py
+++ b/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 .
-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
diff --git a/gpu4pyscf/scf/tests/test_rhf.py b/gpu4pyscf/scf/tests/test_rhf.py
index 39f909d4..b05dc01a 100644
--- a/gpu4pyscf/scf/tests/test_rhf.py
+++ b/gpu4pyscf/scf/tests/test_rhf.py
@@ -16,6 +16,7 @@
# along with this program. If not, see .
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
diff --git a/gpu4pyscf/scf/tests/test_uhf.py b/gpu4pyscf/scf/tests/test_uhf.py
index 3e2fd310..76d4329f 100644
--- a/gpu4pyscf/scf/tests/test_uhf.py
+++ b/gpu4pyscf/scf/tests/test_uhf.py
@@ -16,6 +16,7 @@
# along with this program. If not, see .
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
diff --git a/gpu4pyscf/scf/uhf.py b/gpu4pyscf/scf/uhf.py
index 7928a601..5d376502 100644
--- a/gpu4pyscf/scf/uhf.py
+++ b/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