Skala: Accurate and scalable exchange-correlation with deep learning

Skala is a neural network-based exchange-correlation functional for density functional theory (DFT), developed by Microsoft Research AI for Science. It uses deep learning to predict exchange-correlation energies from electron density features, surpasses state-of-the-art hybrid functionals in accuracy for main group thermochemistry, kinetics and non-covalent interactions, all at a computational cost similar to semi-local DFT.

Trained on a large, diverse dataset — including coupled-cluster atomization energies and public benchmarks — Skala uses scalable message passing and local layers to learn both local and non-local effects. The model has about 385,000 parameters and matches the accuracy of leading hybrid functionals.

The recommended neural functional is skala-1.1, which uses per-atom packed grids, multiple non-local layers, and symmetric contraction. The legacy skala-1.0 traced model is still loadable via load_functional("skala-1.0").

Learn more about Skala in our ArXiv paper.

What's in here

This repository contains two main components:

1. The Python package skala, distributed on PyPI and on conda-forge. It contains a PyTorch implementation of the Skala model and its bindings to the quantum-chemistry packages PySCF, GPU4PySCF, and ASE.
2. Examples of using Skala from compiled code through LibTorch and GauXC:
   • Skala in C++ with libtorch
   • Skala in Fortran with FTorch
   • Skala in C++ with GauXC
   • Skala in C with GauXC
   • Skala in Fortran with GauXC

GauXC development version for PyTorch-based functionals like Skala

GauXC is a CPU/GPU C++ library for XC functionals. A development version with an add-on supporting PyTorch-based functionals like Skala is available in the skala branch of the GauXC repository. GauXC can be used to integrate Skala into other third-party DFT codes. For detailed documentation on using GauXC visit the Skala integration guide.

Getting started: PySCF (CPU)

All information below relates to the Python package skala.

pip install skala works out of the box and pulls every dependency from PyPI. If you don't already have PyTorch installed, install the CPU-only wheel first to avoid pulling a large CUDA build:

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install skala

For a reproducible conda environment, use the provided environment-cpu.yml, which pins CPU-only PyTorch and all runtime dependencies:

mamba env create -n skala -f environment-cpu.yml
mamba activate skala
pip install skala

Run an SCF calculation with Skala for a hydrogen molecule:

from pyscf import gto
from skala.pyscf import SkalaKS

mol = gto.M(
    atom="""H 0 0 0; H 0 0 1.4""",
    basis="def2-tzvp",
)
ks = SkalaKS(mol, xc="skala-1.1")
ks.kernel()

Getting started: GPU4PySCF (GPU)

The GPU install is more involved because gpu4pyscf ships CUDA-version-specific wheels that must match your CUDA toolkit.

To install all dependencies from PyPI, use the GPU specific package with the matching CUDA version, e.g., for CUDA 12:

pip install skala-cuda12x

The skala-cuda11x and skala-cuda13x packages are also available for CUDA 11 and 13, respectively.

The recommended path is the provided environment-gpu.yml, which pins pytorch-gpu, CUDA 12, cuda-nvrtc and matching cuda-cudart-dev headers for CuPy kernel compilation, and cutensor, and installs gpu4pyscf-cuda12x from PyPI:

mamba env create -n skala -f environment-gpu.yml
mamba activate skala
pip install skala

If you are building inside a container without a GPU attached (e.g., CI or a Docker image built on a CPU-only host), set CONDA_OVERRIDE_CUDA so the solver proceeds without a device:

CONDA_OVERRIDE_CUDA=12.0 mamba env create -n skala -f environment-gpu.yml

For CUDA 11 or 13, adjust cuda-version, cuda-nvrtc, cuda-cudart-dev, and the gpu4pyscf-cuda{11,13}x pin in environment-gpu.yml accordingly. Check your driver's maximum supported CUDA version with nvidia-smi.

Run an SCF calculation with Skala for a hydrogen molecule on GPU:

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-1.1")
ks.kernel()

Known issue: multiple visible GPUs

Skala uses a single GPU, but importing gpu4pyscf allocates memory on every visible CUDA device. This can conflict with PyTorch and with other processes sharing those GPUs (e.g. in MPI-parallel workloads).

Restrict CUDA to one device before launching Python:

CUDA_VISIBLE_DEVICES=0 python my_script.py

For MPI-parallel runs, assign one GPU per local rank:

mpirun -np 4 bash -c 'CUDA_VISIBLE_DEVICES=$OMPI_COMM_WORLD_LOCAL_RANK python my_script.py'

Tracked upstream at pyscf/gpu4pyscf#435.

Getting started: ASE calculator

Skala also provides an ASE calculator for energy, force, and geometry optimization workflows:

from ase.build import molecule
from ase.optimize import LBFGSLineSearch
from skala.ase import Skala

atoms = molecule("H2O")
atoms.calc = Skala(xc="skala-1.1", basis="def2-tzvp")

# Single-point energy (eV)
print(atoms.get_potential_energy())

# Geometry optimization
opt = LBFGSLineSearch(atoms)
opt.run(fmax=0.01)

Documentation and examples

See microsoft.github.io/skala for a more detailed installation guide and further examples of how to use the Skala functional with PySCF, GPU4PySCF and ASE.

Security: loading .fun files

Skala model files (.fun) use TorchScript serialization, which can execute arbitrary code when loaded. Never load .fun files from untrusted sources.

When loading the official Skala models via load_functional("skala-1.1") or load_functional("skala-1.0"), file integrity is automatically verified against pinned SHA-256 hashes before deserialization. If you load .fun files directly with TracedFunctional.load(), pass the expected_hash parameter to enable verification:

TracedFunctional.load("model.fun", expected_hash="<sha256-hex-digest>")

Project information

See the following files for more information about contributing, reporting issues, and the code of conduct:

• CONTRIBUTING.md
• LICENSE.txt
• SECURITY.md

Trademarks

This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft trademarks or logos is subject to and must follow Microsoft's Trademark & Brand Guidelines. Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. Any use of third-party trademarks or logos are subject to those third-party's policies.