"""RDF first-peak ``delta_r_ref`` → ε_noise calibration (center-only rattle).
An alternative to the thermal :func:`atom_hifi.training.judge.measure_eps_noise`
route for obtaining the per-element descriptor noise floor ε_noise (``eps_error``).
Instead of a thermal Gaussian rattle at a physics-derived σ, the amplitude is read
from the **database's own** nearest-neighbour RDF first-peak width and applied as a
fixed-magnitude, single-atom ("Dirac") displacement of the focus atom; the induced
embedding shift ‖Δembed‖ is aggregated into a scalar ε per element. The result is
dimensionally identical to the thermal ε_noise and is consumed by exactly the same
downstream code (the c-factor scan bandwidth unit, the r100 ``k_micro`` selection).
Validated by the 2026-06-19 sweep (4 systems × 3600 variants, 0 fallbacks) and the
2026-06-23/24 downstream benchmark. Recommended variant (the locked defaults below):
``xall / gauss / sigma / sqrt2 / bin_width 0.01 / light smoothing / median``.
Pure functions ported verbatim from the validated probe ``rdf_deltar_core.py``;
:func:`center_only_eps` is generalised to accept either an ``embed_op`` matrix or an
``embed`` callable so it matches the ``embed_fn`` convention used by
``measure_eps_noise``. Seed 42; ``n_samples`` default 3000.
"""
import contextlib
import io
import logging
import numpy as np
from ase.data import atomic_numbers
from scipy.ndimage import gaussian_filter1d
from scipy.optimize import curve_fit
from scipy.signal import find_peaks
from sklearn.mixture import GaussianMixture
logger = logging.getLogger(__name__)
# Validated recommended variant. Only ``n_samples`` and ``eps_stat`` are exposed in
# config (``eps_noise.rdf``); the other six are locked here. ``rule`` (r100/r95) is
# NOT here — k-selection stays with ``microstate.reference: r100`` against this ε.
RDF_DEFAULTS = {
'source': 'xall', # xall | pair
'fit': 'gauss', # gauss | gmm
'width': 'sigma', # sigma | fwhm
'conversion': 'sqrt2', # sqrt2 | direct
'bin_width': 0.01, # Å
'smoothing': 'light', # none | light
'eps_stat': 'median', # mean | median | q95 (exposed)
'n_samples': 3000, # center-only rattle samples (exposed)
}
# ── RDF computation ────────────────────────────────────────────────────────────
[docs]
def rdf_first_peak(structs, focus_el, neighbour_el=None, r_max=5.0, bin_width=0.01,
smoothing='none'):
"""Compute g(r) and detect the first peak.
Uses ``get_all_distances(mic=True)`` directly (avoids the ``ase.ga.utilities``
cell-size restriction that fails for small ML training boxes). g(r) uses a
pair-count-weighted density normalisation ``Σ_s n_X^s n_Y^s / V_s`` so bulk
structures dominate over isolated dimers in vacuum boxes. Returns a dict with
``r``, ``g``, ``g_smooth``, ``peak_r/g/i``, ``window_L/R(_i)``, ``pair_counts``,
``n_structs``.
"""
nbins = int(round(r_max / bin_width))
edges = np.linspace(0.0, r_max, nbins + 1)
centers = 0.5 * (edges[:-1] + edges[1:])
Z_x = atomic_numbers[focus_el]
Z_y = atomic_numbers[neighbour_el] if neighbour_el is not None else None
hist_sum = np.zeros(nbins)
pair_counts = {}
n_structs = 0
norm_factor_scalar = 0.0
for s in structs:
nums = s.numbers
syms = np.array(s.get_chemical_symbols())
x_idx = np.where(nums == Z_x)[0]
if len(x_idx) == 0:
continue
V = s.cell.volume
if V < 1.0:
continue
n = len(s)
d = s.get_all_distances(mic=True)
if Z_y is not None:
y_idx = np.where(nums == Z_y)[0]
if len(y_idx) == 0:
continue
n_x = len(x_idx)
n_y = len(y_idx)
n_y_eff = (n_x - 1) if (Z_y == Z_x) else n_y
for i in x_idx:
dists_i = d[i, y_idx]
if Z_y == Z_x:
dists_i = dists_i[dists_i > 1e-6] # exclude self for homo
else:
dists_i = dists_i[dists_i > 0.0]
dists_i = dists_i[dists_i < r_max]
if len(dists_i) == 0:
continue
h, _ = np.histogram(dists_i, bins=edges)
hist_sum += h
pair_counts[neighbour_el] = pair_counts.get(neighbour_el, 0) + len(dists_i)
norm_factor_scalar += n_x * n_y_eff / V
else:
n_x = len(x_idx)
n_others = n - 1
for i in x_idx:
other_idx = [j for j in range(n) if j != i]
dists_i = d[i, other_idx]
dists_i_mask = dists_i < r_max
if not dists_i_mask.any():
continue
h, _ = np.histogram(dists_i[dists_i_mask], bins=edges)
hist_sum += h
for j, ok in zip(other_idx, dists_i_mask):
if ok:
el_j = syms[j]
pair_counts[el_j] = pair_counts.get(el_j, 0) + 1
norm_factor_scalar += n_x * n_others / V
n_structs += 1
if norm_factor_scalar > 0.0:
denom = norm_factor_scalar * 4.0 * np.pi * centers**2 * bin_width
denom = np.where(denom > 0, denom, 1.0)
g = hist_sum / denom
else:
g = np.zeros(nbins)
sigma_out = 3 if smoothing == 'light' else 0
g_smooth = gaussian_filter1d(g, sigma=sigma_out) if sigma_out > 0 else g.copy()
# Peak detection uses a separate always-applied moderate smoothing (sigma=5 bins)
# to suppress noisy spikes from sparse PES-scan structures.
sigma_det = max(5, sigma_out)
g_det = gaussian_filter1d(g, sigma=sigma_det)
start_i = max(1, np.searchsorted(centers, 1.5)) # skip r<1.5 Å ghost bonds
peaks, _ = find_peaks(g_det[start_i:], height=0.1, distance=10,
prominence=max(0.1, 0.001 * g_det.max()))
if len(peaks) == 0:
pk_i = start_i + int(np.argmax(g_det[start_i:]))
else:
pk_i = peaks[0] + start_i
valleys, _ = find_peaks(-g_det[start_i:])
left_v = [v for v in valleys if v + start_i < pk_i]
right_v = [v for v in valleys if v + start_i > pk_i]
L_i = (left_v[-1] + start_i) if left_v else start_i
R_i = (right_v[0] + start_i) if right_v else min(pk_i + 50, nbins - 1)
return {
'r': centers,
'g': g,
'g_smooth': g_smooth,
'peak_r': float(centers[pk_i]),
'peak_g': float(g_smooth[pk_i]),
'peak_i': int(pk_i),
'window_L': float(centers[L_i]),
'window_R': float(centers[R_i]),
'window_L_i': int(L_i),
'window_R_i': int(R_i),
'pair_counts': dict(pair_counts),
'n_structs': n_structs,
}
# ── Peak fitting ────────────────────────────────────────────────────────────────
def _gauss(r, A, mu, sigma):
"""Single Gaussian: A * exp(-0.5 * ((r - mu) / sigma)^2)."""
return A * np.exp(-0.5 * ((r - mu) / sigma) ** 2)
[docs]
def fit_peak(r, g, window, mode='gauss', gmm_weight_min=0.0, sigma_floor=0.0):
"""Fit the first RDF peak within ``window`` → σ / FWHM.
Returns a dict with ``status`` ('ok'|'poor_fit'|'failed'), ``sigma_peak``,
``fwhm``, ``center``, ``residual``, ``mode`` (and GMM extras for mode='gmm').
"""
L, R = window
mask = (r >= L) & (r <= R)
r_w = r[mask]
g_w = g[mask]
_fail = {'status': 'failed', 'sigma_peak': np.nan, 'fwhm': np.nan,
'center': np.nan, 'residual': np.nan, 'mode': mode}
if len(r_w) < 5 or g_w.max() < 1e-10:
return _fail
if mode == 'gauss':
try:
mu0 = float(r_w[np.argmax(g_w)])
A0 = float(g_w.max())
s0 = (R - L) / 4.0
popt, _ = curve_fit(
_gauss, r_w, g_w,
p0=[A0, mu0, abs(s0)],
bounds=([0.0, float(L), 1e-6], [np.inf, float(R), float(R - L)]),
maxfev=5000,
)
A_fit, mu_fit, sigma_fit = popt
g_pred = _gauss(r_w, *popt)
residual = float(np.sqrt(np.mean((g_w - g_pred) ** 2)) / (g_w.max() + 1e-10))
return {
'status': 'ok' if residual < 0.3 else 'poor_fit',
'sigma_peak': float(abs(sigma_fit)),
'fwhm': float(2.3548 * abs(sigma_fit)),
'center': float(mu_fit),
'residual': residual,
'mode': 'gauss',
}
except Exception as e:
return {**_fail, 'err': str(e)}
elif mode == 'gmm':
g_pos = np.maximum(g_w - g_w.min(), 0.0)
total = g_pos.sum()
if total < 1e-10:
return _fail
probs = g_pos / total
rng = np.random.default_rng(42)
n_samp = max(500, len(r_w) * 20)
samples = rng.choice(r_w, size=n_samp, p=probs, replace=True).reshape(-1, 1)
best = None
for n_comp in [2, 3, 4]:
try:
gmm = GaussianMixture(n_components=n_comp, random_state=42, max_iter=200)
gmm.fit(samples)
# Reject sub-bin (σ < sigma_floor) components: physically
# unresolvable near-delta spikes the GMM would overfit.
valid = [
(i,
float(gmm.covariances_[i, 0, 0] ** 0.5),
float(gmm.means_[i, 0]),
float(gmm.weights_[i]))
for i in range(n_comp)
if float(gmm.weights_[i]) >= gmm_weight_min
and float(gmm.covariances_[i, 0, 0] ** 0.5) >= sigma_floor
]
if not valid:
continue
valid.sort(key=lambda x: x[1])
_, sig, mu, w = valid[0]
g_pred = np.zeros_like(g_w, dtype=float)
for ii in range(n_comp):
wi = float(gmm.weights_[ii])
mui = float(gmm.means_[ii, 0])
si = float(gmm.covariances_[ii, 0, 0] ** 0.5)
g_pred += wi * np.exp(-0.5 * ((r_w - mui) / si) ** 2) / (si * np.sqrt(2 * np.pi))
scale = g_w.sum() / (g_pred.sum() + 1e-10)
residual = float(np.sqrt(np.mean((g_w - g_pred * scale) ** 2)) / (g_w.max() + 1e-10))
if best is None or sig < best['sigma_peak']:
best = {
'status': 'ok' if residual < 0.5 else 'poor_fit',
'sigma_peak': sig,
'fwhm': 2.3548 * sig,
'center': mu,
'gmm_weight': w,
'n_components': n_comp,
'residual': residual,
'mode': 'gmm',
}
except Exception:
continue
return best if best is not None else _fail
else:
raise ValueError(f'Unknown fit mode: {mode!r}. Use "gauss" or "gmm".')
# ── delta_r_ref aggregation ────────────────────────────────────────────────────
[docs]
def delta_r_ref(system_inputs, focus_el, source='xall', agg='min',
conversion='direct', fit_mode='gauss', gmm_weight_min=0.0,
width_metric='sigma', bin_width=0.01, smoothing='none'):
"""RDF first-peak width → reference displacement amplitude (Å).
Fallback ladder: pair-resolved → aggregate → if all pairs fail try xall → if
xall fails return ``(None, provenance)`` (caller falls back). See module
docstring for the recommended variant.
"""
structs = system_inputs['structs']
if 'neighbour_els' in system_inputs:
neighbour_els = list(system_inputs['neighbour_els'])
else:
all_syms = set()
for s in structs:
all_syms.update(s.get_chemical_symbols())
neighbour_els = sorted(all_syms)
prov = {
'source': source, 'agg': agg, 'conversion': conversion,
'fit_mode': fit_mode, 'gmm_weight_min': gmm_weight_min,
'width_metric': width_metric, 'bin_width': bin_width, 'smoothing': smoothing,
'pair_results': {},
}
def _extract_width(fit_res):
if fit_res.get('status', 'failed') == 'failed':
return None
key = 'sigma_peak' if width_metric == 'sigma' else 'fwhm'
v = fit_res.get(key, np.nan)
return float(v) if (v is not None and np.isfinite(v) and v > 0) else None
def _apply_conversion(w):
if w is None:
return None
return w / np.sqrt(2.0) if conversion == 'sqrt2' else float(w)
if source == 'pair':
pair_widths = {}
pair_counts_total = {}
for y_el in neighbour_els:
try:
rdf_res = rdf_first_peak(structs, focus_el, neighbour_el=y_el,
r_max=5.0, bin_width=bin_width,
smoothing=smoothing)
win = (rdf_res['window_L'], rdf_res['window_R'])
fit_res = fit_peak(rdf_res['r'], rdf_res['g_smooth'], win,
mode=fit_mode, gmm_weight_min=gmm_weight_min,
sigma_floor=bin_width)
w = _extract_width(fit_res)
n_pairs = sum(rdf_res['pair_counts'].values())
prov['pair_results'][y_el] = {
'peak_r': rdf_res['peak_r'], 'width': w,
'fit_status': fit_res.get('status'), 'n_pairs': n_pairs,
}
if w is not None and n_pairs > 0:
pair_widths[y_el] = w
pair_counts_total[y_el] = n_pairs
except Exception as e:
prov['pair_results'][y_el] = {'error': str(e)}
if pair_widths:
vals = [pair_widths[y] for y in pair_widths]
counts = [pair_counts_total.get(y, 1) for y in pair_widths]
if agg == 'min':
width_agg = float(np.min(vals))
elif agg == 'q10':
width_agg = float(np.percentile(vals, 10))
elif agg == 'q25':
width_agg = float(np.percentile(vals, 25))
elif agg == 'wmean':
width_agg = float(np.average(vals, weights=counts))
else:
raise ValueError(f'Unknown agg: {agg!r}')
value = _apply_conversion(width_agg)
prov.update({'pair_agg_result': width_agg, 'n_valid_pairs': len(pair_widths),
'delta_r_before_convert': width_agg, 'delta_r_final': value})
return value, prov
prov['pair_fallback'] = 'all_pairs_failed → xall'
# xall path (primary if source='xall', fallback if all pairs failed)
try:
rdf_res = rdf_first_peak(structs, focus_el, neighbour_el=None,
r_max=5.0, bin_width=bin_width, smoothing=smoothing)
win = (rdf_res['window_L'], rdf_res['window_R'])
fit_res = fit_peak(rdf_res['r'], rdf_res['g_smooth'], win,
mode=fit_mode, gmm_weight_min=gmm_weight_min,
sigma_floor=bin_width)
w = _extract_width(fit_res)
prov['xall_result'] = {'peak_r': rdf_res['peak_r'], 'width': w,
'fit_status': fit_res.get('status')}
if w is not None:
value = _apply_conversion(w)
prov['delta_r_final'] = value
return value, prov
except Exception as e:
prov['xall_error'] = str(e)
prov['delta_r_final'] = None
return None, prov
# ── Center-only rattle ─────────────────────────────────────────────────────────
def _apply_embed(embed, d_rows):
"""Project descriptor rows (k, n_feat) → (k, embed_dim).
``embed`` may be a linear ``embed_op`` matrix (n_feat, embed_dim) or a callable
matching the ``embed_fn`` convention of ``measure_eps_noise``
(callable((k, n_feat)) → (k, embed_dim)).
"""
if callable(embed):
return np.asarray(embed(d_rows))
return np.asarray(d_rows) @ embed
[docs]
def center_only_eps(lib, focus_el, embed, descriptor_fn, delta_r,
n_samples=3000, seed=42):
"""Center-only fixed-sphere ("Dirac") rattle producing embed-shift samples.
Move ONLY the focus atom by a fixed displacement ``delta_r`` in a uniformly random
direction on the unit sphere; the cage is frozen. ``embed`` is a matrix or an
embed callable (see :func:`_apply_embed`). Returns a 1-D array of per-sample
embedding-shift norms.
"""
rng = np.random.default_rng(seed)
pool = [
(si, ai)
for si, s in enumerate(lib)
for ai, sym in enumerate(s.get_chemical_symbols())
if sym == focus_el
]
if not pool:
raise ValueError(f'No {focus_el} atoms found in lib ({len(lib)} structs).')
M = max(1, int(np.ceil(n_samples / len(pool))))
M = min(M, 50)
n_atoms_draw = min(len(pool), max(1, int(np.ceil(n_samples / M))))
atom_draw = rng.choice(len(pool), size=n_atoms_draw, replace=False)
sampled = [pool[i] for i in atom_draw]
eps_vals = []
for si, ai in sampled:
host = lib[si]
pos0 = host.get_positions().copy()
with contextlib.redirect_stdout(io.StringIO()):
desc0 = descriptor_fn([host]) # (n_atoms, n_feat)
e0 = _apply_embed(embed, desc0[ai:ai + 1])[0]
for _ in range(M):
u = rng.normal(0.0, 1.0, 3)
u /= np.linalg.norm(u)
pos_new = pos0.copy()
pos_new[ai] += delta_r * u
p = host.copy()
p.set_positions(pos_new)
with contextlib.redirect_stdout(io.StringIO()):
desc1 = descriptor_fn([p])
e1 = _apply_embed(embed, desc1[ai:ai + 1])[0]
eps_vals.append(float(np.linalg.norm(e1 - e0)))
return np.array(eps_vals, dtype=float)
# ── eps_tol aggregation ────────────────────────────────────────────────────────
[docs]
def eps_tol(eps_samples, stat='mean'):
"""Aggregate ‖Δembed‖ samples to a scalar ε ('mean' | 'median' | 'q95')."""
a = np.asarray(eps_samples, dtype=float)
if stat == 'mean':
return float(np.mean(a))
elif stat == 'median':
return float(np.median(a))
elif stat == 'q95':
return float(np.percentile(a, 95))
else:
raise ValueError(f'Unknown stat: {stat!r}. Use "mean", "median", or "q95".')
# ── ε_noise calibration entry point (parallel to measure_eps_noise) ─────────────
[docs]
def measure_eps_noise_rdf(atoms_list, descriptor_fn, embed_fn, species,
rdf_cfg=None, *, seed=42):
"""Per-element ε_noise from the RDF ``delta_r_ref`` center-only rattle.
Drop-in alternative to :func:`atom_hifi.training.judge.measure_eps_noise`:
returns a per-element scalar ε that feeds the c-factor scan bandwidth and the
r100 ``k_micro`` selection.
Parameters
----------
atoms_list : list[ase.Atoms]
descriptor_fn : callable(list[Atoms]) → (n_atoms, n_feat)
embed_fn : dict{el: callable|matrix} or a single callable/matrix — the
SAME embedding used to build the microstate covers.
species : list[str]
rdf_cfg : dict — exposed knobs (``n_samples``, ``eps_stat``) override the
locked ``RDF_DEFAULTS``; ``None`` → all defaults.
seed : int
Returns
-------
dict {el: {'eps': float|None, 'delta_r': float|None, 'n_eps_samples': int,
'fit_status': str|None, 'fallback': bool}}
``eps`` is ``None`` (``fallback=True``) when no RDF peak is found; the
caller substitutes the mean ε of the elements that did peak.
"""
cfg = {**RDF_DEFAULTS, **(rdf_cfg or {})}
results = {}
for el in species:
lib_el = [s for s in atoms_list if el in s.get_chemical_symbols()]
if not lib_el:
logger.warning(' eps_noise[rdf] %s: no structures contain %s — fallback.',
el, el)
results[el] = {'eps': None, 'delta_r': None, 'n_eps_samples': 0,
'fit_status': None, 'fallback': True}
continue
dr, prov = delta_r_ref(
{'structs': lib_el}, el,
source=cfg['source'], conversion=cfg['conversion'],
fit_mode=cfg['fit'], width_metric=cfg['width'],
bin_width=cfg['bin_width'], smoothing=cfg['smoothing'])
fit_status = (prov.get('xall_result') or {}).get('fit_status')
if dr is None:
logger.warning(' eps_noise[rdf] %s: no usable RDF peak (delta_r=None) '
'— borrowing peaked-element mean ε.', el)
results[el] = {'eps': None, 'delta_r': None, 'n_eps_samples': 0,
'fit_status': fit_status, 'fallback': True}
continue
embed = embed_fn.get(el) if isinstance(embed_fn, dict) else embed_fn
if embed is None:
raise RuntimeError(f'embed_fn has no entry for element {el!r}')
eps_samples = center_only_eps(
lib_el, el, embed, descriptor_fn, dr,
n_samples=int(cfg['n_samples']), seed=seed)
eps = eps_tol(eps_samples, stat=cfg['eps_stat'])
logger.info(' eps_noise[rdf] %s: delta_r=%.4f Å eps=%.5f '
'(fit=%s, %d samples)',
el, dr, eps, fit_status, len(eps_samples))
results[el] = {'eps': float(eps), 'delta_r': float(dr),
'n_eps_samples': int(len(eps_samples)),
'fit_status': fit_status, 'fallback': False}
return results