diff --git a/docs/ase.ipynb b/docs/ase.ipynb index 8aaf4df..c5d41e5 100644 --- a/docs/ase.ipynb +++ b/docs/ase.ipynb @@ -262,6 +262,51 @@ "print(f\"Maximum force: {max_force:.6f} eV/Å\")" ] }, + { + "cell_type": "markdown", + "id": "a87cc9da", + "metadata": {}, + "source": [ + "## Using Skala on GPU\n", + "\n", + "In case you have access to a compatible GPU, you can enable GPU acceleration by setting the `device` parameter to `\"cuda\"` when initializing the calculator.\n", + "This can significantly speed up calculations for larger systems." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "37835b66", + "metadata": {}, + "outputs": [], + "source": [ + "calc = Skala(\n", + " xc=\"skala-1.1\",\n", + " basis=\"def2-tzvp\",\n", + " with_density_fit=True,\n", + " auxbasis=\"def2-universal-jkfit\",\n", + " device=\"cuda\", # Enable GPU acceleration\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "a3df794c", + "metadata": {}, + "source": [ + "The device can be changed at any time, causing the calculator to reset and use the new device for subsequent calculations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f9cdd62", + "metadata": {}, + "outputs": [], + "source": [ + "calc.set(device=\"cpu\")" + ] + }, { "cell_type": "markdown", "id": "8dd192b6", diff --git a/src/skala/ase/calculator.py b/src/skala/ase/calculator.py index 667dc90..c6d66d2 100644 --- a/src/skala/ase/calculator.py +++ b/src/skala/ase/calculator.py @@ -14,10 +14,16 @@ from ase.units import Bohr, Debye, Hartree from pyscf import grad, gto +import skala.pyscf as skala_cpu from skala.functional.base import ExcFunctionalBase -from skala.pyscf import SkalaKS from skala.pyscf.retry import retry_scf +try: + import skala.gpu4pyscf as skala_gpu +except ImportError as e: + skala_gpu = None # type: ignore[assignment] + gpu4pyscf_import_error = e + class Skala(Calculator): """ @@ -49,6 +55,7 @@ class Skala(Calculator): "multiplicity": None, "verbose": 0, "ks_config": None, + "device": "cpu", } _mol: gto.Mole | None = None @@ -94,6 +101,7 @@ def set(self, **kwargs: Any) -> dict[str, Any]: or "auxbasis" in changed_parameters or "with_newton" in changed_parameters or "with_dftd3" in changed_parameters + or "device" in changed_parameters ): self._ks = None self.reset() @@ -166,16 +174,28 @@ def calculate( dm0 = None if not isinstance(xc_param := self.parameters.xc, (ExcFunctionalBase, str)): # type: ignore raise InputError("XC functional must be a string or ExcFunctionalBase.") - grad_method = SkalaKS( - self._mol, + device = self.parameters.device # type: ignore + _kwargs = dict( + mol=self._mol, xc=xc_param, with_density_fit=bool(self.parameters.with_density_fit), # type: ignore auxbasis=self.parameters.auxbasis, # type: ignore with_newton=bool(self.parameters.with_newton), # type: ignore with_dftd3=bool(self.parameters.with_dftd3), # type: ignore ks_config=self.parameters.ks_config, # type: ignore - ).nuc_grad_method() - self._ks = grad_method + ) + if device == "cuda": + if skala_gpu is None: + raise ImportError( + "gpu4pyscf is not available. Please install gpu4pyscf to use GPU acceleration." + ) from gpu4pyscf_import_error + ks = skala_gpu.SkalaKS(**_kwargs) + elif device == "cpu": + ks = skala_cpu.SkalaKS(**_kwargs) + else: + raise InputError(f"Unsupported device type: {device}") + + self._ks = ks.nuc_grad_method() else: # Mimic PySCF's SCF_Scanner (hf.py:1569-1588): convert old MOs # into a density-matrix guess, wipe mo_coeff so Newton won't