From 177fb055c90a21430b42b11832bb6c294ad46317 Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Fri, 27 Dec 2024 22:05:17 -0800 Subject: [PATCH] Add pickle serialization (#294) * Add pickle serialization (fix #267) * syntax error * Fix DFHF serialization tests --------- Co-authored-by: Qiming Sun --- gpu4pyscf/df/df.py | 6 +++++- gpu4pyscf/df/df_jk.py | 3 +-- gpu4pyscf/df/tests/test_df_rhf.py | 12 +++++++++++- gpu4pyscf/dft/numint.py | 3 +++ gpu4pyscf/dft/rks.py | 10 ++++++++-- gpu4pyscf/dft/tests/test_rks.py | 10 +++++++++- gpu4pyscf/scf/hf.py | 3 +++ 7 files changed, 40 insertions(+), 7 deletions(-) diff --git a/gpu4pyscf/df/df.py b/gpu4pyscf/df/df.py index 52b0ecf8..48e0e8e1 100644 --- a/gpu4pyscf/df/df.py +++ b/gpu4pyscf/df/df.py @@ -36,7 +36,7 @@ class DF(lib.StreamObject): from gpu4pyscf.lib.utils import to_gpu, device - _keys = {'intopt', 'mol', 'auxmol', 'use_gpu_memory'} + _keys = {'intopt', 'nao', 'naux', 'cd_low', 'mol', 'auxmol', 'use_gpu_memory'} def __init__(self, mol, auxbasis=None): self.mol = mol @@ -52,8 +52,12 @@ def __init__(self, mol, auxbasis=None): self.naux = None self.cd_low = None self._cderi = None + self._vjopt = None self._rsh_df = {} + __getstate__, __setstate__ = lib.generate_pickle_methods( + excludes=('cd_low', 'intopt', '_cderi', '_vjopt')) + @property def auxbasis(self): return self._auxbasis diff --git a/gpu4pyscf/df/df_jk.py b/gpu4pyscf/df/df_jk.py index d2083f41..5561cf9c 100644 --- a/gpu4pyscf/df/df_jk.py +++ b/gpu4pyscf/df/df_jk.py @@ -122,7 +122,7 @@ class _DFHF: to_gpu = utils.to_gpu device = utils.device __name_mixin__ = 'DF' - _keys = {'rhoj', 'rhok', 'disp', 'screen_tol'} + _keys = {'rhoj', 'rhok', 'disp', 'screen_tol', 'with_df', 'only_dfj'} def __init__(self, mf, dfobj, only_dfj): self.__dict__.update(mf.__dict__) @@ -132,7 +132,6 @@ def __init__(self, mf, dfobj, only_dfj): self.direct_scf = False self.with_df = dfobj self.only_dfj = only_dfj - self._keys = mf._keys.union(['with_df', 'only_dfj']) def undo_df(self): '''Remove the DFHF Mixin''' diff --git a/gpu4pyscf/df/tests/test_df_rhf.py b/gpu4pyscf/df/tests/test_df_rhf.py index e724015a..c2f3caa9 100644 --- a/gpu4pyscf/df/tests/test_df_rhf.py +++ b/gpu4pyscf/df/tests/test_df_rhf.py @@ -13,12 +13,17 @@ # limitations under the License. import unittest +import pickle import numpy as np import pyscf from pyscf import scf as cpu_scf from pyscf.df import df_jk as cpu_df_jk from gpu4pyscf.df import df_jk as gpu_df_jk from gpu4pyscf import scf as gpu_scf +try: + import cloudpickle +except ImportError: + cloudpickle = None atom = ''' O 0.0000000000 -0.0000000000 0.1174000000 @@ -48,12 +53,17 @@ class KnownValues(unittest.TestCase): ''' def test_rhf(self): print('------- RHF -----------------') - mf = gpu_scf.RHF(mol_sph).density_fit(auxbasis='def2-tzvpp-jkfit') + mf = mol_sph.RHF().density_fit(auxbasis='def2-tzvpp-jkfit').to_gpu() e_tot = mf.kernel() e_qchem = -76.0624582299 print(f'diff from qchem {e_tot - e_qchem}') assert np.abs(e_tot - e_qchem) < 1e-5 + # test serialization + if cloudpickle is not None: + mf1 = pickle.loads(cloudpickle.dumps(mf)) + assert mf1.e_tot == e_tot + def test_cart(self): print('------- RHF Cart -----------------') mf = gpu_scf.RHF(mol_cart).density_fit(auxbasis='def2-tzvpp-jkfit') diff --git a/gpu4pyscf/dft/numint.py b/gpu4pyscf/dft/numint.py index 85662146..a40c976f 100644 --- a/gpu4pyscf/dft/numint.py +++ b/gpu4pyscf/dft/numint.py @@ -1792,6 +1792,9 @@ class NumInt(lib.StreamObject, LibXCMixin): screen_index = None xcfuns = None # can be multiple xc functionals + __getstate__, __setstate__ = lib.generate_pickle_methods( + excludes=('gdftopt',)) + def build(self, mol, coords): self.gdftopt = _GDFTOpt.from_mol(mol) self.grid_blksize = None diff --git a/gpu4pyscf/dft/rks.py b/gpu4pyscf/dft/rks.py index d512caa5..496abfa3 100644 --- a/gpu4pyscf/dft/rks.py +++ b/gpu4pyscf/dft/rks.py @@ -13,9 +13,9 @@ # limitations under the License. # modified by Xiaojie Wu (wxj6000@gmail.com) + import cupy from pyscf.dft import rks - from gpu4pyscf.lib import logger from gpu4pyscf.dft import numint, gen_grid from gpu4pyscf.scf import hf @@ -257,6 +257,7 @@ def __init__(self, xc='LDA,VWN'): ################################################## # don't modify the following attributes, they are not input options self._numint = numint.NumInt() + @property def omega(self): return self._numint.omega @@ -291,8 +292,13 @@ def reset(self, mol=None): hf.SCF.reset(self, mol) self.grids.reset(mol) self.nlcgrids.reset(mol) - self.cphf_grids.reset(mol) self._numint.reset() + # The cphf_grids attribute is not available in the PySCF CPU version. + # In PySCF's to_gpu() function, this attribute is not properly + # initialized. mol of the KS object must be used for initialization. + if mol is None: + mol = self.mol + self.cphf_grids.reset(mol) return self def nuc_grad_method(self): diff --git a/gpu4pyscf/dft/tests/test_rks.py b/gpu4pyscf/dft/tests/test_rks.py index d1bf278d..4bae05ca 100644 --- a/gpu4pyscf/dft/tests/test_rks.py +++ b/gpu4pyscf/dft/tests/test_rks.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pickle import numpy as np import unittest import pyscf @@ -64,11 +65,18 @@ class KnownValues(unittest.TestCase): ''' def test_rks_lda(self): print('------- LDA ----------------') - e_tot = run_dft("LDA, vwn5", mol_sph) + mf = mol_sph.RKS(xc='LDA,vwn5').to_gpu() + mf.grids.level = grids_level + mf.nlcgrids.level = nlcgrids_level + e_tot = mf.kernel() e_ref = -75.9046410402 print('| CPU - GPU |:', e_tot - e_ref) assert np.abs(e_tot - e_ref) < 1e-5 + # test serialization + mf1 = pickle.loads(pickle.dumps(mf)) + assert mf1.e_tot == e_tot + def test_rks_pbe(self): print('------- PBE ----------------') e_tot = run_dft('PBE', mol_sph) diff --git a/gpu4pyscf/scf/hf.py b/gpu4pyscf/scf/hf.py index 28d76b57..3a0497ff 100644 --- a/gpu4pyscf/scf/hf.py +++ b/gpu4pyscf/scf/hf.py @@ -391,6 +391,9 @@ def __init__(self, mol): self._opt_gpu = {None: None} self._eri = None # Note: self._eri requires large amount of memory + __getstate__, __setstate__ = pyscf_lib.generate_pickle_methods( + excludes=('_opt_gpu', '_eri', '_numint')) + def check_sanity(self): s1e = self.get_ovlp() if isinstance(s1e, cupy.ndarray) and s1e.ndim == 2: