omnibias.jax.bo_derivatives
JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.
import omnibias.jax.bo_derivativesModule contract & mathematical conventions
Born-Oppenheimer energy gradient and Hessian estimators.
This module computes nuclear-coordinate derivatives of the Born-Oppenheimer energy
E_{\mathrm{BO}}(R) = \bigl\langle\psi(r;R)\bigm|\hat H(R)\bigm|\psi(r;R)\bigr\rangle /\,\bigl\langle\psi\bigm|\psi\bigr\rangle
as Monte Carlo estimators over walkers r_w drawn from
\rho(r;R) = |\psi(r;R)|^{2}/Z. The unique value of doing
this through omnibias is that the *inner* derivatives -- the
\nabla_r and \nabla_r^{2} of \log|\psi|
that drive the local energy -- are evaluated in closed form via
omnibias.jax.laplacian and the chain rule of
omnibias.jax.ferminet_integration, while
the *outer* derivatives w.r.t. R (low-dimensional;
3N \le 30 for the molecules we care about) are computed
with ordinary jax.grad / jax.hessian.
This gives the first analytic-second-derivative path for a
neural-VMC energy and unlocks Hessian-class observables
(vibrational frequencies, instanton rates, response properties)
that have been infeasible for neural-VMC stacks. See
H2O_ANALYTIC_HESSIAN for the chemistry headline (private
benchmark archive) and AUTOGRAD_PHASE_TRANSITION for the ICLR 2027
paper that frames the autograd compile bomb this side-steps.
Estimators (Casula--Sorella, eqs. 3, 6 of *Mol. Phys.* 109, 2473, 2011):
F_\alpha = -\,\Bigl\langle \partial_{R_\alpha} E_{\mathrm{loc}} + 2\bigl(E_{\mathrm{loc}} - \bar E\bigr)\, \partial_{R_\alpha}\log|\psi| \Bigr\rangle,
K_{\alpha\beta} = \langle\partial^{2}_{R_\alpha R_\beta} E_{\mathrm{loc}}\rangle + 2\bigl\langle (\partial_{R_\alpha} E_{\mathrm{loc}}) (\partial_{R_\beta} \log|\psi|) + (\partial_{R_\beta} E_{\mathrm{loc}}) (\partial_{R_\alpha} \log|\psi|)\bigr\rangle + 2\bigl\langle (E_{\mathrm{loc}} - \bar E) \partial^{2}_{R_\alpha R_\beta} \log|\psi| + 2\,(E_{\mathrm{loc}} - \bar E) (\partial_{R_\alpha} \log|\psi|) (\partial_{R_\beta} \log|\psi|)\bigr\rangle.
All terms cancel in expectation when \psi is the exact
ground state (zero-variance principle). With a fixed, non-exact
ansatz they do not, and the estimator above is the unbiased
form used throughout the QMC force literature.
API
make_local_energy(psi_fn, potential_fn, n_e, ndim=3) Returns e_local(R, r_flat) -> scalar using the omnibias closed-form Laplacian on log|psi| and an explicit potential_fn(R, r_flat) -> scalar for the nucleus-electron / electron-electron / nucleus-nucleus interaction. Pluggable so the same kernel works for hand- rolled wavefunctions, FermiNet, etc.
make_bo_force(psi_fn, potential_fn, n_e) Returns force(R, walkers, mean_energy) -> (3N,) -- the Casula--Sorella force estimator on a batch of walkers.
make_bo_hessian(psi_fn, potential_fn, n_e) Returns hessian(R, walkers, mean_energy) -> (3N, 3N) -- the full Casula--Sorella Hessian estimator. Symmetric by construction.
vibrational_frequencies(hessian, masses, R) Mass-weights and diagonalises the 3N x 3N nuclear Hessian, projects out the 6 (or 5 for diatomics) translation / rotation modes, and returns the remaining eigenvalues in inverse centimetres (negative for imaginary modes).
Convention
Ris the nuclear positions as a flat(3N,)array; callers reshape to(N, 3)as needed.walkersis a batch of electron positions, shape(B, n_e * 3).psi_fn(R, r_flat) -> (sign, log_abs)is the FermiNetLike convention (matches the Tier-1 plumbing inomnibias.jax.make_omnibias_envelope_local_kinetic_energy).potential_fn(R, r_flat) -> scalaris the bare Coulomb potential energy at the given (R, r); the caller assembles this from atomic charges. We exposecoulomb_potentialas the default.
Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.
AMU_TO_ME
omnibias.jax.bo_derivatives.AMU_TO_MESource ↗AMU_TO_ME = 1822.888486209atomic mass unit to electron-mass units (m_e); needed for vibrational eigenvalues since the kinetic operator has 1/m factors in atomic units.
Regression example · vibrational frequencies diatomic harmonic
This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_vibrational_frequencies_diatomic_harmonic():
r"""Toy diatomic with a known harmonic curvature.
Build K_nuclear = k * (e_R outer e_R), with e_R the bond-axis
direction, and check the unique vibrational frequency comes out
to the textbook value
nu = (1 / (2 pi c)) sqrt(k / mu)
where mu is the reduced mass and the factor of 2 pi c is absorbed
in HARTREE_TO_CM after converting eigenvalues to angular
frequencies in atomic units.
"""
m1, m2 = 1.0, 16.0 # H and O masses, in amu (toy values)
R = jnp.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 1.8], dtype=jnp.float64)
k = 0.5 # Ha / Bohr^2
# Place k between the two atoms along z.
e_R = jnp.asarray([0.0, 0.0, -1.0, 0.0, 0.0, 1.0], dtype=jnp.float64)
K = k * jnp.outer(e_R, e_R)
masses = jnp.asarray([m1, m2], dtype=jnp.float64)
freqs, _ = vibrational_frequencies(
K,
masses,
R,
project_tr=True,
diatomic=True,
)
nonzero = jnp.sort(jnp.abs(freqs))[::-1][:1]
from omnibias.jax.bo_derivatives import AMU_TO_ME
# eigenvalue of the mass-weighted Hessian for a single mode along the
# bond axis is k * (1/m1 + 1/m2) = k / mu_me with mu = m1 m2/(m1+m2).
# omega_au = sqrt(eigval); nu_cm = omega_au * HARTREE_TO_CM.
expected_cm = float(
np.sqrt(k * (1.0 / (m1 * AMU_TO_ME) + 1.0 / (m2 * AMU_TO_ME))) * HARTREE_TO_CM
)
err = float(jnp.abs(nonzero[0] - expected_cm))
assert err / expected_cm < 1e-6, (
f"diatomic vib freq: got {float(nonzero[0]):.4f}, "
f"expected {expected_cm:.4f} cm^-1, rel err = {err / expected_cm:.3e}"
)python -m pytest packages/omnibias-jax/tests/test_bo_derivatives.py -k test_vibrational_frequencies_diatomic_harmonic -qRead the implementation
AMU_TO_ME = 1822.888486209BOHR_TO_ANGSTROM
omnibias.jax.bo_derivatives.BOHR_TO_ANGSTROMSource ↗BOHR_TO_ANGSTROM = 0.52917721067Bohr to angstrom
Read the implementation
BOHR_TO_ANGSTROM = 0.52917721067HARTREE_TO_CM
omnibias.jax.bo_derivatives.HARTREE_TO_CMSource ↗HARTREE_TO_CM = 219474.6313632Hartree to cm^-1: E_h / (h c) in cm^-1
Regression example · vibrational frequencies diatomic harmonic
This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_vibrational_frequencies_diatomic_harmonic():
r"""Toy diatomic with a known harmonic curvature.
Build K_nuclear = k * (e_R outer e_R), with e_R the bond-axis
direction, and check the unique vibrational frequency comes out
to the textbook value
nu = (1 / (2 pi c)) sqrt(k / mu)
where mu is the reduced mass and the factor of 2 pi c is absorbed
in HARTREE_TO_CM after converting eigenvalues to angular
frequencies in atomic units.
"""
m1, m2 = 1.0, 16.0 # H and O masses, in amu (toy values)
R = jnp.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 1.8], dtype=jnp.float64)
k = 0.5 # Ha / Bohr^2
# Place k between the two atoms along z.
e_R = jnp.asarray([0.0, 0.0, -1.0, 0.0, 0.0, 1.0], dtype=jnp.float64)
K = k * jnp.outer(e_R, e_R)
masses = jnp.asarray([m1, m2], dtype=jnp.float64)
freqs, _ = vibrational_frequencies(
K,
masses,
R,
project_tr=True,
diatomic=True,
)
nonzero = jnp.sort(jnp.abs(freqs))[::-1][:1]
from omnibias.jax.bo_derivatives import AMU_TO_ME
# eigenvalue of the mass-weighted Hessian for a single mode along the
# bond axis is k * (1/m1 + 1/m2) = k / mu_me with mu = m1 m2/(m1+m2).
# omega_au = sqrt(eigval); nu_cm = omega_au * HARTREE_TO_CM.
expected_cm = float(
np.sqrt(k * (1.0 / (m1 * AMU_TO_ME) + 1.0 / (m2 * AMU_TO_ME))) * HARTREE_TO_CM
)
err = float(jnp.abs(nonzero[0] - expected_cm))
assert err / expected_cm < 1e-6, (
f"diatomic vib freq: got {float(nonzero[0]):.4f}, "
f"expected {expected_cm:.4f} cm^-1, rel err = {err / expected_cm:.3e}"
)python -m pytest packages/omnibias-jax/tests/test_bo_derivatives.py -k test_vibrational_frequencies_diatomic_harmonic -qRead the implementation
HARTREE_TO_CM = 219474.6313632coulomb_potential
omnibias.jax.coulomb_potentialSource ↗coulomb_potential(R: Array, r_flat: Array, charges: Array, n_e: int) -> ArrayBare Coulomb potential energy for fixed nuclear positions.
V(R, r) = -\sum_{ja} \frac{Z_a}{\lVert r_j - R_a\rVert} + \sum_{j < k} \frac{1}{\lVert r_j - r_k\rVert} + \sum_{a < b} \frac{Z_a Z_b}{\lVert R_a - R_b\rVert}.
All in atomic units. The small epsilon regularisation
that appears below is for AD stability (the gradient of
1/r is undefined at the singularity); MC walkers never
sit exactly on the singularity but the AD graph still
evaluates at sample points where the gradient is needed.
Parameters
| Name | Type & default | Contract |
|---|---|---|
Rpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
r_flatpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
chargespositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
n_epositional-or-keyword | intRequired | See the operation contract and implementation below. |
ArrayExecuted regression · translation invariant hessian is zero
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_translation_invariant_hessian_is_zero():
r"""For psi(r-R), the BO Hessian also vanishes (translation invariance)."""
alpha = 0.6
Z = 1.0
R = jnp.asarray([0.1, 0.2, -0.3], dtype=jnp.float64)
psi_fn = _gaussian_psi(alpha)
def potential_fn(R_, r_):
return coulomb_potential(
R_,
r_,
jnp.asarray([Z], dtype=jnp.float64),
n_e=1,
)
walkers = _gaussian_h_walkers(R, alpha, n_walkers=4000)
e_bar = _bo_energy_estimate_array(psi_fn, potential_fn, R, walkers)
bo_hess = make_bo_hessian(psi_fn, potential_fn)
K = bo_hess(R, walkers, e_bar)
assert K.shape == (3, 3)
max_abs = float(jnp.max(jnp.abs(K)))
assert max_abs < 5.0, f"max|K| = {max_abs:.3e}, translation-invariant ansatz should be 0."python -m pytest "packages/omnibias-jax/tests/test_bo_derivatives.py::test_translation_invariant_hessian_is_zero" -qRegression example · translation invariant force is zero
This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_translation_invariant_force_is_zero():
r"""For psi(r-R), the BO force vanishes (zero-variance principle).
We use enough walkers (symmetric around R) for the MC variance to
be well below the assertion threshold.
"""
alpha = 0.5
Z = 1.0
R = jnp.asarray([0.3, -0.7, 0.4], dtype=jnp.float64)
psi_fn = _gaussian_psi(alpha)
def potential_fn(R_, r_):
return coulomb_potential(
R_,
r_,
jnp.asarray([Z], dtype=jnp.float64),
n_e=1,
)
walkers = _gaussian_h_walkers(R, alpha, n_walkers=4000)
e_bar = _bo_energy_estimate_array(psi_fn, potential_fn, R, walkers)
bo_force = make_bo_force(psi_fn, potential_fn)
F = bo_force(R, walkers, e_bar)
assert F.shape == (3,)
max_abs = float(jnp.max(jnp.abs(F)))
# MC noise floor for the centred-gaussian estimator with 4000 walkers
# sits ~ few mHa/Bohr; the systematic part must be << that.
assert max_abs < 5e-2, f"|F|_inf = {max_abs:.3e}, expected near MC zero (~few mHa/Bohr)."python -m pytest packages/omnibias-jax/tests/test_bo_derivatives.py -k test_translation_invariant_force_is_zero -qRead the implementation
def coulomb_potential(
R: Array, # (n_atoms * 3,)
r_flat: Array, # (n_e * 3,)
charges: Array, # (n_atoms,)
n_e: int,
) -> Array:
r"""Bare Coulomb potential energy for fixed nuclear positions.
.. math::
V(R, r) =
-\sum_{ja} \frac{Z_a}{\lVert r_j - R_a\rVert}
+ \sum_{j < k} \frac{1}{\lVert r_j - r_k\rVert}
+ \sum_{a < b} \frac{Z_a Z_b}{\lVert R_a - R_b\rVert}.
All in atomic units. The small ``epsilon`` regularisation
that appears below is for AD stability (the gradient of
``1/r`` is undefined at the singularity); MC walkers never
sit exactly on the singularity but the AD graph still
evaluates at sample points where the gradient is needed.
"""
eps2 = 1e-30
n_atoms = charges.shape[0]
R_atoms = R.reshape((n_atoms, 3))
r = r_flat.reshape((n_e, 3))
# electron-nucleus
d_en = jnp.sqrt(
jnp.sum((r[:, None, :] - R_atoms[None, :, :]) ** 2, axis=-1) + eps2
) # (n_e, n_atoms)
V_en = -jnp.sum(charges[None, :] / d_en)
# electron-electron
if n_e > 1:
d_ee = jnp.sqrt(jnp.sum((r[:, None, :] - r[None, :, :]) ** 2, axis=-1) + eps2) # (n_e, n_e)
tri = jnp.triu(1.0 / d_ee, k=1)
V_ee = jnp.sum(tri)
else:
V_ee = jnp.asarray(0.0, dtype=r_flat.dtype)
# nucleus-nucleus
if n_atoms > 1:
d_nn = jnp.sqrt(
jnp.sum(
(R_atoms[:, None, :] - R_atoms[None, :, :]) ** 2,
axis=-1,
)
+ eps2
)
ZZ = charges[:, None] * charges[None, :]
tri = jnp.triu(ZZ / d_nn, k=1)
V_nn = jnp.sum(tri)
else:
V_nn = jnp.asarray(0.0, dtype=r_flat.dtype)
return V_en + V_ee + V_nnOther public import paths
omnibias.jax.coulomb_potentialomnibias.jax.bo_derivatives.coulomb_potential
make_bo_force
omnibias.jax.make_bo_forceSource ↗make_bo_force(psi_fn: Callable[[Array, Array], tuple[Array, Array]], potential_fn: Callable[[Array, Array], Array]) -> Callable[[Array, Array, Array], Array]Return force(R, walkers, mean_energy) -> (3N,).
Estimator (single walker, then average):
f_\alpha(r;R) = -\,\partial_{R_\alpha} E_{\mathrm{loc}}(R, r) -\,2\bigl(E_{\mathrm{loc}}(R, r) - \bar E\bigr) \partial_{R_\alpha} \log|\psi(R, r)|.
The minus sign convention matches F = -dE/dR. Caller
supplies mean_energy (typically a separate MC estimate of
\bar E from the same walkers).
Parameters
| Name | Type & default | Contract |
|---|---|---|
psi_fnpositional-or-keyword | Callable[[Array, Array], tuple[Array, Array]]Required | See the operation contract and implementation below. |
potential_fnpositional-or-keyword | Callable[[Array, Array], Array]Required | See the operation contract and implementation below. |
Callable[[Array, Array, Array], Array]Executed regression · translation invariant force is zero
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_translation_invariant_force_is_zero():
r"""For psi(r-R), the BO force vanishes (zero-variance principle).
We use enough walkers (symmetric around R) for the MC variance to
be well below the assertion threshold.
"""
alpha = 0.5
Z = 1.0
R = jnp.asarray([0.3, -0.7, 0.4], dtype=jnp.float64)
psi_fn = _gaussian_psi(alpha)
def potential_fn(R_, r_):
return coulomb_potential(
R_,
r_,
jnp.asarray([Z], dtype=jnp.float64),
n_e=1,
)
walkers = _gaussian_h_walkers(R, alpha, n_walkers=4000)
e_bar = _bo_energy_estimate_array(psi_fn, potential_fn, R, walkers)
bo_force = make_bo_force(psi_fn, potential_fn)
F = bo_force(R, walkers, e_bar)
assert F.shape == (3,)
max_abs = float(jnp.max(jnp.abs(F)))
# MC noise floor for the centred-gaussian estimator with 4000 walkers
# sits ~ few mHa/Bohr; the systematic part must be << that.
assert max_abs < 5e-2, f"|F|_inf = {max_abs:.3e}, expected near MC zero (~few mHa/Bohr)."python -m pytest "packages/omnibias-jax/tests/test_bo_derivatives.py::test_translation_invariant_force_is_zero" -qRegression example · force matches finite difference fixed psi
This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_force_matches_finite_difference_fixed_psi():
r"""When psi is fixed and V(R, r) depends on R, dE/dR is finite-difference
of the MC energy estimator. With ~10k walkers, the FD-vs-analytic
relative error sits in the 1% band (MC noise dominated)."""
alpha = 0.7
Z = 1.5
R = jnp.asarray([0.15, -0.20, 0.30], dtype=jnp.float64)
R0 = R + 0.05 # psi is centred slightly off from the nucleus
psi_fn = _fixed_gaussian_psi(R0, alpha)
def potential_fn(R_, r_):
return coulomb_potential(
R_,
r_,
jnp.asarray([Z], dtype=jnp.float64),
n_e=1,
)
walkers = _gaussian_h_walkers(R0, alpha, n_walkers=8000)
e_bar = _bo_energy_estimate_array(psi_fn, potential_fn, R, walkers)
bo_force = make_bo_force(psi_fn, potential_fn)
F_analytic = bo_force(R, walkers, e_bar)
# Finite-difference of E_BO(R) along each direction.
h = 1e-3
F_fd = np.zeros(3, dtype=np.float64)
for k in range(3):
R_plus = R.at[k].add(h)
R_minus = R.at[k].add(-h)
e_p = _bo_energy_estimate_array(psi_fn, potential_fn, R_plus, walkers)
e_m = _bo_energy_estimate_array(psi_fn, potential_fn, R_minus, walkers)
F_fd[k] = -(float(e_p) - float(e_m)) / (2.0 * h)
F_fd = jnp.asarray(F_fd, dtype=jnp.float64)
err = float(jnp.max(jnp.abs(F_analytic - F_fd)))
scale = float(jnp.max(jnp.abs(F_fd)) + 1e-3)
rel_err = err / scale
assert rel_err < 0.02, f"F_analytic = {F_analytic}, F_fd = {F_fd}, rel_err = {rel_err:.3e}"python -m pytest packages/omnibias-jax/tests/test_bo_derivatives.py -k test_force_matches_finite_difference_fixed_psi -qRead the implementation
def make_bo_force(
psi_fn: Callable[[Array, Array], tuple[Array, Array]],
potential_fn: Callable[[Array, Array], Array],
) -> Callable[[Array, Array, Array], Array]:
r"""Return ``force(R, walkers, mean_energy) -> (3N,)``.
Estimator (single walker, then average):
.. math::
f_\alpha(r;R)
= -\,\partial_{R_\alpha} E_{\mathrm{loc}}(R, r)
-\,2\bigl(E_{\mathrm{loc}}(R, r) - \bar E\bigr)
\partial_{R_\alpha} \log|\psi(R, r)|.
The minus sign convention matches ``F = -dE/dR``. Caller
supplies ``mean_energy`` (typically a separate MC estimate of
:math:`\bar E` from the same walkers).
"""
e_local = make_local_energy(psi_fn, potential_fn)
def log_abs_fn(R: Array, r: Array) -> Array:
return psi_fn(R, r)[1]
def force_per_walker(R: Array, r_flat: Array, e_bar: Array) -> Array:
e_l = e_local(R, r_flat)
grad_e_l = jax.grad(e_local, argnums=0)(R, r_flat)
grad_log = jax.grad(log_abs_fn, argnums=0)(R, r_flat)
out: Array = -(grad_e_l + 2.0 * (e_l - e_bar) * grad_log)
return out
def force(R: Array, walkers: Array, mean_energy: Array) -> Array:
per_walker = jax.vmap(force_per_walker, in_axes=(None, 0, None))(
R,
walkers,
mean_energy,
)
return jnp.mean(per_walker, axis=0)
return forceOther public import paths
omnibias.jax.make_bo_forceomnibias.jax.bo_derivatives.make_bo_force
make_bo_hessian
omnibias.jax.make_bo_hessianSource ↗make_bo_hessian(psi_fn: Callable[[Array, Array], tuple[Array, Array]], potential_fn: Callable[[Array, Array], Array]) -> Callable[[Array, Array, Array], Array]Return hessian(R, walkers, mean_energy) -> (3N, 3N).
The per-walker integrand is
K_{\alpha\beta}(r;R) = \partial^{2}_{R_\alpha R_\beta} E_{\mathrm{loc}}(R, r) + 2 (\partial_{R_\alpha} E_{\mathrm{loc}})(\partial_{R_\beta}\log|\psi|) + 2 (\partial_{R_\beta} E_{\mathrm{loc}})(\partial_{R_\alpha}\log|\psi|) + 2 (E_{\mathrm{loc}} - \bar E)\, \partial^{2}_{R_\alpha R_\beta} \log|\psi| + 4 (E_{\mathrm{loc}} - \bar E) (\partial_{R_\alpha} \log|\psi|) (\partial_{R_\beta} \log|\psi|).
Then averaged over walkers. The result is symmetric by construction; we explicitly symmetrise the cross-derivative term to guard against MC asymmetry in float64.
Parameters
| Name | Type & default | Contract |
|---|---|---|
psi_fnpositional-or-keyword | Callable[[Array, Array], tuple[Array, Array]]Required | See the operation contract and implementation below. |
potential_fnpositional-or-keyword | Callable[[Array, Array], Array]Required | See the operation contract and implementation below. |
Callable[[Array, Array, Array], Array]Executed regression · translation invariant hessian is zero
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_translation_invariant_hessian_is_zero():
r"""For psi(r-R), the BO Hessian also vanishes (translation invariance)."""
alpha = 0.6
Z = 1.0
R = jnp.asarray([0.1, 0.2, -0.3], dtype=jnp.float64)
psi_fn = _gaussian_psi(alpha)
def potential_fn(R_, r_):
return coulomb_potential(
R_,
r_,
jnp.asarray([Z], dtype=jnp.float64),
n_e=1,
)
walkers = _gaussian_h_walkers(R, alpha, n_walkers=4000)
e_bar = _bo_energy_estimate_array(psi_fn, potential_fn, R, walkers)
bo_hess = make_bo_hessian(psi_fn, potential_fn)
K = bo_hess(R, walkers, e_bar)
assert K.shape == (3, 3)
max_abs = float(jnp.max(jnp.abs(K)))
assert max_abs < 5.0, f"max|K| = {max_abs:.3e}, translation-invariant ansatz should be 0."python -m pytest "packages/omnibias-jax/tests/test_bo_derivatives.py::test_translation_invariant_hessian_is_zero" -qRegression example · vibrational frequencies diatomic harmonic
This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_vibrational_frequencies_diatomic_harmonic():
r"""Toy diatomic with a known harmonic curvature.
Build K_nuclear = k * (e_R outer e_R), with e_R the bond-axis
direction, and check the unique vibrational frequency comes out
to the textbook value
nu = (1 / (2 pi c)) sqrt(k / mu)
where mu is the reduced mass and the factor of 2 pi c is absorbed
in HARTREE_TO_CM after converting eigenvalues to angular
frequencies in atomic units.
"""
m1, m2 = 1.0, 16.0 # H and O masses, in amu (toy values)
R = jnp.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 1.8], dtype=jnp.float64)
k = 0.5 # Ha / Bohr^2
# Place k between the two atoms along z.
e_R = jnp.asarray([0.0, 0.0, -1.0, 0.0, 0.0, 1.0], dtype=jnp.float64)
K = k * jnp.outer(e_R, e_R)
masses = jnp.asarray([m1, m2], dtype=jnp.float64)
freqs, _ = vibrational_frequencies(
K,
masses,
R,
project_tr=True,
diatomic=True,
)
nonzero = jnp.sort(jnp.abs(freqs))[::-1][:1]
from omnibias.jax.bo_derivatives import AMU_TO_ME
# eigenvalue of the mass-weighted Hessian for a single mode along the
# bond axis is k * (1/m1 + 1/m2) = k / mu_me with mu = m1 m2/(m1+m2).
# omega_au = sqrt(eigval); nu_cm = omega_au * HARTREE_TO_CM.
expected_cm = float(
np.sqrt(k * (1.0 / (m1 * AMU_TO_ME) + 1.0 / (m2 * AMU_TO_ME))) * HARTREE_TO_CM
)
err = float(jnp.abs(nonzero[0] - expected_cm))
assert err / expected_cm < 1e-6, (
f"diatomic vib freq: got {float(nonzero[0]):.4f}, "
f"expected {expected_cm:.4f} cm^-1, rel err = {err / expected_cm:.3e}"
)python -m pytest packages/omnibias-jax/tests/test_bo_derivatives.py -k test_vibrational_frequencies_diatomic_harmonic -qRead the implementation
def make_bo_hessian(
psi_fn: Callable[[Array, Array], tuple[Array, Array]],
potential_fn: Callable[[Array, Array], Array],
) -> Callable[[Array, Array, Array], Array]:
r"""Return ``hessian(R, walkers, mean_energy) -> (3N, 3N)``.
The per-walker integrand is
.. math::
K_{\alpha\beta}(r;R)
= \partial^{2}_{R_\alpha R_\beta} E_{\mathrm{loc}}(R, r)
+ 2 (\partial_{R_\alpha} E_{\mathrm{loc}})(\partial_{R_\beta}\log|\psi|)
+ 2 (\partial_{R_\beta} E_{\mathrm{loc}})(\partial_{R_\alpha}\log|\psi|)
+ 2 (E_{\mathrm{loc}} - \bar E)\,
\partial^{2}_{R_\alpha R_\beta} \log|\psi|
+ 4 (E_{\mathrm{loc}} - \bar E)
(\partial_{R_\alpha} \log|\psi|)
(\partial_{R_\beta} \log|\psi|).
Then averaged over walkers. The result is symmetric by
construction; we explicitly symmetrise the cross-derivative
term to guard against MC asymmetry in float64.
"""
e_local = make_local_energy(psi_fn, potential_fn)
def log_abs_fn(R: Array, r: Array) -> Array:
return psi_fn(R, r)[1]
def hess_per_walker(R: Array, r_flat: Array, e_bar: Array) -> Array:
e_l = e_local(R, r_flat)
grad_e_l = jax.grad(e_local, argnums=0)(R, r_flat)
hess_e_l = jax.hessian(e_local, argnums=0)(R, r_flat)
grad_log = jax.grad(log_abs_fn, argnums=0)(R, r_flat)
hess_log = jax.hessian(log_abs_fn, argnums=0)(R, r_flat)
cross = jnp.outer(grad_e_l, grad_log)
cross_sym = cross + cross.T
delta_e = e_l - e_bar
out: Array = (
hess_e_l
+ 2.0 * cross_sym
+ 2.0 * delta_e * hess_log
+ 4.0 * delta_e * jnp.outer(grad_log, grad_log)
)
return out
def hessian(
R: Array,
walkers: Array,
mean_energy: Array,
) -> Array:
per_walker = jax.vmap(hess_per_walker, in_axes=(None, 0, None))(
R,
walkers,
mean_energy,
)
K = jnp.mean(per_walker, axis=0)
return 0.5 * (K + K.T)
return hessianOther public import paths
omnibias.jax.make_bo_hessianomnibias.jax.bo_derivatives.make_bo_hessian
make_local_energy
omnibias.jax.make_local_energySource ↗make_local_energy(psi_fn: Callable[[Array, Array], tuple[Array, Array]], potential_fn: Callable[[Array, Array], Array]) -> Callable[[Array, Array], Array]Factory returning e_local(R, r_flat) -> scalar.
Parameters
psi_fn (R, r_flat) -> (sign, log_abs). Matches the FermiNetLike convention.
potential_fn (R, r_flat) -> scalar -- the bare Coulomb potential at (R, r); see coulomb_potential for the default.
Parameters
| Name | Type & default | Contract |
|---|---|---|
psi_fnpositional-or-keyword | Callable[[Array, Array], tuple[Array, Array]]Required | (R, r_flat) -> (sign, log_abs). Matches the FermiNetLike convention. |
potential_fnpositional-or-keyword | Callable[[Array, Array], Array]Required | (R, r_flat) -> scalar -- the bare Coulomb potential at (R, r); see coulomb_potential for the default. |
Callable[[Array, Array], Array]Executed regression · translation invariant hessian is zero
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_translation_invariant_hessian_is_zero():
r"""For psi(r-R), the BO Hessian also vanishes (translation invariance)."""
alpha = 0.6
Z = 1.0
R = jnp.asarray([0.1, 0.2, -0.3], dtype=jnp.float64)
psi_fn = _gaussian_psi(alpha)
def potential_fn(R_, r_):
return coulomb_potential(
R_,
r_,
jnp.asarray([Z], dtype=jnp.float64),
n_e=1,
)
walkers = _gaussian_h_walkers(R, alpha, n_walkers=4000)
e_bar = _bo_energy_estimate_array(psi_fn, potential_fn, R, walkers)
bo_hess = make_bo_hessian(psi_fn, potential_fn)
K = bo_hess(R, walkers, e_bar)
assert K.shape == (3, 3)
max_abs = float(jnp.max(jnp.abs(K)))
assert max_abs < 5.0, f"max|K| = {max_abs:.3e}, translation-invariant ansatz should be 0."python -m pytest "packages/omnibias-jax/tests/test_bo_derivatives.py::test_translation_invariant_hessian_is_zero" -qRead the implementation
def make_local_energy(
psi_fn: Callable[[Array, Array], tuple[Array, Array]],
potential_fn: Callable[[Array, Array], Array],
) -> Callable[[Array, Array], Array]:
r"""Factory returning ``e_local(R, r_flat) -> scalar``.
Parameters
----------
psi_fn
``(R, r_flat) -> (sign, log_abs)``. Matches the
FermiNetLike convention.
potential_fn
``(R, r_flat) -> scalar`` -- the bare Coulomb potential
at ``(R, r)``; see :func:`coulomb_potential` for the
default.
"""
def log_abs_fn(R: Array, r: Array) -> Array:
return psi_fn(R, r)[1]
def e_local(R: Array, r_flat: Array) -> Array:
T = _kinetic_energy_density(log_abs_fn, R, r_flat)
V = potential_fn(R, r_flat)
return T + V
return e_localOther public import paths
omnibias.jax.make_local_energyomnibias.jax.bo_derivatives.make_local_energy
vibrational_frequencies
omnibias.jax.vibrational_frequenciesSource ↗vibrational_frequencies(K_nuclear: Array, masses_amu: Array, R: Array, project_tr: bool=True, diatomic: bool=False) -> tuple[Array, Array]Diagonalise the mass-weighted nuclear Hessian.
Parameters
K_nuclear (3N, 3N) symmetric matrix of \partial^{2} E_{\mathrm{BO}} / \partial R_\alpha \partial R_\beta in atomic units (Hartree / Bohr^2).
masses_amu (N,) atomic masses, in atomic mass units (amu).
R (3N,) equilibrium nuclear positions in Bohr. Used to construct the translation and rotation projectors.
project_tr If True (default), project out the six translation and rotation modes before diagonalising. For a diatomic molecule pass diatomic=True so we project out five modes (one rotation about the bond axis is degenerate).
diatomic See above.
Returns
freqs_cm Vibrational frequencies in inverse centimetres. Negative values indicate an imaginary mode (saddle point / instability).
eigvecs Mass-weighted Cartesian eigenvectors of the projected Hessian, shape (3N, 3N).
Parameters
| Name | Type & default | Contract |
|---|---|---|
K_nuclearpositional-or-keyword | ArrayRequired | (3N, 3N) symmetric matrix of \partial^{2} E_{\mathrm{BO}} / \partial R_\alpha \partial R_\beta in atomic units (Hartree / Bohr^2). |
masses_amupositional-or-keyword | ArrayRequired | (N,) atomic masses, in atomic mass units (amu). |
Rpositional-or-keyword | ArrayRequired | (3N,) equilibrium nuclear positions in Bohr. Used to construct the translation and rotation projectors. |
project_trpositional-or-keyword | boolDefault: True | If True (default), project out the six translation and rotation modes before diagonalising. For a diatomic molecule pass diatomic=True so we project out five modes (one rotation about the bond axis is degenerate). |
diatomicpositional-or-keyword | boolDefault: False | See above. |
tuple[Array, Array]Executed regression · vibrational frequencies diatomic harmonic
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_vibrational_frequencies_diatomic_harmonic():
r"""Toy diatomic with a known harmonic curvature.
Build K_nuclear = k * (e_R outer e_R), with e_R the bond-axis
direction, and check the unique vibrational frequency comes out
to the textbook value
nu = (1 / (2 pi c)) sqrt(k / mu)
where mu is the reduced mass and the factor of 2 pi c is absorbed
in HARTREE_TO_CM after converting eigenvalues to angular
frequencies in atomic units.
"""
m1, m2 = 1.0, 16.0 # H and O masses, in amu (toy values)
R = jnp.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 1.8], dtype=jnp.float64)
k = 0.5 # Ha / Bohr^2
# Place k between the two atoms along z.
e_R = jnp.asarray([0.0, 0.0, -1.0, 0.0, 0.0, 1.0], dtype=jnp.float64)
K = k * jnp.outer(e_R, e_R)
masses = jnp.asarray([m1, m2], dtype=jnp.float64)
freqs, _ = vibrational_frequencies(
K,
masses,
R,
project_tr=True,
diatomic=True,
)
nonzero = jnp.sort(jnp.abs(freqs))[::-1][:1]
from omnibias.jax.bo_derivatives import AMU_TO_ME
# eigenvalue of the mass-weighted Hessian for a single mode along the
# bond axis is k * (1/m1 + 1/m2) = k / mu_me with mu = m1 m2/(m1+m2).
# omega_au = sqrt(eigval); nu_cm = omega_au * HARTREE_TO_CM.
expected_cm = float(
np.sqrt(k * (1.0 / (m1 * AMU_TO_ME) + 1.0 / (m2 * AMU_TO_ME))) * HARTREE_TO_CM
)
err = float(jnp.abs(nonzero[0] - expected_cm))
assert err / expected_cm < 1e-6, (
f"diatomic vib freq: got {float(nonzero[0]):.4f}, "
f"expected {expected_cm:.4f} cm^-1, rel err = {err / expected_cm:.3e}"
)python -m pytest "packages/omnibias-jax/tests/test_bo_derivatives.py::test_vibrational_frequencies_diatomic_harmonic" -qRead the implementation
def vibrational_frequencies(
K_nuclear: Array, # (3N, 3N), Hartree / Bohr^2, symmetric
masses_amu: Array, # (N,) in atomic mass units
R: Array, # (3N,) equilibrium nuclear positions, atomic units
project_tr: bool = True,
diatomic: bool = False,
) -> tuple[Array, Array]:
r"""Diagonalise the mass-weighted nuclear Hessian.
Parameters
----------
K_nuclear
``(3N, 3N)`` symmetric matrix of :math:`\partial^{2}
E_{\mathrm{BO}} / \partial R_\alpha \partial R_\beta` in
atomic units (Hartree / Bohr^2).
masses_amu
``(N,)`` atomic masses, in atomic mass units (amu).
R
``(3N,)`` equilibrium nuclear positions in Bohr. Used to
construct the translation and rotation projectors.
project_tr
If ``True`` (default), project out the six translation
and rotation modes before diagonalising. For a diatomic
molecule pass ``diatomic=True`` so we project out five
modes (one rotation about the bond axis is degenerate).
diatomic
See above.
Returns
-------
freqs_cm
Vibrational frequencies in inverse centimetres. Negative
values indicate an imaginary mode (saddle point /
instability).
eigvecs
Mass-weighted Cartesian eigenvectors of the projected
Hessian, shape ``(3N, 3N)``.
"""
masses_me = masses_amu * AMU_TO_ME
inv_sqrt_m = jnp.repeat(1.0 / jnp.sqrt(masses_me), 3) # (3N,)
K_mw = (inv_sqrt_m[:, None] * K_nuclear) * inv_sqrt_m[None, :]
if project_tr:
K_mw = _project_translations_rotations(
K_mw,
masses_me,
R,
diatomic=diatomic,
)
eigvals, eigvecs = jnp.linalg.eigh(K_mw)
# eigenvalues are in (Hartree / Bohr^2) / m_e = Hartree / (Bohr^2 m_e).
# The corresponding angular frequency is sqrt(eigval) in atomic units
# (where omega_au = 1 / t_au, t_au = a0 sqrt(m_e/E_h)).
# 1 au of angular frequency = 219474.63... cm^-1 / (2 pi); we convert
# via the same constant since omega -> wavenumber uses
# nu_cm = omega_au * HARTREE_TO_CM / (2 pi). Actually the standard
# convention used in vibrational spectroscopy is to report
# nu_cm = sqrt(eigval [Ha/Bohr^2/m_e]) * HARTREE_TO_CM
# because the relation E = hbar omega becomes nu_tilde = omega/(2 pi c)
# and in atomic units omega = sqrt(eigval) (since hbar = 1, m_e = 1),
# so nu_tilde_cm = sqrt(eigval) * HARTREE_TO_CM.
freqs_cm = jnp.sign(eigvals) * jnp.sqrt(jnp.abs(eigvals)) * HARTREE_TO_CM
return freqs_cm, eigvecsOther public import paths
omnibias.jax.vibrational_frequenciesomnibias.jax.bo_derivatives.vibrational_frequencies