"""Config handling for the Atom-HiFi CLI and Python API.
The commented YAML in ``_TEMPLATE`` below is the **single source of truth** for the
public CLI configuration: it drives
* ``atom-hifi init`` — written verbatim as the starter ``config.yaml``,
* the **YAML defaults** — parsed from it and merged under any user config, and
* the **allowed YAML keys** — strict validation rejects keys absent from the template.
Some expert controls are intentionally Python-API-only. They are not accepted in
YAML, but :func:`atom_hifi.runner.run` still receives stable hidden defaults via
``fill_defaults``.
Note on the ``custom`` descriptor: a custom descriptor is a Python callable and
cannot be expressed in YAML. ``descriptor.kind`` therefore accepts only ``eds`` or
``ace`` (MACE ACE) here; to use a custom descriptor, call :func:`atom_hifi.runner.run` from
Python with ``config['descriptor']['custom_fn']`` set to your callable.
The same applies to a custom cluster backend: ``cluster.backend`` accepts only
the built-in names in YAML; to plug your own :class:`ClusterBackend`, call
:func:`atom_hifi.runner.run` with ``cluster.backend='custom'`` and
``cluster.custom_backend`` set to your instance.
"""
from __future__ import annotations
import copy
import os
import re
import warnings
import yaml
# (1 + sqrt(3)) / 2 — short-range SOAP cutoff default, in units of lengthscale.
_S_CUT_DEFAULT = (1 + 3 ** 0.5) / 2
# `microstate.reference: r<NUM>` — user-defined strictness for the ε-radius
# criterion: NUM is the within-cluster-distance percentile (r100 = the maximum,
# today's `r100`; rNUM < 100 is looser). r0/r101/r10x/R95 are rejected here,
# so every downstream consumer can treat any accepted r-string as valid.
R_NUM_RE = re.compile(r'^r([1-9]\d?|100)$')
[docs]
def is_rnum(value) -> bool:
"""True iff ``value`` is an ``r<NUM>`` reference string, NUM in [1, 100]."""
return isinstance(value, str) and R_NUM_RE.match(value) is not None
_TEMPLATE = f"""\
# ─────────────────────────────────────────────────────────────────────────────
# Atom-HiFi configuration. Edit values below, then run: atom-hifi run config.yaml
# Unknown keys are rejected. Lists use [a, b]; null disables optional fields.
# ─────────────────────────────────────────────────────────────────────────────
# ── Paths ────────────────────────────────────────────────────────────────────
paths:
lib_path: train_structs.xyz # ASE-readable structure library
focus_elements: [O, H] # elements to cluster on (null = all present)
output_dir: results
# ── Descriptor ───────────────────────────────────────────────────────────────
# kind: 'eds' or 'ace' (needs torch+mace); custom descriptor → Python API only.
descriptor:
kind: eds
# ED-SOAP settings (used when kind: eds)
eds:
species: null # null = auto-detect from library
lengthscale: 1.2 # Å; typical nearest-neighbour bond length
r_cut: 3.0 # Å; ε_noise rattling cutoff (≈ shortest descriptor cutoff)
# Short-range channel (cutoffs in units of lengthscale)
s_cut: {_S_CUT_DEFAULT!r} # ≈ 1.366 = (1+√3)/2; between 1st and 2nd shells
s_nmax: 8
s_lmax: 4
# Long-range channel
l_cut: 2.5 # ≈ 3rd/4th shell midpoint
l_nmax: 4
l_lmax: 3
periodic: true # apply PBC per structure; false = always isolated
custom_r_cut: 5.0 # Å; ε_noise rattling cutoff for a custom descriptor
# MACE ACE settings (used when kind: ace)
ace:
model_path: MACE_model_stagetwo.model
device: cpu # 'cuda' for GPU
r_cut: 5.0 # Å; must match the MACE model cutoff
# ── Clusterer backend (ADR-0009) ─────────────────────────────────────────────
# decaf (default; supports loop/deployable OOD via Python API) |
# kmeans (fastest, smallest covers; scans a k-fraction axis, not c) |
# hdbscan (honest noise channel + dedup-noise forced covers).
cluster:
backend: decaf
min_cluster_size: 2 # hdbscan only
# ── Post-selection method ────────────────────────────────────────────────────
# baseline | mu_tiebreak (recommended) | atom_aware
selection:
method: mu_tiebreak
# ── Parallel workers ─────────────────────────────────────────────────────────
n_jobs: 16
# ── Console logging ──────────────────────────────────────────────────────────
# debug | info | warning | error (CLI: -v = debug, -q = warning)
log_level: info
# ── Reproducibility ──────────────────────────────────────────────────────────
seed: 42
# ── L/R scan tuning ──────────────────────────────────────────────────────────
scan:
l_tol: 0.95 # minimum acceptable likeness (L); null → no L floor, auto-pick = best-fidelity (max F) point
c_factor_range: null # per-element {{El: [lo, hi]}}; null → [0.1, 5.0]
n_coarse: 60 # coarse-scan grid points
n_fine: 60 # fine-scan grid points
cache_file: null # path to el_fit_cache.npz from a prior run; pre-warms scan to skip DECAF
# ── Microstate reference ────────────────────────────────────────────────────
microstate:
reference: n2 # n2 → k=N//2 (default);
# r100 → KMeans ε_noise-radius criterion (bisected);
# rNUM → same criterion with user strictness NUM (1-100);
# f_tol → force-resolution ε_F criterion (opt-in)
# ── Descriptor computation ───────────────────────────────────────────────────
compute:
n_jobs_desc: 60 # parallel workers for descriptor computation
n_jobs_eps: null # workers for ε_noise calibration; null → top-level n_jobs
# ── ε_noise (descriptor noise floor) ─────────────────────────────────────────
eps_noise:
method: thermal # thermal | rdf | precomputed
temperature: 300 # K; used by method: thermal
precomputed: # final ε_noise values; skips calibration
per_species: null # e.g. {{Ni: 0.012, O: 0.030}}
thermal: # center-only σ_thermal rattle calibration
sigma_source: qho_table # qho_table | mass_only
sigma_override: null # exact run σ in Å, e.g. {{H: 0.12}}
rdf: # used only when method: rdf (other knobs locked to the validated variant)
n_samples: 3000 # center-only rattle samples
eps_stat: median # mean | median | q95
force: # used only when microstate.reference: f_tol
f_tol: 0.05 # eV/Å per atom; force-resolution tolerance
# ── Optional active-learning loop ────────────────────────────────────────────
loop:
run: false # true → run loop (Stages 2+3); false → skip to Stage 3
initial_sample: 500 # initial pool size
batch_size: 500 # structures added per batch
skip_fine_scan: true # coarse-only inner iters; fine scan runs once in Stage 3
resume_from_iter: null # null → full library; N → load iter-N checkpoint; mutually exclusive with run: true
# ── NSGA-II optimizer ────────────────────────────────────────────────────────
nsga2:
run: false # default multi-element c-optimisation (N_el ≥ 2)
n_gen: 50
pop_size: 40 # ≥ 10 × N_focus_elements, min 40
inner_n_gen: 5 # cheap per-loop budget when loop.run=true
inner_pop_size: null # null → max(8, 4 × N_focus_elements)
restore: false # true → reuse pareto_front.csv from output_dir
# ── Piecing (super-large structures) ─────────────────────────────────────────
# See README § "Piecing for super-large structures" for rationale and tuning.
piecing:
enabled: false # explicit opt-in; no auto-trigger
target_atoms: 64 # approx. atoms per voxel; tune down until no element collapses
# ── Output control ───────────────────────────────────────────────────────────
output:
delta_pick: null # null → auto-select L≥l_tol optimum from fine scan; non-null overrides scan/NSGA2 auto-pick
extract_surroundings: false # Stage 6: extract per-group representative environments
save_judge: true # write judge_final.pkl (the fitted Judge) after the final fit
# ── Surroundings extraction (Stage 6) ────────────────────────────────────────
surroundings:
output_mode: full_structure # 'sphere' (isolated cluster) or 'full_structure' (original cell)
r_cut: null # Å; null → descriptor r_cut
n_per_group: 1 # representatives per DECAF group
# ── Re-fit from a saved run ───────────────────────────────────────────────────
refit:
delta: null # e.g. 1.20 (c_factor units)
"""
_ALLOWED_DESCRIPTOR_KINDS = ('eds', 'ace')
_API_ONLY_DEFAULTS = {
'paths': {
'patient_path': None,
},
'cluster': {
# A custom clusterer is a Python object and cannot be expressed in
# YAML, exactly like descriptor.custom_fn. load_config rejects
# cluster.backend: custom with a pointer to the Python API.
'custom_backend': None,
},
'grid': {
'run': False,
'n_coarse': 20,
'n_fine': 20,
'restore': False,
},
'patient': {
'run_ood': False,
},
'refit': {
'grid_point': None,
},
}
[docs]
def default_config_yaml() -> str:
"""Return the canonical commented config template (verbatim ``init`` output)."""
return _TEMPLATE
def _load_defaults() -> dict:
"""Parse the template into the default config dict."""
return yaml.safe_load(_TEMPLATE)
def _load_runtime_defaults() -> dict:
"""Return defaults consumed by ``runner.run``, including API-only keys."""
return _deep_merge(_load_defaults(), _API_ONLY_DEFAULTS)
def _deep_merge(base: dict, override: dict) -> dict:
"""Recursively merge ``override`` into a copy of ``base`` (override wins)."""
out = copy.deepcopy(base)
for key, val in override.items():
if isinstance(val, dict) and isinstance(out.get(key), dict):
out[key] = _deep_merge(out[key], val)
else:
out[key] = copy.deepcopy(val)
return out
def _warn_legacy(message: str) -> None:
warnings.warn(message, DeprecationWarning, stacklevel=3)
def _normalize_legacy_config(config: dict) -> dict:
"""Map one-release legacy keys to the clean config schema."""
out = copy.deepcopy(config or {})
eps = out.get('eps_noise')
if isinstance(eps, dict) and 'per_species' in eps:
if any(k in eps for k in ('precomputed', 'thermal', 'rdf')) or eps.get('method') is not None:
raise ValueError(
"Legacy eps_noise.per_species cannot be mixed with the new "
"eps_noise.method/precomputed/thermal/rdf schema.")
_warn_legacy(
"eps_noise.per_species is deprecated; use "
"eps_noise.method: precomputed and eps_noise.precomputed.per_species.")
val = eps.pop('per_species')
eps['method'] = 'precomputed'
eps['precomputed'] = {'per_species': val}
scan = out.get('scan')
if isinstance(scan, dict) and 'k_micro_ratio' in scan:
if 'microstate' in out:
raise ValueError(
"Legacy scan.k_micro_ratio cannot be mixed with microstate.reference.")
val = scan.pop('k_micro_ratio')
if val is None:
ref = 'n2'
elif val == 'r100' or is_rnum(val):
ref = val
else:
raise ValueError(
"Legacy scan.k_micro_ratio integer ratios are no longer exposed "
"in the CLI; use microstate.reference: n2, r100, r<NUM>, or f_tol.")
_warn_legacy(
"scan.k_micro_ratio is deprecated; use microstate.reference.")
out['microstate'] = {'reference': ref}
return out
def _reject_unknown(user: dict, allowed: dict, path: str = '') -> None:
"""Raise ValueError on any key in ``user`` absent from the ``allowed`` tree."""
for key, val in user.items():
here = f'{path}.{key}' if path else str(key)
if key not in allowed:
raise ValueError(
f'Unknown config key: {here!r}. '
f'Run `atom-hifi init` to see the full set of valid keys.')
if isinstance(val, dict) and isinstance(allowed.get(key), dict):
_reject_unknown(val, allowed[key], here)
[docs]
def fill_defaults(config: dict) -> dict:
"""Return ``config`` deep-merged under runtime defaults (no YAML validation).
Used by :func:`atom_hifi.runner.run` so callers may pass a partial dict.
Python API callers may use API-only controls such as ``grid`` and
``patient``; those keys are deliberately rejected by :func:`load_config`.
"""
user = _normalize_legacy_config(config or {})
_validate_config_overrides(user, reject_noop_combos=False)
return _deep_merge(_load_runtime_defaults(), user)
def _validate_config_overrides(user: dict, *, reject_noop_combos: bool = True) -> None:
_eps = user.get('eps_noise') or {}
method = _eps.get('method')
if method is not None and method not in ('thermal', 'rdf', 'precomputed'):
raise ValueError(
f"eps_noise.method must be 'thermal', 'rdf', or 'precomputed', got {method!r}.")
eps_stat = (_eps.get('rdf') or {}).get('eps_stat')
if eps_stat is not None and eps_stat not in ('mean', 'median', 'q95'):
raise ValueError(
f"eps_noise.rdf.eps_stat must be 'mean', 'median', or 'q95', got {eps_stat!r}.")
sigma_source = (_eps.get('thermal') or {}).get('sigma_source')
if sigma_source is not None and sigma_source not in ('qho_table', 'mass_only'):
raise ValueError(
"eps_noise.thermal.sigma_source must be 'qho_table' or 'mass_only', "
f"got {sigma_source!r}.")
f_tol = (_eps.get('force') or {}).get('f_tol')
if f_tol is not None and float(f_tol) <= 0:
raise ValueError(
f"eps_noise.force.f_tol must be positive, got {f_tol!r}.")
if reject_noop_combos and method == 'rdf' and 'thermal' in _eps:
raise ValueError(
"eps_noise.method: rdf cannot be combined with eps_noise.thermal settings.")
if method == 'precomputed':
if reject_noop_combos and ('thermal' in _eps or 'rdf' in _eps):
raise ValueError(
"eps_noise.method: precomputed cannot be combined with calibration settings.")
per = (_eps.get('precomputed') or {}).get('per_species')
if reject_noop_combos and not per:
raise ValueError(
"eps_noise.method: precomputed requires eps_noise.precomputed.per_species.")
ref = (user.get('microstate') or {}).get('reference')
if ref is not None and not (ref == 'n2' or ref == 'f_tol' or is_rnum(ref)):
raise ValueError(
f"microstate.reference must be 'n2', 'r100', 'f_tol', or "
f"'r<NUM>' with NUM in [1, 100], got {ref!r}.")
# Internal force-resolution flag for r<NUM> references (never YAML-settable:
# _reject_unknown rejects keys absent from the template). A naive
# `'force' in cfg['eps_noise']` check is vacuous — fill_defaults always
# injects the template's force: {f_tol: 0.05} — so presence is read off the
# RAW user dict here. Respect-if-present is load-bearing: the CLI mutates
# the already-merged config (--f-tol sets _force_active=True after
# load_config), then runner.run re-enters fill_defaults; recomputing from
# the merged dict would clobber that True with the template-injected force.
micro = user.setdefault('microstate', {})
if '_force_active' not in micro:
micro['_force_active'] = 'force' in (user.get('eps_noise') or {})
# Deliberately NOT gated on reject_noop_combos, unlike the eps_noise
# checks above (issue #72). fill_defaults — the path runner.run takes —
# passes reject_noop_combos=False, so gating this would leave the primary
# entry point doing exactly what the bug report describes: accepting
# resume_from_iter and silently discarding it. The eps_noise combos are
# gated because there a no-op is merely redundant; here it means the run
# starts from scratch when the user asked to resume, and the checkpoint
# they meant to continue is never read.
loop = user.get('loop') or {}
resume_from_iter = loop.get('resume_from_iter')
if (resume_from_iter is not None
and (type(resume_from_iter) is not int or resume_from_iter < 1)):
raise ValueError(
'loop.resume_from_iter must be a positive integer or null; '
f'got {resume_from_iter!r}.')
if loop.get('run') and resume_from_iter is not None:
raise ValueError(
"loop.run=true cannot be combined with "
"loop.resume_from_iter; choose an active loop or a checkpoint "
"resume run.")
# cluster.backend (ADR-0009): loop and patient OOD require a deployable
# classifier (all built-ins); grid scan remains ε_noise-anchored and has
# no count-kind analogue. NSGA-II optimises kmeans' k-fraction axis directly.
backend = (user.get('cluster') or {}).get('backend')
if backend is not None:
# Imported here, not at module scope: config.py is otherwise free of the
# scientific stack, and `from atom_hifi import load_config` must not
# drag in decaf/sklearn. These are capability questions, so they read
# the backend classes instead of re-hardcoding a name list (issue #69).
from atom_hifi.training.cluster_backends import BACKEND_REGISTRY
if backend not in BACKEND_REGISTRY:
names = ', '.join(repr(n) for n in BACKEND_REGISTRY)
raise ValueError(
f'cluster.backend must be one of {names}; got {backend!r}.')
caps = BACKEND_REGISTRY[backend]
if not caps.has_classifier:
if (user.get('loop') or {}).get('run') is True:
raise ValueError(
f"loop.run: true requires a backend with a deployable "
f"classifier — the loop's OOD early-stop has no {backend} "
f"analogue (ADR-0009).")
if (user.get('patient') or {}).get('run_ood') is True:
raise ValueError(
f"patient.run_ood: true requires a backend with a "
f"deployable classifier — {backend} has none (ADR-0009).")
if caps.granularity_kind == 'count':
if (user.get('grid') or {}).get('run') is True:
raise ValueError(
f"grid.run: true is not supported with cluster.backend: "
f"{backend} (count-kind backend; the grid scan is "
f"ε_noise-anchored).")
[docs]
def load_config(path: str) -> dict:
"""Load a user YAML config, merge it under the template defaults, and validate.
Strict: any key not present in the template tree raises ``ValueError``; a
``custom`` descriptor kind is rejected with a pointer to the Python API.
"""
if not os.path.exists(path):
raise FileNotFoundError(f'Config file not found: {path}')
with open(path) as fh:
user = yaml.safe_load(fh) or {}
if not isinstance(user, dict):
raise ValueError(f'Config root must be a mapping, got {type(user).__name__}.')
user = _normalize_legacy_config(user)
defaults = _load_defaults()
_reject_unknown(user, defaults)
kind = (user.get('descriptor') or {}).get('kind')
if kind == 'custom':
raise ValueError(
"descriptor.kind: 'custom' is not supported in YAML — a custom "
"descriptor is a Python callable. Use the Python API instead:\n"
" from atom_hifi.runner import run\n"
" run({'descriptor': {'kind': 'custom', 'custom_fn': my_fn}, ...})")
if (user.get('cluster') or {}).get('backend') == 'custom':
raise ValueError(
"cluster.backend: 'custom' is not supported in YAML — a custom "
"backend is a Python object. Use the Python API instead:\n"
" from atom_hifi.runner import run\n"
" run({'cluster': {'backend': 'custom',\n"
" 'custom_backend': MyBackend()}, ...})")
if kind is not None and kind not in _ALLOWED_DESCRIPTOR_KINDS:
raise ValueError(
f"descriptor.kind must be one of {_ALLOWED_DESCRIPTOR_KINDS}, got {kind!r}.")
_validate_config_overrides(user)
return _deep_merge(_load_runtime_defaults(), user)