Source code for atom_hifi.training.qho_sigma

"""QHO-backed sigma_thermal lookup for runtime ε_noise calibration."""
from __future__ import annotations

import json
import logging
import math
from functools import lru_cache
from importlib import resources

from ase.data import atomic_masses, atomic_numbers

from atom_hifi.training import calibration as _calibration

logger = logging.getLogger(__name__)

_AMU_KG = 1.66053906660e-27
_ANGSTROM_M = 1e-10
_EV_J = 1.602176634e-19
_HBAR = 1.054571817e-34
_KB = 1.380649e-23
_THIN_N_MATERIALS = 10
_NO_OVERRIDE = object()




@lru_cache(maxsize=1)
def load_qho_table() -> dict:
    """Load packaged QHO frequency table.

    Default is the msq/kappa_stiff (stretch-mode, second-moment) estimator,
    not the centroid (full-DOS average) estimator — see ADR-0010: clean
    chemical-motif identity under centroid does not imply MLIP-training-
    adequate granularity, so the default now shares the stiffer moment used
    by the opt-in force-truth (`microstate.reference: f_tol`) path.
    """
    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 _omega_eff_meV(entry: dict) -> float:
    """Effective phonon energy (meV) from a table entry, estimator-agnostic."""
    omega = entry.get("omega_eff_meV", entry.get("omega_stiff_meV"))
    if omega is None:
        raise KeyError("table entry has neither 'omega_eff_meV' nor 'omega_stiff_meV'")
    return float(omega)




def qho_table_version() -> str:
    """Version tag of the packaged QHO table (``_meta.version``; ``'unknown'`` if absent).

    Saved in ``eps_noise_raw.npz`` and compared by the runner's ε-cache check
    when ``sigma_source`` is ``qho_table``.
    """
    return str(load_qho_table().get("_meta", {}).get("version", "unknown"))





def qho_sigma_from_omega_mev(element: str, omega_eff_meV: float, T: float) -> float:
    """Return QHO RMS displacement σ (Å) for an element/frequency/temperature."""
    if element not in atomic_numbers:
        raise ValueError(f"Unknown element symbol: {element!r}")
    if omega_eff_meV <= 0:
        raise ValueError(f"omega_eff_meV must be positive, got {omega_eff_meV!r}")
    if T <= 0:
        raise ValueError(f"Temperature must be positive, got {T!r}")
    m_kg = atomic_masses[atomic_numbers[element]] * _AMU_KG
    energy_j = float(omega_eff_meV) * 1e-3 * _EV_J
    omega = energy_j / _HBAR
    x = energy_j / (2.0 * _KB * float(T))
    coth = 1.0 / math.tanh(x)
    sigma_m = math.sqrt(_HBAR / (2.0 * m_kg * omega) * coth)
    return float(sigma_m / _ANGSTROM_M)





def qho_sigma_thermal(element: str, T: float = 300.0) -> float:
    """QHO-table RMS displacement σ_thermal (Å)."""
    table = load_qho_table()
    entry = table.get("elements", {}).get(element)
    if entry is None:
        raise ValueError(
            f"No QHO sigma_thermal table entry for {element!r}. "
            "Use eps_noise.thermal.sigma_source: mass_only or set "
            "eps_noise.thermal.sigma_override for this element.")
    return qho_sigma_from_omega_mev(element, _omega_eff_meV(entry), T)





def resolve_sigma_thermal_targets(
    species,
    *,
    T: float,
    source: str = "qho_table",
    config_override: dict | None = None,
) -> dict[str, float]:
    """Resolve per-element σ targets for thermal ε calibration.

    Precedence (highest first): ``config_override``
    (``eps_noise.thermal.sigma_override``), then the deprecated
    ``calibration.SIGMA_THERMAL_OVERRIDE`` (its warning is logged even when a
    config pin shadows it), then the ``source`` model.  Pin values are exact Å
    for the current run temperature (not rescaled with ``T``).  Explicit pins
    skip the QHO-table lookup, so pinned species need no table entry.  Every
    target must be finite and positive.
    """
    if source not in ("qho_table", "mass_only"):
        raise ValueError(
            f"eps_noise.thermal.sigma_source must be 'qho_table' or 'mass_only', "
            f"got {source!r}.")
    config_override = dict(config_override or {})
    table = None
    out = {}
    for el in species:
        if el not in atomic_numbers:
            raise ValueError(f"Unknown element symbol: {el!r}")
        override = _calibration.SIGMA_THERMAL_OVERRIDE.get(el, _NO_OVERRIDE)
        if override is not _NO_OVERRIDE:
            logger.warning(
                "SIGMA_THERMAL_OVERRIDE for %s is deprecated; prefer "
                "eps_noise.thermal.sigma_override.", el)
        if el in config_override:
            sigma = float(config_override[el])
        elif override is not _NO_OVERRIDE:
            sigma = float(override)
        elif source == "qho_table":
            if table is None:
                table = load_qho_table()
            entry = table.get("elements", {}).get(el)
            if entry is None:
                raise ValueError(
                    f"No QHO sigma_thermal table entry for {el!r}. "
                    "Use eps_noise.thermal.sigma_source: mass_only or set "
                    "eps_noise.thermal.sigma_override for this element.")
            n_materials = int(entry.get("n_materials", 0))
            if n_materials < _THIN_N_MATERIALS:
                logger.warning(
                    "QHO sigma_thermal table entry for %s has thin provenance "
                    "(n_materials=%d < %d); using it anyway.",
                    el, n_materials, _THIN_N_MATERIALS)
            sigma = qho_sigma_from_omega_mev(el, _omega_eff_meV(entry), T)
        else:
            sigma = _calibration.sigma_thermal(el, T)
        if not math.isfinite(sigma) or sigma <= 0:
            raise ValueError(f"sigma_thermal target for {el!r} must be positive and finite.")
        out[el] = float(sigma)
    return out