Refactoring HF and DFT modules; Add CONTRIBUTING.md (#91)

* Refactoring HF and DFT modules; Add CONTRIBUTING.md

* Mute old CI job in github_actions.yml

* flake8 errors

.github/workflows/github_actions.yml

@@ -1,6 +1,6 @@
 name: CI
 
-on: [push]
+on: #[push]
 
 jobs:
   build:

CONTRIBUTING.md

@@ -0,0 +1,52 @@
+# Contributing Guidelines
+
+We welcome contributions from everyone. Please follow the guidelines below when
+adding features to the library for PySCF functions running on GPU devices.
+
+## Principles
+
+* Functions in GPU4PySCF are meant to be performance-optimized versions of
+  the features implemented in PySCF. The primary focus of this library is on
+  performance. It is understandable and acceptable to have some API differences
+  and incompatibilities with the PySCF CPU code.
+
+* GPU4PySCF functions may process input data or objects created by PySCF. It is
+  important to consider this possibility and perform necessary data type
+  conversions in GPU4PySCF functions. However, the data or objects produced by
+  GPU4PySCF do not need to be directly consumed by PySCF functions. In other
+  words, you only need to ensure that GPU4PySCF supports PySCF functions and not
+  the other way around.
+
+* GPU4PySCF uses CuPy arrays as the default array container. Please ensure that
+  functions in GPU4PySCF can handle any possible mixing of CuPy and NumPy
+  arrays, as many NumPy ufunc operations do not support this mixing.
+
+## Naming Conventions
+
+* If the GPU module has a corresponding one in PySCF, please consider using
+  the same function names and similar function signatures.
+
+## Additional Suggestions
+
+* If applicable, Please consider providing the `to_cpu` method to convert the
+  GPU4PySCF object to the corresponding PySCF object, and also providing the
+  `to_gpu` method in the corresponding classes in PySCF.
+
+* Inheriting classes from the corresponding classes in PySCF is not required.
+  If you choose to inherit classes from PySCF, please mute the unsupported
+  methods. You can overwrite the unsupported methods with a dummy method that
+  raise a `NotImplementedError`. Alternatively, you can assign the Python
+  keyword `None` or `NotImplemented` to these methods.
+
+* When writing unit tests, please consider including tests:
+  - To verify the values returned by the functions against the corresponding values in PySCF.
+  - To ensure proper handling of data conversions for those produced by PySCF CPU code.
+
+* While examples or documentation are not mandatory, it is highly recommended to
+  include examples of how to invoke the new module.
+
+* CUDA compute capability 70 (sm_70) is required. Please avoid using features
+  that are only available on CUDA compute capability 80 or newer. The CUDA code
+  should be compiled and run using CUDA 11 and CUDA 12 toolkits.
+
+Thank you for your contributions!

gpu4pyscf/df/df.py

@@ -178,6 +178,9 @@ def reset(self, mol=None):
         self._cderi = None
         return self
 
+    get_ao_eri = get_eri = NotImplemented
+    get_mo_eri = ao2mo = NotImplemented
+
 def cholesky_eri_gpu(intopt, mol, auxmol, cd_low, omega=None, sr_only=False):
     '''
     Returns:

gpu4pyscf/dft/gks.py

@@ -17,13 +17,16 @@
 
 from pyscf.dft import gks
 from gpu4pyscf.dft import numint
+from gpu4pyscf.dft import rks
 from gpu4pyscf.scf.ghf import GHF
 
-class GKS(gks.GKS):
+class GKS(gks.GKS, GHF):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
 
     def __init__(self, mol, xc='LDA,VWN'):
         raise NotImplementedError
 
-    get_jk = GHF.get_jk
-    _eigh = GHF._eigh
+    energy_elec = rks.RKS.energy_elec
+    get_veff = NotImplemented
+    nuc_grad_method = NotImplemented
+    to_hf = NotImplemented

gpu4pyscf/dft/rks.py

@@ -22,12 +22,15 @@
 from pyscf import lib
 from pyscf.dft import rks
 
-from gpu4pyscf import scf
 from gpu4pyscf.lib import logger
 from gpu4pyscf.dft import numint, gen_grid
 from gpu4pyscf.scf.hf import RHF
 from gpu4pyscf.lib.cupy_helper import load_library, tag_array
 
+__all__ = [
+    'get_veff', 'RKS'
+]
+
 libcupy_helper = load_library('libcupy_helper')
 
 LINEAR_DEP_THR = 1e-12
@@ -74,9 +77,8 @@ def initialize_grids(ks, mol=None, dm=None):
         #ks.grids.build(with_non0tab=True)
         ks.grids.weights = cupy.asarray(ks.grids.weights)
         ks.grids.coords = cupy.asarray(ks.grids.coords)
-        if (ks.small_rho_cutoff > 1e-20 and
-            # dm.ndim == 2 indicates ground state
-            isinstance(dm, cupy.ndarray) and dm.ndim == 2):
+        ground_state = getattr(dm, 'ndim', 0) == 2
+        if ks.small_rho_cutoff > 1e-20 and ground_state:
             # Filter grids the first time setup grids
             ks.grids = prune_small_rho_grids_(ks, ks.mol, dm, ks.grids)
         t0 = logger.timer_debug1(ks, 'setting up grids', *t0)
@@ -88,9 +90,7 @@ def initialize_grids(ks, mol=None, dm=None):
                 ks.nlcgrids.build()
                 ks.nlcgrids.weights = cupy.asarray(ks.nlcgrids.weights)
                 ks.nlcgrids.coords = cupy.asarray(ks.nlcgrids.coords)
-                if (ks.small_rho_cutoff > 1e-20 and
-                    # dm.ndim == 2 indicates ground state
-                    isinstance(dm, cupy.ndarray) and dm.ndim == 2):
+                if ks.small_rho_cutoff > 1e-20 and ground_state:
                     # Filter grids the first time setup grids
                     ks.nlcgrids = prune_small_rho_grids_(ks, ks.mol, dm, ks.nlcgrids)
                 t0 = logger.timer_debug1(ks, 'setting up nlc grids', *t0)
@@ -129,7 +129,7 @@ def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
 
     if hasattr(ks, 'screen_tol') and ks.screen_tol is not None:
         ks.direct_scf_tol = ks.screen_tol
-    ground_state = (isinstance(dm, cupy.ndarray) and dm.ndim == 2)
+    ground_state = getattr(dm, 'ndim', 0) == 2
 
     ni = ks._numint
     if hermi == 2:  # because rho = 0
@@ -224,7 +224,7 @@ def energy_elec(ks, dm=None, h1e=None, vhf=None):
     logger.debug(ks, 'E1 = %s  Ecoul = %s  Exc = %s', e1, ecoul, exc)
     return e1+e2, e2
 
-class RKS(scf.hf.RHF, rks.RKS):
+class RKS(rks.RKS, RHF):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
 
     _keys = {'disp'}
@@ -271,6 +271,5 @@ def nuc_grad_method(self):
         return rks_grad.Gradients(self)
 
     energy_elec = energy_elec
-    get_jk = RHF.get_jk
     get_veff = get_veff
-    _eigh = RHF._eigh
+    to_hf = NotImplemented

gpu4pyscf/dft/roks.py

@@ -15,16 +15,39 @@
 # 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 numpy as np
 from pyscf.dft import roks
 from gpu4pyscf.dft import numint
 from gpu4pyscf.scf.rohf import ROHF
+from gpu4pyscf.dft import uks, gen_grid
+from gpu4pyscf.lib.cupy_helper import tag_array
 
-class ROKS(roks.ROKS):
+class ROKS(roks.ROKS, ROHF):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
 
     def __init__(self, mol, xc='LDA,VWN'):
         super().__init__(mol, xc)
         self._numint = numint.NumInt()
+        self.disp = None
+        self.screen_tol = 1e-14
 
-    get_jk = ROHF.get_jk
-    _eigh = ROHF._eigh
+        grids_level = self.grids.level
+        self.grids = gen_grid.Grids(mol)
+        self.grids.level = grids_level
+
+        nlcgrids_level = self.nlcgrids.level
+        self.nlcgrids = gen_grid.Grids(mol)
+        self.nlcgrids.level = nlcgrids_level
+
+    def get_veff(self, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
+        if getattr(dm, 'mo_coeff', None) is not None:
+            mo_coeff = dm.mo_coeff
+            mo_occ_a = (dm.mo_occ > 0).astype(np.double)
+            mo_occ_b = (dm.mo_occ ==2).astype(np.double)
+            dm = tag_array(dm, mo_coeff=(mo_coeff,mo_coeff),
+                           mo_occ=(mo_occ_a,mo_occ_b))
+        return uks.get_veff(self, mol, dm, dm_last, vhf_last, hermi)
+
+    energy_elec = uks.UKS.energy_elec
+    nuc_grad_method = NotImplemented
+    to_hf = NotImplemented

gpu4pyscf/dft/uks.py

@@ -18,9 +18,9 @@
 import cupy
 from pyscf.dft import uks
 from pyscf import lib
-from gpu4pyscf import scf
 from gpu4pyscf.lib import logger
 from gpu4pyscf.dft import numint, gen_grid, rks
+from gpu4pyscf.scf.uhf import UHF
 from gpu4pyscf.lib.cupy_helper import tag_array
 
 
@@ -35,7 +35,7 @@ def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
 
     if hasattr(ks, 'screen_tol') and ks.screen_tol is not None:
         ks.direct_scf_tol = ks.screen_tol
-    ground_state = (isinstance(dm, cupy.ndarray) and dm.ndim == 3)
+    ground_state = getattr(dm, 'ndim', 0) == 3
 
     ni = ks._numint
     if hermi == 2:  # because rho = 0
@@ -112,7 +112,7 @@ def energy_elec(ks, dm=None, h1e=None, vhf=None):
     return rks.energy_elec(ks, dm, h1e, vhf)
 
 
-class UKS(scf.uhf.UHF, uks.UKS):
+class UKS(uks.UKS, UHF):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
     _keys = {'disp', 'screen_tol'}
 
@@ -132,4 +132,4 @@ def __init__(self, mol, xc='LDA,VWN', disp=None):
 
     energy_elec = energy_elec
     get_veff = get_veff
-
+    to_hf = NotImplemented

gpu4pyscf/lib/CMakeLists.txt

@@ -28,6 +28,7 @@ if(DEFINED CUDA_ARCHITECTURES)
 else()
   set(CMAKE_CUDA_ARCHITECTURES "70-real;80-real")
 endif()
+message("CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}")
 
 if (NOT CMAKE_BUILD_TYPE)
   set(CMAKE_BUILD_TYPE RELWITHDEBINFO)
@@ -172,4 +173,4 @@ if(BUILD_DFTD4)
             -DCMAKE_INSTALL_LIBDIR:PATH=lib
   )
   add_dependencies(dftd4 dftd4_static)
-endif()
+endif()

gpu4pyscf/lib/cupy_helper.py

@@ -18,9 +18,9 @@
 import os
 import sys
 import functools
+import ctypes
 import numpy as np
 import cupy
-import ctypes
 from pyscf import lib
 from gpu4pyscf.lib import logger
 from gpu4pyscf.gto import mole
@@ -114,6 +114,25 @@ def tag_array(a, **kwargs):
     t.__dict__.update(kwargs)
     return t
 
+def to_cupy(a):
+    '''Converts a numpy (and subclass) object to a cupy object'''
+    if isinstance(a, lib.NPArrayWithTag):
+        attrs = {k: to_cupy(v) for k, v in a.__dict__.items()}
+        return tag_array(cupy.asarray(a), **attrs)
+    if isinstance(a, np.ndarray):
+        return cupy.asarray(a)
+    return a
+
+def return_cupy_array(fn):
+    '''Ensure that arrays in returns are cupy objects'''
+    @functools.wraps(fn)
+    def filter_ret(*args, **kwargs):
+        ret = fn(*args, **kwargs)
+        if isinstance(ret, tuple):
+            return tuple(to_cupy(x) for x in ret)
+        return to_cupy(ret)
+    return filter_ret
+
 def unpack_tril(cderi_tril, cderi, stream=None):
     nao = cderi.shape[1]
     count = cderi_tril.shape[0]

gpu4pyscf/scf/ghf.py

@@ -17,12 +17,34 @@
 
 import cupy
 from pyscf.scf import ghf
-from gpu4pyscf.scf.hf import _get_jk as _get_jk_nr
-from gpu4pyscf.scf.hf import eigh
+from gpu4pyscf.scf import hf
 
 class GHF(ghf.GHF):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
 
+    _eigh = hf.RHF._eigh
+    scf = kernel = hf.RHF.kernel
+    get_hcore = hf.return_cupy_array(ghf.GHF.get_hcore)
+    get_ovlp = hf.return_cupy_array(ghf.GHF.get_ovlp)
+    get_init_guess = hf.RHF.get_init_guess
+    make_rdm2 = NotImplemented
+    dump_chk = NotImplemented
+    newton = NotImplemented
+    x2c = x2c1e = sfx2c1e = NotImplemented
+    to_rhf = NotImplemented
+    to_uhf = NotImplemented
+    to_ghf = NotImplemented
+    to_rks = NotImplemented
+    to_uks = NotImplemented
+    to_gks = NotImplemented
+    to_ks = NotImplemented
+    canonicalize = NotImplemented
+    # TODO: Enable followings after testing
+    analyze = NotImplemented
+    stability = NotImplemented
+    mulliken_pop = NotImplemented
+    mulliken_meta = NotImplemented
+
     def get_jk(self, mol=None, dm=None, hermi=0, with_j=True, with_k=True,
                omega=None):
         if mol is None: mol = self.mol
@@ -31,12 +53,10 @@ def get_jk(self, mol=None, dm=None, hermi=0, with_j=True, with_k=True,
         dm = cupy.asarray(dm)
 
         def jkbuild(mol, dm, hermi, with_j, with_k, omega=None):
-            return _get_jk_nr(self, mol, dm, hermi, with_j, with_k, omega)
+            return hf._get_jk(self, mol, dm, hermi, with_j, with_k, omega)
 
         if nao == dm.shape[-1]:
             vj, vk = jkbuild(mol, dm, hermi, with_j, with_k, omega)
         else:  # GHF density matrix, shape (2N,2N)
             vj, vk = ghf.get_jk(mol, dm, hermi, with_j, with_k, jkbuild, omega)
         return vj, vk
-
-    _eigh = staticmethod(eigh)