diff --git a/pysisyphus/calculators/PySCF.py b/pysisyphus/calculators/PySCF.py index c3ec9e2a3..6193463f6 100644 --- a/pysisyphus/calculators/PySCF.py +++ b/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