Refactor classes (pyscf#132)

* loose cupy requirements

* support third-row transition metals

* reverse grad/rhf

* fixed an issue in unit test

* reduce class inheritance from pyscf

* linter

* fixed a bug in unit test

* move high-angular momentum to CPU

* delete patch && add coverage

* flake8

* be compatible with pyscf

* replace numpy with cupy

* skip mp2 to_gpu

* refactor CCSDBase

* update unittest workflow

* install pytest cov

Author: wxj6000

.coveragerc

@@ -0,0 +1,31 @@
+# .coveragerc to control coverage.py
+[run]
+branch = True
+omit = */tests/*
+disable_warnings = include-ignored
+
+[report]
+# Regexes for lines to exclude from consideration
+exclude_lines =
+    # Have to re-enable the standard pragma
+    pragma: no cover
+
+    # Don't complain about missing debug-only code:
+    def __repr__
+
+    # Don't complain if tests don't hit defensive assertion code:
+    raise RuntimeError
+    raise NotImplementedError
+    pass
+
+    # Don't complain if non-runnable code isn't run:
+    if 0:
+    if __name__ == .__main__.:
+    if sys.version_info.*:
+    if DEBUG:
+    except ImportError:
+
+ignore_errors = True
+
+[html]
+directory = coverage_html_report

.flake8

@@ -15,7 +15,7 @@ ignore =
         E701, E731, E741, E275,
         F401, C901, W391, W503, W504, W291, W292, W293
 
-exclude = test, tests, .git, __pycache__, build, dist, __init__.py .eggs, *.egg
+exclude = test, tests, .git, __pycache__, build, dist, __init__.py .eggs, *.egg, .coveragerc
 max-line-length = 160
 
 per-file-ignores =

.github/workflows/unittest.yml

@@ -22,7 +22,7 @@ jobs:
     - name: Install dependencies
       run: |
         python3 -m pip install --upgrade pip
-        pip3 install flake8 pytest coverage gpu4pyscf-libxc-cuda11x
+        pip3 install flake8 pytest coverage gpu4pyscf-libxc-cuda11x pytest-cov
     - name: Build GPU4PySCF
       run: |
         export CUDA_HOME=/usr/local/cuda
@@ -33,4 +33,4 @@ jobs:
       run: |
         echo $GITHUB_WORKSPACE
         export PYTHONPATH="${PYTHONPATH}:$GITHUB_WORKSPACE"
-        coverage run -m pytest
+        pytest --cov=$GITHUB_WORKSPACE

.gitignore

@@ -18,6 +18,7 @@ output/
 *.log
 *.h5
 **/qchem_input.in
+.coverage
 #*
 dockerfiles/manylinux/wheelhouse
 gpu4pyscf/lib/bin

gpu4pyscf/__init__.py

@@ -5,8 +5,3 @@
 # monkey patch libxc reference due to a bug in nvcc
 from pyscf.dft import libxc
 libxc.__reference__ = 'unable to decode the reference due to https://github.com/NVIDIA/cuda-python/issues/29'
-
-from gpu4pyscf.lib.utils import patch_cpu_kernel
-from gpu4pyscf.lib.cupy_helper import tag_array
-from pyscf import lib
-lib.tag_array = tag_array

gpu4pyscf/cc/ccsd_incore.py

@@ -529,8 +529,83 @@ def _make_eris_incore(mycc, mo_coeff=None):
         cupy.get_default_memory_pool().free_all_blocks()
     return eris
 
-class CCSD(ccsd.CCSD):
-    from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
+class CCSDBase(lib.StreamObject):
+    # attributes
+    _keys                  = ccsd.CCSDBase._keys
+    max_cycle              = ccsd.CCSDBase.max_cycle
+    conv_tol               = ccsd.CCSDBase.conv_tol
+    iterative_damping      = ccsd.CCSDBase.iterative_damping
+    conv_tol_normt         = ccsd.CCSDBase.conv_tol_normt
+
+    diis                   = ccsd.CCSDBase.diis
+    diis_space             = ccsd.CCSDBase.diis_space
+    diis_file              = None
+    diis_start_cycle       = ccsd.CCSDBase.diis_start_cycle
+    diis_start_energy_diff = ccsd.CCSDBase.diis_start_energy_diff
+
+    direct                 = ccsd.CCSDBase.direct
+    async_io               = None
+    incore_complete        = ccsd.CCSDBase.incore_complete
+    cc2                    = ccsd.CCSDBase.cc2
+    callback               = None
+
+    # functions
+    __init__           = ccsd.CCSDBase.__init__
+    ecc                = ccsd.CCSDBase.ecc
+    e_tot              = ccsd.CCSDBase.e_tot
+    nocc               = ccsd.CCSDBase.nocc
+    nmo                = ccsd.CCSDBase.nmo
+    reset              = ccsd.CCSDBase.reset
+    get_nocc           = ccsd.CCSDBase.get_nocc
+    get_nmo            = ccsd.CCSDBase.get_nmo
+    get_frozen_mask    = ccsd.CCSDBase.get_frozen_mask
+    get_e_hf           = ccsd.CCSDBase.get_e_hf
+    set_frozen         = ccsd.CCSDBase.set_frozen
+    dump_flags         = ccsd.CCSDBase.dump_flags
+    get_init_guess     = ccsd.CCSDBase.get_init_guess
+    init_amps          = ccsd.CCSDBase.init_amps
+    energy             = ccsd.CCSDBase.energy
+    _add_vvvv          = ccsd.CCSDBase._add_vvvv
+    update_amps        = update_amps
+    kernel             = ccsd.CCSDBase.kernel
+    _finalize          = ccsd.CCSDBase._finalize
+    as_scanner         = ccsd.CCSDBase.as_scanner
+    restore_from_diis_ = ccsd.CCSDBase.restore_from_diis_
+
+    solve_lambda         = NotImplemented
+    ccsd_t               = NotImplemented
+    ipccsd               = NotImplemented
+    eaccsd               = NotImplemented
+    eeccsd               = NotImplemented
+    eomee_ccsd_singlet   = NotImplemented
+    eomee_ccsd_triplet   = NotImplemented
+    eomsf_ccsd           = NotImplemented
+    eomip_method         = NotImplemented
+    eomea_method         = NotImplemented
+    eomee_method         = NotImplemented
+    make_rdm1            = NotImplemented
+    make_rdm2            = NotImplemented
+    ao2mo                = _make_eris_incore
+    run_diis             = ccsd.CCSDBase.run_diis
+    amplitudes_to_vector = ccsd.CCSDBase.amplitudes_to_vector
+    vector_to_amplitudes = ccsd.CCSDBase.vector_to_amplitudes
+    dump_chk             = None
+    density_fit          = NotImplemented
+    nuc_grad_method      = NotImplemented
+
+    # to_cpu can be reused only when __init__ still takes mf
+    def to_cpu(self):
+        mf = self._scf.to_cpu()
+        from importlib import import_module
+        mod = import_module(self.__module__.replace('gpu4pyscf', 'pyscf'))
+        cls = getattr(mod, self.__class__.__name__)
+        obj = cls(mf)
+        return obj
+
+CCSDBase.ccsd = ccsd.CCSDBase.ccsd
+
+class CCSD(CCSDBase):
+    from gpu4pyscf.lib.utils import to_gpu, device
 
     def __init__(self, mf, *args, **kwargs):
         if hasattr(mf, 'to_cpu'):
@@ -539,5 +614,3 @@ def __init__(self, mf, *args, **kwargs):
             lib.logger.warn(mf.mol, 'DF-CCSD not available. Run the standard CCSD.')
         ccsd.CCSD.__init__(self, mf, *args, **kwargs)
 
-    update_amps = update_amps
-    ao2mo = _make_eris_incore

gpu4pyscf/cc/tests/test_ccsd.py

@@ -53,6 +53,15 @@ def test_ccsd_incore_kernel(self):
         self.assertAlmostEqual(abs(mcc.t1 - ref.t1).max(), 0, 6)
         self.assertAlmostEqual(abs(mcc.t2 - ref.t2).max(), 0, 6)
 
+    def test_to_gpu(self):
+        mcc = ccsd_incore.CCSD(mf.to_gpu()).run()
+        mcc = mcc.run()
+        e_gpu = mcc.e_tot
+        mcc = mcc.to_gpu()
+        mcc = mcc.run()
+        e_cpu = mcc.e_tot
+        assert (e_cpu - e_gpu) < 1e-6
+
 
 if __name__ == '__main__':
     print("Full Tests for CCSD")

gpu4pyscf/df/df.py

@@ -32,24 +32,38 @@
 ALIGNED = getattr(__config__, 'ao_aligned', 32)
 LINEAR_DEP_TOL = 1e-7
 
-class DF(df.DF):
+class DF(lib.StreamObject):
     from gpu4pyscf.lib.utils import to_gpu, device
 
-    _keys = {'intopt'}
+    _keys = {'intopt', 'mol', 'auxmol'}
 
     def __init__(self, mol, auxbasis=None):
-        super().__init__(mol, auxbasis)
+        self.mol = mol
+        self.stdout = mol.stdout
+        self.verbose = mol.verbose
+        self.max_memory = mol.max_memory
+        self._auxbasis = auxbasis
+
         self.auxmol = None
         self.intopt = None
         self.nao = None
         self.naux = None
         self.cd_low = None
         self._cderi = None
+        self._rsh_df = {}
+
+    @property
+    def auxbasis(self):
+        return self._auxbasis
+    @auxbasis.setter
+    def auxbasis(self, x):
+        if self._auxbasis != x:
+            self.reset()
+            self._auxbasis = x
 
     def to_cpu(self):
         from gpu4pyscf.lib.utils import to_cpu
         obj = to_cpu(self)
-        del obj.intopt, obj.cd_low, obj.nao, obj.naux
         return obj.reset()
 
     def build(self, direct_scf_tol=1e-14, omega=None):
@@ -167,15 +181,17 @@ def loop(self, blksize=None, unpack=True):
                 buf = buf_prefetch
 
     def reset(self, mol=None):
-        '''
-        reset object for scanner
-        '''
-        super().reset(mol)
+        '''Reset mol and clean up relevant attributes for scanner mode'''
+        if mol is not None:
+            self.mol = mol
+            self.auxmol = None
+        self._cderi = None
+        self._vjopt = None
+        self._rsh_df = {}
         self.intopt = None
         self.nao = None
         self.naux = None
         self.cd_low = None
-        self._cderi = None
         return self
 
     get_ao_eri = get_eri = NotImplemented

gpu4pyscf/df/df_jk.py

@@ -85,7 +85,7 @@ def _density_fit(mf, auxbasis=None, with_df=None, only_dfj=False):
     Examples:
     '''
 
-    assert isinstance(mf, scf.hf.SCF)
+    assert isinstance(mf, hf.SCF)
 
     if with_df is None:
         if isinstance(mf, dhf.UHF):
@@ -110,7 +110,8 @@ def _density_fit(mf, auxbasis=None, with_df=None, only_dfj=False):
     dfmf = _DFHF(mf, with_df, only_dfj)
     return lib.set_class(dfmf, (_DFHF, mf.__class__))
 
-class _DFHF(df_jk._DFHF):
+from gpu4pyscf.lib import utils
+class _DFHF:
     '''
     Density fitting SCF class
     Attributes for density-fitting SCF:
@@ -121,8 +122,8 @@ class _DFHF(df_jk._DFHF):
         with_df : DF object
             Set mf.with_df = None to switch off density fitting mode.
     '''
-
-    from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
+    to_gpu = utils.to_gpu
+    device = utils.device
 
     _keys = {'rhoj', 'rhok', 'disp', 'screen_tol'}
 
@@ -181,20 +182,23 @@ def nuc_grad_method(self):
         raise NotImplementedError()
 
     def Hessian(self):
-        from pyscf.dft.rks import KohnShamDFT
-        if isinstance(self, scf.rhf.RHF):
-            from gpu4pyscf.df.hessian import rhf, rks
+        from gpu4pyscf.dft.rks import KohnShamDFT
+        if isinstance(self, hf.RHF):
             if isinstance(self, KohnShamDFT):
-                return rks.Hessian(self)
+                from gpu4pyscf.df.hessian import rks as rks_hess
+                return rks_hess.Hessian(self)
             else:
-                return rhf.Hessian(self)
-        else:
-            from gpu4pyscf.df.hessian import uhf, uks
+                from gpu4pyscf.df.hessian import rhf as rhf_hess
+                return rhf_hess.Hessian(self)
+        elif isinstance(self, uhf.UHF):
             if isinstance(self, KohnShamDFT):
-                return uks.Hessian(self)
+                from gpu4pyscf.df.hessian import uks as uks_hess
+                return uks_hess.Hessian(self)
             else:
-                return uhf.Hessian(self)
-
+                from gpu4pyscf.df.hessian import uhf as uhf_hess
+                return uhf_hess.Hessian(self)
+        else:
+            raise NotImplementedError
     @property
     def auxbasis(self):
         return getattr(self.with_df, 'auxbasis', None)
@@ -207,7 +211,7 @@ def get_veff(self, mol=None, dm=None, dm_last=None, vhf_last=0, hermi=1):
         if dm is None: dm = self.make_rdm1()
 
         # for DFT
-        if isinstance(self, scf.hf.KohnShamDFT):
+        if isinstance(self, rks.KohnShamDFT):
             if dm.ndim == 2:
                 return rks.get_veff(self, dm=dm)
             elif dm.ndim == 3:
@@ -234,33 +238,11 @@ def get_veff(self, mol=None, dm=None, dm_last=None, vhf_last=0, hermi=1):
         else:
             raise NotImplementedError("Please check the dimension of the density matrix, it should not reach here.")
 
-    """
-    def energy_tot(self, dm=None, h1e=None, vhf=None):
-        '''
-        compute tot energy
-        '''
-        nuc = self.energy_nuc()
-        e_tot = self.energy_elec(dm, h1e, vhf)[0] + nuc
-        self.scf_summary['nuc'] = nuc.real
-        return e_tot
-    """
-    '''
     def to_cpu(self):
         obj = self.undo_df().to_cpu().density_fit()
-        keys = dir(obj)
-        obj.__dict__.update(self.__dict__)
-        for key in set(dir(self)).difference(keys):
-            print(key)
-            delattr(obj, key)
-
-        for key in keys:
-            val = getattr(obj, key)
-            if isinstance(val, cupy.ndarray):
-                setattr(obj, key, cupy.asnumpy(val))
-            elif hasattr(val, 'to_cpu'):
-                setattr(obj, key, val.to_cpu())
-        return obj
-    '''
+        print(type(obj))
+        return utils.to_cpu(self, obj)
+
 
 def get_jk(dfobj, dms_tag, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-14, omega=None):
     '''

gpu4pyscf/df/grad/rhf.py

@@ -21,7 +21,7 @@
 from pyscf import lib, scf, gto
 from gpu4pyscf.df import int3c2e
 from gpu4pyscf.lib.cupy_helper import print_mem_info, tag_array, unpack_tril, contract, load_library, take_last2d
-from gpu4pyscf.grad.rhf import grad_elec
+from gpu4pyscf.grad import rhf as rhf_grad
 from gpu4pyscf import __config__
 from gpu4pyscf.lib import logger
 
@@ -237,11 +237,19 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
     return vj, vk, vjaux, vkaux
 
 
-class Gradients(rhf.Gradients):
-    from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
+class Gradients(rhf_grad.Gradients):
+    from gpu4pyscf.lib.utils import to_gpu, device
+
+    _keys = {'with_df', 'auxbasis_response'}
+    def __init__(self, mf):
+        # Whether to include the response of DF auxiliary basis when computing
+        # nuclear gradients of J/K matrices
+        self.auxbasis_response = True
+        rhf_grad.Gradients.__init__(self, mf)
 
     get_jk = get_jk
-    grad_elec = grad_elec
+    grad_elec = rhf_grad.grad_elec
+    check_sanity = NotImplemented
 
     def get_j(self, mol=None, dm=None, hermi=0):
         vj, _, vjaux, _ = self.get_jk(mol, dm, with_k=False)
@@ -267,3 +275,5 @@ def extra_force(self, atom_id, envs):
             return envs['dvhf'].aux[atom_id]
         else:
             return 0
+
+Grad = Gradients