Add gpu4pyscf support

* Adds support for using gpu4pyscf instead of pyscf for running SCF

(cherry picked from commit 7c1ce3bce4b84bf05d6b5765c914f71f4ced98d5)

Author: DerkKooi

.github/workflows/docs.yml

@@ -24,7 +24,7 @@ jobs:
       - name: Setup micromamba
         uses: mamba-org/setup-micromamba@v1
         with:
-          environment-file: environment.yml
+          environment-file: environment-cpu.yml
           environment-name: skala
           cache-environment: true
           cache-downloads: true

.github/workflows/test.yml

@@ -23,7 +23,7 @@ jobs:
     - name: Setup micromamba
       uses: mamba-org/setup-micromamba@v1
       with:
-        environment-file: environment.yml
+        environment-file: environment-cpu.yml
         environment-name: skala
         cache-environment: true
         cache-downloads: true

README.md

@@ -43,9 +43,31 @@ mol = gto.M(
 ks = SkalaKS(mol, xc="skala")
 ks.kernel()
 ```
-
 Go to [microsoft.github.io/skala](https://microsoft.github.io/skala) for a more detailed installation guide and further examples of how to use Skala functional with PySCF and ASE and in [Azure Foundry](https://ai.azure.com/catalog/models/Skala).
 
+## Getting started (GPU support)
+
+Install using Pip:
+```bash
+cu_version=128 #or 126 or 130 depending on your CUDA version
+pip install torch cupy --extra-index-url "https://download.pytorch.org/whl/cu${cu_version}"
+pip install --no-deps "gpu4pyscf-cuda${cu_version:0:2}x>=1.0,<2" "gpu4pyscf-libxc-cuda${cu_version:0:2}x>=0.4,<1"
+pip install microsoft-skala
+```
+Run an SCF calculation with Skala for a hydrogen molecule on GPU:
+
+```python
+from pyscf import gto
+from skala.gpu4pyscf import SkalaKS
+
+mol = gto.M(
+    atom="""H 0 0 0; H 0 0 1.4""",
+    basis="def2-tzvp",
+)
+ks = SkalaKS(mol, xc="skala")
+ks.kernel()
+```
+
 ## Project information
 
 See the following files for more information about contributing, reporting issues, and the code of conduct:

docs/installation.rst

@@ -33,10 +33,16 @@ If you prefer to install Skala from the source code, you can clone the repositor
 
    git clone https://github.com/microsoft/skala
    cd skala
-   mamba env create -n skala -f environment.yml
+   mamba env create -n skala -f environment-cpu.yml
    mamba activate skala
    pip install -e .
 
+where `environment-cpu.yml` can be replaced for `environment-gpu.yml` for gpu support (CUDA 12) with gpu4pyscf, in which case gpu4pyscf needs to be separately installed *after creating the environment* via
+
+.. code-block:: bash
+
+   pip install --no-deps 'gpu4pyscf-cuda12x>=1.0,<2' 'gpu4pyscf-libxc-cuda12x>=0.4,<1'
+
 To install the development dependencies, you can run:
 
 .. code-block:: bash

docs/pyscf/scf_settings.ipynb

@@ -11,19 +11,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "dd93a68a",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/giulialuise/miniconda3/envs/skala/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from pyscf import gto\n",
     "\n",

docs/pyscf/singlepoint.ipynb

@@ -104,11 +104,42 @@
     "\n",
     "print(ks.dump_scf_summary())"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fd56964d",
+   "metadata": {},
+   "source": [
+    "# Using the Skala functional in GPU4PySCF\n",
+    "\n",
+    "The Skala functional can also be used in GPU4PySCF with an appropriate PyTorch CUDA version by creating a new Kohn-Sham calculator based on the `SkalaKS` constructor from the `skala.gpu4pyscf`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cec1b5f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pyscf import gto\n",
+    "\n",
+    "from skala.gpu4pyscf import SkalaKS\n",
+    "\n",
+    "mol = gto.M(\n",
+    "    atom=\"\"\"H 0 0 0; H 0 0 1.4\"\"\",\n",
+    "    basis=\"def2-tzvp\",\n",
+    ")\n",
+    "ks = SkalaKS(mol, xc=\"skala\")\n",
+    "ks.kernel()\n",
+    "\n",
+    "print(ks.dump_scf_summary())"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "skala",
+   "display_name": "skala_test",
    "language": "python",
    "name": "python3"
   },
@@ -122,7 +153,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.5"
+   "version": "3.14.2"
   }
  },
  "nbformat": 4,

environment-cpu.yml

@@ -0,0 +1,21 @@
+name: skala-cpu
+channels:
+  - conda-forge
+dependencies:
+  - ase
+  - azure-core
+  - azure-identity
+  - dftd3-python
+  - e3nn
+  - numpy
+  - opt_einsum_fx
+  - pyscf<2.10.0
+  - python
+  - pytorch * cpu_*
+  - qcelemental
+# Testing and development
+  - pre-commit
+  - pytest
+  - pytest-cov
+  - pip:
+    - huggingface_hub

environment-gpu.yml

@@ -0,0 +1,25 @@
+name: skala-gpu
+channels:
+  - conda-forge
+dependencies:
+  - ase
+  - azure-core
+  - azure-identity
+  - dftd3-python
+  - e3nn
+  - numpy
+  - opt_einsum_fx
+  - pyscf<2.10.0
+  - python
+  - pytorch-gpu
+  - qcelemental
+  - cuda-toolkit
+  - cupy
+  - cutensor
+  - cuda-version ==12.*
+  # Testing and development
+  - pre-commit
+  - pytest
+  - pytest-cov
+  - pip:
+    - huggingface_hub

environment.yml

@@ -8,7 +8,7 @@ Set up the conda environment using the provided environment file:
 
 ```bash
 cd examples/cpp/cpp_integration
-conda env create -n skala_cpp_integration -f environment.yml
+conda env create -n skala_cpp_integration -f environment-cpu.yml
 conda activate skala_cpp_integration
 ```