Source code for atom_hifi.training.force_truth

"""Force-resolution helpers for ε_F microstate references."""
from __future__ import annotations

import json
from functools import lru_cache
from importlib import resources

import numpy as np
from ase.data import atomic_numbers




@lru_cache(maxsize=1)
def load_kappa_table() -> dict:
    """Load packaged QHO force-stiffness table."""
    path = resources.files("atom_hifi").joinpath(
        "data/qho_freq_table_v1_kappa_msq.json")
    with path.open(encoding="utf-8") as fh:
        return json.load(fh)





def kappa_table_version() -> str:
    return str(load_kappa_table().get("_meta", {}).get("version", "unknown"))





def qho_force_kappa(element: str) -> float:
    """Return QHO κ = m<ω²> for force-resolution scaling (eV/Ų)."""
    if element not in atomic_numbers:
        raise ValueError(f"Unknown element symbol: {element!r}")
    table = load_kappa_table()
    entry = table.get("elements", {}).get(element)
    if entry is None:
        raise ValueError(f"No QHO force κ table entry for {element!r}.")
    kappa = float(entry["kappa_stiff_eV_A2"])
    if kappa <= 0:
        raise ValueError(f"QHO force κ for {element!r} must be positive.")
    return kappa





def force_eps_from_curve(
    element: str,
    *,
    f_tol: float,
    sigma_rattle,
    mean_eps,
) -> tuple[float, float]:
    """Interpolate ε_F from cached ε(σ) calibration curve.

    Returns ``(eps_F, u_star)`` where ``u_star = F_tol / κ``.
    """
    f_tol = float(f_tol)
    if f_tol <= 0:
        raise ValueError(f"f_tol must be positive, got {f_tol!r}.")
    sigma = np.asarray(sigma_rattle, dtype=float)
    eps = np.asarray(mean_eps, dtype=float)
    if sigma.ndim != 1 or eps.ndim != 1 or len(sigma) != len(eps):
        raise ValueError("sigma_rattle and mean_eps must be 1-D arrays of equal length.")
    if len(sigma) < 2:
        raise ValueError("ε_F interpolation requires at least two ε(σ) points.")
    order = np.argsort(sigma)
    sigma = sigma[order]
    eps = eps[order]
    if np.any(np.diff(sigma) <= 0):
        raise ValueError("sigma_rattle values must be unique.")

    u_star = f_tol / qho_force_kappa(element)
    if u_star < sigma[0] or u_star > sigma[-1]:
        raise ValueError(
            f"u*={u_star:.6g} Å for {element} lies outside cached ε(σ) grid "
            f"[{sigma[0]:.6g}, {sigma[-1]:.6g}] Å; rerun thermal ε calibration "
            "with a wider sigma_rattle grid.")
    return float(np.interp(u_star, sigma, eps)), float(u_star)