Local kinetic energy for neural VMC

Variational Monte Carlo with neural-network wavefunctions spends most of its compute on the local kinetic energy — a Laplacian of the log-wavefunction at every walker. omnibias provides a drop-in, bit-identical replacement.

Prerequisites

pip install omnibias-ferminet     # pulls in jax + core

The bottleneck

The local kinetic energy is:

E_kin(x) = -1/2 · ( ∇² logψ(x) + ‖∇ logψ(x)‖² )

The Laplacian term is the expensive one. The textbook autodiff routes — a dense Hessian trace or a forward-Laplacian — are correct but costly, and the dense path is quadratic in the electron-coordinate dimension.

A drop-in adapter

from omnibias.ferminet.folx_compat import forward_laplacian
from omnibias.ferminet.integration import (
    make_omnibias_envelope_local_kinetic_energy,
)

# Drop-in for FermiNet's laplacian_method = "default".
local_kinetic = make_omnibias_envelope_local_kinetic_energy(
    network=ansatz,
    # ... envelope and system configuration ...
)

E_kin = local_kinetic(params, walkers)

Bit-identical to the method you trust

On closed-shell systems this is bit-identical to FermiNet's default autograd Laplacian — verified to rel-err 0.0 (ULP) on a Beryllium 8-determinant body and ≤ 5.07e-15 versus folx. The energy curve does not move; only the wall-clock and memory ceiling change.

Why bit-identical matters here

Swapping a numerical kernel in the inner loop of a physics calculation is risky: if the new kernel returns subtly different numbers, you cannot distinguish a speedup from a regression. "Bit-identical to the trusted method" removes that risk entirely — the optimizer sees the same landscape.

Relativistic corrections

Mass–velocity corrections need iterated Laplacians Δᵏψ. This is where the closed form shines: it stays flat in k instead of exploding, because each order is one tower evaluation, not a nested graph.

from omnibias.jax import neural_field_polylaplacian
# Δᵏ of the log-wavefunction stays O(1) per order.

Performance shape

Next

• Performance guide — measure it on your system.
• FermiNet API — the full bridge surface.
• Benchmarks — the published numbers.