funxc

Differentiable exchange–correlation functionals in JAX: a clean-room reimplementation of libxc, aiming for numeric compatibility with libxc 7.0.0 while being composable, autodifferentiable, and hardware-agnostic.

Why

1. libxc defines its mathematics in Maple (non-FOSS) and ships generated C; funxc defines each functional family once as a pure JAX function.
2. JAX provides a hardware-agnostic acceleration layer (CPU/GPU/TPU) and exact derivatives to any order via autodiff, rather than generated derivative code.
3. Differentiable functionals open the path from SCF-on-a-BO-surface toward end-to-end gradient methods.

Most of libxc's 600+ functionals are a small set of analytic kernels with different parameters. funxc implements the kernels as code and the parameters as a TOML database: adding a parameter-only variant is one TOML entry.

Usage

JAX-native:

import funxc
import jax, jax.numpy as jnp

f = funxc.functional("GGA_X_PBE")            # polarized by default
zk = f.exc(rho, sigma)                       # rho (N,2), sigma (N,3)
out = f.exc_vxc(rho, sigma)                  # zk, vrho, vsigma via autodiff

# per-point functions compose with any JAX transform
grad_e = jax.grad(f.energy_density, argnums=(0, 1))

pylibxc drop-in shim:

from funxc.libxc_compat import LibXCFunctional
f = LibXCFunctional("gga_x_pbe", "polarized")
out = f.compute({"rho": rho, "sigma": sigma})   # zk, vrho, vsigma arrays

Compatibility & testing

Every registered functional is validated against libxc 7.0.0's own regression data (energies and first derivatives, polarized and unpolarized) to a tolerance of 2e-9 + 1e-7·|ref| for a given |ref| per functional; typical agreement is machine precision. Kernels are guarded so that values and gradients (including second derivatives) stay finite at ρ→0 tails, σ=0, ζ=±1, s→∞ and τ→0.

uv sync
uv run pytest

Status

Stage 1 complete (39 functionals): Slater exchange, PW92 + VWN local correlation, the PBE rational-exchange family (PBE, PBEsol, revPBE, xPBE, APBE, ...), RPBE, B88, the PBE H(t) correlation family, LYP, and the B97 power-series family (B97-D, B97-3c, B97-GGA1, the HCTH functionals, HLE16). Meta-GGA plumbing (τ and ∇²ρ inputs, vtau/vlapl via autodiff) is in place and validated with the LTA family. Stage 2 kernels (TPSS, SCAN/r2SCAN) and hybrid/range-separated recipes are next; see mapOfFunctionals.md for the roadmap.