Learnable Exchange-Correlation Models in JAX

This repository contains the reference implementation for egxc, di-loss, and SAIL alongside scripts, model parameters and configuration files used to train and evaluate exchange-correlation (XC) functionals. The code combines classical density functional theory with neural network models implemented in JAX.

The methods implemented here are described in detail in:

• Transferable SCF-Acceleration through Solver-Aligned Initialization Learning, preprint. It introduces SAIL, which trains SCF initial guesses by backpropagating through the solver to reduce wall-clock time.
• Derivative Informed Learning of Exchange-Correlation Functionals, presented at ICML 2026. It introduces a new loss function for training learnable exchange-correlation models that is based on the derivative of the exchange-correlation energy with respect to the density.
• Learning Equivariant Non-Local Electron Density Functionals, presented as a Spotlight at ICLR 2025. It introduces non-local learnable exchange-correlation models built on modern machine learning techniques.

@misc{eberhardTransferableSCFAccelerationSolverAligned2026,
  title = {Transferable {{SCF-Acceleration}} through {{Solver-Aligned Initialization Learning}}},
  author = {Eberhard, Eike S. and Kotsev, Viktor and G{\"u}thle, Timm and G{\"u}nnemann, Stephan},
  year = 2026,
  url = {https://arxiv.org/abs/2604.21657}
}

@inproceedings{eberhardDerivativeInformedLearning2026,
  title = {Derivative {{Informed Learning}} of {{Exchange-Correlation Functionals}}},
  booktitle = {Proceedings of the {{43rd International Conference}} on {{Machine Learning}}, Seoul, South Korea},
  author = {Eberhard, Eike S. and Thiede, Luca A. and Aldossary, Abdul and Burger, Andreas and Gao, Nicholas and Bhethanabotla, Vignesh and {Aspuru-Guzik}, Al{\'a}n and G{\"u}nnemann, Stephan},
  year = 2026,
  url = {https://arxiv.org/abs/2606.04279}
}

@inproceedings{gaoLearningEquivariantNonLocal2024,
  title = {Learning {{Equivariant Non-Local Electron Density Functionals}}},
  booktitle = {The {{Thirteenth International Conference}} on {{Learning Representations}}},
  author = {Gao, Nicholas and Eberhard, Eike and G{\"u}nnemann, Stephan},
  year = 2024,
  url = {https://openreview.net/forum?id=FhBT596F1X}
}

Repository Structure

• src/egxc – Original EG-XC implementation and associated modules.
  • dataloading – dataset abstractions and loaders.
  • discretization – quadrature grids, GTO basis evaluation, and a pure‑JAX GPU integral engine.
    • integrals/md.py – class‑specialized (angular‑momentum‑bucketed) McMurchie–Davidson integral engine, architecture inspired by gpu4pyscf's gint kernels but written in plain JAX, so integrals are natively differentiable wrt nuclear positions with no custom VJP (module docstring md.py:1). Public API re‑exported from integrals/__init__.py: build_integral_layout (host‑side shell structure → IntegralLayout pytree, only the class‑shape meta is jit‑static), one_electron_integrals (overlap S, kinetic T, nuclear‑attraction V), two_center_coulomb (DF Coulomb metric (P|Q)), three_center_coulomb (DF (ij|P)), three_center_j_vector (contracted DF Coulomb vector d_Q = Σ_ij (ij|Q) P_ij, never materializing the full (B,B,Q) tensor), plus IntegralLayout, MDPadSpec, pad_from_compositions, merge_pad_specs, worst_case_pad, MAX_AUX_L. The engine deliberately builds only DF 3‑center/2‑center Coulomb integrals and never forms a dense 4‑center ERI (there is no four_center symbol). The 2c/3c Coulomb kernels support aux DF bases up to g shells (MAX_AUX_L = 4); overlap/kinetic do not (primitives.py:29). Conventions reproduce PySCF/libcint exactly (primitives.py:1). See src/egxc/discretization/README.md for details.
    • gto/containers.py – GTO basis containers: GTOBasis, PreloadedGTOBasis, GTOShellIndexing, RadialPrimitives, and the get_gto_preloader factory.
    • gto/grid_eval.py – AO‑on‑grid evaluation: get_gto_grid_eval_fn, compute_radial_primitives, compute_angular_components (plus gto/constants.py).
    • grids/ – quadrature grid construction (atomic.py, lebedev.py, quadrature.py; get_grid_fn re‑exported from the package root).
    • __init__.py re‑exports GTOBasis, PreloadedGTOBasis, get_gto_grid_eval_fn, get_grid_fn; the integral engine lives under discretization.integrals.
  • systems – immutable containers for molecular structures and precomputed tensors.
    • base.py – immutable containers System, Grid, FockTensors, plus nuclear_energy_fn and nuclear_energy_and_force.
    • preload.py – PySCF‑backed CPU preloading of Fock tensors / grids: preload_fock_tensors_using_pyscf, preload_system_using_pyscf, get_minao_initial_density_matrices, get_aux_basis. The dense 4‑center ERI is built here via compute_electron_repulsion_tensor → mol.intor('int2e') (preload.py:106), and the DF Cholesky factor (naux, nao, nao) via _cholesky_decomposed_eri (preload.py:74, with a truncated‑eigh fallback for rank‑deficient long‑range metrics).
    • observables.py – homo_lumo_gap_fn.
  • solver – self‑consistent field solver with convergence acceleration.
    • scf/scf.py – SelfConsistentFieldSolver (DIIS/Momentum/Vanilla acceleration, DF and restricted‑open‑shell options).
    • scf/diis.py – jittable Pulay/commutator DIIS: compute_residual, solve_pulay_equation, DiisState, diis_update.
    • fock.py – FockMatrix (nn.Module), mean_field_energy, roothaan_effective_fock (ROHF effective Fock), and get_coulomb_matrix_fn. The latter contracts a precomputed full ERI into the Coulomb matrix J when no DF is used (jnp.einsum('ijkl,ij->kl', eri, P), fock.py:82) or the DF factor otherwise (fock.py:84).
    • base.py – Solver base nn.Module.
  • xc_energy – classical and learnable XC functionals and density features.
  • training – training loop, optimizer utilities and loss functions.
  • utils – common helpers: physical/grid constants (constants.py: BOHR_TO_ANGSTROM/ANGSTROM_TO_BOHR, HATREE_TO_KCAL_PER_MOL, HARTREE_TO_EV, the Z‑keyed ATOMIC_GROUND_STATE_SPIN, BRAGG_RADII, RAD_GRIDS, ANG_ORDER, TREUTLER_AHLRICHS_XI, L_MAX, DENSITY_FLOOR, ZETA_GUARD), shape‑specific tensor aliases (typing.py: e.g. FloatBxB, FloatAx3, FloatOxV, FloatQxBxB), plus linalg.py, pad.py, checkpointing.py, and logging.
  • visualization – interactive Plotly utilities.
  • data_generation – scripts for producing reference data used in training.
• src/dixc – Derivative‑Informed eXchange‑Correlation (DI‑XC) plugin.
  • data_generation – generation of DI‑XC reference energies and forces.
  • dataset.py – dataset interface for DI‑XC training data.
  • orbital_transforms.py – jitted, jnp.einsum‑based GPU JAX AO↔MO transforms and orbital‑rotation parameterization: ao_to_mo / mo_to_ao (project an AO response tensor into the occupied‑virtual MO block and back), expand_occupied_virtual_block_to_full_basis, orbital_rotation / first_order_orbital_rotation (rotate MO coefficients via the matrix exponential of the occ‑virt generator), dm_gradient_to_orbital_rotation_gradient, ao_density_perturbation_from_occupied_virtual_rotation, and orthonormalize_delta_densities.
  • scf.py – DI‑XC self‑consistent field driver.
  • training – end‑to‑end DI‑XC training (train.py, loss functions, utils).
• src/utils – shared evaluation utilities (eval_scf.py).
• configs – Example experiment configurations (YAML) for seml, organized by method:
  • egxc/ – EG-XC training runs.
  • dixc/ – DI-loss / DI-XC training, data generation, and evaluation.
  • sail/ – SAIL data generation.
• scripts – Entry points for running experiments, including:
  • egxc_main.py – EG-XC training.
  • dixc_main.py – DI-XC / DI-loss training.
  • evaluate.py – standalone checkpoint evaluation.
  • gen_solver_trajectory_data.py, end2end_scf_acceleration_with_distillates.py – SAIL reference data generation and end-to-end SCF acceleration.
• doc – Technical notes (density_features.md, orbital_rotations.md, diis_backprop.md, padding.md, and others).
• notebooks – Jupyter notebooks with examples and analysis.
• figures – Plot assets for papers and presentations.
• evaluations – experiment results and metadata.
• test – Unit tests covering basis functions, solvers and XC modules.

Where to find the electronic-structure primitives

A concept → file → key‑symbols index for the DFT / electronic‑structure machinery. All integral, grid and AO↔MO code is pure JAX (GPU‑capable, differentiable); PySCF appears only in the CPU preload step. Paths are under src/.

Concept	File	Key symbols
GPU McMurchie–Davidson integral engine (one‑electron + DF Coulomb, no dense 4‑center ERI)	egxc/discretization/integrals/md.py (API via integrals/__init__.py)	build_integral_layout, one_electron_integrals, two_center_coulomb, three_center_coulomb, three_center_j_vector, IntegralLayout, MDPadSpec
Per‑element Cartesian GTO shell data (libcint conventions)	egxc/discretization/integrals/primitives.py	cartesian_components, MAX_AUX_L
GTO basis containers & AO‑on‑grid evaluation	egxc/discretization/gto/containers.py, gto/grid_eval.py	GTOBasis, PreloadedGTOBasis, get_gto_preloader, get_gto_grid_eval_fn
Quadrature grids	egxc/discretization/grids/	get_grid_fn (quadrature.py), atomic.py, lebedev.py
Dense 4‑center ERI (CPU, PySCF)	egxc/systems/preload.py	compute_electron_repulsion_tensor → mol.intor('int2e') (line 106)
DF Cholesky factor (naux, nao, nao) (CPU, PySCF)	egxc/systems/preload.py	_cholesky_decomposed_eri (line 74)
Fock build (J from ERI/DF factor)	egxc/solver/fock.py	FockMatrix, get_coulomb_matrix_fn (line 73; J einsum at line 82), mean_field_energy, roothaan_effective_fock
SCF solver + DIIS	egxc/solver/scf/scf.py, scf/diis.py	SelfConsistentFieldSolver, compute_residual, solve_pulay_equation, diis_update
AO↔MO transforms & orbital rotations (GPU JAX)	dixc/orbital_transforms.py	ao_to_mo, mo_to_ao, orbital_rotation, expand_occupied_virtual_block_to_full_basis
System / Fock‑tensor containers	egxc/systems/base.py	System, Grid, FockTensors, nuclear_energy_fn, nuclear_energy_and_force
Physical / grid constants	egxc/utils/constants.py	HATREE_TO_KCAL_PER_MOL, ANGSTROM_TO_BOHR, ATOMIC_GROUND_STATE_SPIN, BRAGG_RADII, L_MAX
Shape‑specific tensor type aliases	egxc/utils/typing.py	FloatBxB, FloatAx3, FloatOxV, FloatQxBxB

Getting Started

1. Install uv and create the virtual environment defined by pyproject.toml and uv.lock:

   curl -Ls https://astral.sh/uv/install.sh | bash
   uv sync

2. Explore example configuration files in configs/ and run experiments via seml or directly:
   • EG-XC training: python scripts/egxc_main.py with configs/egxc/*
   • DI-XC / DI-loss training: python scripts/dixc_main.py with configs/dixc/* (data generation via configs/dixc/gen_*.yaml, evaluation via configs/dixc/evaluate_batch.yaml and scripts/evaluate.py)
   • SAIL: python scripts/gen_solver_trajectory_data.py with configs/sail/* for solver-trajectory data, then python scripts/end2end_scf_acceleration_with_distillates.py for end-to-end SCF acceleration with distilled initial guesses
3. Review the notes in doc/ and the notebooks under notebooks/ for additional explanations and analysis.
4. Run the test suite (see Testing and Checks for the recommended invocations and dataset requirements).

Key implementations:

• EG-XC functional: src/egxc/xc_energy/functionals/learnable/egxc.py
• DI-XC functional and DI-loss training: src/egxc/xc_energy/functionals/learnable/dixc.py and src/dixc/

Core APIs

Common entry points when working with the codebase:

from egxc.systems import examples                  # built-in molecules via examples.get(...)
from egxc.dataloading import MD17, QM9, QM40       # dataset loaders
from egxc.solver import scf                        # SCF solver
from egxc.xc_energy.functionals.classical import get_lda, get_gga, get_mgga, get_hybrid
from egxc.xc_energy.functionals.learnable import (
    DIXC, EGXC, Nagai2020, Nagai2022, NNmGGA_PP, Skala, XCDiff,
)
from egxc.training import run                      # EG-XC training loop

from dixc.dataset import DIXCDataset
from dixc.scf import DerivativeInformedSelfConsistentFieldSolver

Development Guidelines

• Use JAX and Flax best practices for performance and numerical stability: apply @jax.jit to performance‑critical functions and prefer jax.vmap and jax.lax.scan over Python loops.
• Prefer static_argnames over static_argnums for JIT‑compiled functions with static arguments.
• Use explicit, shape‑specific tensor aliases (for example FloatAx3) instead of bare jnp.ndarray or jax.Array type hints.
• For offline or demo runs, start from the example molecules in egxc.systems.examples (examples.get(...)).
• When the SCF solver does not converge, adjust tolerances or mixing parameters explicitly rather than silently accepting non‑converged solutions.
• The project uses uv to manage the virtual environment in .venv; you can activate it manually (for example source .venv/bin/activate) if needed.

Testing and Checks

• The test suite is tiered by runtime markers; select a tier with -m:

  uv run pytest -m quick          # <= 1 min per test
  uv run pytest -m intermediate   # 1-3 min per test
  uv run pytest -m slow           # > 3 min per test

  data and modelling markers additionally tag the data-pipeline and ML-modelling test subsets.
• Two test modules (test/test_dataloading.py and test/test_dixc_components.py) require pre-generated datasets and are skipped automatically when the dataset root is missing. The root defaults to data/datasets; point the EGXC_DATA_DIR environment variable at an existing dataset directory to run them (test_dataloading.py will download the raw QM9/MD17/3BPA archives on first use, ~800 MB; test_dixc_components.py additionally needs reference densities generated via configs/dixc/gen_*.yaml). To exclude them explicitly, mirror the CI invocation:

  uv run pytest -m quick --ignore=test/test_dataloading.py --ignore=test/test_dixc_components.py

• Where applicable, validate energies against established reference implementations such as PySCF using small example molecules from egxc.systems.examples with basis sets like def2-SVP.
• Run uv run pytest for the tests that are affected by your changes before committing.
• Do not relax test strictness to make checks pass; instead, fix the underlying issues and rerun the checks.
• Run pre-commit run --files <changed files> on any modified files before committing.