"""Atom-HiFi workflow runner.
Single source of truth for the L/R representative-selection pipeline. The CLI
(``atom-hifi run``), the tutorial wrapper, and the tests all call :func:`run` with
a config dict — the same orchestration that previously lived inline in
``hifi_workflow_tutorial.py``.
Config shape is the nested mapping produced by :func:`atom_hifi.config.load_config`
(see that module / ``atom-hifi init`` for the full key set). Partial dicts are
accepted: :func:`run` merges its argument under the template defaults. A custom
descriptor is supported here (and only here) via
``config['descriptor']['custom_fn']`` set to a callable ``fn(List[Atoms]) -> ndarray``.
"""
from __future__ import annotations
import hashlib
import os
import contextlib
import io
import json
import logging
import pickle as pkl
import tempfile
import time
from itertools import chain
logger = logging.getLogger(__name__)
import numpy as np
import pandas as pd
from ase.io import read, write as ase_write
from ase.io.jsonio import MyEncoder as ASEJSONEncoder
from joblib import Parallel, delayed
from atom_hifi.config import fill_defaults, is_rnum
from atom_hifi.corpus import (
filter_atoms as _filter_atoms,
focus_arrays as _focus_arrays,
validate_focus_species as _validate_focus_species,
)
from atom_hifi.piecing import (
assign_back,
make_pieces,
original_representatives,
)
from atom_hifi.soap_ML.soap import generate_ED_SOAP
from atom_hifi.training.hifi_workflow import HiFiWorkflow
from atom_hifi.training.qho_sigma import (
qho_table_version,
resolve_sigma_thermal_targets,
)
from atom_hifi.training.rdf_eps import RDF_DEFAULTS
[docs]
def collapsed_elements(n_groups_per_el, min_groups=1):
"""Focus elements whose clustering collapsed (<= ``min_groups`` groups).
A collapsed element resolves no environment diversity at the operating point
— every atom of that element shares one group — so the representative set
captures nothing distinctive for it even when ``|S|`` looks healthy.
Counts **density-supported groups only**. Promoted rare-environment groups
(issue #63) are one-atom outliers, not resolved diversity, so including them
would mask a genuine collapse: PtOH_big O would report 64 groups while still
being a single density blob.
"""
return [el for el, n in n_groups_per_el.items() if n <= min_groups]
def _post_selection_metrics(wf, rep_idx, struct_map, selection_method,
focus_elements=None):
"""Compute L/R/F for the selected subset; log a summary; return the metric dict."""
from scipy.stats import entropy as scipy_entropy
rep_idx = sorted(rep_idx)
mu_labels = wf.mu_labels
k_micro = wf.k_micro
N_L = len(mu_labels)
all_counts = np.bincount(mu_labels, minlength=k_micro)
k_occ_L = int((all_counts > 0).sum())
# Labels use disjoint element offsets: normalize both entropies over the
# joint population, as Judge does. wf.h_mu is a legacy per-element mean.
H_L = float(scipy_entropy(all_counts[all_counts > 0], base=np.e))
atom_mask = np.isin(struct_map, rep_idx)
sub_mu = mu_labels[atom_mask]
n_atoms_S = int(atom_mask.sum())
counts_S = np.bincount(sub_mu, minlength=k_micro)
occ_S = counts_S[counts_S > 0]
H_S = float(scipy_entropy(occ_S, base=np.e)) if len(occ_S) else 0.0
L = min(H_S / H_L, 1.0) if H_L > 0 else 0.0
k_occ_S = int((counts_S > 0).sum())
R = (n_atoms_S / k_occ_S) / (N_L / k_occ_L) if k_occ_S > 0 else float('inf')
logger.info(f' Method: {selection_method or "baseline"}')
logger.info(f' |S|: {len(rep_idx)}')
logger.info(f' atoms: {n_atoms_S}')
logger.info(f' k_occ_S: {k_occ_S} / k_micro={k_micro}')
logger.info(f' L: {L:.4f}')
logger.info(f' R: {R:.4f}')
logger.info(f' F (L/R): {L/R:.4f}')
if focus_elements is not None and wf.mu_labels_per_el is not None:
logger.info('\n Per-element breakdown:')
for el in focus_elements:
if el in (wf.k_micro_per_el or {}):
logger.info(f' {el}: k_micro={wf.k_micro_per_el[el]} '
f'ε_noise={wf.eps_noise_per_species[el]:.5f}')
return {'L': L, 'R': R, 'F': L / R, 'n_structs': len(rep_idx), 'n_atoms': n_atoms_S}
def _build_descriptor_fn(desc_cfg, descriptor_fn_override, lib_for_species):
"""Return full/center descriptor callbacks, r_cut, and descriptor species."""
descriptor = desc_cfg['kind']
if descriptor == 'eds':
eds = dict(desc_cfg['eds'])
species = (np.unique(np.concatenate(
[s.numbers for s in lib_for_species])).astype(int).tolist()
if eds['species'] is None else list(eds['species']))
lengthscale = eds['lengthscale']
logger.info(f' ED-SOAP species: {species} lengthscale={lengthscale} Å')
logger.info(f' Short: cut={eds["s_cut"]:.4f}×ls={eds["s_cut"]*lengthscale:.3f}Å '
f'nmax={eds["s_nmax"]} lmax={eds["s_lmax"]}')
logger.info(f' Long: cut={eds["l_cut"]:.4f}×ls={eds["l_cut"]*lengthscale:.3f}Å '
f'nmax={eds["l_nmax"]} lmax={eds["l_lmax"]} periodic={eds["periodic"]}')
def center_descriptor_fn(structs, atom_ids_per_structure):
with contextlib.redirect_stdout(io.StringIO()):
return generate_ED_SOAP(
structs,
species=np.array(species),
lengthscale=lengthscale,
SOAP_S_cut=eds['s_cut'],
SOAP_S_nmax=eds['s_nmax'],
SOAP_S_lmax=eds['s_lmax'],
SOAP_L_cut=eds['l_cut'],
SOAP_L_nmax=eds['l_nmax'],
SOAP_L_lmax=eds['l_lmax'],
periodic=eds['periodic'],
atom_ids_per_structure=atom_ids_per_structure,
)
def descriptor_fn(structs):
return center_descriptor_fn(structs, None)
descriptor_fn._cache_settings = {'kind': 'eds', **eds, 'species': species}
return descriptor_fn, center_descriptor_fn, eds['r_cut'], species
if descriptor == 'ace':
mace_cfg = desc_cfg['ace']
try:
from mace.calculators import MACECalculator
except ImportError:
raise ImportError(
"descriptor.kind='ace' requires mace-torch. "
"Install with: pip install mace-torch")
try:
model_hash = _file_sha256(mace_cfg['model_path'])
except (OSError, TypeError):
model_hash = None # Remote/opaque models cannot be verified locally.
calc = MACECalculator(model_paths=mace_cfg['model_path'],
device=mace_cfg['device'], default_dtype='float32')
def descriptor_fn(structs):
rows = []
for atoms in structs:
# MACE needs to be the active calculator while descriptors are
# evaluated, but the library structures are retained through the
# final selection. Leaving ``calc`` on each Atoms would keep the
# model alive after this callback is released (and overwrite a
# caller-provided calculator). Restore the exact prior state on
# success and failure.
prior_calc = atoms.calc
try:
atoms.calc = calc
desc = calc.get_descriptors(atoms) # (n_atoms, D)
finally:
atoms.calc = prior_calc
rows.append(desc)
return np.vstack(rows)
descriptor_fn._cache_settings = (
{'kind': 'ace', **mace_cfg, 'default_dtype': 'float32',
'model_sha256': model_hash} if model_hash else None)
return descriptor_fn, None, mace_cfg['r_cut'], None
if descriptor == 'custom':
if descriptor_fn_override is None:
raise ValueError(
"descriptor.kind='custom' but descriptor.custom_fn is None. "
"Pass a callable fn(List[Atoms]) -> ndarray via the Python API.")
return descriptor_fn_override, None, desc_cfg['custom_r_cut'], None
raise ValueError(f"Unknown descriptor.kind={descriptor!r}. "
"Choose 'eds', 'ace', or 'custom' (custom via Python API).")
def _load_library(lib_path, focus_elements):
logger.info('=' * 70)
logger.info('Step 1: load structure library')
logger.info('=' * 70)
lib_all = read(lib_path, index=':')
logger.info(f' Loaded {len(lib_all)} structures from {lib_path}')
focus_explicit = bool(focus_elements)
if not focus_elements:
focus_elements = sorted({sym for a in lib_all
for sym in a.get_chemical_symbols()})
logger.info(f' focus_elements=null → clustering all elements present: {focus_elements}')
lib = [a for a in lib_all if any(el in set(a.get_chemical_symbols())
for el in focus_elements)]
logger.info(f' Retained {len(lib)} structures containing {focus_elements}')
return lib_all, lib, focus_elements, focus_explicit
def _stream_sha256(stream):
digest = hashlib.sha256()
for chunk in iter(lambda: stream.read(1024 * 1024), b''):
digest.update(chunk)
return digest.hexdigest()
def _file_sha256(path):
with open(path, 'rb') as stream:
return _stream_sha256(stream)
def _descriptor_cache_key(lib, descriptor_fn, descriptor_kind):
"""Bind built-in settings and ordered ASE contents; opaque callbacks bypass reuse."""
settings = getattr(descriptor_fn, '_cache_settings', None)
if descriptor_kind == 'custom' or settings is None:
return None
digest = hashlib.sha256()
try:
# Hash one structure at a time to avoid duplicating a large library.
for payload in chain((settings,), (atoms.todict() for atoms in lib)):
encoded = json.dumps(payload, cls=ASEJSONEncoder, sort_keys=True,
separators=(',', ':'), allow_nan=False)
digest.update(hashlib.sha256(encoded.encode()).digest())
except (TypeError, ValueError):
logger.info(' Descriptor inputs cannot be fingerprinted; bypassing cache reuse.')
return None
return digest.hexdigest()
def _atomic_cache_write(path, write):
"""Replace one cache file without leaving a partially written result."""
with tempfile.NamedTemporaryFile(dir=os.path.dirname(os.path.abspath(path)),
delete=False) as stream:
temporary = stream.name
try:
write(stream)
stream.flush()
stream.seek(0)
digest = _stream_sha256(stream)
stream.close()
os.replace(temporary, path)
return digest
finally:
if os.path.exists(temporary):
os.unlink(temporary)
def _compute_descriptors(desc_cache, lib, descriptor_fn, descriptor_kind,
n_jobs_desc):
"""Load or compute per-structure descriptors for the library."""
def _one_desc(atoms):
return descriptor_fn([atoms])
desc_list = None
cache_key = _descriptor_cache_key(lib, descriptor_fn, descriptor_kind)
manifest_path = os.fspath(desc_cache) + '.manifest.json'
if cache_key is not None and os.path.exists(desc_cache):
try:
with open(manifest_path) as f:
manifest = json.load(f)
if manifest.get('schema') != 1 or manifest.get('input_sha256') != cache_key:
raise ValueError('descriptor cache identity mismatch')
with open(desc_cache, 'rb') as f:
if manifest.get('cache_sha256') != _stream_sha256(f):
raise ValueError('descriptor cache/manifest pairing mismatch')
f.seek(0)
cached = pkl.load(f)
if (len(cached) == len(lib)
and all(c.ndim == 2 and c.shape[0] == len(atoms)
for c, atoms in zip(cached, lib))):
desc_list = cached
logger.info(f' Loaded cached descriptors from {desc_cache}')
except Exception as exc:
logger.warning(f' Descriptor cache {desc_cache} unreadable ({exc}); recomputing.')
if desc_list is None:
logger.info(f' Computing descriptors for {len(lib)} structures '
f'(n_jobs={n_jobs_desc}) …')
t0 = time.time()
prefer = 'threads' if descriptor_kind == 'eds' else 'processes'
desc_list = Parallel(n_jobs=n_jobs_desc, prefer=prefer)(
delayed(_one_desc)(atoms) for atoms in lib
)
logger.info(f' Done in {time.time() - t0:.1f}s')
cache_hash = _atomic_cache_write(desc_cache, lambda f: pkl.dump(
desc_list, f, protocol=pkl.HIGHEST_PROTOCOL))
if cache_key is not None:
manifest = {'schema': 1, 'input_sha256': cache_key,
'cache_sha256': cache_hash}
_atomic_cache_write(manifest_path, lambda f: f.write(
json.dumps(manifest, sort_keys=True).encode()))
elif os.path.exists(manifest_path):
os.unlink(manifest_path)
logger.info(f' Saved {desc_cache}')
return desc_list
def _load_patient_descriptors(patient_path, run_patient_ood, pat_cache,
descriptor_fn, descriptor_kind, n_jobs_desc,
focus_elements):
"""Return patient focus descriptors, or None when Stage 4 is disabled."""
if patient_path is None or not run_patient_ood:
return None
logger.info('\n' + '=' * 70)
logger.info('Step 4: patient (OOD) structures')
logger.info('=' * 70)
patient_lib = read(patient_path, index=':')
logger.info(f' Loaded {len(patient_lib)} patient structures from {patient_path}')
pat_list = _compute_descriptors(
pat_cache, patient_lib, descriptor_fn, descriptor_kind, n_jobs_desc)
pat_focus = {el: [] for el in (focus_elements or [])}
n_feat = pat_list[0].shape[1] if pat_list else 0
for atoms, desc in zip(patient_lib, pat_list):
_, foc, foc_atoms = _filter_atoms(atoms, desc, focus_elements)
syms = np.array(foc_atoms.get_chemical_symbols())
for el in (focus_elements or []):
rows = foc[syms == el]
if len(rows):
pat_focus[el].append(rows)
patient_desc_per_el = {
el: (np.vstack(rows) if rows else np.empty((0, n_feat)))
for el, rows in pat_focus.items()
}
total = sum(len(v) for v in patient_desc_per_el.values())
logger.info(f' Patient focus-atom descriptors: total={total} '
f'per_element={ {el: v.shape for el, v in patient_desc_per_el.items()} }')
return patient_desc_per_el
def _eps_cache_fingerprint_ok(d, fingerprint):
"""True if a cached eps_noise_raw.npz matches the requested calibration.
Runner fingerprints also bind descriptor/corpus identity and sampling policy.
Missing identities and opaque callbacks force recalibration. Legacy method
metadata remains readable for callers that do not request an input identity.
"""
if fingerprint is None:
return True
if 'calibration_input_sha256' in fingerprint:
requested = fingerprint['calibration_input_sha256']
if (requested is None or 'calibration_input_sha256' not in d
or str(d['calibration_input_sha256']) != requested):
logger.info(' ε_noise calibration inputs changed or unverifiable; recomputing.')
return False
cached_method = str(d['eps_method']) if 'eps_method' in d else 'thermal'
if cached_method != fingerprint['method']:
logger.warning(f' ε_noise cache method={cached_method!r} ≠ requested '
f'{fingerprint["method"]!r}; recomputing.')
return False
if cached_method == 'rdf':
cn = int(d['rdf_n_samples']) if 'rdf_n_samples' in d else None
cs = str(d['rdf_eps_stat']) if 'rdf_eps_stat' in d else None
if cn != int(fingerprint['n_samples']) or cs != str(fingerprint['eps_stat']):
logger.warning(f' ε_noise[rdf] cache variant '
f'(n_samples={cn}, eps_stat={cs}) ≠ requested '
f'(n_samples={fingerprint["n_samples"]}, '
f'eps_stat={fingerprint["eps_stat"]}); recomputing.')
return False
if cached_method == 'thermal':
cached_source = str(d['sigma_source']) if 'sigma_source' in d else 'mass_only'
requested_source = str(fingerprint.get('sigma_source'))
if cached_source != requested_source:
logger.warning(f' ε_noise[thermal] cache sigma_source={cached_source!r} '
f'≠ requested {requested_source!r}; recomputing.')
return False
cached_table = str(d['qho_table_version']) if 'qho_table_version' in d else None
requested_table = fingerprint.get('qho_table_version')
if requested_source == 'qho_table' and cached_table != requested_table:
logger.warning(f' ε_noise[thermal] cache QHO table={cached_table!r} '
f'≠ requested {requested_table!r}; recomputing.')
return False
for el, target in (fingerprint.get('sigma_targets') or {}).items():
key = f'sigma_target_{el}'
if key not in d:
logger.warning(f' ε_noise[thermal] cache lacks {key}; recomputing.')
return False
if not np.isclose(float(d[key]), float(target), rtol=1e-10, atol=1e-12):
logger.warning(f' ε_noise[thermal] cache target σ[{el}]='
f'{float(d[key]):.8f} ≠ requested '
f'{float(target):.8f}; recomputing.')
return False
return True
def _load_eps_noise_cache(eps_cache, focus_elements, fingerprint=None,
*, require_curve=False):
eps_noise_global = None
eps_per_species_cache = None
fallback_elements = set()
eps_curve_cache = None
if os.path.exists(eps_cache):
with np.load(eps_cache, allow_pickle=True) as d:
if not _eps_cache_fingerprint_ok(d, fingerprint):
return None, None, set(), None
eps_noise_global = float(d['eps_noise_global'])
loaded_per_el = {el: float(d[f'eps_{el}'])
for el in (focus_elements or []) if f'eps_{el}' in d}
if len(loaded_per_el) == len(focus_elements or []) and loaded_per_el:
eps_per_species_cache = loaded_per_el
elif loaded_per_el:
logger.warning(f'per-element ε_noise cache incomplete '
f'({list(loaded_per_el.keys())} found, {focus_elements} needed); '
f'falling back to global ε_noise={eps_noise_global:.5f}')
fallback_elements = {el for el in (focus_elements or [])
if f'fallback_{el}' in d and bool(d[f'fallback_{el}'])}
if 'sigma_rattle' in d:
mean_per_el = {
el: np.array(d[f'mean_{el}'], dtype=float)
for el in (focus_elements or []) if f'mean_{el}' in d
}
if mean_per_el:
eps_curve_cache = {
'sigma_rattle': np.array(d['sigma_rattle'], dtype=float),
'mean_per_el': mean_per_el,
}
if require_curve and (
eps_curve_cache is None
or set(eps_curve_cache['mean_per_el']) != set(focus_elements or [])
):
logger.warning(
f' ε_noise cache {eps_cache} lacks full ε(σ) curve required by '
'microstate.reference=f_tol / r<NUM> force-resolution; recomputing.')
return None, None, set(), None
if eps_per_species_cache is not None:
logger.info(f' Loaded per-element ε_noise from {eps_cache}: '
f'{ {el: f"{v:.5f}" for el, v in eps_per_species_cache.items()} }')
if fallback_elements:
logger.info(f' RDF no-peak (N//2) elements from cache: {sorted(fallback_elements)}')
elif eps_noise_global is not None:
logger.info(f' Loaded ε_noise from {eps_cache}: {eps_noise_global:.5f}')
return eps_noise_global, eps_per_species_cache, fallback_elements, eps_curve_cache
def _calibrate_eps_noise_if_needed(wf, lib, descriptor_fn,
center_descriptor_fn, lib_focus,
focus_desc_list, initial_sample, seed):
if wf.eps_noise_global is not None:
return
# Count-kind backends (kmeans) scan a k-fraction axis: ε_noise is only
# needed for the r100/r<NUM>/f_tol microstate references (ADR-0009).
if (getattr(wf, 'granularity_kind', 'distance') == 'count'
and not wf._backend_caps.emits_noise
and wf.k_micro_ratio not in ('r100', 'f_tol')
and not is_rnum(wf.k_micro_ratio)):
logger.info('Stage 1 skipped: %s backend (count-kind) with the N//2 '
'microstate reference never consumes ε_noise.',
wf.cluster_backend)
return
logger.info('\n' + '=' * 70)
logger.info('Stage 1: ε_noise calibration')
logger.info('=' * 70)
n_init = min(initial_sample, len(lib))
rng = np.random.default_rng(seed)
init_idx = rng.choice(len(lib), size=n_init, replace=False)
wf.calibrate_eps_noise(
atoms_list = [lib[int(i)] for i in init_idx],
descriptor_fn = descriptor_fn,
center_descriptor_fn = center_descriptor_fn,
lib_focus = [lib_focus[int(i)] for i in init_idx],
desc_array = np.vstack([focus_desc_list[int(i)] for i in init_idx]),
K = 200,
M = 20,
)
if wf.eps_noise_per_species is not None:
logger.info(f' Per-element ε_noise: { {el: f"{v:.5f}" for el, v in wf.eps_noise_per_species.items()} }')
def _score_remaining_ood(wf, judger, remaining, lib_focus, focus_desc_list):
"""Score remaining structures by OOD focus atoms against a fitted Judge."""
rows = []
for i in remaining:
i = int(i)
desc_i = focus_desc_list[i]
n_focus_atoms = int(len(desc_i))
ood_count = 0
if n_focus_atoms:
smap_i = np.zeros(n_focus_atoms, dtype=int)
_, desc_per_el_i, _ = wf._build_el_data(
[lib_focus[i]], desc_i, smap_i)
for el in wf.species:
clu = judger.cluster_per_el.get(el)
if clu is None or el not in desc_per_el_i:
continue
embedded = clu.embedder.embed(desc_per_el_i[el])
gID_distinct, _ = clu.classify(embedded)
ood_count += int(len(gID_distinct[-1]))
rows.append({
'index': i,
'ood_count': ood_count,
'n_focus_atoms': n_focus_atoms,
'ood_frac': ood_count / n_focus_atoms if n_focus_atoms else 0.0,
})
return rows
def _pick_ood_ranked_batch(scored_rows, batch_size):
"""Return top OOD-positive structures by frac, count, stable index."""
positive = [r for r in scored_rows if int(r['ood_count']) > 0]
positive.sort(key=lambda r: (-float(r['ood_frac']),
-int(r['ood_count']),
int(r['index'])))
return positive[:int(batch_size)]
def _structure_fingerprint(atoms):
"""Stable content fingerprint used to bind loop checkpoints to a library."""
try:
payload = json.dumps(
atoms, cls=ASEJSONEncoder, sort_keys=True,
separators=(',', ':'), ensure_ascii=True,
).encode('utf-8')
except Exception as exc:
raise RuntimeError(
'loop checkpointing requires JSON-serializable ASE structure '
'metadata so source identity can be verified on resume.') from exc
return hashlib.sha256(payload).hexdigest()
def _build_final_corpus(wf, run_loop, resume_from_iter, initial_sample,
loop_batch_size, loop_skip_fine_scan, lib, lib_focus,
focus_desc_list, descriptor_array, struct_map,
output_path_fn, seed, run_nsga2=False,
inner_nsga2_n_gen=5, inner_nsga2_pop_size=None):
"""Build the final training corpus from loop, resume, or full library."""
if run_loop:
logger.info('\n' + '=' * 70)
logger.info('Stage 2: active-learning loop')
logger.info('=' * 70)
n_init = min(initial_sample, len(lib))
rng = np.random.default_rng(seed)
init_idx = rng.choice(len(lib), size=n_init, replace=False)
sampled = set(int(i) for i in init_idx)
sampled_order = [int(i) for i in init_idx]
lib_focus_init = [lib_focus[int(i)] for i in init_idx]
desc_init = [focus_desc_list[int(i)] for i in init_idx]
full_init = [lib[int(i)] for i in init_idx]
def sample_next():
remaining = [i for i in range(len(lib)) if i not in sampled]
if not remaining:
return None
n_next = min(loop_batch_size, len(remaining))
rng_next = np.random.default_rng(seed + len(sampled))
picks = rng_next.choice(remaining, size=n_next, replace=False)
nf, na, ns = [], [], []
for i in picks:
nf.append(lib_focus[int(i)])
na.append(focus_desc_list[int(i)])
ns.append(lib[int(i)])
sampled.add(int(i))
sampled_order.append(int(i))
return nf, na, ns
def sample_next_ood(best_judger):
remaining = [i for i in range(len(lib)) if i not in sampled]
picked = _pick_ood_ranked_batch(
_score_remaining_ood(wf, best_judger, remaining,
lib_focus, focus_desc_list),
loop_batch_size,
)
if not picked:
return None
nf, na, ns = [], [], []
ood_next = 0
n_atoms_next = 0
for row in picked:
i = int(row['index'])
nf.append(lib_focus[i])
na.append(focus_desc_list[i])
ns.append(lib[i])
sampled.add(i)
sampled_order.append(i)
ood_next += int(row['ood_count'])
n_atoms_next += int(row['n_focus_atoms'])
ood_frac = ood_next / n_atoms_next if n_atoms_next else 0.0
return nf, na, ns, ood_next, ood_frac
inner_enabled = bool(run_nsga2 and len(getattr(wf, 'species', [])) > 1)
wf.run_loop(
lib_focus_init = lib_focus_init,
desc_focus_init = desc_init,
full_struct_lib_init = full_init,
sample_next_fn = sample_next,
sample_next_ood_fn = sample_next_ood if inner_enabled else None,
skip_inner_fine = loop_skip_fine_scan,
inner_nsga2 = inner_enabled,
inner_nsga2_n_gen = inner_nsga2_n_gen,
inner_nsga2_pop_size = inner_nsga2_pop_size,
checkpoint_source_indices_fn = lambda: np.asarray(
sampled_order, dtype=np.int64),
checkpoint_source_fingerprints_fn = lambda: np.asarray([
_structure_fingerprint(lib[i]) for i in sampled_order
]),
)
lib_focus_final = wf._loop_final_lib_focus
desc_final = np.vstack(wf._loop_final_desc_list)
full_final = wf._loop_final_full_lib
struct_map_final = np.repeat(
np.arange(len(wf._loop_final_desc_list)),
[len(x) for x in wf._loop_final_desc_list],
)
return lib_focus_final, desc_final, full_final, struct_map_final
if resume_from_iter is not None:
npy = output_path_fn(f'descp_iter{resume_from_iter:02d}.npy')
smap = output_path_fn(f'descp_iter{resume_from_iter:02d}_struct_map.npy')
source_idx_path = output_path_fn(
f'descp_iter{resume_from_iter:02d}_source_idx.npy')
source_hash_path = output_path_fn(
f'descp_iter{resume_from_iter:02d}_source_hash.npy')
if (not os.path.exists(npy) or not os.path.exists(smap)
or not os.path.exists(source_idx_path)
or not os.path.exists(source_hash_path)):
raise FileNotFoundError(
f'loop.resume_from_iter={resume_from_iter}: checkpoint not found:\n'
f' {npy}\n {smap}\n {source_idx_path}\n {source_hash_path}\n'
'Source identity metadata is required for exact reconstruction; '
'legacy checkpoints that only saved descriptors cannot safely '
'resume OOD-ranked loops.')
desc_final = np.load(npy)
struct_map_final = np.load(smap)
raw_picked = np.load(source_idx_path)
if (raw_picked.ndim != 1
or not np.issubdtype(raw_picked.dtype, np.integer)):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: source indices must '
'be a one-dimensional integer sequence.')
picked = np.asarray(raw_picked, dtype=np.int64)
if len(np.unique(picked)) != len(picked):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: source indices must '
'be unique.')
if len(picked) and (picked.min() < 0 or picked.max() >= len(lib)):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: source indices are '
f'outside the current library of {len(lib)} structures.')
saved_hashes = np.asarray(np.load(source_hash_path))
current_hashes = np.asarray([
_structure_fingerprint(lib[int(i)]) for i in picked
])
if (saved_hashes.ndim != 1 or len(saved_hashes) != len(picked)
or not np.array_equal(saved_hashes.astype(str), current_hashes)):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: recorded source '
'structures do not match the current library at the saved '
'indices; restore the original library ordering and contents.')
lib_focus_final = [lib_focus[int(i)] for i in picked]
full_final = [lib[int(i)] for i in picked]
logger.info(f' Iter-{resume_from_iter} corpus: {len(lib_focus_final)} structs '
f'({desc_final.shape[0]} focus atoms) loaded from checkpoint')
n_unique_smap = len(np.unique(struct_map_final))
if n_unique_smap != len(lib_focus_final):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: struct count mismatch — '
f'{n_unique_smap} unique structs in saved smap vs {len(lib_focus_final)} '
'recorded source structures; the checkpoint is inconsistent.')
raw_struct_map = np.asarray(struct_map_final)
if (raw_struct_map.ndim != 1
or not np.issubdtype(raw_struct_map.dtype, np.integer)
or (len(raw_struct_map) and (
raw_struct_map.min() < 0
or raw_struct_map.max() >= len(picked)))):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: saved structure map '
'must be a one-dimensional integer array bounded by the '
'recorded source-structure count.')
if desc_final.ndim != 2 or desc_final.shape[0] != len(raw_struct_map):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: descriptor rows '
f'({desc_final.shape[0] if desc_final.ndim else 0}) do not '
f'match structure-map rows ({len(raw_struct_map)}).')
saved_counts = np.bincount(
raw_struct_map.astype(int, copy=False), minlength=len(picked))
expected_counts = np.asarray(
[len(lib_focus[int(i)]) for i in picked], dtype=int)
if (len(saved_counts) != len(expected_counts)
or not np.array_equal(saved_counts, expected_counts)):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: saved descriptor '
'rows do not match the focus-atom counts of the recorded source '
'structures; the library or checkpoint is incompatible.')
offset = 0
for i, source_index in enumerate(picked):
rows = focus_desc_list[int(source_index)]
count = len(rows)
if (desc_final.dtype != rows.dtype
or not np.array_equal(struct_map_final[offset:offset + count],
np.full(count, i))
or not np.array_equal(desc_final[offset:offset + count], rows)):
raise RuntimeError(
f'loop.resume_from_iter={resume_from_iter}: saved descriptor '
'values or row ordering do not match the current selected '
'source structures; use the original descriptor inputs or '
'start a new run.')
offset += count
return lib_focus_final, desc_final, full_final, struct_map_final
return lib_focus, descriptor_array, lib, struct_map
def _apply_piecing(target_atoms, full_final, desc_list, focus_elements):
"""Substitute voxel pieces for structures."""
logger.info('\n' + '=' * 70)
logger.info(f'Piecing: voxel substitution (target_atoms={target_atoms})')
logger.info('=' * 70)
pieces_lib, pieces_desc_list, piece_map = make_pieces(
full_final, desc_list, target_atoms)
sizes = np.array([len(p) for p in pieces_lib])
logger.info(f' {len(full_final)} structures → {len(pieces_lib)} pieces '
f'(atoms/piece: min={sizes.min()} median={int(np.median(sizes))} '
f'max={sizes.max()})')
logger.info(
' NOTE: |S| and fidelity reports count PIECES; representative '
'files contain unique original structures.')
lib_focus_final, desc_final, struct_map_final, _ = _focus_arrays(
pieces_lib, pieces_desc_list, focus_elements)
return piece_map, pieces_lib, lib_focus_final, desc_final, struct_map_final
def _run_final_scan(wf, lib_focus_final, desc_final, full_final, struct_map_final,
l_tol):
"""Run the 1-D final L/R scan and log the fine-scan table."""
logger.info('\n' + '=' * 70)
logger.info('Stage 3: final L/R scan')
logger.info('=' * 70)
wf.run_final(
lib_focus = lib_focus_final,
desc_array = desc_final,
full_struct_lib = full_final,
struct_map = struct_map_final,
)
logger.debug('')
scan_col = 'c_factor'
logger.info(f"{'c':>10} {'L':>7} {'R':>7} {'F':>7}"
f" {'|S|':>6} {'atoms':>8} {'★':2}")
logger.info('-' * 60)
for _, row in wf.df_fine.iterrows():
star = '★' if (wf.delta_star_L095 is not None
and abs(row[scan_col] - wf.delta_star_L095) < 1e-9) else ''
logger.info(f"{row[scan_col]:>10.5f} {min(row['L'], 1.0):>7.4f} "
f"{row['R']:>7.4f} {row['F']:>7.4f} "
f"{int(row['n_structs_S']):>6} {int(row['n_atoms_S']):>8} {star}")
logger.debug('')
lbl = 'c*'
if l_tol is None:
logger.info(f" {lbl} (L floor disabled — l_tol: null; auto-pick = max F)")
else:
logger.info(f" {lbl} (L≥{l_tol}): {wf.delta_star_L095}")
logger.info(f" {lbl} (max F): {wf.delta_star_F}")
def _run_optional_cfactor_optimizations(
wf, run_grid_scan, n_coarse_grid, n_fine_grid, restore_grid,
run_nsga2, n_nsga2_gen, n_nsga2_pop, restore_nsga2,
focus_elements, l_tol, seed, output_dir, output_path_fn):
"""Run optional per-element grid/NSGA c-factor optimizers."""
best_c_per_el = None
run_grid = run_grid_scan and wf.mu_labels_per_el is not None
if run_grid:
logger.info('\n' + '=' * 70)
logger.info('Stage 3b: ND grid scan (cached per-element DECAF)')
logger.info('=' * 70)
c_star_common = wf.delta_star_L095 or wf.delta_star_F
grid_csv = output_path_fn('scan_grid.csv')
if restore_grid and os.path.isfile(grid_csv):
df_grid = pd.read_csv(grid_csv)
wf.df_grid = df_grid
logger.info(f' Restored {len(df_grid)} scan points from {grid_csv}')
else:
if restore_grid:
logger.info(f' grid.restore=true but {grid_csv} not found — running scan.')
df_grid = wf.run_grid_scan(
n_coarse_per_el = n_coarse_grid,
n_fine_per_el = n_fine_grid,
l_tol = l_tol,
save_dir = output_dir,
)
df_pareto = wf.report_cfactor_grid(df_grid, c_star_1d=c_star_common,
save_dir=output_dir, l_tol=l_tol)
best_grid = df_grid.loc[df_grid['F'].idxmax()]
fr_1d = float(wf.df_fine['F'].max())
s_1d = int(wf.df_fine.loc[wf.df_fine['F'].idxmax(), 'n_structs_S'])
fr_grid = float(best_grid['F'])
s_grid = int(best_grid['n_structs_S'])
f95 = (df_pareto[df_pareto['L'] >= l_tol]
if l_tol is not None else df_pareto.iloc[0:0])
logger.debug('')
logger.info(f' 1-D baseline: c*={c_star_common:.5f} F={fr_1d:.4f} |S|={s_1d}')
cgrid_str = ' '.join(
f'c_{el}={best_grid[f"c_{el}"]:.4f}' for el in (focus_elements or []))
logger.info(f' Grid max F: {cgrid_str} F={fr_grid:.4f} |S|={s_grid}')
logger.info(f' ΔF={fr_grid-fr_1d:+.4f} Δ|S|={s_grid-s_1d:+d}')
best_c_per_el = wf._best_grid_operating_point(df_grid)
if l_tol is None:
logger.info(f' Recommended (max F, L floor disabled): {cgrid_str} '
f'F={fr_grid:.4f} |S|={s_grid}')
elif len(f95):
rec = f95.loc[f95['F'].idxmax()]
rec_str = ' '.join(
f'c_{el}={rec[f"c_{el}"]:.4f}' for el in (focus_elements or []))
logger.info(f' Recommended (L≥{l_tol}): {rec_str} '
f'F={rec["F"]:.4f} |S|={int(rec["n_structs_S"])}')
else:
logger.info(f' No grid point meets L≥{l_tol} — lower l_tol or run wider grid.')
logger.info(f' Pareto: {len(df_pareto)} pts → scan_grid_operating_points.out + scan_grid_report.png')
run_nsga2_enabled = run_nsga2 and wf.mu_labels_per_el is not None
if run_nsga2_enabled:
logger.info('\n' + '=' * 70)
logger.info('Stage 3c: NSGA-II Pareto optimisation of per-element c-factors')
logger.info('=' * 70)
pareto_csv = output_path_fn('pareto_front.csv')
if restore_nsga2 and os.path.isfile(pareto_csv):
df_pareto_nsga2 = pd.read_csv(pareto_csv)
logger.info(f' Restored {len(df_pareto_nsga2)}-point Pareto front from {pareto_csv}')
if 'fidelity_best' in df_pareto_nsga2.columns and df_pareto_nsga2['fidelity_best'].any():
fb = df_pareto_nsga2[df_pareto_nsga2['fidelity_best']].iloc[0]
else:
from atom_hifi.training.optimizer import fidelity_envelope
env = fidelity_envelope(df_pareto_nsga2)
fb = env.loc[env['F'].idxmax()]
if l_tol is None:
rec_nsga2 = fb # no L floor — auto-pick = best-fidelity point
else:
q = df_pareto_nsga2[df_pareto_nsga2['L_min'] >= l_tol]
rec_nsga2 = q.iloc[0] if not q.empty else df_pareto_nsga2.loc[df_pareto_nsga2['L_min'].idxmax()]
best_c_per_el = {el: float(rec_nsga2[f'c_{el}']) for el in (focus_elements or [])}
c_str = ' '.join(f'c_{el}={best_c_per_el[el]:.4f}' for el in (focus_elements or []))
logger.info(f' Auto-picked: {c_str} L_min={rec_nsga2["L_min"]:.4f} |S|={int(rec_nsga2["n_structs_S"])}')
fb_str = ' '.join(f'c_{el}={float(fb[f"c_{el}"]):.4f}' for el in (focus_elements or []))
logger.info(f' Best fidelity (global): {fb_str} F={fb["F"]:.4f} L_min={fb["L_min"]:.4f} |S|={int(fb["n_structs_S"])}')
else:
if restore_nsga2:
logger.info(f' nsga2.restore=true but {pareto_csv} not found — running optimizer.')
_, df_pareto_nsga2, _, best_c_per_el = wf.run_nsga2(
n_gen = n_nsga2_gen,
pop_size = n_nsga2_pop,
c_star_1d = wf.delta_star_L095 or wf.delta_star_F,
l_tol = l_tol,
seed = seed,
save_dir = output_dir,
)
return best_c_per_el
def _select_final_representatives(wf, best_c_per_el, delta_pick,
patient_desc_array, selection_method,
struct_map_final, focus_elements,
full_final, output_path_fn,
save_representatives=True):
"""Fit at the selected operating point and return its representatives.
Optionally saves representatives.xyz; returns (structures, indices,
metrics, path).
"""
logger.info('\n' + '=' * 70)
logger.info('Stage 5: final fit and representative-set selection')
logger.info(f' selection_method = {selection_method!r}')
logger.info('=' * 70)
if delta_pick and best_c_per_el:
logger.info(
f' user delta_pick={delta_pick!r} overrides auto-picked '
f'best_c_per_el={best_c_per_el!r}'
)
fit_kwargs = dict(
delta_pick = delta_pick or best_c_per_el,
patient_descriptors = patient_desc_array,
selection_method = selection_method,
)
if not save_representatives:
fit_kwargs['save_representatives'] = False
rep_structs, rep_idx = wf.plot_final_fit(**fit_kwargs)
logger.info('\n' + '=' * 70)
logger.info('Post-selection metrics')
logger.info('=' * 70)
metrics = _post_selection_metrics(
wf, rep_idx, struct_map_final, selection_method, focus_elements or [])
out_xyz = output_path_fn('representatives.xyz')
if save_representatives:
ase_write(out_xyz, rep_structs)
logger.info(
f'\n Saved {len(rep_structs)} representative structures → {out_xyz}')
return rep_structs, rep_idx, metrics, out_xyz
def _meta_by_group(judge):
"""Return final global group id → meta-group id."""
result = {}
for element, order in getattr(judge, 'group_order_per_el', {}).items():
offset = int(judge.group_id_offset[element])
for local_id, meta_id in enumerate(order.meta_of_local):
result[offset + local_id] = int(meta_id)
return result
def _noise_by_group(judge):
"""Return final global group id → was it promoted from the noise slot.
Mirrors :func:`_meta_by_group`: refined ids (issue #40) permute the
backend-native order, so the flag is mapped through ``old_of_local``.
"""
result = {}
for element, order in getattr(judge, 'group_order_per_el', {}).items():
fit = (getattr(judge, 'cluster_per_el', None) or {}).get(element)
flags = getattr(fit, 'noise_derived', None)
if flags is None:
continue
flags = np.asarray(flags)
offset = int(judge.group_id_offset[element])
for local_id, old_id in enumerate(order.old_of_local):
if 0 <= int(old_id) < len(flags):
result[offset + local_id] = bool(flags[int(old_id)])
return result
def _piecing_representatives(piece_map, lib, judge, focus_elements, rep_idx):
"""Build user-facing full-parent representatives from complete fit labels."""
if judge is None or not hasattr(judge, 'per_el_group_libs'):
raise RuntimeError(
'Piecing representative output requires final per-piece group labels')
return original_representatives(
piece_map,
lib,
selected_piece_idx=rep_idx,
per_el_group_libs=judge.per_el_group_libs,
species=focus_elements or [],
meta_by_group=_meta_by_group(judge),
noise_by_group=_noise_by_group(judge),
)
def _save_piecing_diagnostics(piece_map, lib, judge, focus_elements, rep_idx,
output_path_fn):
"""Write piece provenance and complete group labels onto original parents.
Callers pass a PieceMap and a Judge already checked by
``_piecing_representatives``.
"""
diagnosed = assign_back(
piece_map, lib, judge.per_el_group_libs,
focus_elements or [], rep_idx,
ids_are_global=getattr(judge, '_refined_ids_applied', False))
diag_xyz = output_path_fn('diagnosed.xyz')
ase_write(diag_xyz, diagnosed)
np.savez(
output_path_fn('diagnosed.npz'),
motif_id=np.array(
[atoms.get_array('motif_id') for atoms in diagnosed], dtype=object),
piece_id=np.array(
[atoms.get_array('piece_id') for atoms in diagnosed], dtype=object),
in_selection=np.array(
[atoms.get_array('in_selection') for atoms in diagnosed],
dtype=object),
)
logger.info(
f' Piecing: mapped group ids onto {len(diagnosed)} original '
f'structures → {diag_xyz}')
def _register_custom_backend(config):
"""Register ``cluster.custom_backend`` and rewrite the name, before validation.
Returns a shallow-copied config whose ``cluster.backend`` is the registered
name, so ``_validate_config_overrides`` sees a name it knows. The live
instance stays in the registry; clearing it from the normalized config
prevents ``fill_defaults`` from deep-copying configured backend state.
"""
_cl = (config or {}).get('cluster') or {}
custom = _cl.get('custom_backend')
if _cl.get('backend') != 'custom' and custom is None:
return config
if custom is None:
raise ValueError(
"cluster.backend='custom' but cluster.custom_backend is None. "
"Pass a ClusterBackend instance (Python API only — issue #69).")
backend = _cl.get('backend')
if backend not in (None, 'custom'):
raise ValueError(
f"cluster.backend={backend!r} cannot be combined with "
'cluster.custom_backend; use backend=\'custom\'.')
from atom_hifi.training.cluster_backends import register_backend
config = dict(config)
config['cluster'] = dict(
_cl, backend=register_backend(custom), custom_backend=None)
return config
[docs]
def run(config: dict) -> dict:
"""Execute the full L/R workflow described by ``config``.
Returns a small summary dict (metrics + output paths). REFIT modes return
early after re-fitting at the saved state.
"""
# Register a custom backend BEFORE fill_defaults: fill_defaults runs
# _validate_config_overrides, which resolves cluster.backend against
# BACKEND_REGISTRY. Registering afterwards made the documented
# backend='custom' call raise before it ever got here (issue #69 review).
config = _register_custom_backend(config)
cfg = fill_defaults(config)
# ── Unpack config into the flat names used by the pipeline ────────────────
_paths = cfg['paths']
LIB_PATH = _paths['lib_path']
PATIENT_PATH = _paths['patient_path']
FOCUS_ELEMENTS = _paths['focus_elements']
OUTPUT_DIR = _paths['output_dir']
_desc = cfg['descriptor']
DESCRIPTOR = _desc['kind']
DESCRIPTOR_FN = _desc.get('custom_fn') # Python-API only
SELECTION_METHOD = cfg['selection']['method']
SEED = int(cfg['seed'])
N_JOBS = cfg['n_jobs']
_scan = cfg['scan']
L_TOL = _scan['l_tol']
C_FACTOR_RANGE = _scan['c_factor_range']
N_COARSE = _scan['n_coarse']
N_FINE = _scan['n_fine']
MICROSTATE_REFERENCE = cfg['microstate']['reference']
K_MICRO_RATIO = None if MICROSTATE_REFERENCE == 'n2' else MICROSTATE_REFERENCE
SCAN_CACHE_FILE = _scan['cache_file']
N_JOBS_DESC = cfg['compute']['n_jobs_desc']
_n_jobs_eps = cfg['compute'].get('n_jobs_eps')
N_JOBS_EPS = N_JOBS if _n_jobs_eps is None else int(_n_jobs_eps)
_eps_cfg = cfg['eps_noise']
EPS_METHOD = _eps_cfg.get('method', 'thermal')
# r<NUM> + explicit force-resolution block (raw-config `eps_noise.force`,
# or `--f-tol` on the CLI — see config._validate_config_overrides) routes
# the reference against ε_F instead of ε_noise. r100 is excluded: its
# dispatch branch ignores eps_force entirely (hifi_workflow's r100 branch
# precedes the r<NUM> one), so the f_tol-style gates below must not fire
# for r100 + an unused force block (that combo used to run fine with
# precomputed/rdf ε — keep it byte-identical).
_rnum_force = (is_rnum(K_MICRO_RATIO) and K_MICRO_RATIO != 'r100'
and bool((cfg['microstate'] or {}).get('_force_active', False)))
EPS_TEMPERATURE = _eps_cfg['temperature']
EPS_NOISE_PER_SPECIES = (
(_eps_cfg.get('precomputed') or {}).get('per_species')
if EPS_METHOD == 'precomputed' else None
)
EPS_THERMAL_CFG = _eps_cfg.get('thermal') or {}
EPS_SIGMA_SOURCE = EPS_THERMAL_CFG.get('sigma_source', 'qho_table')
EPS_SIGMA_OVERRIDE = EPS_THERMAL_CFG.get('sigma_override') or {}
EPS_RDF_CFG = _eps_cfg.get('rdf') or {}
_loop = cfg['loop']
RUN_LOOP = _loop['run']
INITIAL_SAMPLE = _loop['initial_sample']
LOOP_BATCH_SIZE = _loop['batch_size']
LOOP_SKIP_FINE_SCAN = _loop['skip_fine_scan']
RESUME_FROM_ITER = _loop['resume_from_iter']
_grid = cfg['grid']
RUN_GRID_SCAN = _grid['run']
N_COARSE_GRID = _grid['n_coarse']
N_FINE_GRID = _grid['n_fine']
RESTORE_GRID = _grid['restore']
_nsga2 = cfg['nsga2']
RUN_NSGA2 = _nsga2['run']
N_NSGA2_GEN = _nsga2['n_gen']
N_NSGA2_POP = _nsga2['pop_size']
INNER_NSGA2_GEN = _nsga2['inner_n_gen']
INNER_NSGA2_POP = _nsga2['inner_pop_size']
RESTORE_NSGA2 = _nsga2['restore']
RUN_PATIENT_OOD = cfg['patient']['run_ood'] is True
_cluster = cfg.get('cluster') or {}
# Never a live backend object: _register_custom_backend registered any
# cluster.custom_backend under its name and cleared it before fill_defaults.
# HiFiWorkflow resolves a registered custom name to its executable spec,
# which it threads to process workers (issue #69).
CLUSTER_BACKEND = _cluster.get('backend', 'decaf')
MIN_CLUSTER_SIZE = _cluster.get('min_cluster_size', 2)
_piecing = cfg['piecing']
PIECING_ENABLED = _piecing['enabled']
PIECING_TARGET_ATOMS = _piecing['target_atoms']
_out = cfg['output']
DELTA_PICK = _out['delta_pick']
EXTRACT_SURROUNDINGS = _out['extract_surroundings']
SAVE_JUDGE = _out['save_judge']
_surr = cfg['surroundings']
SURR_KW = dict(output_mode=_surr['output_mode'],
r_cut=_surr['r_cut'],
n_per_group=_surr['n_per_group'])
_refit = cfg['refit']
REFIT_DELTA = _refit['delta']
REFIT_GRID_POINT = _refit['grid_point']
# ── Config validation ─────────────────────────────────────────────────────
if REFIT_GRID_POINT is not None and REFIT_DELTA is not None:
raise ValueError("Set only one of refit.grid_point or refit.delta, not both.")
if RUN_PATIENT_OOD and PATIENT_PATH is None:
raise ValueError("patient.run_ood=true requires paths.patient_path to be set.")
if RUN_LOOP and (REFIT_DELTA is not None or REFIT_GRID_POINT is not None):
logger.warning("loop.run=true with refit.* set — loop will run but refit "
"will then restore a saved checkpoint, discarding loop results.")
if not RUN_LOOP and RESUME_FROM_ITER is not None and (REFIT_DELTA is not None or REFIT_GRID_POINT is not None):
logger.warning("refit.* mode calls restore_state() and bypasses "
"loop.resume_from_iter corpus reconstruction — it has no effect.")
if RUN_GRID_SCAN and RUN_NSGA2:
logger.warning("Both grid.run and nsga2.run are true — "
"grid-scan results will be overwritten by NSGA-II.")
if PIECING_ENABLED and (RUN_LOOP or RESUME_FROM_ITER is not None):
raise ValueError(
"piecing.enabled=true is not supported with the active-learning loop "
"(loop.run) or loop.resume_from_iter — these modes build the corpus "
"incrementally from whole structures. Run piecing on a full-library "
"scan (loop.run=false, loop.resume_from_iter=null).")
if (K_MICRO_RATIO == 'f_tol' or _rnum_force) and EPS_METHOD != 'thermal':
raise ValueError(
"microstate.reference: f_tol or an r<NUM> force-resolution "
"reference requires eps_noise.method: thermal because ε_F is "
"interpolated from the cached thermal ε(σ) curve.")
os.makedirs(OUTPUT_DIR, exist_ok=True)
_p = lambda fname: os.path.join(OUTPUT_DIR, fname) # noqa: E731
# Scan-specific outputs go to a subdirectory when resume_from_iter is active so
# they never overwrite full-library run artifacts. Cache files (desc_lib.pkl,
# eps_noise_raw.npz, descp_iter*.npy) always live in the parent OUTPUT_DIR.
_SCAN_OUT = (os.path.join(OUTPUT_DIR, f'iter{RESUME_FROM_ITER:02d}')
if (not RUN_LOOP and RESUME_FROM_ITER is not None)
else OUTPUT_DIR)
os.makedirs(_SCAN_OUT, exist_ok=True)
_ps = lambda fname: os.path.join(_SCAN_OUT, fname) # noqa: E731
# ── Step 1: load library ──────────────────────────────────────────────────
_, lib, FOCUS_ELEMENTS, _focus_explicit = _load_library(
LIB_PATH, FOCUS_ELEMENTS)
# ── Step 2: build descriptor function ─────────────────────────────────────
logger.info('\n' + '=' * 70)
logger.info(f'Step 2: descriptor setup [descriptor.kind={DESCRIPTOR!r}]')
logger.info('=' * 70)
descriptor_fn, center_descriptor_fn, r_cut, _ = _build_descriptor_fn(
_desc, DESCRIPTOR_FN, lib)
logger.info(f' r_cut for ε_noise calibration: {r_cut} Å')
# ── Step 3: compute / load descriptors ────────────────────────────────────
logger.info('\n' + '=' * 70)
logger.info('Step 3: compute descriptors')
logger.info('=' * 70)
desc_list = _compute_descriptors(
_p('desc_lib.pkl'), lib, descriptor_fn, DESCRIPTOR, N_JOBS_DESC)
lib_focus, descriptor_array, struct_map, focus_desc_list = _focus_arrays(
lib, desc_list, FOCUS_ELEMENTS)
logger.info(f' Focus-atom descriptor: {descriptor_array.shape} '
f'({len(lib_focus)} structures)')
FOCUS_ELEMENTS = _validate_focus_species(
lib, lib_focus, FOCUS_ELEMENTS, _focus_explicit)
logger.info(f' Per-element focus species: {FOCUS_ELEMENTS}')
# The full per-structure descriptors are only needed to derive the focus
# arrays, except when voxel piecing later reuses whole-structure rows.
# Releasing them before calibration avoids overlap with its large rattle
# descriptor allocation.
if not PIECING_ENABLED:
del desc_list
# ── Step 4: patient structures (optional) ─────────────────────────────────
if REFIT_GRID_POINT is None:
patient_desc_array = _load_patient_descriptors(
PATIENT_PATH, RUN_PATIENT_OOD, _p('desc_patient.pkl'), descriptor_fn,
DESCRIPTOR, N_JOBS_DESC, FOCUS_ELEMENTS)
else:
# Grid-point refit uses only the restored fit and final corpus.
patient_desc_array = None
# ── Step 5: HiFiWorkflow initialisation ───────────────────────────────────
logger.info('\n' + '=' * 70)
logger.info('Step 5: HiFiWorkflow initialisation')
logger.info('=' * 70)
if EPS_METHOD == 'rdf' and not is_rnum(K_MICRO_RATIO):
logger.info(" eps_noise.method=rdf sets ε, but microstate.reference is not "
"an r… reference (r100/r<NUM>) → k_micro stays N//2 "
"(RDF ε still scales the scan bandwidth). Set "
"microstate.reference: r100 or keep the default N//2 "
"reference.")
_sigma_targets = {}
if EPS_METHOD == 'thermal':
_sigma_targets = resolve_sigma_thermal_targets(
FOCUS_ELEMENTS,
T=EPS_TEMPERATURE,
source=EPS_SIGMA_SOURCE,
config_override=EPS_SIGMA_OVERRIDE,
)
_eps_fingerprint = {'method': EPS_METHOD,
'n_samples': (EPS_RDF_CFG or {}).get('n_samples', 3000),
'eps_stat': (EPS_RDF_CFG or {}).get('eps_stat', 'median'),
'sigma_source': EPS_SIGMA_SOURCE,
'qho_table_version': (
qho_table_version() if EPS_SIGMA_SOURCE == 'qho_table' else None),
'sigma_targets': _sigma_targets}
descriptor_key = (_descriptor_cache_key(lib, descriptor_fn, DESCRIPTOR)
if EPS_METHOD != 'precomputed' else None)
calibration_key = None
if descriptor_key is not None:
calibration_inputs = {
'schema': 1, 'descriptor_sha256': descriptor_key,
'species': FOCUS_ELEMENTS, 'seed': SEED,
'initial_sample': int(min(INITIAL_SAMPLE, len(lib))),
'calibration': _eps_fingerprint,
}
if EPS_METHOD == 'rdf':
calibration_inputs['rdf'] = {**RDF_DEFAULTS, **EPS_RDF_CFG}
calibration_key = hashlib.sha256(json.dumps(
calibration_inputs, cls=ASEJSONEncoder,
sort_keys=True, separators=(',', ':')).encode()).hexdigest()
_eps_fingerprint['calibration_input_sha256'] = calibration_key
if EPS_METHOD == 'precomputed':
eps_noise_global, _eps_per_species_cache, _eps_fallback_cache, _eps_curve_cache = None, None, set(), None
else:
eps_noise_global, _eps_per_species_cache, _eps_fallback_cache, _eps_curve_cache = _load_eps_noise_cache(
_p('eps_noise_raw.npz'), FOCUS_ELEMENTS, fingerprint=_eps_fingerprint,
require_curve=(K_MICRO_RATIO == 'f_tol' or _rnum_force))
_c_lo = {el: v[0] for el, v in C_FACTOR_RANGE.items()} if C_FACTOR_RANGE else 0.1
_c_hi = {el: v[1] for el, v in C_FACTOR_RANGE.items()} if C_FACTOR_RANGE else 5.0
_eps_per_species = EPS_NOISE_PER_SPECIES or _eps_per_species_cache
wf = HiFiWorkflow(
species = FOCUS_ELEMENTS or [],
lib_path = LIB_PATH,
sigma_rattle = (0.001, 0.003, 0.01, 0.03, 0.05, 0.08, 0.1),
eps_r_cut = r_cut,
delta_lo_factor = _c_lo,
delta_hi_factor = _c_hi,
n_coarse = N_COARSE,
n_fine = N_FINE,
k_micro_ratio = K_MICRO_RATIO,
l_tol = L_TOL,
n_jobs = N_JOBS,
n_jobs_eps = N_JOBS_EPS,
seed = SEED,
eps_noise_global = None if _eps_per_species else eps_noise_global,
eps_noise_per_species = _eps_per_species,
eps_temperature = EPS_TEMPERATURE,
eps_noise_method = EPS_METHOD,
eps_rdf_cfg = EPS_RDF_CFG,
eps_sigma_source = EPS_SIGMA_SOURCE,
eps_sigma_override = EPS_SIGMA_OVERRIDE,
eps_force_f_tol = float((_eps_cfg.get('force') or {}).get('f_tol', 0.05)),
k_micro_force_active = _rnum_force,
cluster_backend = CLUSTER_BACKEND,
min_cluster_size = MIN_CLUSTER_SIZE,
output_dir = _SCAN_OUT,
save_judge = SAVE_JUDGE,
)
# Save the runner's input binding beside the existing calibration values and
# curves. Direct workflow callers without a verifiable identity remain valid,
# but their raw artifacts cannot be trusted for automatic runner reuse.
wf._eps_calibration_input_sha256 = calibration_key
# Restore RDF no-peak (N//2) elements when ε came from cache (calibration
# skipped). Skip when per_species is explicit — the cache is not the ε source.
if not EPS_NOISE_PER_SPECIES and _eps_per_species_cache and _eps_fallback_cache:
wf.eps_fallback_elements = set(_eps_fallback_cache)
if _eps_curve_cache:
wf.eps_curve_sigma = _eps_curve_cache['sigma_rattle']
wf.eps_curve_mean_per_el = _eps_curve_cache['mean_per_el']
# ── Stage 1: ε_noise calibration ──────────────────────────────────────────
_calibrate_eps_noise_if_needed(
wf, lib, descriptor_fn, center_descriptor_fn, lib_focus,
focus_desc_list, INITIAL_SAMPLE, SEED)
# All later stages consume precomputed descriptors. Drop the runner's local
# callback references now; a caller-provided custom callback may still be
# retained by its config or caller.
del descriptor_fn, center_descriptor_fn
# ── Scan cache pre-warm (optional) ────────────────────────────────────────
if SCAN_CACHE_FILE:
if RUN_LOOP:
logger.warning(
' scan.cache_file is set with loop.run=true; the pre-loaded cache '
'will be cleared at each loop iteration boundary (corpus changes → '
'covers stale). Cache pre-warm only benefits the final run_final call.')
logger.info(f' Warming scan cache from {SCAN_CACHE_FILE}')
wf.warmup_scan_cache(SCAN_CACHE_FILE)
# ── Stage 2: active-learning loop (optional) ──────────────────────────────
_lib_focus_final, _desc_final, _full_final, _struct_map_final = (
_build_final_corpus(
wf, RUN_LOOP, RESUME_FROM_ITER, INITIAL_SAMPLE, LOOP_BATCH_SIZE,
LOOP_SKIP_FINE_SCAN, lib, lib_focus, focus_desc_list,
descriptor_array, struct_map, _p, SEED,
run_nsga2=RUN_NSGA2,
inner_nsga2_n_gen=INNER_NSGA2_GEN,
inner_nsga2_pop_size=INNER_NSGA2_POP))
if RUN_LOOP and hasattr(wf, '_loop_final_desc_list'):
# _build_final_corpus stacked this list into _desc_final.
del wf._loop_final_desc_list
if RUN_LOOP or RESUME_FROM_ITER is not None:
# Loop/resume creates an independent final corpus. Piecing is forbidden
# for both modes, so nothing below needs the original full library.
del lib
# Final-corpus aliases above now own every downstream input. The original
# per-structure focus rows otherwise duplicate the stacked descriptor array.
del lib_focus, focus_desc_list, descriptor_array, struct_map
# ── Piecing: substitute pieces for structures (container boundary) ─────────
# Cut large structures into disjoint voxel pieces and rerun the focus-atom
# reduction on them. ε_noise was already calibrated on the original
# structures above; everything downstream sees pieces as ordinary small
# structures. Deterministic (no RNG), so it is safe under refit re-runs.
piece_map = None
if PIECING_ENABLED:
(piece_map, _full_final, _lib_focus_final, _desc_final,
_struct_map_final) = _apply_piecing(
PIECING_TARGET_ATOMS, _full_final, desc_list, FOCUS_ELEMENTS)
# Piecing is the only post-calibration consumer of full rows.
del desc_list
logger.info(f'\nFinal training corpus: {len(_lib_focus_final)} structures, '
f'{_desc_final.shape[0]} focus atoms')
# ── REFIT mode (skip L/R scan) ────────────────────────────────────────────
if REFIT_GRID_POINT is not None:
logger.info('\n' + '=' * 70)
logger.info(f'REFIT_GRID MODE: loading saved state and fitting at {REFIT_GRID_POINT}')
logger.info('=' * 70)
wf.restore_state(
npz_path = _ps('hifi_final.npz'),
lib_focus = _lib_focus_final,
desc_array = _desc_final,
struct_map = _struct_map_final,
full_struct_lib = _full_final,
)
# Refit skips every scan/optimizer consumer of a prewarmed fit cache.
wf._el_fit_cache.clear()
# This refit mode never consumes patient descriptors.
del patient_desc_array
fit_kwargs = {
'delta_pick': REFIT_GRID_POINT,
'selection_method': 'mu_tiebreak',
}
if piece_map is not None:
fit_kwargs['save_representatives'] = False
rep_structs_grid, rep_idx_grid = wf.plot_final_fit(**fit_kwargs)
judge = getattr(wf, '_last_final_judge', None)
if piece_map is not None:
rep_structs_grid = _piecing_representatives(
piece_map, lib, judge, FOCUS_ELEMENTS, rep_idx_grid)
out_xyz_grid = _ps('representatives_grid.xyz')
ase_write(out_xyz_grid, rep_structs_grid)
logger.info(f'\n Saved {len(rep_structs_grid)} representatives → {out_xyz_grid}')
if piece_map is not None:
_save_piecing_diagnostics(
piece_map, lib, judge, FOCUS_ELEMENTS, rep_idx_grid, _ps)
wf.plot_learning_curve()
if EXTRACT_SURROUNDINGS:
logger.info('\n' + '=' * 70)
logger.info('Stage 6: extract_rep_surroundings')
logger.info('=' * 70)
surroundings = wf.extract_rep_surroundings(save_dir=OUTPUT_DIR, **SURR_KW)
for el, groups in surroundings.items():
logger.info(f' {el}: {len(groups)} groups → {OUTPUT_DIR}/surroundings_{el}.xyz')
logger.info('\n' + '=' * 70)
logger.info(f'Done (refit grid). All outputs in: {OUTPUT_DIR}')
logger.info('=' * 70)
return {'mode': 'refit_grid', 'output_dir': OUTPUT_DIR,
'representatives': out_xyz_grid, 'n_representatives': len(rep_structs_grid)}
if REFIT_DELTA is not None:
logger.info('\n' + '=' * 70)
logger.info(f'REFIT MODE: loading saved state and fitting at c={REFIT_DELTA}')
logger.info('=' * 70)
wf.restore_state(
npz_path = _ps('hifi_final.npz'),
lib_focus = _lib_focus_final,
desc_array = _desc_final,
struct_map = _struct_map_final,
full_struct_lib = _full_final,
)
# Refit skips every scan/optimizer consumer of a prewarmed fit cache.
wf._el_fit_cache.clear()
fit_kwargs = {
'delta_pick': REFIT_DELTA,
'patient_descriptors': patient_desc_array,
'selection_method': SELECTION_METHOD,
}
if piece_map is not None:
fit_kwargs['save_representatives'] = False
rep_structs_r, rep_idx_r = wf.plot_final_fit(**fit_kwargs)
# plot_final_fit is this refit path's last patient-descriptor consumer.
del patient_desc_array
# ── Post-selection metrics ────────────────────────────────────────────
logger.debug('')
_m_r = _post_selection_metrics(wf, rep_idx_r, _struct_map_final, SELECTION_METHOD)
L_r, R_r = _m_r['L'], _m_r['R']
# ── Save representatives ───────────────────────────────────────────────
judge = getattr(wf, '_last_final_judge', None)
if piece_map is not None:
rep_structs_r = _piecing_representatives(
piece_map, lib, judge, FOCUS_ELEMENTS, rep_idx_r)
out_xyz_r = _ps(f'representatives_delta{REFIT_DELTA:.5f}.xyz')
ase_write(out_xyz_r, rep_structs_r)
logger.info(
f'\n Saved {len(rep_structs_r)} representative structures → {out_xyz_r}')
if piece_map is not None:
_save_piecing_diagnostics(
piece_map, lib, judge, FOCUS_ELEMENTS, rep_idx_r, _ps)
# ── Learning curve (no-op if no AL loop records) ───────────────────────
wf.plot_learning_curve()
logger.info('\n' + '=' * 70)
logger.info(f'Done (refit). All outputs in: {OUTPUT_DIR}')
logger.info('=' * 70)
return {'mode': 'refit', 'output_dir': OUTPUT_DIR,
'representatives': out_xyz_r,
'n_representatives': len(rep_structs_r),
'L': L_r, 'R': R_r, 'F': _m_r['F']}
# ── Stage 3: final L/R scan ───────────────────────────────────────────────
_run_final_scan(
wf, _lib_focus_final, _desc_final, _full_final, _struct_map_final, L_TOL)
# ── Stage 3b/3c: optional per-element c-factor optimizers ─────────────────
best_c_per_el = _run_optional_cfactor_optimizations(
wf, RUN_GRID_SCAN, N_COARSE_GRID, N_FINE_GRID, RESTORE_GRID,
RUN_NSGA2, N_NSGA2_GEN, N_NSGA2_POP, RESTORE_NSGA2,
FOCUS_ELEMENTS, L_TOL, SEED, OUTPUT_DIR, _p)
# Per-element fitted-group cache only supports the optional optimizers.
# Stage 4 onward refits its own final Judge, so release cached DECAF state.
wf._el_fit_cache.clear()
# ── Stage 4: patient OOD (optional) ───────────────────────────────────────
if patient_desc_array is not None:
logger.info('\n' + '=' * 70)
logger.info('Stage 4: patient OOD per fine δ')
logger.info('=' * 70)
wf.compute_patient_ood(patient_desc_array)
# ── Stage 5: final fit at δ* ──────────────────────────────────────────────
select_kwargs = (
{'save_representatives': False} if piece_map is not None else {})
rep_structs, rep_idx, _m, out_xyz = _select_final_representatives(
wf, best_c_per_el, DELTA_PICK, patient_desc_array, SELECTION_METHOD,
_struct_map_final, FOCUS_ELEMENTS, _full_final, _ps, **select_kwargs)
judge = getattr(wf, '_last_final_judge', None)
if piece_map is not None:
rep_structs = _piecing_representatives(
piece_map, lib, judge, FOCUS_ELEMENTS, rep_idx)
ase_write(out_xyz, rep_structs)
logger.info(
f' Piecing: saved {len(rep_structs)} unique original '
f'representatives → {out_xyz}')
L, R = _m['L'], _m['R']
# The final fit is the last consumer; do not retain patient descriptor rows
# through piecing, reporting, or optional surroundings extraction.
del patient_desc_array
# ── Piecing: assign group ids back onto the ORIGINAL structures ────────────
if piece_map is not None:
_save_piecing_diagnostics(
piece_map, lib, judge, FOCUS_ELEMENTS, rep_idx, _ps)
# ── Learning curve ─────────────────────────────────────────────────────────
if RUN_LOOP:
wf.plot_learning_curve()
logger.info('\n' + '=' * 70)
logger.info(f'Done. All outputs in: {OUTPUT_DIR}')
logger.info('=' * 70)
# ── Stage 6: representative environment extraction (optional) ──────────────
if EXTRACT_SURROUNDINGS and piece_map is not None:
logger.warning('\nStage 6 (extract_surroundings) skipped: not supported with '
'piecing — surroundings would be carved from voxel pieces '
'(truncated neighbourhoods). Use diagnosed.xyz instead.')
elif EXTRACT_SURROUNDINGS:
logger.info('\n' + '=' * 70)
logger.info('Stage 6: representative environment extraction')
logger.info('=' * 70)
surroundings = wf.extract_rep_surroundings(save_dir=OUTPUT_DIR, **SURR_KW)
for focus_el, groups in surroundings.items():
n_groups = len(groups)
n_clusters = sum(len(v) for v in groups.values())
logger.info(f' {focus_el}: {n_groups} groups, {n_clusters} surroundings'
f' → {OUTPUT_DIR}/surroundings_{focus_el}.xyz')
# Per-element density-supported group count of the final fit at the
# operating point (any backend): distinct local atomic states per element,
# excluding promoted rare-environment groups (issue #63).
n_groups_per_el = {}
n_promoted_per_el = {}
if judge is not None and getattr(judge, 'cluster_per_el', None):
for el in (FOCUS_ELEMENTS or []):
cl = judge.cluster_per_el.get(el)
if cl is not None:
n_total = int(cl.n_groups)
nd = getattr(cl, 'noise_derived', None)
n_noise = 0 if nd is None else int(np.asarray(nd).sum())
n_groups_per_el[el] = n_total - n_noise # density-supported
n_promoted_per_el[el] = n_noise
if n_noise:
logger.info(
f' {el}: {n_total - n_noise} density groups '
f'+ {n_noise} promoted rare-environment groups')
# Degeneracy warning: an element resolving to <=1 DENSITY-supported group at
# c* means its clustering collapsed — every environment of that element is
# treated as identical, so the representative set captures no diversity for
# it. This is deliberately LOUD: the run still writes a complete, healthy-
# looking result (Pareto front with a selected row, representatives,
# group_id_mapping, surroundings), and nothing in those artifacts says they
# are degenerate. Observed on PtOH_big, where |S|=62 came entirely from the
# noise tax while O had a single density group at every c in [0.1, 5.0].
_collapsed = collapsed_elements(n_groups_per_el)
if _collapsed:
_bar = '!' * 78
logger.warning('\n%s', _bar)
logger.warning(
'DEGENERATE CLUSTERING at c*: %s resolved to <=1 density-supported '
'group.', ', '.join(_collapsed))
logger.warning(
'The outputs below are still written and will LOOK healthy — a '
'Pareto front with a selected row, representatives, '
'group_id_mapping, surroundings — but they capture no environment '
'diversity for these elements.')
for el in _collapsed:
n_p = n_promoted_per_el[el]
if n_p:
logger.warning(
' %s: 1 density group + %d promoted rare-environment '
'group(s); |S| for %s is driven by the noise tax, not by '
'cover.', el, n_p, el)
logger.warning(
'Usually c* is too high — lower it, or under piecing reduce '
'target_atoms (issue #18). If NO c in the scan resolves >1 density '
'group, the backend/descriptor cannot resolve this system at all: '
'see docs/handoffs/2026-06-24-kmeans-backend/HANDOFF.md §H3 '
'(hdbscan on a continuous density manifold) and switch backend.')
logger.warning('%s\n', _bar)
return {'mode': 'full', 'output_dir': OUTPUT_DIR,
'representatives': out_xyz, 'n_representatives': len(rep_structs),
'fine_scan': _ps('fine_scan.out'),
'delta_star_L095': wf.delta_star_L095, 'delta_star_F': wf.delta_star_F,
'L': L, 'R': R, 'F': _m['F'], 'n_groups_per_el': n_groups_per_el}