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 leverages deep learning to predict exchange-correlation energies from electron density features, achieving chemical accuracy for atomization energies and strong performance on broad thermochemistry and kinetics benchmarks, 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 276,000 parameters and matches the accuracy of leading hybrid functionals.

Learn more about Skala in our ArXiv paper.

What's in here

This repository contains two main components:

1. The Python package skala, which is also distributed on PyPI and contains a PyTorch implementation of the Skala model, its hookups to quantum chemistry packages PySCF, GPU4PySCF and ASE.
2. An example of using Skala in C++ CPU applications through LibTorch, see examples/cpp/cpp_integration.

Skala in Azure AI Foundry

The Skala model is also served on Azure AI Foundry.

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 is part of the stack that serves Skala in Azure AI Foundry and 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.

Install using Pip:

# Install CPU-only PyTorch (skip if you already have CPU or GPU-enabled PyTorch installed)
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install skala

Or using Conda (Mamba):

mamba install -c conda-forge skala "pytorch=*=cpu*"

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

Getting started: GPU4PySCF (GPU)

These instructions use Mamba and pip to install CUDA toolkit, Torch, and CuPy. It supports CUDA version 11, 12 or 13. You can find the most recent CUDA version that is supported on your system using nvidia-smi.

cu_version=12 #or 11 or 13 depending on your CUDA version
mamba env create -n skala -f environment-gpu.yml  "cuda-version==${cu_version}.*" skala
mamba activate skala
pip install --no-deps "gpu4pyscf-cuda${cu_version}x>=1.0,<2" "gpu4pyscf-libxc-cuda${cu_version}x>=0.4,<1"

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

Documentation and examples

Go to 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 and in Azure AI Foundry.

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.