jxc-python 0.1.0__tar.gz

jxc_python-0.1.0/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: jxc-python
+Version: 0.1.0
+Summary: JAX-based exchange-correlation functionals library
+Requires-Python: <3.14,>=3.11
+Requires-Dist: wheel>=0.37.0
+Requires-Dist: setuptools==68.2.2
+Requires-Dist: absl-py==2.1.0
+Requires-Dist: jaxtyping>=0.2.0
+Requires-Dist: numpy<3.0.0,>=2.0.0
+Requires-Dist: scipy>=1.15.0
+Requires-Dist: ml-dtypes>=0.5.0
+Requires-Dist: opt-einsum>=3.4.0
+Requires-Dist: jinja2>=3.0.0
+Requires-Dist: matplotlib>=3.10.7

jxc_python-0.1.0/README.md

@@ -0,0 +1,320 @@
+# JXC
+
+Translation of LibXC into JAX and Julia.
+
+## Quickstart
+
+### Python Usage
+
+**Preferred API: `get_xc_functional()` (supports composites)**
+
+`jxc.get_xc_functional()` returns a callable that handles both single-component
+and composite functionals without having to import generated Maple modules:
+
+``` python
+import jxc
+import jax.numpy as jnp
+
+# Works with single-component functionals
+lda_func = jxc.get_xc_functional('lda_x', polarized=False)
+rho = jnp.array([0.1, 0.2, 0.3])
+print(lda_func(rho))  # Returns epsilon_xc values
+
+# Works with composite functionals (e.g., B3LYP)
+b3lyp = jxc.get_xc_functional('hyb_gga_xc_b3lyp', polarized=False)
+sigma = jnp.array([0.01, 0.05, 0.1])  # |grad rho|^2
+print(b3lyp(rho, s=sigma))  # Composite of 4 functionals
+
+# Access hybrid parameters
+print(b3lyp.cam_alpha)  # 0.2 for B3LYP (20% HF exchange)
+```
+
+This API automatically handles:
+- **Single-component functionals**: Direct calls to generated Maple translations (272 functionals)
+- **Composite functionals**: Weighted combinations like B3LYP (407 additional functionals)
+- **Hybrid metadata**: Exposes `cam_alpha`, `cam_beta`, `cam_omega` for range-separated hybrids
+- **NLC parameters**: Exposes `nlc_b`, `nlc_C` for non-local correlation
+
+#### High-Level Derivative API (Python)
+
+`get_xc_functional` also returns fully assembled derivative callables:
+
+```python
+import jxc
+import numpy as np
+
+rho = np.linspace(1e-4, 0.5, 32)
+sigma = 0.1 * rho**(4/3)
+
+vxc = jxc.get_xc_functional("gga_x_pbe", polarized=False, order="vxc")
+print("vrho:", vxc(rho, sigma=sigma)["vrho"][:3])
+
+fxc = jxc.get_xc_functional("gga_x_pbe", polarized=False, order="fxc")
+print("v2rho2:", fxc(rho, sigma=sigma)["v2rho2"][:3])
+
+# Higher orders use the same API:
+kxc = jxc.get_xc_functional("gga_x_pbe", polarized=False, order="kxc")
+lxc = jxc.get_xc_functional("gga_x_pbe", polarized=False, order="lxc")
+print("v3rho3 shape:", kxc(rho, sigma=sigma)["v3rho3"].shape)
+print("v4rho4 shape:", lxc(rho, sigma=sigma)["v4rho4"].shape)
+```
+
+All derivative orders (`vxc`, `fxc`, `kxc`, `lxc`) can also accept a
+precomputed parameter object as the first positional argument
+(`f(params, rho, ...)`), but this is optional; if omitted, `get_xc_functional`
+internally calls `get_params` for you. Julia’s API remains EXC-only for now.
+
+#### Listing available functionals
+
+To discover which LibXC functionals JXC can use, call:
+
+```python
+import jxc
+
+all_names = jxc.list_functionals()
+print(len(all_names), "functionals available")
+print([n for n in all_names if "blyp" in n.lower()])
+```
+
+This uses LibXC’s registry via `pylibxc.util.xc_available_functional_names()`
+when available and otherwise falls back to the set of generated Maple modules
+in `jxc.functionals`.
+
+#### Example: constructing BLYP from components
+
+LibXC does not expose a single `gga_xc_blyp` name; “BLYP” is conventionally
+`B88` exchange + `LYP` correlation. You can build it explicitly:
+
+```python
+import jax.numpy as jnp
+import jxc
+
+rho = jnp.linspace(1e-4, 0.5, 32)
+s = 0.1 * rho ** (4 / 3)  # |∇ρ|^2 for GGA
+
+ex_b88 = jxc.get_xc_functional("gga_x_b88", polarized=False)  # exchange only
+ec_lyp = jxc.get_xc_functional("gga_c_lyp", polarized=False)  # correlation only
+
+def blyp(rho, s=None):
+    """Pure BLYP = B88 exchange + LYP correlation."""
+    return ex_b88(rho, s=s) + ec_lyp(rho, s=s)
+
+eps_xc = blyp(rho, s=s)
+print("BLYP ε_xc:", eps_xc[:3])
+```
+
+This BLYP construction is different from “B3LYP without HF”: B3LYP uses
+additional LSDA/VWN mixing and fixed coefficients on B88/LYP, whereas the
+example above is pure B88 + pure LYP.
+
+**Important Note on Hybrid Functionals:**
+JXC returns the DFT component (epsilon_xc) only. For hybrid functionals with exact (Hartree-Fock) exchange:
+- Use `cam_alpha` to determine the HF exchange fraction (e.g., 0.2 for B3LYP = 20% HF exchange)
+- Your upstream code must compute the HF exchange term separately
+- Total energy: E_xc = (1 - cam_alpha) * E_DFT + cam_alpha * E_HF
+
+### Julia Usage (Package: JXC)
+
+We provide a native Julia module named `JXC` under `JXC.jl`.
+
+Start a Julia REPL in the module folder (set `JULIA_PKG_SERVER=""` if your network blocks the regional mirrors):
+```bash
+pushd JXC.jl
+export JULIA_PKG_SERVER=""   # optional but avoids geo-routed mirrors such as in.pkg.julialang.org
+julia --project
+```
+
+Then in Julia (option A: use as a package):
+```julia
+using Pkg
+Pkg.instantiate()
+using JXC
+
+# Python-style EXC API parity
+lda = JXC.get_xc_functional("lda_x"; polarized=false)
+rho = fill(0.3, 4)
+p = JXC.get_params("lda_x", JXC.XC_UNPOLARIZED)
+println("LDA ε_xc: ", lda(p, rho))
+
+# Composite example (B3LYP)
+b3lyp = JXC.get_xc_functional("hyb_gga_xc_b3lyp"; polarized=false)
+sigma = fill(0.01, length(rho))
+println("B3LYP ε_xc: ", b3lyp(rho; s=sigma))
+println("B3LYP cam_alpha: ", b3lyp.cam_alpha)
+```
+
+Or (option B: without package registration): `include("src/JXC.jl"); using .JXC`
+If you are activating a different project and want to make this checkout available globally,
+run `Pkg.develop(path="/path/to/JXC.jl")` from that outer environment.
+
+Notes:
+- Python Maple outputs live under `jxc/functionals/` (tracked in git).
+- Julia Maple outputs live under `JXC.jl/src/functionals/` (tracked in git).
+- Julia callables expose the same hybrid/NLC metadata as Python (`cam_alpha`, `cam_beta`, `cam_omega`, `nlc_b`, `nlc_C`) so you can inspect HF fractions directly from `get_xc_functional`.
+- To (re)generate locally, run `make convert-maple` and/or `make convert-maple-julia`, then `make pregenerate-commit` to commit the results.
+
+## Testing & Parity
+
+| Command | Scope | Notes |
+| --- | --- | --- |
+| `bazel test //tests:jax_parity_test_*` | Python energy parity | auto-generated targets covering the full LibXC catalogue |
+| `bazel test //tests/derivatives/vxc:all_vxc_functionals` | Python VXC parity | compares `jxc.derivatives.ad_derivs` derivatives against LibXC |
+| `bazel test //tests:julia_parity_test` | Julia energy smoke | mirrors the structure of `jax_parity_test.py` for a small representative set (`lda_x`, `lda_c_pw`, `gga_x_pbe`, `gga_c_pbe`) |
+| `bazel test //tests:julia_vxc_parity_test` | Julia VXC smoke | ensures the PythonCall-backed Laplacian helpers are wired correctly |
+| `make julia-test` | Convenience wrapper | runs the two Julia Bazel targets above |
+
+The parity suites are the canonical source of truth—consult `tests/jax_parity_test.py`, `tests/derivatives/vxc/`, and `tests/julia_parity_test.jl` for the exact tolerances in use. Expand the Julia list as additional functionals stabilise.
+
+## Prerequisites
+
+### Required Software
+
+1. **Maple 18** (optional) – Only needed if you regenerate functionals
+   - If not available, pre-generated functionals in `JXC.jl/src/functionals` are used
+
+2. **Python 3.12** - Required Python version (project uses `uv` for env management)
+
+3. **Build tools**:
+   - CMake (>= 3.5)
+   - C/C++ compiler (gcc/clang)
+   - wget
+   - tar
+   - patch
+
+### System Dependencies
+`Arch Linux`
+```bash
+sudo pacman -S cmake gcc wget tar patch
+```
+
+`Ubuntu/Debian`
+```bash
+sudo apt install cmake build-essential wget tar patch
+```
+
+
+## Local Build
+
+All commands below assume you are inside `projects/jxc/`. The repository uses
+[`uv`](https://docs.astral.sh/uv/) so you can stay in a regular shell—prefix
+build steps with `uv run` instead of activating the virtualenv.
+
+```bash
+# create venv, build and install both pylibxc + jxc (default python version 3.13)
+make install
+
+# switch to python 3.12
+PYVER=3.12 make install
+
+# Launch an interactive shell
+uv run ipython
+
+# Discover other handy targets
+make help
+```
+
+`make build-wheel` automatically runs `git submodule update --init --recursive`,
+regenerates the helper glue, builds (or reuses) the LibXC core under
+`.libxc-core/`, compiles the pybind11 helper against that installation, bundles
+the Python sources, and drops a wheel into `dist/`. Use `make install` to
+install the freshly built `pylibxc` and `jxc` wheels into the selected `uv`
+environment. The first phase (`scripts/build_libxc_core.sh`) only
+runs when `.libxc-core` is absent; subsequent builds for new Python versions
+reuse the same C artifacts.
+
+The Makefile always targets `.venv`; when you set `PYVER=3.12` (or similar)
+the `venv` prerequisite will recreate `.venv` with that interpreter before the
+rest of the build runs. `make build` remains an alias for the same pipeline.
+
+### Julia (Native Package: JXC)
+
+The Julia package now lives under `projects/jxc/JXC.jl`.
+
+> **Heads-up:** `JXC.get_params` defers to the Python helper via
+> [PythonCall.jl](https://github.com/JuliaPy/PythonCall.jl). Build the helper
+> (`make build-wheel` / `make install`) so that `pylibxc` and the `jxc.helper`
+> extension are importable, or point Julia at your existing virtualenv with
+> `JULIA_CONDAPKG_BACKEND=Null JULIA_PYTHONCALL_EXE=@venv julia --project`.
+> If those prerequisites are missing, the Julia API will fall back to a minimal
+> parameter stub and emit a warning.
+
+> Use `make sync-julia-wheels` (and optionally `PYVER=3.12 make sync-julia-wheels`)
+> to copy freshly built `pylibxc`/`jxc` wheels, the matching helper, and a cached
+> NumPy wheel into `JXC.jl/python_wheels`. Run `make package-julia` to produce a
+> tarball under `artifacts/` that GitHub Actions can publish as the Julia bundle.
+
+If you are preparing a distributable Julia package, run `make sync-julia-wheels`
+after building so the latest `pylibxc`/`jxc` wheels are copied into
+`JXC.jl/python_wheels/`. The Julia runtime will install from that cache on first
+use when the helper is absent.
+
+```bash
+# Start Julia REPL in project
+pushd projects/jxc/JXC.jl && julia --project
+```
+
+```julia
+using Pkg; Pkg.instantiate(); using JXC
+p = JXC.get_params("lda_x", JXC.XC_UNPOLARIZED); println("LDA unpol:", JXC.lda_x.unpol(p, 0.3))
+pp = JXC.get_params("gga_x_pbe", JXC.XC_UNPOLARIZED); println("PBE unpol:", JXC.gga_x_pbe.unpol(pp, 0.3))
+```
+
+## Performance Benchmarks
+
+JXC provides significant performance improvements over pylibxc, especially for larger batch sizes. Below are benchmark results comparing JXC against pylibxc on both CPU and GPU backends.
+
+### CPU Performance
+
+![CPU Benchmark](docs/bench_cpu.png)
+
+**CPU Speedup Distribution:**
+- **Batch=1,000**: Median 7.6x faster, Mean 9.5x faster (range: 0.3x - 120x)
+- **Batch=100,000**: Median 439x faster, Mean 616x faster (range: 3.1x - 9,120x)
+
+### GPU Performance
+
+![GPU Benchmark](docs/bench_gpu.png)
+
+**GPU Speedup Distribution:**
+- **Batch=1,000**: Median 2.7x faster, Mean 3.3x faster (range: 0.0x - 14.9x)
+- **Batch=100,000**: Median 194x faster, Mean 243x faster (range: 0.0x - 1,413x)
+
+The histograms show that JXC provides substantial speedups across most functionals, with performance gains increasing significantly for larger batch sizes. The CPU backend shows particularly impressive speedups, with many functionals achieving >100x performance improvements at batch size 100,000.
+
+### Derivative Performance (VXC)
+
+![VXC Derivative Benchmark](docs/bench_deriv_vxc.png)
+
+**VXC Derivative Speedup Highlights:**
+- **Maple-generated code**: Up to 130x faster than pylibxc for functionals like `mgga_x_r2scan` and `hyb_gga_xc_b3lyp`
+- **JAX AD**: Comparable performance to Maple code, with most functionals showing 20-100x speedup
+- Median speedup across all functionals: ~25x for both Maple and AD approaches
+
+The chart compares first derivative (VXC) computation performance for a batch size of 100 points. Blue bars show speedup from Maple-generated analytical derivatives, while orange bars show JAX automatic differentiation performance. Both approaches significantly outperform pylibxc's implementation.
+
+**Benchmark Details:**
+- Tested on 620 functionals (626 total, 6 excluded due to extreme compilation times)
+- Each functional tested with polarized/unpolarized variants
+- Batch sizes: 1,000 and 100,000 grid points
+- Speedup = pylibxc_time / jxc_time (higher is better)
+
+## Notes on Coverage
+
+- The Bazel parity suites (see table above) are the authoritative source for supported functionals and tolerances.
+- The Julia harness is intentionally limited to a smoke subset while the native codegen evolves.
+- Deprecated tables that previously tracked unsupported or numerically delicate functionals have been removed; refer to the test logs instead.
+
+## Troubleshooting
+
+### Maple Not Found
+Maple is only required for local code generation (`make convert-maple*`). CI never runs Maple.
+If you need to regenerate locally:
+1. Ensure Maple 18 is installed
+2. Add to PATH: `export PATH="$HOME/maple18/bin:$PATH"`
+3. Or set: `export MAPLE_PATH="$HOME/maple18/bin"`
+
+### CMake Errors
+If CMake fails:
+- Ensure CMake >= 3.5 is installed
+- Check that the patch file is applied correctly
+- Run `make distclean` and try again

jxc_python-0.1.0/jxc/__init__.py

@@ -0,0 +1,28 @@
+"""
+JAX-XC: JAX implementations of exchange-correlation functionals.
+This file should be copied to jxc/__init__.py during build.
+"""
+
+from . import functionals
+
+# Import main functions from get_params if it exists
+try:
+    from .get_params import (
+        XC_POLARIZED,
+        XC_UNPOLARIZED,
+        get_params,
+        get_xc_functional,
+        list_functionals,
+    )
+
+    __all__ = [
+        "functionals",
+        "get_params",
+        "get_xc_functional",
+        "list_functionals",
+        "XC_UNPOLARIZED",
+        "XC_POLARIZED",
+    ]
+except ImportError:
+    # If get_params doesn't exist, just export functionals
+    __all__ = ["functionals"]

jxc_python-0.1.0/jxc/composite.py

@@ -0,0 +1,175 @@
+"""
+Builder for composite (mixed) functionals.
+
+Composite functionals are linear combinations of component functionals,
+such as B3LYP which combines LDA exchange, B88 GGA exchange, VWN correlation, and LYP correlation.
+
+All functionals are expected to have Maple-derived Python implementations;
+if a module is missing, the runtime raises an explicit error so it can be
+fixed via codegen rather than silently falling back to LibXC.
+"""
+
+import importlib
+from typing import Callable, Any
+
+from .get_params import get_params as _get_params
+
+
+def make_epsilon_xc(p, deo_functional=None) -> Callable:
+    """
+    Build an epsilon_xc function from a params namedtuple.
+
+    Parameters
+    ----------
+    p : namedtuple
+        Parameters from get_params(), containing fields:
+        - maple_name: str (empty for composite, "DEORBITALIZE" for deorbitalized, or actual name)
+        - func_aux: list of child params (for composite/deorbitalize)
+        - mix_coef: list of float (weights for composite)
+        - cam_alpha, cam_beta, cam_omega: float (range-separated hybrid params)
+        - nlc_b, nlc_C: float (non-local correlation params)
+        - nspin: int (1 for unpolarized, 2 for polarized)
+    deo_functional : Callable, optional
+        For DEORBITALIZE functionals, the deorbitalized functional to use.
+
+    Returns
+    -------
+    Callable
+        A function that computes epsilon_xc given input densities.
+        For composites, also has attributes: cam_alpha, cam_beta, cam_omega, nlc_b, nlc_C
+    """
+    if p.maple_name == "DEORBITALIZE":
+        # Deorbitalize pattern: p.func_aux = [p0, p1]
+        # Build p1 first, then build p0 using p1 as deo_functional
+        # TODO: Full implementation requires threading mo (molecular orbitals) through inputs
+        p0, p1 = (p.func_aux[0], p.func_aux[1])
+        epsilon_xc_p1 = make_epsilon_xc(p1)
+        epsilon_xc_p0 = make_epsilon_xc(p0, epsilon_xc_p1)
+        return epsilon_xc_p0
+
+    elif hasattr(p, 'func_aux') and p.func_aux and len(p.func_aux) > 0:
+        # Composite: linear mix of func_aux with mix_coef weights
+        # Detect via presence of func_aux (not maple_name=="")
+        # Rehydrate child params if helper provided uninitialized params (common in LibXC composites)
+        def _nonempty_params(cp) -> bool:
+            try:
+                return bool(getattr(cp, 'params', None) and cp.params._asdict())
+            except Exception:
+                return False
+        children_ps = []
+        for child_p in p.func_aux:
+            if _nonempty_params(child_p):
+                children_ps.append(child_p)
+            else:
+                child_name = getattr(child_p, 'name', getattr(child_p, 'maple_name', ''))
+                if not child_name:
+                    children_ps.append(child_p)
+                else:
+                    try:
+                        rebuilt = _get_params(child_name, polarized=(p.nspin == 2))
+                        children_ps.append(rebuilt)
+                    except Exception:
+                        children_ps.append(child_p)
+        # Build children outside closure to avoid recompiling on every call
+        # Preserve child-specific CAM params when rehydrating children.
+        # If a child entry from helper already carries cam_* fields, keep them.
+        # Only rebuild minimal child metadata via get_params when the helper left
+        # params empty, and carry over any cam_* present on the original child.
+        patched_children = []
+        for orig_cp, cp in zip(p.func_aux, children_ps):
+            rebuilt = cp
+            # If orig child has explicit cam params, prefer them over any defaults
+            for attr in ("cam_alpha", "cam_beta", "cam_omega"):
+                if hasattr(orig_cp, attr):
+                    try:
+                        rebuilt = rebuilt._replace(**{attr: getattr(orig_cp, attr)})
+                    except Exception:
+                        # Not a namedtuple or field missing; ignore
+                        pass
+            patched_children.append(rebuilt)
+
+        children = [make_epsilon_xc(cp) for cp in patched_children]
+
+        EXCLUDE_BASE_ADDITION = {
+            'hyb_gga_xc_lb07',
+            'hyb_gga_xc_apbe0',  # defined in zvpbeloc.c; base is not part of mix
+        }
+
+        def epsilon_xc(*args, **kwargs):
+            # Specialization: LibXC mixes b0kcis as (1 - alpha) * Ex_B88 + 2 * Ec_KCIS.
+            if p.name == 'hyb_mgga_xc_b0kcis':
+                try:
+                    ex = children[0](*args, **kwargs)
+                    corr = children[1](*args, **kwargs)
+                    return (1.0 - getattr(p, 'cam_alpha', 0.25)) * ex + 2.0 * corr
+                except Exception:
+                    pass
+            # Weighted sum - children already built
+            result = sum(coef * child(*args, **kwargs)
+                         for coef, child in zip(p.mix_coef, children))
+
+            # For some LibXC hybrids, a base functional is implied in addition
+            # to func_aux; re-add it if not already included, except for known
+            # cases (e.g., LB07) where this would double-count.
+            if p.maple_name and (p.name not in EXCLUDE_BASE_ADDITION):
+                base_in_aux = any(child_p.name == p.maple_name or child_p.maple_name == p.maple_name
+                                   for child_p in p.func_aux)
+                if not base_in_aux:
+                    try:
+                        base_module = importlib.import_module(f"jxc.functionals.{p.maple_name}")
+                        base_func = base_module.unpol if p.nspin == 1 else base_module.pol
+                        # Rebuild correct params for the base functional, and propagate CAM params
+                        base_p = _get_params(p.maple_name, polarized=(p.nspin == 2))
+                        try:
+                            base_p = base_p._replace(
+                                cam_alpha=getattr(p, 'cam_alpha', 0.0),
+                                cam_beta=getattr(p, 'cam_beta', 0.0),
+                                cam_omega=getattr(p, 'cam_omega', 0.0),
+                            )
+                        except Exception:
+                            pass
+                        result = result + base_func(base_p, *args, **kwargs)
+                    except (ImportError, AttributeError):
+                        pass
+            return result
+
+        # Attach hybrid/nlc metadata as attributes
+        epsilon_xc.cam_alpha = p.cam_alpha
+        epsilon_xc.cam_beta = p.cam_beta
+        epsilon_xc.cam_omega = p.cam_omega
+        epsilon_xc.nlc_b = p.nlc_b
+        epsilon_xc.nlc_C = p.nlc_C
+        return epsilon_xc
+
+    else:
+        # Single maple-based functional: import the generated module
+        module_name = p.maple_name or p.name
+
+        try:
+            module = importlib.import_module(f"jxc.functionals.{module_name}")
+        except ImportError as exc:
+            raise ImportError(
+                f"No Maple implementation found for '{p.name}' (expected module '{module_name}'). "
+                "Run the Maple codegen pipeline to generate it."
+            ) from exc
+
+        # Select pol/unpol based on nspin
+        if p.nspin == 1:
+            func = module.unpol
+        elif p.nspin == 2:
+            func = module.pol
+        else:
+            raise ValueError(f"Unsupported nspin={p.nspin} for {p.maple_name}")
+
+        # Wrap to pass full p object
+        def epsilon_xc(*args, **kwargs):
+            return func(p, *args, **kwargs)
+
+        # Attach hybrid/nlc metadata
+        epsilon_xc.cam_alpha = p.cam_alpha
+        epsilon_xc.cam_beta = p.cam_beta
+        epsilon_xc.cam_omega = p.cam_omega
+        epsilon_xc.nlc_b = p.nlc_b
+        epsilon_xc.nlc_C = p.nlc_C
+
+        return epsilon_xc

jxc_python-0.1.0/jxc/functionals/__init__.py

@@ -0,0 +1,9 @@
+"""Maple-generated functionals and helpers.
+
+This package contains EXC energy implementations under this directory and
+first-derivative VXC implementations under the subpackage `vxc` when generated
+via `make convert-maple`.
+"""
+
+__all__ = []
+

jxc_python-0.1.0/jxc/functionals/_libxc_laplacian.py

@@ -0,0 +1,119 @@
+"""LibXC-backed wrappers for Laplacian-sensitive MGGA XC functionals."""
+
+from __future__ import annotations
+
+import functools
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pylibxc
+
+
+@functools.lru_cache(maxsize=None)
+def _get_libxc(name: str, spin: int) -> pylibxc.LibXCFunctional:
+  return pylibxc.LibXCFunctional(name, spin)
+
+
+def _prepare_unpolarized(rho: jnp.ndarray, sigma: jnp.ndarray, lapl: jnp.ndarray):
+  rho_arr = jnp.asarray(rho, dtype=jnp.float64).reshape(-1, 1)
+  sigma_arr = jnp.asarray(sigma, dtype=jnp.float64).reshape(-1, 1)
+  lapl_arr = jnp.asarray(lapl, dtype=jnp.float64).reshape(-1, 1)
+  return rho_arr, sigma_arr, lapl_arr
+
+
+def _prepare_polarized(
+  rho: Tuple[jnp.ndarray, jnp.ndarray],
+  sigma: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
+  lapl: Tuple[jnp.ndarray, jnp.ndarray],
+):
+  rho_arr = jnp.stack([jnp.asarray(rho[0], dtype=jnp.float64),
+                       jnp.asarray(rho[1], dtype=jnp.float64)], axis=1)
+  sigma_arr = jnp.stack(
+      [jnp.asarray(sigma[0], dtype=jnp.float64),
+       jnp.asarray(sigma[1], dtype=jnp.float64),
+       jnp.asarray(sigma[2], dtype=jnp.float64)],
+      axis=1,
+  )
+  lapl_arr = jnp.stack([jnp.asarray(lapl[0], dtype=jnp.float64),
+                        jnp.asarray(lapl[1], dtype=jnp.float64)], axis=1)
+  return rho_arr, sigma_arr, lapl_arr
+
+
+def _libxc_callback(name: str, spin: int, rho, sigma, lapl):
+  rho_np = np.asarray(rho, dtype=np.float64)
+  sigma_np = np.asarray(sigma, dtype=np.float64)
+  lapl_np = np.asarray(lapl, dtype=np.float64)
+  inputs = {
+      "rho": rho_np,
+      "sigma": sigma_np,
+      "lapl": lapl_np,
+  }
+  out = _get_libxc(name, spin).compute(inputs, do_vxc=True)
+  energy = np.asarray(out["zk"], dtype=np.float64).reshape(rho_np.shape[0])
+  vrho = np.asarray(out["vrho"], dtype=np.float64)
+  vsigma = np.asarray(out["vsigma"], dtype=np.float64)
+  vlapl = np.asarray(out["vlapl"], dtype=np.float64)
+  return energy, vrho, vsigma, vlapl
+
+
+def _libxc_energy_core(name: str, spin: int, rho, sigma, lapl):
+  shapes = (
+      jax.ShapeDtypeStruct((rho.shape[0],), jnp.float64),
+      jax.ShapeDtypeStruct((rho.shape[0], rho.shape[1]), jnp.float64),
+      jax.ShapeDtypeStruct((sigma.shape[0], sigma.shape[1]), jnp.float64),
+      jax.ShapeDtypeStruct((lapl.shape[0], lapl.shape[1]), jnp.float64),
+  )
+  return jax.pure_callback(
+      lambda rho_cb, sigma_cb, lapl_cb: _libxc_callback(name, spin, rho_cb, sigma_cb, lapl_cb),
+      shapes,
+      rho,
+      sigma,
+      lapl,
+  )
+
+
+@functools.lru_cache(maxsize=None)
+def _make_energy_fn(name: str, spin: int):
+  def core(rho, sigma, lapl):
+    return _libxc_energy_core(name, spin, rho, sigma, lapl)
+
+  @jax.custom_vjp
+  def energy(rho, sigma, lapl):
+    energy_val, _, _, _ = core(rho, sigma, lapl)
+    return energy_val
+
+  def fwd(rho, sigma, lapl):
+    energy_val, vrho, vsigma, vlapl = core(rho, sigma, lapl)
+    return energy_val, (vrho, vsigma, vlapl)
+
+  def bwd(residual, g):
+    vrho, vsigma, vlapl = residual
+    g_arr = jnp.asarray(g, dtype=jnp.float64).reshape(-1, 1)
+    grad_rho = g_arr * jnp.asarray(vrho, dtype=jnp.float64)
+    grad_sigma = g_arr * jnp.asarray(vsigma, dtype=jnp.float64)
+    grad_lapl = g_arr * jnp.asarray(vlapl, dtype=jnp.float64)
+    return grad_rho, grad_sigma, grad_lapl
+
+  energy.defvjp(fwd, bwd)
+  return energy
+
+
+def eval_unpolarized(name: str, rho, sigma, lapl):
+  rho_arr, sigma_arr, lapl_arr = _prepare_unpolarized(rho, sigma, lapl)
+  energy_fn = _make_energy_fn(name, 1)
+  energy = energy_fn(rho_arr, sigma_arr, lapl_arr)
+  return energy.reshape(-1)
+
+
+def eval_polarized(
+  name: str,
+  rho: Tuple[jnp.ndarray, jnp.ndarray],
+  sigma: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
+  lapl: Tuple[jnp.ndarray, jnp.ndarray],
+):
+  rho_arr, sigma_arr, lapl_arr = _prepare_polarized(rho, sigma, lapl)
+  energy_fn = _make_energy_fn(name, 2)
+  energy = energy_fn(rho_arr, sigma_arr, lapl_arr)
+  return energy.reshape(-1)