101 lines
3.8 KiB
Python
101 lines
3.8 KiB
Python
"""NaN-safety and differentiability tests for the guard layer.
|
|
|
|
The regression tests establish numeric parity; these establish that funxc
|
|
stays finite (values *and* gradients) at the singular corners libxc's grids
|
|
avoid: rho -> 0 tails, sigma = 0, full spin polarization, and combinations.
|
|
"""
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import funxc
|
|
|
|
GGA_IDS = [fid for fid in funxc.REGISTRY if funxc.REGISTRY[fid].family == "gga"]
|
|
LDA_IDS = [fid for fid in funxc.REGISTRY if funxc.REGISTRY[fid].family == "lda"]
|
|
|
|
# Adversarial polarized points:
|
|
# (rho_up, rho_dn, sigma_uu, sigma_ud, sigma_dd, tau_up, tau_dn)
|
|
EDGE_POINTS = [
|
|
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # vacuum
|
|
(1e-300, 0.0, 0.0, 0.0, 0.0, 1e-300, 0.0), # denormal tail
|
|
(0.3, 0.0, 7e-3, 0.0, 0.0, 2.9e-3, 0.0), # fully polarized, tau ~ tau_W
|
|
(0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0), # sigma = tau = 0 exactly
|
|
(0.3, 1e-16, 1e-2, 0.0, 1e-30, 5e-2, 1e-30), # one channel below thr
|
|
(1e-14, 1e-14, 1e-25, -1e-25, 1e-25, 1e-25, 1e-25), # negative sigma_ud
|
|
(1e3, 1e3, 1e8, 1e8, 1e8, 1e8, 1e8), # high density, huge gradient
|
|
(0.5, 0.5, 1e12, 0.0, 1e12, 1e-5, 1e-5), # s -> inf, tau << tau_W (FHC)
|
|
]
|
|
|
|
|
|
def _edge_args(spec, f):
|
|
rho = np.array([[p[0], p[1]] for p in EDGE_POINTS])
|
|
sigma = np.array([[p[2], p[3], p[4]] for p in EDGE_POINTS])
|
|
tau = np.array([[p[5], p[6]] for p in EDGE_POINTS])
|
|
lapl = np.zeros_like(tau)
|
|
if spec.family == "lda":
|
|
return (rho,)
|
|
if spec.family == "gga":
|
|
return (rho, sigma)
|
|
return (rho, sigma, lapl, tau)
|
|
|
|
|
|
@pytest.mark.parametrize("fid", sorted(funxc.REGISTRY))
|
|
def test_values_and_gradients_finite_at_edges(fid):
|
|
spec = funxc.REGISTRY[fid]
|
|
f = funxc.functional(fid, polarized=True)
|
|
out = f.exc_vxc(*_edge_args(spec, f))
|
|
for key, val in out.items():
|
|
assert np.all(np.isfinite(val)), f"{fid}: non-finite {key} at edge points"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"fid",
|
|
["LDA_X", "LDA_C_PW", "GGA_X_PBE", "GGA_C_PBE", "GGA_C_LYP",
|
|
"GGA_XC_B97_D", "MGGA_X_LTA"],
|
|
)
|
|
def test_second_derivatives_finite(fid):
|
|
"""fxc-level derivatives must also stay finite (autodiff twice)."""
|
|
spec = funxc.REGISTRY[fid]
|
|
f = funxc.functional(fid, polarized=True)
|
|
if spec.family == "lda":
|
|
hess = jax.hessian(f.energy_density, argnums=0)
|
|
args = lambda p: (jnp.array(p[:2]),) # noqa: E731
|
|
elif spec.family == "gga":
|
|
hess = jax.hessian(f.energy_density, argnums=(0, 1))
|
|
args = lambda p: (jnp.array(p[:2]), jnp.array(p[2:5])) # noqa: E731
|
|
else:
|
|
hess = jax.hessian(f.energy_density, argnums=(0, 1, 2, 3))
|
|
args = lambda p: ( # noqa: E731
|
|
jnp.array(p[:2]), jnp.array(p[2:5]), jnp.zeros(2), jnp.array(p[5:7])
|
|
)
|
|
for p in EDGE_POINTS:
|
|
h = hess(*args(p))
|
|
leaves = jax.tree.leaves(h)
|
|
assert all(jnp.all(jnp.isfinite(x)) for x in leaves), (
|
|
f"{fid}: non-finite second derivative at {p}"
|
|
)
|
|
|
|
|
|
def test_vacuum_outputs_are_zero():
|
|
f = funxc.functional("GGA_X_PBE", polarized=True)
|
|
out = f.exc_vxc(np.zeros((1, 2)), np.zeros((1, 3)))
|
|
for key, val in out.items():
|
|
assert np.all(np.asarray(val) == 0.0), f"nonzero {key} in vacuum"
|
|
|
|
|
|
def test_jit_and_vmap_compose():
|
|
"""The per-point eps must survive user-side jit/vmap/grad composition."""
|
|
f = funxc.functional("GGA_C_PBE", polarized=True)
|
|
g = jax.jit(jax.vmap(jax.grad(f.energy_density, argnums=(0, 1))))
|
|
rho = jnp.array([[0.3, 0.2], [1.0, 1.0]])
|
|
sigma = jnp.array([[0.01, 0.0, 0.02], [0.1, 0.05, 0.1]])
|
|
vrho, vsigma = g(rho, sigma)
|
|
assert vrho.shape == (2, 2) and vsigma.shape == (2, 3)
|
|
assert jnp.all(jnp.isfinite(vrho)) and jnp.all(jnp.isfinite(vsigma))
|
|
|
|
|
|
def test_float64_active():
|
|
assert jnp.asarray(1.0).dtype == jnp.float64
|