Merge pull request eljost#313 from bytedance/gpu4pyscf_support

Add gpu4pyscf support, fix bug with pyscf density fitting

pysisyphus/calculators/PySCF.py

@@ -3,7 +3,8 @@
 
 import numpy as np
 import pyscf
-from pyscf import gto, grad, lib, hessian, tddft, qmmm
+from pyscf import gto, lib, qmmm
+from pyscf import __all__ # ensure all modules are accessible under the pyscf namespace
 
 from pysisyphus.calculators.OverlapCalculator import OverlapCalculator
 from pysisyphus.helpers import geom_loader
@@ -47,6 +48,7 @@ def __init__(
         unrestricted=None,
         grid_level=3,
         pruning="nwchem",
+        use_gpu=False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -79,6 +81,8 @@ def __init__(
         self.chkfile = None
         self.out_fn = "pyscf.out"
 
+        self.use_gpu = use_gpu
+
         lib.num_threads(self.pal)
 
     @staticmethod
@@ -98,7 +102,10 @@ def build_grid(self, mf):
 
     def prepare_mf(self, mf):
         # Method can be overriden in a subclass to modify the mf object.
-        return mf
+        if self.use_gpu:
+            return mf.to_gpu()
+        else:
+            return mf
 
     def get_driver(self, step, mol=None, mf=None):
         def _get_driver():
@@ -244,7 +251,7 @@ def run(self, mol, point_charges=None):
                         f"Using '{self.chkfile}' as initial guess for {step} calculation."
                     )
                 if self.auxbasis:
-                    mf.density_fit(auxbasis=self.auxbasis)
+                    mf = mf.density_fit(auxbasis=self.auxbasis)
                     self.log(f"Using density fitting with auxbasis {self.auxbasis}.")
 
                 if point_charges is not None:
@@ -270,7 +277,7 @@ def run(self, mol, point_charges=None):
 
         # Keep mf and dump mol
         # save_mol(mol, self.make_fn("mol.chk"))
-        self.mf = mf
+        self.mf = mf.reset() # release integrals and other temporary intermediates.
         self.calc_counter += 1
 
         return mf