70 lines
2.5 KiB
Markdown
70 lines
2.5 KiB
Markdown
# funxc
|
||
|
||
Differentiable exchange–correlation functionals in JAX: a clean-room
|
||
reimplementation of [libxc](https://libxc.gitlab.io/), 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:
|
||
|
||
```python
|
||
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:
|
||
|
||
```python
|
||
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.
|