From f5b836c8f1efe435f81a1ca664358c5e7c395ac8 Mon Sep 17 00:00:00 2001 From: Sebastian Ehlert Date: Mon, 4 May 2026 15:16:23 +0200 Subject: [PATCH 1/4] Allow using ASE calculator on GPU --- docs/ase.ipynb | 45 +++++++++++++++++++++++++++++++++++++ src/skala/ase/calculator.py | 29 +++++++++++++++++++----- 2 files changed, 69 insertions(+), 5 deletions(-) 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..bf4111a 100644 --- a/src/skala/ase/calculator.py +++ b/src/skala/ase/calculator.py @@ -15,9 +15,15 @@ from pyscf import grad, gto from skala.functional.base import ExcFunctionalBase -from skala.pyscf import SkalaKS +from skala.pyscf import SkalaKS as SkalaKSCPU from skala.pyscf.retry import retry_scf +try: + from skala.gpu4pyscf import SkalaKS as SkalaKSGPU +except ImportError as e: + SkalaKSGPU = None + 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,27 @@ 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, + _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 str(self.parameters.device) == "cuda": + if SkalaKSGPU is None: + raise ImportError( + "gpu4pyscf is not available. Please install gpu4pyscf to use GPU acceleration." + ) from gpu4pyscf_import_error + ks = SkalaKSGPU(**_kwargs) + elif str(self.parameters.device) == "cpu": + ks = SkalaKSCPU(**_kwargs) + else: + raise InputError(f"Unsupported device type: {self.parameters.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 From d89242589a32a32eaec394a505aa0a64206d5cc6 Mon Sep 17 00:00:00 2001 From: Sebastian Ehlert Date: Mon, 4 May 2026 15:36:11 +0200 Subject: [PATCH 2/4] Typing --- src/skala/ase/calculator.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/skala/ase/calculator.py b/src/skala/ase/calculator.py index bf4111a..c451e56 100644 --- a/src/skala/ase/calculator.py +++ b/src/skala/ase/calculator.py @@ -15,13 +15,13 @@ from pyscf import grad, gto from skala.functional.base import ExcFunctionalBase -from skala.pyscf import SkalaKS as SkalaKSCPU +import skala.pyscf as skala_cpu from skala.pyscf.retry import retry_scf try: - from skala.gpu4pyscf import SkalaKS as SkalaKSGPU + import skala.gpu4pyscf as skala_gpu except ImportError as e: - SkalaKSGPU = None + skala_gpu = None # type: ignore[assignment] gpu4pyscf_import_error = e @@ -184,13 +184,13 @@ def calculate( ks_config=self.parameters.ks_config, # type: ignore ) if str(self.parameters.device) == "cuda": - if SkalaKSGPU is None: + if skala_gpu is None: raise ImportError( "gpu4pyscf is not available. Please install gpu4pyscf to use GPU acceleration." ) from gpu4pyscf_import_error - ks = SkalaKSGPU(**_kwargs) + ks = skala_gpu.SkalaKS(**_kwargs) elif str(self.parameters.device) == "cpu": - ks = SkalaKSCPU(**_kwargs) + ks = skala_cpu.SkalaKS(**_kwargs) else: raise InputError(f"Unsupported device type: {self.parameters.device}") From 042be11338cbdd373c17d31e504299dba3136188 Mon Sep 17 00:00:00 2001 From: Sebastian Ehlert Date: Mon, 4 May 2026 15:43:23 +0200 Subject: [PATCH 3/4] Lint --- src/skala/ase/calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/skala/ase/calculator.py b/src/skala/ase/calculator.py index c451e56..05c6662 100644 --- a/src/skala/ase/calculator.py +++ b/src/skala/ase/calculator.py @@ -14,8 +14,8 @@ from ase.units import Bohr, Debye, Hartree from pyscf import grad, gto -from skala.functional.base import ExcFunctionalBase import skala.pyscf as skala_cpu +from skala.functional.base import ExcFunctionalBase from skala.pyscf.retry import retry_scf try: From 8cda843628c7e0a70777241c1881bcae81e4c2dc Mon Sep 17 00:00:00 2001 From: Sebastian Ehlert Date: Mon, 4 May 2026 15:57:44 +0200 Subject: [PATCH 4/4] Typing --- src/skala/ase/calculator.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/skala/ase/calculator.py b/src/skala/ase/calculator.py index 05c6662..c6d66d2 100644 --- a/src/skala/ase/calculator.py +++ b/src/skala/ase/calculator.py @@ -174,6 +174,7 @@ 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.") + device = self.parameters.device # type: ignore _kwargs = dict( mol=self._mol, xc=xc_param, @@ -183,16 +184,16 @@ def calculate( with_dftd3=bool(self.parameters.with_dftd3), # type: ignore ks_config=self.parameters.ks_config, # type: ignore ) - if str(self.parameters.device) == "cuda": + 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 str(self.parameters.device) == "cpu": + elif device == "cpu": ks = skala_cpu.SkalaKS(**_kwargs) else: - raise InputError(f"Unsupported device type: {self.parameters.device}") + raise InputError(f"Unsupported device type: {device}") self._ks = ks.nuc_grad_method() else: