"""
atom_hifi/training/hifi_workflow.py — L/R bandwidth-selection workflow and parallel workers.
Public API
----------
HiFiWorkflow
End-to-end L/R bandwidth selection and visualisation. Stages:
1. calibrate_eps_noise — ε_noise calibration (optional; distance-kind runs
need ε from here or from eps_noise_per_species).
2. run_loop — active-learning loop (optional).
3. run_final — final coarse + fine L/R scan (mandatory). Optional
per-element optimisers on its state: run_grid_scan / scan_cfactor_grid
(distance-kind backends only) and run_nsga2.
4. compute_patient_ood — patient OOD per fine scan value (optional).
5. plot_final_fit — final fit and representative set (mandatory).
extract_rep_surroundings optionally extracts per-group environments at the
chosen operating point.
Private helpers
~~~~~~~~~~~~~~~
_fit_cfactor_worker
The only process-pool worker defined in this module: one scan point per
call, called in-process when n_jobs == 1 and sent to loky processes
otherwise. It must stay module-level so joblib pickles it by reference;
optimizer imports it by name. _fit_cfactor_worker_body is its in-process
body.
_ood_worker, _fit_decaf_single_el
Dispatched with prefer='threads' (patient-OOD refits; grid Phase-1 fits
via a closure).
_fit_judge_at_delta
Shared helper that fits a fresh Judge.
"""
# Standard library
import contextlib
import json
import logging
from pathlib import Path
# Scientific stack
import numpy as np
import pandas as pd
from scipy.stats import entropy as _scipy_entropy
# ASE
from ase.io import write
# Descriptors / clustering
import decaf
# Internal
from atom_hifi.config import is_rnum
from atom_hifi.training.judge import (
Judge,
measure_eps_noise,
fit_microstate_reference,
atom_aware_greedy,
)
from atom_hifi.training.cluster_base import ClusterBackend
from atom_hifi.training.cluster_backends import (
BACKEND_REGISTRY,
BUILTIN_BACKENDS,
PCAEmbedder,
fit_cluster_backend,
register_backend,
)
from atom_hifi.training.qho_sigma import (
qho_table_version,
resolve_sigma_thermal_targets,
)
from atom_hifi.training.force_truth import force_eps_from_curve
from atom_hifi.training.rdf_eps import RDF_DEFAULTS, measure_eps_noise_rdf
from atom_hifi.training.plotting import _PlottingMixin
from atom_hifi.training.state_io import _StateIOMixin
logger = logging.getLogger(__name__)
# ── Parallel workers ──────────────────────────────────────────────────────────
def _strip_lib_for_workers(lib):
"""Return symbol-only Atoms stubs for IPC-efficient parallel dispatch.
Worker-side consumers (Judge.fit_decaf, _fit_decaf_single_el, the scan
row's atom count) read only len(atoms) and per-atom symbols —
positions/cell/pbc are never accessed. Stripping to symbols reduces
pickle size ~15× (3 KB → 200 bytes per structure), cutting joblib IPC
from ~8 GB → ~500 MB for a 400-point 2-D scan on a 6542-structure library.
"""
from ase import Atoms as _A
return [_A(a.get_chemical_symbols()) for a in lib]
def _build_calibration_embed_fns(species, lib_focus, lib_path, desc_array, seed):
"""Fit element-local PCA, then retain only its compact embedding operator.
The temporary Judge/ClusterFit also owns raw descriptors, embeddings, and
DECAF clustering state. None of that is needed once ``embed_op`` has been
copied: ε-noise calibration applies only ``rows @ embed_op``. Keeping the
bound ``ClusterFit.embed`` method would pin the full temporary fit while the
much larger rattle-result buffer is active.
"""
el_descs = {el: [] for el in species}
el_libs = {el: [] for el in species}
offset = 0
for atoms in lib_focus:
n_atoms = len(atoms)
desc = desc_array[offset:offset + n_atoms]
symbols = np.array(atoms.get_chemical_symbols())
for el in species:
mask = symbols == el
if mask.any():
el_descs[el].append(desc[mask])
el_libs[el].append(atoms[np.where(mask)[0]])
offset += n_atoms
embed_fns = {}
for el in species:
if not el_descs[el]:
continue
desc_el = np.vstack(el_descs[el])
judge = _fit_judge_at_delta(
[el], el_libs[el], lib_path, desc_el, delta=0,
restart_seed=seed,
)
cluster = judge.cluster_per_el.get(el)
if cluster is None:
continue
embed_op = np.array(cluster.embedder.embed_op, copy=True)
def embed(rows, op=embed_op):
return np.asarray(rows) @ op
embed_fns[el] = embed
# δ_auto is a HaarWavelet reference only. The operational bandwidth is
# c·ε_noise; the temporary clustering result has no downstream use.
logger.info(
' %s: embed_op built (δ_auto reference=%.5f)',
el, cluster.granularity,
)
return embed_fns
def _to_shared_arrays(per_el_dict):
"""Copy dict[el -> ndarray] into shared memory; return (shm_list, meta_dict).
The caller must close + unlink every SharedMemory in shm_list after all
workers have finished. Workers only attach (create=False) and close, never
unlink.
"""
from multiprocessing.shared_memory import SharedMemory
shm_list, shm_meta = [], {}
try:
for el, arr in per_el_dict.items():
arr = np.ascontiguousarray(arr)
shm = SharedMemory(create=True, size=max(arr.nbytes, 1))
# Register before copying: a copy failure must not leak this block.
shm_list.append(shm)
np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf)[:] = arr
shm_meta[el] = (shm.name, arr.shape, arr.dtype.str)
except Exception:
for shm in shm_list:
shm.close()
shm.unlink()
raise
return shm_list, shm_meta
def _from_shared_arrays(shm_meta):
"""Attach to existing shared memory blocks; return (read-only views, handles).
Workers call shm.close() on each handle when done — never shm.unlink();
the main process owns the lifecycle. On POSIX, attaching also registers
the name with multiprocessing's resource tracker (always on CPython ≤3.12;
on ≥3.13 because track defaults to True and is not passed here). joblib's
loky workers share the parent's tracker, and the parent's unlink()
unregisters the name, so no worker-side unregister is needed.
On partial failure, already-opened handles are closed before re-raising.
"""
from multiprocessing.shared_memory import SharedMemory
per_el, handles = {}, []
try:
for el, (name, shape, dtype_str) in shm_meta.items():
shm = SharedMemory(name=name, create=False)
handles.append(shm)
arr = np.ndarray(shape, dtype=np.dtype(dtype_str), buffer=shm.buf)
arr.flags.writeable = False
per_el[el] = arr
except Exception:
for shm in handles:
shm.close()
raise
return per_el, handles
@contextlib.contextmanager
def _shm_scope(desc_per_el, species, pre_embedded_per_el=None):
"""Precompute the PCA embedding; serve it from SharedMemory.
Yields (shm_meta_desc, shm_meta_emb) — opaque dicts passed to
_fit_cfactor_worker as _shm_meta_* kwargs. Guarantees cleanup (close+unlink)
on both normal exit and exception.
pre_embedded_per_el optionally supplies the iteration's already-fitted
embedding, avoiding another PCA pass before the SHM copy.
Was a 3-tuple until issue #69: the third slot held an O(N²) pairwise-distance
block, retired in #68 once decaf's fit-path self-classification (its only
consumer, and the identity) was short-circuited. #68 left the empty slot in
place to bound its own diff; nothing ever read it.
"""
t_emb_per_el = {}
for el in species:
pre_embedded = (None if pre_embedded_per_el is None
else pre_embedded_per_el.get(el))
if pre_embedded is None:
_tmp = decaf.embed_cluster(desc_per_el[el], embed_str='PCA', cluster_str='MSC')
_tmp.get_embed(dim=0)
pre_embedded = np.asarray(_tmp.t_embedded)
del _tmp
# Serve the embedding at the dtype the serial path actually clusters on.
# The historical float32 cast here WAS the bug: workers received a
# truncated copy while in-process fits kept decaf's own precision, so
# borderline KMeans assignments diverged with n_jobs at a fixed seed.
#
# Read the dtype off the embedding, not off the descriptors.
_t_emb = np.asarray(pre_embedded)
t_emb_per_el[el] = np.array(_t_emb, dtype=_t_emb.dtype)
# Append each completed allocation immediately: if a later allocation
# fails, the outer finally still closes and unlinks the earlier blocks.
shm_all = []
try:
shm_list_desc, shm_meta_desc = _to_shared_arrays(desc_per_el)
shm_all.extend(shm_list_desc)
# The descriptor SHM now owns its copy. This generator remains alive
# through the scan, so keeping its argument would otherwise pin the
# source N×D arrays for the entire parallel section.
desc_per_el = None
shm_list_emb, shm_meta_emb = _to_shared_arrays(t_emb_per_el)
shm_all.extend(shm_list_emb)
del t_emb_per_el
yield shm_meta_desc, shm_meta_emb
finally:
for shm in shm_all:
shm.close()
shm.unlink()
def _fit_decaf_single_el(el, desc_el, delta_el, lib_len, lib_stub,
t_emb_el=None,
cluster_backend='decaf', min_cluster_size=2,
eps_noise_el=None, seed=42):
"""Run one backend fit for one element at one granularity. No cover, no L/R.
Used by ``scan_cfactor_grid`` for on-demand Phase 1 fits of c values not
already in ``_el_fit_cache``.
``fit_cluster_backend`` has already deduped the trailing noise slot and
promoted each survivor to a real group (issue #63), so ``pegl``/``all_gid``
contain rare-environment groups alongside density ones and no atom is
excluded. ``eps_noise_el`` is the dedup radius and is required whenever the
backend produces a trailing slot.
Returns
-------
pegl : list[list[int]]
Per-structure group ID lists, shape (lib_len, variable).
all_gid : set[int]
All group IDs encountered across structures.
granularity : float
``float(delta_el)``, the granularity this fit used. Stored as the
cache entry's third element; scan workers reuse the cover only when it
serves the point (equal δ/k, or, for distance kind, a request for the
cover's own cache key; #99).
"""
fit = fit_cluster_backend(cluster_backend, desc_el, delta_el,
t_embedded=t_emb_el,
min_cluster_size=min_cluster_size,
eps_noise=eps_noise_el, seed=seed)
atomic_tags = fit.labels
# Per-structure atom counts
el_sizes = np.array([sum(1 for a in s if a.symbol == el) for s in lib_stub])
# Build per-structure lists. Noise reps were already deduped and promoted
# to real groups inside fit_cluster_backend (issue #63), so this path has no
# dedup of its own and no forced_idx side-channel to keep in sync.
pegl = [[] for _ in range(lib_len)]
all_gid = set()
N = 0
for i_full, n_el in enumerate(el_sizes):
gids = [g for g in atomic_tags[N:N + n_el].tolist() if g >= 0]
pegl[i_full] = gids
all_gid.update(gids)
N += n_el
return pegl, all_gid, float(delta_el)
def _make_el_fit_cache_entry(judge, el, lib_len):
"""Build the cover pair; the caller attaches its resolved granularity.
Promoted rare-environment groups are ordinary members of ``pegl``
(issue #63), so cover semantics survive the cache without the retired
``forced_idx`` third element.
"""
pegl = [
[int(gid) for gid in gids if int(gid) >= 0]
for gids in judge.per_el_group_libs[el]
]
all_gid = set(gid for gids in pegl for gid in gids)
return pegl, all_gid
def _c_cache_slice(el_fit_cache, c_factor, species):
"""Return only the cache entries needed for this single c_factor value.
Avoids shipping the full _el_fit_cache (n_c_points × N_atoms) to every
worker; each worker receives only its own c-point slice (N_atoms entries).
Cache keys are (el, rounded_c) tuples, same structure as _el_fit_cache.
"""
if not el_fit_cache:
return None
if isinstance(c_factor, dict):
c_vals = {el: float(c_factor[el]) for el in species if el in c_factor}
else:
c_vals = {el: float(c_factor) for el in species}
result = {(el, round(c, 10)): el_fit_cache[(el, round(c, 10))]
for el, c in c_vals.items()
if (el, round(c, 10)) in el_fit_cache}
return result if result else None
def _fit_cfactor_worker(c_factor, species, lib, lib_path,
desc_per_el, struct_map_per_el,
mu_per_el, counts_L_per_el, k_occ_L_per_el, H_L_per_el,
eps_noise_per_el, embed_op_per_el,
pre_embedded_per_el=None,
max_training_size=None,
seed=42,
worker_subprocess: bool = False,
_shm_meta_desc=None,
_shm_meta_emb=None,
cluster_cache_per_el=None,
cluster_backend='decaf',
min_cluster_size=2):
"""Per-element clustering-backend scan worker for one scan-axis value.
Fits one multi-element Judge at this point's granularity (distance-kind:
δ = c·ε_noise; count-kind: k = max(2, round(f·k_micro)) capped at
max(1, n_el−1)), with one backend fit per element unless the caller
supplies a cached cover whose stored granularity serves this point (equal
δ/k or, for distance kind, a request for the cover's own cache key). Then
runs a joint set-cover over (el, gid) pairs and computes joint L/R.
Returns a row dict or None on failure.
``k_occ_L_per_el``, ``H_L_per_el`` and ``embed_op_per_el`` are accepted for
caller compatibility and unused (#103).
"""
import os as _os
# Only override thread counts when running as a subprocess (n_jobs > 1).
# In the main process (n_jobs=1), respect the caller's env vars.
if worker_subprocess:
_os.environ['OPENBLAS_NUM_THREADS'] = '1'
_os.environ['OMP_NUM_THREADS'] = '1'
try:
from threadpoolctl import threadpool_limits
threadpool_limits(limits=1)
except ImportError:
pass
# Attach shared-memory descriptor arrays when dispatched by _scan_cfactor
# with n_jobs > 1. Falls back to direct arrays when metadata is None.
_shm_handles = []
_shm_t_emb = None
if _shm_meta_desc is not None:
desc_per_el, _h = _from_shared_arrays(_shm_meta_desc)
_shm_handles.extend(_h)
if _shm_meta_emb is not None:
_shm_t_emb, _h = _from_shared_arrays(_shm_meta_emb)
_shm_handles.extend(_h)
try:
return _fit_cfactor_worker_body(
c_factor, species, lib, lib_path,
desc_per_el, struct_map_per_el,
mu_per_el, counts_L_per_el,
eps_noise_per_el, seed,
pre_embedded_per_el, max_training_size,
_shm_t_emb,
cluster_cache_per_el=cluster_cache_per_el,
cluster_backend=cluster_backend,
min_cluster_size=min_cluster_size,
)
finally:
for _shm in _shm_handles:
_shm.close()
def _count_kind_k_from_fraction(f, k_micro, *, floor=2, cap=None) -> int:
"""Realize a count-kind scan-axis fraction into an integer cluster count.
k = max(floor, round(f*k_micro)), then capped to `cap` if given. Without
a cap, this is the nominal k — KMeansBackend.fit separately clips to
n_el-1 regardless, so an uncapped caller's actual fit is unaffected, but
any value the caller reports/caches as "the k used" will be wrong past
that ceiling unless it also caps (issue #82 codex-review finding: the
cache key and the reported delta_{el} column must reflect the same k the
backend will actually realize, not the pre-clip nominal value).
"""
k = max(floor, int(round(float(f) * k_micro)))
if cap is not None:
k = min(k, int(cap))
return k
def _fit_cfactor_worker_body(c_factor, species, lib, lib_path,
desc_per_el, struct_map_per_el,
mu_per_el, counts_L_per_el,
eps_noise_per_el, seed,
pre_embedded_per_el, max_training_size,
shm_t_emb,
cluster_cache_per_el=None,
cluster_backend='decaf',
min_cluster_size=2):
"""Inner body of _fit_cfactor_worker, extracted for clean try/finally."""
# Subsample library if larger than max_training_size
n_lib = len(lib)
if max_training_size and n_lib > max_training_size:
rng = np.random.default_rng(seed=seed)
keep_idx = np.sort(rng.choice(n_lib, max_training_size, replace=False))
lib = [lib[i] for i in keep_idx]
desc_per_el = {el: desc_per_el[el][np.isin(struct_map_per_el[el], keep_idx)]
for el in species}
# Remap struct indices to 0..max_training_size-1
idx_map = {old: new for new, old in enumerate(keep_idx)}
struct_map_per_el = {
el: np.array([idx_map[s] for s in struct_map_per_el[el]
if s in idx_map])
for el in species
}
# Build per-element delta dict and c_vals_dict (needed for cache key lookup)
if isinstance(c_factor, dict):
c_vals_dict = {el: float(c_factor[el]) for el in species}
else:
c_vals_dict = {el: float(c_factor) for el in species}
backend_caps = (BACKEND_REGISTRY[cluster_backend]
if isinstance(cluster_backend, str)
else cluster_backend)
if backend_caps.granularity_kind == 'count':
# scan axis is the k-fraction f: k_el = round(f · k_micro_el), capped
# at n_el-1 to match KMeansBackend.fit's own clip — the reported/
# cached delta_{el} must reflect the k actually realized, not the
# pre-clip nominal value (issue #82 codex-review finding).
delta_per_el = {
el: float(_count_kind_k_from_fraction(
c_vals_dict[el], len(counts_L_per_el[el]),
cap=max(1, len(desc_per_el[el]) - 1)))
for el in species
}
else:
delta_per_el = {el: float(c_vals_dict[el] * eps_noise_per_el[el])
for el in species}
# A rounded scan-axis key alone cannot prove the fitted bandwidth/count.
# Reuse needs the cover's stored granularity: equal to this point's δ/k,
# or (distance kind) a request for the cover's own cache key, as on
# scan_cfactor_grid's default and run_grid_scan's coarse axes, when
# δ_stored/ε rounds to that key (|Δc| ≤ 5e-11; the row reports δ = key·ε).
# Covers without this evidence are recomputed, never rescored as a
# different clustering. Input identity is checked before worker dispatch.
def _serves_request(key, entry):
if len(entry) < 3 or not np.isscalar(entry[2]):
return False
el, c_key = key
fitted = float(entry[2])
if fitted == delta_per_el.get(el):
return True
eps_el = (eps_noise_per_el or {}).get(el)
return (backend_caps.granularity_kind != 'count' and bool(eps_el)
and c_vals_dict.get(el) == c_key
and round(fitted / eps_el, 10) == c_key)
cluster_cache_per_el = {
key: entry for key, entry in (cluster_cache_per_el or {}).items()
if _serves_request(key, entry)
}
t_emb_per_el = (shm_t_emb if shm_t_emb is not None
else (pre_embedded_per_el or {}))
if max_training_size and len(lib) < n_lib:
t_emb_per_el = {}
try:
j = _fit_judge_at_delta(
species, lib, lib_path,
desc_array = desc_per_el,
delta = delta_per_el,
mu_per_el = mu_per_el,
counts_L_per_el = counts_L_per_el,
t_embedded = t_emb_per_el,
cluster_cache_per_el = cluster_cache_per_el,
c_vals_dict = c_vals_dict,
restart_seed = seed,
cluster_backend = cluster_backend,
min_cluster_size = min_cluster_size,
eps_noise_per_el = eps_noise_per_el,
)
except Exception as exc:
_c_str = (str({k: f'{v:.4f}' for k, v in c_factor.items()})
if isinstance(c_factor, dict) else f'{c_factor:.4f}')
logger.info(f' c={_c_str} skipped: {exc}')
return None
n_structs_S = len(j.sub_lib_idx)
if isinstance(c_factor, dict):
row = {f'c_{el}': float(v) for el, v in c_factor.items()}
row['n_structs_S'] = n_structs_S
else:
row = dict(c_factor=float(c_factor), n_structs_S=n_structs_S)
for el in species:
row[f'delta_{el}'] = delta_per_el[el]
row['n_atoms_S'] = sum(len(lib[i]) for i in j.sub_lib_idx)
row['L'] = float(j.L) if j.L is not None else float('nan')
row['R'] = float(j.R) if j.R is not None else float('nan')
row['F'] = float(j.F) if j.F is not None else float('nan')
row['N_groups'] = j.n_groups
row['sub_lib_idx'] = j.sub_lib_idx
# Collect freshly-fitted per-el results for _el_fit_cache population in
# _scan_cfactor. Cache entries contain all coverable groups, including
# promoted rare-environment groups.
row['_el_fits'] = {
el: (*_make_el_fit_cache_entry(j, el, len(lib)), delta_per_el[el])
for el in species
if j.cluster_per_el.get(el) is not None # skip cache-hit elements
}
_c_str = (str({el: f'{v:.4f}' for el, v in c_factor.items()})
if isinstance(c_factor, dict) else f'{c_factor:.4f}')
logger.debug(f' c={_c_str} L={row["L"]:.4f} R={row["R"]:.4f} '
f'F={row["F"]:.4f} |S|={n_structs_S}')
return row
def _fit_judge_at_delta(species, lib_focus, lib_path, desc_array, delta,
mu_per_el=None, counts_L_per_el=None,
t_embedded=None,
cluster_cache_per_el=None, c_vals_dict=None,
use_mu_tiebreak=None, restart_seed=None,
cluster_backend='decaf', min_cluster_size=2,
eps_noise_per_el=None):
"""Fit a fresh Judge at the given bandwidth (no file output).
Parameters
----------
species : list[str]
Focus elements.
lib_focus : list[Atoms]
Full training library (all species). Judge filters per element internally.
lib_path : str
Path to the library file (metadata only).
desc_array : ndarray or dict[str, ndarray]
Bare ndarray → single-element mode.
``{el: ndarray}`` dict → per-element mode.
delta : float or dict[str, float]
DECAF bandwidth. Scalar for single-element, dict for per-element.
mu_per_el : dict[str, ndarray], optional
Microstate labels per element; enables L/R computation on the Judge.
counts_L_per_el : dict[str, ndarray], optional
Precomputed microstate distribution; required alongside ``mu_per_el``.
t_embedded : ndarray or dict[str, ndarray], optional
Pre-computed PCA embedding for SharedMemory fast-path.
cluster_cache_per_el : dict, optional
Per-element DECAF cache keyed by ``(el, round(c_val, 10))``.
c_vals_dict : dict[str, float], optional
Raw c-factor values per element (before eps_noise scaling), used as cache keys.
use_mu_tiebreak : bool or None, optional
Forwarded to Judge.fit_decaf.
restart_seed : int, optional
Seed for Judge.fit_decaf shuffled restarts.
"""
j = Judge(species, lib_focus, lib_path, descriptor='ED_SOAP',
cluster_backend=cluster_backend,
min_cluster_size=min_cluster_size)
j.fit_decaf(
desc_array = desc_array,
delta = delta,
mu_per_el = mu_per_el,
counts_L_per_el = counts_L_per_el,
t_embedded = t_embedded,
save = False,
cluster_cache_per_el = cluster_cache_per_el,
c_vals_dict = c_vals_dict,
use_mu_tiebreak = use_mu_tiebreak,
restart_seed = restart_seed,
eps_noise_per_el = eps_noise_per_el,
)
return j
def _normalise_patient_descriptors(patient_descriptors, species):
"""Return {element: descriptors} for patient OOD classification."""
if patient_descriptors is None:
return {}
if isinstance(patient_descriptors, dict):
return {
el: np.asarray(patient_descriptors.get(el, np.empty((0, 0))))
for el in species
}
if len(species) == 1:
return {species[0]: np.asarray(patient_descriptors)}
raise ValueError(
'multi-element patient_descriptors must be a dict {element: descriptors}; '
'the runner builds this automatically for patient.run_ood=true.')
def _ood_len(gid_distinct):
"""Return number of atoms in a classifier's trailing OOD slot."""
if hasattr(gid_distinct, 'get'):
return len(gid_distinct.get(-1, []))
try:
return len(gid_distinct[-1])
except (KeyError, IndexError, TypeError):
return 0
def _classify_patient_ood(judge, patient_desc_per_el, species):
"""Classify patient descriptors with each element's fitted cluster."""
per_el = {}
total_atoms = 0
for el in species:
desc = np.asarray(patient_desc_per_el.get(el, np.empty((0, 0))))
total_atoms += len(desc)
if len(desc) == 0 or el not in judge.cluster_per_el:
per_el[el] = 0
continue
cluster = judge.cluster_per_el[el]
embedded = cluster.embedder.embed(desc)
gid_distinct, _ = cluster.classify(embedded)
per_el[el] = int(_ood_len(gid_distinct))
return int(sum(per_el.values())), total_atoms, per_el
def _ood_worker(delta, delta_fit, species, lib_focus, lib_path, desc_array,
patient_desc, cluster_backend='decaf', min_cluster_size=2,
eps_noise_per_el=None):
"""Refit Judge at δ on full corpus, return patient OOD count.
Picklable joblib worker for `HiFiWorkflow.compute_patient_ood`.
"""
import os as _os
_os.environ['OPENBLAS_NUM_THREADS'] = '4'
_os.environ['OMP_NUM_THREADS'] = '4'
j = Judge(species, lib_focus, lib_path, descriptor='ED_SOAP',
cluster_backend=cluster_backend,
min_cluster_size=min_cluster_size)
try:
j.fit_decaf(desc_array=desc_array, delta=delta_fit, save=False,
eps_noise_per_el=eps_noise_per_el)
except Exception as exc:
logger.info(f' δ={delta:.5f} skipped: {exc}')
return None
pat_ood, pat_atoms, pat_ood_per_el = _classify_patient_ood(
j, patient_desc, species)
return dict(delta=float(delta),
pat_ood=pat_ood,
pat_atoms=pat_atoms,
pat_ood_per_el=pat_ood_per_el,
n_structs_S=int(len(j.sub_lib_idx)),
n_groups=int(j.n_groups))
# ── Representative tagging ────────────────────────────────────────────────────
def _build_reps_with_group_id(rep_idx, judge, full_struct_lib, species):
"""Return copies of full structures with per-atom 'group_id' array.
Focus-element atoms receive their group id from
``judge.per_el_group_libs``; all other atoms receive -1. When the judge
carries refined/sorted ids (issue #40), a 'meta_group_id' array is added
from the per-element GroupOrder (non-focus atoms stay -1).
A 'noise_derived' array records provenance (issue #63): -1 marks an atom
whose group was promoted from the trailing noise slot, 0 a density-supported
group. Every focus atom now has a real group id, so this flag is the only
surviving trace of "this was a rare environment".
"""
meta_of_global = {}
for el, order in getattr(judge, 'group_order_per_el', {}).items():
off = judge.group_id_offset[el]
for loc, m in enumerate(order.meta_of_local):
meta_of_global[off + loc] = int(m)
# global gid -> was it promoted from noise? Refined ids (issue #40) permute
# the backend-native order, so map through old_of_local when present.
noise_of_global = {}
for el, fit in (getattr(judge, 'cluster_per_el', None) or {}).items():
nd = getattr(fit, 'noise_derived', None)
if fit is None or nd is None:
continue
nd = np.asarray(nd)
order = getattr(judge, 'group_order_per_el', {}).get(el)
off = (judge.group_id_offset or {}).get(el, 0) if order is not None else 0
if order is not None:
for loc, old in enumerate(order.old_of_local):
if int(old) < len(nd):
noise_of_global[off + loc] = bool(nd[int(old)])
else:
for g in range(len(nd)):
noise_of_global[g] = bool(nd[g])
result = []
for i in rep_idx:
full = full_struct_lib[i].copy()
syms = np.array(full.get_chemical_symbols())
gid_arr = np.full(len(full), -1, dtype=int)
is_focus = (full.arrays['is_focus'].astype(bool)
if 'is_focus' in full.arrays else None)
for el in species:
el_gids = judge.per_el_group_libs.get(el, [])
if i < len(el_gids) and len(el_gids[i]) > 0:
# gids cover only the FOCUS atoms of `el` (is_focus=1); the judge
# clusters lib_focus, not the full cell. Mask must match: element
# el AND is_focus, else an is_focus=0 atom of a focus element
# (e.g. an interface Pd) makes |mask| > |gids| → assignment error.
mask = syms == el
if is_focus is not None:
mask &= is_focus
gid_arr[mask] = el_gids[i]
full.set_array('group_id', gid_arr)
if meta_of_global:
full.set_array('meta_group_id', np.array(
[meta_of_global.get(int(g), -1) for g in gid_arr], dtype=int))
if noise_of_global:
full.set_array('noise_derived', np.array(
[-1 if noise_of_global.get(int(g), False) else 0
for g in gid_arr], dtype=int))
result.append(full)
return result
# ── HiFiWorkflow ────────────────────────────────────────────────────────────────
[docs]
class HiFiWorkflow(_PlottingMixin, _StateIOMixin):
"""End-to-end L/R-based bandwidth selection + visualisation.
See module docstring for the five-stage pipeline. Mandatory stages are
`run_final` and `plot_final_fit`; the others are optional.
Parameters
----------
species : list[str]
Element symbols clustering focuses on (e.g. ['Pd']).
lib_path : str
Path to the original library file (passed through to Judge as metadata).
sigma_rattle : tuple[float]
Rattling amplitudes (Å) for the thermal ε(σ) curve
(``eps_noise_method='thermal'``). Each element's σ_thermal target —
``resolve_sigma_thermal_targets`` at ``eps_temperature`` from
``eps_sigma_source``, overridden by ``eps_sigma_override`` — is added to
this grid, and ε_noise is read at exactly that σ.
eps_temperature : float
Temperature (K) at which σ_thermal is evaluated (default 300).
eps_r_cut : float
Neighbour cutoff (Å) for rattling. Use the descriptor's r_cut.
delta_lo_factor, delta_hi_factor : float or dict[str, float]
Coarse L/R scan range = [factor_lo, factor_hi] × ε_noise.
Pass a dict ``{el: value}`` for per-element overrides; missing elements
fall back to the scalar default (0.1 / 5.0).
k_micro_ratio : int, 'r100', 'f_tol', 'r<NUM>', or None
Microstate granularity for fit_microstate_reference. None (default):
conservative N//2. Integer ratio: explicit k = N // ratio
(ratio=2 is also N//2). ``'r100'``: opt-in smallest k whose max
within-cluster radius is ≤ ε_noise; useful for known-rigid systems.
``'r<NUM>'``: same criterion at user percentile NUM (1–100); for
NUM < 100 it is evaluated against ε_F instead of ε_noise when
``k_micro_force_active`` is True (``'r100'`` always bisects on
ε_noise).
``'f_tol'``: opt-in force-resolution ε_F criterion for microstate
granularity only.
k_micro_force_active : bool
True when the run carries an explicit force-resolution block
(``eps_noise.force`` in the raw config, or ``--f-tol`` on the CLI): an
``'r<NUM>'`` reference then bisects against ε_F (ceiling N_atoms−1)
instead of ε_noise (also ceiling N_atoms−1). Default False; irrelevant for
``'r100'``/``'f_tol'``.
n_coarse, n_fine : int
Number of δ points in coarse / fine scan.
l_tol : float or None
Likeness floor l_tol for the L≥l_tol constrained optimum. None
disables the L floor; the auto-pick becomes the best-fidelity
(max F) point.
n_jobs : int
Parallel processes for L/R scans.
n_jobs_eps : int
joblib process workers for ε_noise calibration (loky; -1 → all
cores). Config key compute.n_jobs_eps; null falls back to the
top-level n_jobs. See measure_eps_noise's n_jobs note.
seed : int
Reproducibility seed for workflow sampling, ε_noise sampling/rattles,
microstate k-means, grid subsampling, and NSGA-II defaults.
eps_noise_global : float, optional
Pre-computed scalar ε_noise. Promoted to ``eps_noise_per_species`` only
for a single focus element; when ``eps_noise_per_species`` is given,
this is replaced by the per-species mean. Multi-element runs that
consume ε (distance-kind or noise-emitting backends, or
r100/f_tol/r<NUM> references) need ``eps_noise_per_species`` or
``calibrate_eps_noise()`` before ``run_final``.
output_dir : str
Where to write artifacts (default '.').
save_judge : bool, keyword-only
Write the fitted Judge to ``judge_final.pkl`` at the end of
``plot_final_fit`` (config key output.save_judge). The pickle drops the
raw descriptor block — see ``Judge._PICKLE_DROP``.
retain_eps_samples : bool
Keep raw per-rattle distance samples in ``eps_results`` after thermal
calibration. Disabled by default because later workflow stages use only
the compact mean/std curves. Enable only for Python-side diagnostics.
"""
# ── Construction ────────────────────────────────────────────────────────
def __init__(self,
species,
lib_path,
sigma_rattle=(0.001, 0.003, 0.01, 0.03, 0.05, 0.08, 0.1),
eps_r_cut=5.0,
delta_lo_factor=0.1,
delta_hi_factor=5.0,
k_micro_ratio=None,
n_coarse=80,
n_fine=60,
l_tol=0.95,
n_jobs=1,
n_jobs_eps=1,
seed=42,
eps_noise_global=None,
eps_noise_per_species=None,
eps_temperature=300.0,
eps_noise_method='thermal',
eps_rdf_cfg=None,
eps_sigma_source='qho_table',
eps_sigma_override=None,
eps_force_f_tol=0.05,
cluster_backend='decaf',
min_cluster_size=2,
output_dir='.',
retain_eps_samples=False,
*,
save_judge=True,
k_micro_force_active=False):
# Two different things, deliberately kept apart (issue #69):
# self.cluster_backend — the NAME. Cache tag; state_io enforces
# equality on it, artifacts record it. Always a str.
# self._backend_spec — the EXECUTABLE. A name for built-ins (cheap,
# resolved in-process), the instance/class for a custom plug. This is
# what goes into worker kwargs: a child process re-imports the module
# and its registry holds only the built-ins, so a custom backend must
# travel with the work, not be looked up by name on arrival.
if isinstance(cluster_backend, ClusterBackend):
register_backend(cluster_backend)
self._backend_spec = cluster_backend
self.cluster_backend = cluster_backend.name
else:
if cluster_backend not in BACKEND_REGISTRY:
raise ValueError(
f'cluster_backend must be one of {tuple(BACKEND_REGISTRY)}; '
f'got {cluster_backend!r}. Register a custom ClusterBackend '
'first (Python API only — issue #69).')
entry = BACKEND_REGISTRY[cluster_backend]
self.cluster_backend = cluster_backend
self._backend_spec = (cluster_backend
if cluster_backend in BUILTIN_BACKENDS
else entry)
self.min_cluster_size = int(min_cluster_size)
self._backend_caps = BACKEND_REGISTRY[self.cluster_backend]
self.granularity_kind = self._backend_caps.granularity_kind
self.species = list(species)
self.lib_path = str(lib_path)
self.sigma_rattle = list(sigma_rattle)
self.eps_r_cut = float(eps_r_cut)
self.delta_lo_factor = delta_lo_factor # float or {el: float}
self.delta_hi_factor = delta_hi_factor # float or {el: float}
if isinstance(k_micro_ratio, str):
if k_micro_ratio not in ('r100', 'f_tol') and not is_rnum(k_micro_ratio):
raise ValueError(
f"k_micro_ratio must be >= 1, None, 'r100', 'r<NUM>', "
f"or 'f_tol'; got {k_micro_ratio!r}")
self.k_micro_ratio = k_micro_ratio
elif k_micro_ratio is not None:
if int(k_micro_ratio) < 1:
raise ValueError(
f"k_micro_ratio must be >= 1, None, 'r100', 'r<NUM>', "
f"or 'f_tol'; got {k_micro_ratio!r}")
self.k_micro_ratio = int(k_micro_ratio)
else:
self.k_micro_ratio = None
self.n_coarse = int(n_coarse)
self.n_fine = int(n_fine)
self.l_tol = None if l_tol is None else float(l_tol)
self.n_jobs = int(n_jobs)
self.n_jobs_eps = int(n_jobs_eps)
self.seed = int(seed)
self.eps_temperature = float(eps_temperature)
if eps_noise_method not in ('thermal', 'rdf', 'precomputed'):
raise ValueError(
f"eps_noise_method must be 'thermal', 'rdf', or 'precomputed'; "
f"got {eps_noise_method!r}")
self.eps_noise_method = eps_noise_method
self.eps_rdf_cfg = dict(eps_rdf_cfg) if eps_rdf_cfg else {}
self.eps_sigma_source = str(eps_sigma_source)
self.eps_sigma_override = dict(eps_sigma_override or {})
self.eps_force_f_tol = float(eps_force_f_tol)
self.k_micro_force_active = bool(k_micro_force_active)
self.retain_eps_samples = bool(retain_eps_samples)
if self.eps_force_f_tol <= 0:
raise ValueError(
f"eps_force_f_tol must be positive, got {eps_force_f_tol!r}.")
self.sigma_thermal_per_el = {}
self.eps_curve_sigma = None
self.eps_curve_mean_per_el = {}
# Elements whose RDF ε calibration found no peak (filled by
# calibrate_eps_noise in method='rdf'); their k_micro is forced to N//2.
self.eps_fallback_elements = set()
self.output_dir = Path(output_dir)
self.output_dir.mkdir(parents=True, exist_ok=True)
self.save_judge = bool(save_judge)
# Resolve eps_noise: per_species dict takes priority.
# eps_noise_global alone + single element → promote to dict.
# Both None → must call calibrate_eps_noise later.
if eps_noise_per_species is not None:
self.eps_noise_per_species = dict(eps_noise_per_species)
self.eps_noise_global = float(
np.mean(list(self.eps_noise_per_species.values())))
elif eps_noise_global is not None:
g = float(eps_noise_global)
if len(species) == 1:
self.eps_noise_per_species = {species[0]: g}
else:
self.eps_noise_per_species = None
self.eps_noise_global = g
else:
self.eps_noise_per_species = None
self.eps_noise_global = None
# State filled by stage methods
self.eps_results = None
self.records = []
self.lib = None
self.desc_array = None
self.struct_map = None
self.full_struct_lib = None
self.judge_emb = None
self.delta_auto = None
self.mu_labels = None
self.k_micro = None
self.h_mu = None
self.mu_labels_per_el = None
self.k_micro_per_el = None
self.h_mu_per_el = None
self.embed_op_per_el = None
self.df_coarse = None
self.df_fine = None
self.df_grid = None
# (el, round(c, 10)) → (pegl, all_gid, granularity); legacy-loaded
# entries may lack granularity and are refit, not reused.
self._el_fit_cache = {}
self.delta_star_F = None
self.delta_star_L095 = None
self.delta_star_per_el = None # per-element c_factors from plot_final_fit(dict)
self.fine_pat_ood = None
[docs]
def force_eps_for_element(self, el):
"""Return (ε_F, u*) for one element from cached thermal ε(σ)."""
if self.eps_curve_sigma is None or el not in self.eps_curve_mean_per_el:
raise RuntimeError(
"microstate.reference: f_tol or a force-resolution r<NUM> "
"reference requires the cached thermal ε(σ) curve; run "
"eps_noise.method: thermal calibration instead of precomputed/RDF.")
return force_eps_from_curve(
el,
f_tol=self.eps_force_f_tol,
sigma_rattle=self.eps_curve_sigma,
mean_eps=self.eps_curve_mean_per_el[el],
)
@staticmethod
def _build_el_data(lib_focus, desc_array, struct_map):
"""Split desc_array and struct_map by element symbol.
Returns
-------
el_symbols : (N_focus,) array of element symbols
desc_per_el : dict el → (N_el, D) array
smap_per_el : dict el → (N_el,) int array of struct indices
"""
symbols = []
for atoms in lib_focus:
symbols.extend(atoms.get_chemical_symbols())
el_symbols = np.array(symbols)
desc_per_el = {el: desc_array[el_symbols == el] for el in np.unique(el_symbols)}
smap_per_el = {el: struct_map[el_symbols == el] for el in np.unique(el_symbols)}
return el_symbols, desc_per_el, smap_per_el
def _get_cfactor_range(self, el=None):
if self.granularity_kind == 'count':
if el is None:
raise RuntimeError(
"_get_cfactor_range(el=None) is not defined for a count-kind "
"backend — bounds are per-element (f_el), pass el explicitly.")
k_micro_el = (self.k_micro_per_el or {}).get(el)
if k_micro_el is None:
raise RuntimeError(
f'k_micro for {el!r} unknown — call run_final or restore_state first.')
lo = min(10.0 / max(1, k_micro_el), 1.25)
hi = 1.25
if lo >= hi:
realized_k = _count_kind_k_from_fraction(hi, k_micro_el)
logger.warning(
f'[{el}] k_micro={k_micro_el} is too small for the standard '
f'NSGA-II k-fraction search domain [10, 1.25·k_micro]: fixing '
f'f_el = {hi} (realized k = {realized_k}). The nominal k ≥ 10 '
f'search floor is INFEASIBLE for this element — the fixed '
f'point does not satisfy it.')
return lo, hi
lo, hi = self.delta_lo_factor, self.delta_hi_factor
if isinstance(lo, dict):
lo = lo.get(el, 0.1) if el is not None else min(lo.values(), default=0.1)
if isinstance(hi, dict):
hi = hi.get(el, 5.0) if el is not None else max(hi.values(), default=5.0)
return float(lo), float(hi)
def _scan_value_to_delta(self, scan_val, el):
"""Convert one scan-axis value → the backend granularity for ``el``.
distance-kind: δ = c · ε_noise_el; count-kind: k = round(f · k_micro_el).
"""
if self.granularity_kind == 'count':
k_micro_el = (self.k_micro_per_el or {}).get(el)
if k_micro_el is None:
raise RuntimeError(
'count-kind granularity conversion needs k_micro_per_el — '
'call run_final (or restore_state) first.')
# #82's helper, not main's inline max(2, round(...)): the two are
# identical at the default floor=2/no-cap, and the helper is what
# the cap-aware callers share.
return float(_count_kind_k_from_fraction(
float(scan_val), k_micro_el))
# main's STRICT lookup (#75) replaces this branch's permissive
# `.get(el, self.eps_noise_global)`: per-element ε values differ by
# design, so a missing entry must fail rather than silently borrow the
# global one.
per_el = self.eps_noise_per_species
if per_el:
if el not in per_el:
raise RuntimeError(
f'eps_noise_per_species is set but has no entry for {el!r}; '
'refusing to substitute eps_noise_global — per-element ε values '
'differ by design (ADR-0004).')
eps_el = per_el[el]
else:
eps_el = self.eps_noise_global
if eps_el is None:
raise RuntimeError(
f'eps_noise for {el!r} unknown — call calibrate_eps_noise or '
'pass eps_noise_per_species.')
return float(scan_val) * float(eps_el)
def _judge_backend_kwargs(self):
"""kwargs forwarded to _fit_judge_at_delta for this workflow's backend."""
return dict(cluster_backend=self._backend_spec,
min_cluster_size=self.min_cluster_size,
eps_noise_per_el=self.eps_noise_per_species)
# ── Stage 1: ε_noise calibration ─────────────────────────────────────────
[docs]
def calibrate_eps_noise(self,
atoms_list,
descriptor_fn,
embed_fn=None,
lib_focus=None,
desc_array=None,
K=200, M=20,
save=True,
center_descriptor_fn=None):
"""Calibrate per-species ε_noise by ``eps_noise_method``. Optional.
``'thermal'``: rattling over ``sigma_rattle`` plus each element's
σ_thermal target, with ε read at exactly that target. ``'rdf'``: RDF
first-peak width; no-peak elements borrow the peaked mean.
``'precomputed'``: returns the constructor's ``eps_noise_per_species``
unchanged (ValueError if None); nothing is fitted or written.
For thermal/rdf: if `embed_fn` is None, temporary per-element Judges
are fitted from (lib_focus, desc_array) at auto-bandwidth only to build
embedding callbacks; ``eps_noise_per_species`` and ``eps_noise_global``
(their mean) are set; eps_noise_raw.npz is written when save=True.
Returns the per-species ε dict.
"""
if self.eps_noise_method == 'precomputed':
if self.eps_noise_per_species is None:
raise ValueError(
"eps_noise_method='precomputed' requires eps_noise_per_species.")
return self.eps_noise_per_species
if self.eps_noise_method == 'thermal':
self.sigma_thermal_per_el = resolve_sigma_thermal_targets(
self.species,
T=self.eps_temperature,
source=self.eps_sigma_source,
config_override=self.eps_sigma_override,
)
if embed_fn is None:
if lib_focus is None or desc_array is None:
raise ValueError(
'calibrate_eps_noise: either pass embed_fn or '
'(lib_focus, desc_array) so a temp Judge can be fit.')
logger.info('Fitting temporary Judge for embed_fn …')
embed_fn = _build_calibration_embed_fns(
self.species,
lib_focus,
self.lib_path,
desc_array,
self.seed,
)
if self.eps_noise_method == 'rdf':
eps_per_species = self._calibrate_eps_noise_rdf(
atoms_list, descriptor_fn, embed_fn, save=save)
else:
eps_per_species = self._calibrate_eps_noise_thermal(
atoms_list, descriptor_fn, embed_fn, self.sigma_thermal_per_el,
K=K, M=M, save=save,
center_descriptor_fn=center_descriptor_fn)
self.eps_noise_per_species = eps_per_species
self.eps_noise_global = float(np.mean(list(eps_per_species.values())))
logger.info(f' ε_noise per species [{self.eps_noise_method}]: '
f'{ {el: round(v, 5) for el, v in eps_per_species.items()} }')
logger.info(f' ε_noise global: {self.eps_noise_global:.5f}')
return self.eps_noise_per_species
def _calibrate_eps_noise_thermal(self, atoms_list, descriptor_fn, embed_fn,
sigma_thermal_per_el, K=200, M=20, save=True,
center_descriptor_fn=None):
"""Thermal-rattle ε_noise (the default). Returns eps_per_species."""
# Raw per-rattle distance samples are an intermediate. Compact means,
# stds, and ε(σ) curves below are the complete downstream state.
self.eps_results = None
sigma_rattle = sorted({
float(s) for s in self.sigma_rattle
} | {
float(sigma_thermal_per_el[el]) for el in self.species
})
logger.info(f'Running measure_eps_noise (K={K}, M={M}, '
f'r_cut={self.eps_r_cut} Å, n_jobs={self.n_jobs_eps}) …')
eps_results = measure_eps_noise(
atoms_list = atoms_list,
descriptor_fn = descriptor_fn,
embed_fn = embed_fn,
species = self.species,
sigma_rattle_list = sigma_rattle,
K = K,
M = M,
r_cut = self.eps_r_cut,
n_jobs = self.n_jobs_eps,
seed = self.seed,
center_descriptor_fn = center_descriptor_fn,
)
has_full_curve = all(
all(s in eps_results[el] for s in sigma_rattle)
for el in self.species)
if has_full_curve:
self.eps_curve_sigma = np.array(sigma_rattle, dtype=float)
self.eps_curve_mean_per_el = {
el: np.array([eps_results[el][s]['mean'] for s in sigma_rattle],
dtype=float)
for el in self.species
}
else:
self.eps_curve_sigma = None
self.eps_curve_mean_per_el = {}
eps_per_species = {
el: float(eps_results[el][float(sigma_thermal_per_el[el])]['mean'])
for el in self.species
}
if self.retain_eps_samples:
# Explicit diagnostic mode preserves the historical public result
# shape, including the potentially large sample arrays.
self.eps_results = eps_results
else:
# Later workflow stages need only these summaries. Python callers
# that need raw samples can opt in above or call measure_eps_noise.
self.eps_results = {
el: {
sigma: {
'mean': float(stats['mean']),
'std': float(stats['std']),
}
for sigma, stats in eps_results[el].items()
}
for el in self.species
}
# Cache writing below needs only the compact statistics. Release
# raw per-rattle arrays at their exact last consumer, not merely
# when this method returns.
del eps_results
if save:
if not has_full_curve:
raise RuntimeError(
"thermal ε calibration did not return the full sigma_rattle grid.")
out = self.output_dir / 'eps_noise_raw.npz'
np.savez(out,
species = np.array(self.species),
eps_method = 'thermal',
calibration_input_sha256 = getattr(
self, '_eps_calibration_input_sha256', None) or '',
sigma_source = self.eps_sigma_source,
qho_table_version = qho_table_version(),
sigma_rattle = np.array(sigma_rattle),
eps_noise_global = float(np.mean(list(eps_per_species.values()))),
**{f'eps_{el}': np.float64(v)
for el, v in eps_per_species.items()},
**{f'sigma_target_{el}': np.float64(sigma_thermal_per_el[el])
for el in self.species},
**{f'mean_{el}': self.eps_curve_mean_per_el[el]
for el in self.species},
**{f'std_{el}': np.array(
[self.eps_results[el][s]['std']
for s in sigma_rattle])
for el in self.species})
logger.info(f' Saved {out}')
return eps_per_species
def _calibrate_eps_noise_rdf(self, atoms_list, descriptor_fn, embed_fn, save=True):
"""RDF ``delta_r_ref`` center-only ε_noise. Returns eps_per_species.
No-peak elements borrow the mean ε of the elements that did peak (and are
recorded in ``self.eps_fallback_elements`` → k_micro forced to N//2). If no
element yields a peak, raises (RDF method unusable on this library).
"""
cfg = {**RDF_DEFAULTS, **self.eps_rdf_cfg}
logger.info('Running RDF ε_noise calibration '
f'(n_samples={cfg["n_samples"]}, eps_stat={cfg["eps_stat"]}) …')
rdf_res = measure_eps_noise_rdf(
atoms_list, descriptor_fn, embed_fn, self.species,
rdf_cfg=self.eps_rdf_cfg, seed=self.seed)
self.eps_results = rdf_res
peaked = {el: r['eps'] for el, r in rdf_res.items()
if not r['fallback'] and r['eps'] is not None}
if not peaked:
raise RuntimeError(
'RDF ε_noise calibration found no usable RDF first peak for ANY '
f'element {list(self.species)} — cannot calibrate. Use '
"eps_noise.method: thermal or set eps_noise.per_species explicitly.")
borrow = float(np.mean(list(peaked.values())))
self.eps_fallback_elements = {el for el, r in rdf_res.items() if r['fallback']}
eps_per_species = {}
for el in self.species:
if el in self.eps_fallback_elements:
eps_per_species[el] = borrow
logger.warning(' ε_noise[rdf] %s: no peak → borrowed peaked-mean '
'ε=%.5f; k_micro forced to N//2.', el, borrow)
else:
eps_per_species[el] = float(rdf_res[el]['eps'])
if save:
out = self.output_dir / 'eps_noise_raw.npz'
def _dr(el):
v = rdf_res[el]['delta_r']
return np.float64(v if v is not None else np.nan)
np.savez(out,
species = np.array(self.species),
eps_method = 'rdf',
calibration_input_sha256 = getattr(
self, '_eps_calibration_input_sha256', None) or '',
rdf_n_samples = int(cfg['n_samples']),
rdf_eps_stat = str(cfg['eps_stat']),
eps_noise_global = float(np.mean(list(eps_per_species.values()))),
**{f'eps_{el}': np.float64(v)
for el, v in eps_per_species.items()},
**{f'delta_r_{el}': _dr(el) for el in self.species},
**{f'fallback_{el}': np.bool_(rdf_res[el]['fallback'])
for el in self.species})
logger.info(f' Saved {out}')
return eps_per_species
def _require_runnable_backend(self):
"""Refuse to start a scan the configured backend cannot survive.
``fit_cluster_backend`` carries the authoritative check — it is the one
choke point every fit passes through. This is the *second* line, for
message quality only: the per-point ``except Exception`` in
``_fit_cfactor_worker`` swallows fit failures into 'c=... skipped', so
the reason has to be raised before the first point is tried, or it
reaches the user as 'zero valid fits out of N points' with the actual
cause discarded.
Deliberately unconditional, so it also refuses a scan whose every point
is a cache hit — one that would technically succeed, since a cached
``Judge.fit_decaf`` never calls the backend. Refusing is the honest
answer: the fine scan's window is adaptive, so a run that starts fully
cached can still reach a fresh fit and die there instead.
"""
reason = self._backend_caps.unsupported_reason()
if reason is not None:
raise RuntimeError(reason)
# ── Internal: c_factor scan (per-element DECAF) ──────────────────────────
def _scan_cfactor(self, lib_focus, desc_per_el, struct_map_per_el,
mu_per_el, counts_L_per_el, k_occ_L_per_el, H_L_per_el,
embed_op_per_el, c_factor_list,
pre_embedded_per_el=None,
_ext_shm_meta=None):
"""Scan over common resolution factor c; one per-element DECAF fit per point.
_ext_shm_meta : (meta_desc, meta_emb) tuple, optional
Pre-allocated SharedMemory metadata from an outer _shm_scope context.
When provided the internal _shm_scope call is skipped entirely,
avoiding a redundant PCA embedding.
"""
self._require_runnable_backend()
eps_per = self.eps_noise_per_species
# Strip to symbol-only stubs for parallel dispatch: fit_decaf only needs
# len(atoms) and per-atom symbols; positions/cell are never accessed.
lib_workers = lib_focus if self.n_jobs == 1 else _strip_lib_for_workers(lib_focus)
worker_kwargs = dict(
species=self.species, lib=lib_workers,
lib_path=self.lib_path, desc_per_el=desc_per_el,
struct_map_per_el=struct_map_per_el, mu_per_el=mu_per_el,
counts_L_per_el=counts_L_per_el, k_occ_L_per_el=k_occ_L_per_el,
H_L_per_el=H_L_per_el, eps_noise_per_el=eps_per,
embed_op_per_el=embed_op_per_el,
pre_embedded_per_el=pre_embedded_per_el,
seed=self.seed,
cluster_backend=self._backend_spec,
min_cluster_size=self.min_cluster_size,
)
if self.n_jobs == 1:
rows = [_fit_cfactor_worker(
c, **worker_kwargs,
cluster_cache_per_el=self._cached_scan_covers(c))
for c in c_factor_list]
rows = [r for r in rows if r is not None]
else:
from joblib import Parallel, delayed
use_shm = (getattr(self, 'max_training_size', None) is None)
if _ext_shm_meta is not None:
# Caller already owns a live _shm_scope — reuse it, skip allocation
cm = contextlib.nullcontext(_ext_shm_meta)
elif use_shm:
cm = _shm_scope(desc_per_el, self.species,
pre_embedded_per_el=pre_embedded_per_el)
else:
cm = contextlib.nullcontext((None, None))
with cm as (meta_desc, meta_emb):
if meta_desc is not None:
# Workers attach descriptors through metadata below. Drop
# both main-process aliases immediately after the SHM copy;
# otherwise this function and worker_kwargs pin the source
# N×D arrays for the full parallel scan.
desc_per_el = None
worker_kwargs['desc_per_el'] = None
if meta_emb is not None:
# Workers attach the authoritative embedding through SHM.
# Do not also pickle/memmap the same per-element arrays in
# every task as an unused fallback payload.
worker_kwargs['pre_embedded_per_el'] = None
shm_kwargs = {**worker_kwargs,
'_shm_meta_desc': meta_desc,
'_shm_meta_emb': meta_emb}
from tqdm.auto import tqdm
from atom_hifi._progress import bar_disabled, bar_file, tqdm_joblib
with tqdm_joblib(tqdm(total=len(c_factor_list), desc='1-D line scan',
file=bar_file(), disable=bar_disabled())):
results = Parallel(n_jobs=self.n_jobs, prefer='processes')(
delayed(_fit_cfactor_worker)(
c, **shm_kwargs,
cluster_cache_per_el=self._cached_scan_covers(c),
worker_subprocess=True)
for c in c_factor_list
)
rows = sorted((r for r in results if r is not None),
key=lambda r: r['c_factor'])
# Populate per-element DECAF cache from freshly-fitted results
for row in rows:
el_fits = row.pop('_el_fits', {})
c_val = row.get('c_factor') # scalar for 1-D scan
if c_val is not None:
for el, entry in el_fits.items():
self._el_fit_cache[(el, round(float(c_val), 10))] = entry
if not rows:
raise RuntimeError(
f'c_factor scan produced zero valid fits out of '
f'{len(c_factor_list)} points — every point raised inside '
f'_fit_judge_at_delta. Rerun with n_jobs=1 (or log_level: debug) '
f'to see the per-point tracebacks; joblib worker subprocesses '
f'do not inherit the CLI\'s logging configuration, so failures '
f'there are otherwise silent.')
df = pd.DataFrame(rows)
if 'L' in df.columns:
df['L'] = df['L'].clip(upper=1.0)
return df
def _fine_scan_cfactor(self, lib_focus, desc_per_el, struct_map_per_el,
mu_per_el, counts_L_per_el, k_occ_L_per_el, H_L_per_el,
embed_op_per_el, df_coarse,
pre_embedded_per_el=None,
_ext_shm_meta=None):
"""Fine scan around the F peak of df_coarse."""
F_c = df_coarse['F'].values
L_c = df_coarse['L'].values
n_c = len(df_coarse)
peak = int(np.argmax(F_c))
lo = max(0, peak - 3)
hi = min(n_c - 1, peak + 3)
# Extend lo leftward to cover the L>=l_tol region of the coarse scan
# (skipped when l_tol is None — no L floor, keep the default window)
if self.l_tol is not None:
qualifying_idx = np.where(L_c >= self.l_tol)[0]
if len(qualifying_idx):
lo = min(lo, max(0, int(qualifying_idx[0]) - 2))
c_lo = float(df_coarse['c_factor'].iloc[lo])
c_hi = float(df_coarse['c_factor'].iloc[hi])
c_fine = np.linspace(c_lo, c_hi, self.n_fine)
df_fine = self._scan_cfactor(
lib_focus, desc_per_el, struct_map_per_el,
mu_per_el, counts_L_per_el, k_occ_L_per_el, H_L_per_el,
embed_op_per_el, c_fine,
pre_embedded_per_el=pre_embedded_per_el,
_ext_shm_meta=_ext_shm_meta,
)
best_row = df_fine.loc[df_fine['F'].idxmax()]
qualifying = (df_fine[df_fine['L'] >= self.l_tol]
if self.l_tol is not None else df_fine.iloc[0:0])
best_constrained = (qualifying.loc[qualifying['F'].idxmax()]
if not qualifying.empty else None)
return df_fine, best_row, best_constrained
# ── ND grid scan (cached per-element DECAF) ───────────────────────────────
[docs]
def scan_cfactor_grid(self, c_vals_per_el=None, n_points_max=None):
"""ND Cartesian grid scan using cached per-element DECAF fits.
DECAF fits are reused from ``_el_fit_cache`` (populated by the 1-D scan
in ``run_final``). Only the joint set-cover is re-evaluated at each grid
point, so cost is O(n^N_el) covers — not O(n^N_el × N_el) DECAF fits.
Parameters
----------
c_vals_per_el : dict[str, array-like], optional
Per-element c values. If None, uses all cached values per element.
Missing (el, c_val) pairs not in ``_el_fit_cache`` are computed
on-demand via ``_fit_decaf_single_el`` before dispatching covers.
A value that only rounds to a cached key (neither the key itself
nor the cover's own c) reuses no cover and is refit at every grid
point.
n_points_max : int, optional
Cap on total grid points (random subsample if exceeded).
Returns
-------
pd.DataFrame
Columns: ``c_{el}`` for each element, ``L``, ``R``, ``F``,
``n_structs_S``, ``n_atoms_S``.
"""
import itertools
if self.granularity_kind == 'count':
raise RuntimeError(
"scan_cfactor_grid requires a distance-kind backend (decaf/"
f"hdbscan) — {self.cluster_backend!r} (count-kind) has no "
"c-factor axis to sweep (ADR-0009).")
# Own worker dispatch, so it does not inherit _scan_cfactor's check.
self._require_runnable_backend()
if self.mu_labels_per_el is None:
raise RuntimeError(
'scan_cfactor_grid requires per-element path. '
'Call run_final with eps_noise_per_species set.')
if self.lib is None:
raise RuntimeError('scan_cfactor_grid requires run_final to have been called.')
els = list(self.species)
# Axis values remain useful even if the covers themselves cannot be
# certified (e.g. an opaque custom backend setting). Capture them before
# the binding below; if it invalidates the cache, Phase 1 recomputes
# covers at those operating points. Inputs that cannot be certified
# leave the binding key unset, so the workers refit instead.
if c_vals_per_el is None:
c_vals_per_el = {
el: sorted({c for (e, c) in self._el_fit_cache if e == el})
for el in els
}
for el in els:
if not c_vals_per_el[el]:
raise RuntimeError(
f'No cached c values for element {el!r}. '
'Run run_final (1-D scan) first, or pass c_vals_per_el explicitly.')
self._bind_scan_cache_inputs(self.lib, self.desc_array, self.struct_map)
# Phase 1: on-demand fits for c values not yet in cache
_, desc_per_el, smap_per_el = self._build_el_data(
self.lib, self.desc_array, self.struct_map)
lib_len = len(self.lib)
lib_stub = _strip_lib_for_workers(self.lib)
# Collect (el, c_val) pairs missing from cache — need on-demand Phase 1 fit
missing = [
(el, float(c_val))
for el in els
for c_val in c_vals_per_el[el]
if (el, round(float(c_val), 10)) not in self._el_fit_cache
]
if missing:
from joblib import Parallel, delayed
from tqdm.auto import tqdm
from atom_hifi._progress import bar_disabled, bar_file, tqdm_joblib
n_phase1 = min(self.n_jobs if self.n_jobs > 0 else 1, len(missing))
logger.info(f' Phase 1 on-demand: {len(missing)} fits, n_jobs={n_phase1}')
def _phase1_task(el, c_val):
# CONVERTED, never deleted (issue #69). This branch is
# guard-dead, not logic-dead: scan_cfactor_grid raises for
# count-kind backends earlier in this same function, so it
# cannot run today. It goes live the moment that guard is
# relaxed — and deleting it would then send a count-kind
# backend down the else branch computing delta = c * eps_noise,
# a *distance* granularity on a *count* backend: silently
# wrong, no crash.
if self._backend_caps.granularity_kind == 'count':
# self.k_micro_per_el is the requested reference — not
# len(np.bincount(...)), the occupied count, which can be
# smaller when KMeans under-populates a requested cluster
# count (issue #82 codex-review finding).
k_micro_el = self.k_micro_per_el[el]
delta_el = float(_count_kind_k_from_fraction(
c_val, k_micro_el))
else:
delta_el = c_val * self.eps_noise_per_species[el]
return el, c_val, _fit_decaf_single_el(
el, desc_per_el[el], delta_el, lib_len, lib_stub,
cluster_backend=self._backend_spec,
min_cluster_size=self.min_cluster_size,
eps_noise_el=(self.eps_noise_per_species or {}).get(el), seed=self.seed)
with tqdm_joblib(tqdm(total=len(missing), desc='grid Phase-1 fits',
file=bar_file(), disable=bar_disabled())):
results_p1 = Parallel(n_jobs=n_phase1, prefer='threads')(
delayed(_phase1_task)(el, c_val) for el, c_val in missing
)
for el, c_val, entry in results_p1:
self._el_fit_cache[(el, round(c_val, 10))] = entry
# Rebuild mu/counts for workers
mu_per_el = self.mu_labels_per_el
counts_L_per_el = {}
k_occ_L_per_el = {}
H_L_per_el = {}
for el in els:
# minlength pads to the REQUESTED k_micro (self.k_micro_per_el), not
# just the occupied labels — sklearn can under-populate a requested
# cluster count (duplicate points), and len(counts_L_per_el[el]) is
# consumed elsewhere as a k_micro stand-in; without this pad the two
# sources silently diverge (issue #82 codex-review finding).
counts = np.bincount(mu_per_el[el],
minlength=(self.k_micro_per_el or {}).get(el, 0))
counts_L_per_el[el] = counts
k_occ_L_per_el[el] = int((counts > 0).sum())
H_L_per_el[el] = float(_scipy_entropy(
counts[counts > 0] / counts[counts > 0].sum(), base=np.e))
# Cartesian product
c_arrays = [np.asarray(c_vals_per_el[el]) for el in els]
grid_pts = list(itertools.product(*c_arrays))
if n_points_max and len(grid_pts) > n_points_max:
rng = np.random.default_rng(self.seed)
idxs = np.sort(rng.choice(len(grid_pts), n_points_max, replace=False))
grid_pts = [grid_pts[i] for i in idxs]
c_dicts = [{el: float(v) for el, v in zip(els, pt)} for pt in grid_pts]
eps_d = self.eps_noise_per_species or {}
for el in els:
c_arr = np.asarray(c_vals_per_el[el])
eps_el = eps_d.get(el, float('nan'))
logger.info(f' {el}: c∈[{c_arr.min():.3f}, {c_arr.max():.3f}] '
f'δ∈[{c_arr.min()*eps_el:.5f}, {c_arr.max()*eps_el:.5f}] '
f'({len(c_arr)} pts)')
logger.info(f' Grid scan (cached DECAF): '
f'{" × ".join(str(len(c_vals_per_el[el])) for el in els)} = '
f'{len(c_dicts)} points, n_jobs={self.n_jobs}')
# Each call receives only its own per-point cover slice, and only while
# the scan-input binding is verified (_cached_scan_covers). The worker
# reuses a cover whose stored δ equals this point's δ, or, for distance
# kind, a request for the cover's own rounded key; it refits elements
# whose cover lacks granularity evidence or matches neither.
lib_workers = (self.lib if self.n_jobs == 1
else _strip_lib_for_workers(self.lib))
worker_kwargs = dict(
species=self.species, lib=lib_workers,
lib_path=self.lib_path, desc_per_el=desc_per_el,
struct_map_per_el=smap_per_el,
mu_per_el=mu_per_el, counts_L_per_el=counts_L_per_el,
k_occ_L_per_el=k_occ_L_per_el, H_L_per_el=H_L_per_el,
eps_noise_per_el=self.eps_noise_per_species,
embed_op_per_el=self.embed_op_per_el,
cluster_backend=self._backend_spec,
min_cluster_size=self.min_cluster_size,
)
from tqdm.auto import tqdm
from atom_hifi._progress import bar_disabled, bar_file, tqdm_joblib
if self.n_jobs == 1:
rows = [_fit_cfactor_worker(
c, **worker_kwargs,
cluster_cache_per_el=self._cached_scan_covers(c))
for c in tqdm(c_dicts, desc='grid scan', file=bar_file(),
disable=bar_disabled())]
else:
from joblib import Parallel, delayed
with tqdm_joblib(tqdm(total=len(c_dicts), desc='grid scan',
file=bar_file(), disable=bar_disabled())):
rows = Parallel(n_jobs=self.n_jobs, prefer='processes')(
delayed(_fit_cfactor_worker)(
c, **worker_kwargs,
cluster_cache_per_el=self._cached_scan_covers(c),
worker_subprocess=True)
for c in c_dicts
)
rows = [r for r in rows if r is not None]
for row in rows:
row.pop('_el_fits', None) # nothing fresh to cache
df_grid = pd.DataFrame(rows)
if 'L' in df_grid.columns:
df_grid['L'] = df_grid['L'].clip(upper=1.0)
self.df_grid = df_grid
df_grid.to_csv(Path(self.output_dir) / 'scan_grid.csv', index=False)
return df_grid
# ── ND grid scan workflow wrapper ─────────────────────────────────────────
[docs]
def run_grid_scan(self, n_coarse_per_el=None, n_fine_per_el=None,
l_tol=None, save_dir=None):
"""Coarse ND grid scan → per-element optimum → fine ND grid scan.
Uses ``_el_fit_cache`` populated by ``run_final`` — no DECAF
recomputation for c values already visited by the 1-D scan.
Parameters
----------
n_coarse_per_el : int, optional
Number of c values per element for coarse grid. Defaults to all
cached c values (from 1-D scan).
n_fine_per_el : int, optional
Number of c values per element for fine grid. Default 20.
l_tol : float, optional
Likeness floor: the constrained optimum requires L ≥ l_tol. Uses
``self.l_tol`` if None.
save_dir : str or Path, optional
Directory for report output. Uses ``self.output_dir`` if None.
Returns
-------
df_grid : pd.DataFrame
Merged coarse + fine grid with duplicates removed.
"""
l_tol = l_tol if l_tol is not None else self.l_tol
save_dir = Path(save_dir or self.output_dir)
n_fine = n_fine_per_el if n_fine_per_el is not None else 20
els = list(self.species)
# ── Coarse grid ───────────────────────────────────────────────────────
if n_coarse_per_el is not None:
# Subsample cached values evenly per element
coarse_axes = {}
for el in els:
cached_vals = sorted({c for (e, c) in self._el_fit_cache if e == el})
if len(cached_vals) <= n_coarse_per_el:
coarse_axes[el] = cached_vals
else:
idxs = np.linspace(0, len(cached_vals) - 1,
n_coarse_per_el, dtype=int)
coarse_axes[el] = [cached_vals[i] for i in idxs]
else:
coarse_axes = None # scan_cfactor_grid defaults to all cached
logger.info('\n── run_grid_scan: coarse grid ──')
df_coarse = self.scan_cfactor_grid(c_vals_per_el=coarse_axes)
# ── Find per-element coarse optimum ───────────────────────────────────
qualifying = (df_coarse[df_coarse['L'] >= l_tol]
if l_tol is not None else df_coarse.iloc[0:0])
if qualifying.empty:
if l_tol is not None:
logger.warning(f'no coarse grid point has F ≥ {l_tol}; '
'using unconstrained F peak for fine-scan centre.')
best_row = df_coarse.loc[df_coarse['F'].idxmax()]
else:
best_row = qualifying.loc[qualifying['F'].idxmax()]
# ── Fine grid centred on per-element coarse optimum ──────────────────
fine_axes = {}
for el in els:
c_star_el = float(best_row[f'c_{el}'])
lo_el, hi_el = self._get_cfactor_range(el)
# Fine window: ±15 % of the element's full range (30 % total width),
# clamped to valid bounds
half_width = (hi_el - lo_el) * 0.15
c_lo = max(lo_el, c_star_el - half_width)
c_hi = min(hi_el, c_star_el + half_width)
fine_axes[el] = np.linspace(c_lo, c_hi, n_fine).tolist()
logger.info('\n── run_grid_scan: fine grid ──')
df_fine = self.scan_cfactor_grid(c_vals_per_el=fine_axes)
# ── Merge coarse + fine, deduplicate ─────────────────────────────────
c_cols = [f'c_{el}' for el in els]
df_all = pd.concat([df_coarse, df_fine], ignore_index=True)
df_all = df_all.drop_duplicates(subset=c_cols).reset_index(drop=True)
df_all = df_all.sort_values(c_cols).reset_index(drop=True)
if 'L' in df_all.columns:
df_all['L'] = df_all['L'].clip(upper=1.0)
self.df_grid = df_all
df_all.to_csv(Path(self.output_dir) / 'scan_grid.csv', index=False)
logger.info('\n── run_grid_scan: report ──')
self.report_cfactor_grid(df_all, save_dir=save_dir, l_tol=l_tol)
return df_all
# ── NSGA-II multi-objective c-factor optimiser ────────────────────────────
def _run_nsga2_for_state(
self,
lib_focus,
desc_array,
struct_map,
mu_labels_per_el,
embed_op_per_el,
n_gen: int = 50,
pop_size: int = 40,
seed=None,
c_star_1d=None,
c_vec_2d_best=None,
l_tol=None,
save_dir=None,
pareto_filename: str = 'pareto_front.csv',
history_filename: str = 'pareto_history.csv',
plot: bool = True,
):
"""Run NSGA-II against an explicit workflow state."""
# Here rather than in run_nsga2: run_loop's inner NSGA-II calls this
# directly, and the fits happen inside the pymoo workers, past the
# swallow. HDBSCAN can reach this path through its native classifier;
# it must still pass the fitting-environment compatibility guard.
self._require_runnable_backend()
from .optimizer import (
HiFiCFactorProblem,
extract_pareto_front,
optimize_c_factors_nsga2,
plot_nsga2_results,
_print_sanity_summary,
)
seed = self.seed if seed is None else int(seed)
l_tol = l_tol or self.l_tol
save_dir = Path(save_dir) if save_dir is not None else self.output_dir
_, desc_per_el, smap_per_el = self._build_el_data(
lib_focus, desc_array, struct_map)
counts_L_per_el: dict = {}
k_occ_L_per_el: dict = {}
H_L_per_el: dict = {}
for el in self.species:
# minlength: see the matching comment in scan_cfactor_grid — pads
# to the requested k_micro so len(counts_L_per_el[el]) can't
# diverge from self.k_micro_per_el[el] (issue #82 codex-review).
counts = np.bincount(mu_labels_per_el[el],
minlength=(self.k_micro_per_el or {}).get(el, 0))
counts_L_per_el[el] = counts
k_occ_L_per_el[el] = int((counts > 0).sum())
H_L_per_el[el] = float(_scipy_entropy(
counts[counts > 0] / counts[counts > 0].sum(), base=np.e))
eps_d = self.eps_noise_per_species or {}
els = list(self.species)
xl = [self._get_cfactor_range(el)[0] for el in els]
xu = [self._get_cfactor_range(el)[1] for el in els]
evaluator_kwargs = dict(
lib=lib_focus,
lib_path=str(self.lib_path) if self.lib_path else '',
desc_per_el=desc_per_el,
smap_per_el=smap_per_el,
mu_per_el=mu_labels_per_el,
counts_L_per_el=counts_L_per_el,
k_occ_L_per_el=k_occ_L_per_el,
H_L_per_el=H_L_per_el,
eps_noise_per_el=eps_d,
embed_op_per_el=embed_op_per_el,
cluster_backend=self._backend_spec,
min_cluster_size=self.min_cluster_size,
)
problem = HiFiCFactorProblem(
species_order=els,
xl=xl,
xu=xu,
evaluator_kwargs=evaluator_kwargs,
n_jobs=self.n_jobs,
)
# Problem owns the descriptors until its optimiser copies them to SHM.
# This local alias would otherwise survive throughout minimize().
del desc_per_el
res, history, hv_history = optimize_c_factors_nsga2(
problem,
n_gen=n_gen,
pop_size=pop_size,
seed=seed,
c_star_1d=c_star_1d,
c_vec_2d_best=c_vec_2d_best,
n_jobs=self.n_jobs,
)
df_pareto, best_c_per_el = extract_pareto_front(
res, history, els, eps_d,
l_tol=l_tol,
save_dir=save_dir,
pareto_filename=pareto_filename,
history_filename=history_filename,
key_fn=problem.canonical_key,
granularity_kind=self.granularity_kind,
)
if plot:
plot_nsga2_results(
df_pareto, history, hv_history, els,
best_c_per_el, l_tol=l_tol, save_dir=save_dir,
)
_print_sanity_summary(res, df_pareto, hv_history, els,
best_c_per_el, l_tol)
return res, df_pareto, hv_history, best_c_per_el
[docs]
def run_nsga2(
self,
n_gen: int = 50,
pop_size: int = 40,
seed=None,
c_star_1d=None,
c_vec_2d_best=None,
l_tol=None,
save_dir=None,
):
"""Run NSGA-II Pareto optimisation of per-element c-factors.
Replaces the exhaustive grid scan (scan_cfactor_grid / run_grid_scan) as
the default multi-element c-optimisation path. run_final must have
been called first.
Parameters
----------
n_gen, pop_size, seed : NSGA-II budget.
Rule of thumb: pop_size ≥ 10 × N_focus_elements (min 40).
Watch pareto_convergence.png — increase n_gen if HV hasn't plateaued.
c_star_1d : float or None
Common 1-D c* to seed the initial population.
If None, auto-derived from self.delta_star_L095, which is already
stored in dimensionless c-factor units.
c_vec_2d_best : array-like or None
Best c-vector from a prior 2-D grid scan (optional extra seed).
l_tol : float or None
Likeness threshold for auto-picking operating point.
Defaults to self.l_tol.
save_dir : str or Path or None
Directory for CSV + PNG outputs. Defaults to self.output_dir.
Returns
-------
tuple
``(res, df_pareto, hv_history, best_c_per_el)`` where:
* ``res`` is the pymoo Result (``.X``, ``.F`` on the Pareto front).
* ``df_pareto`` is a DataFrame of non-dominated points.
* ``hv_history`` is the list of HV per generation.
* ``best_c_per_el`` is the auto-picked ``{el: c_factor}`` mapping;
pass it to ``plot_final_fit(delta_pick=best_c_per_el)``.
"""
if self.mu_labels_per_el is None:
raise RuntimeError(
'run_nsga2 requires per-element path. '
'Call run_final with eps_noise_per_species set.')
if self.lib is None:
raise RuntimeError('run_nsga2 requires run_final to have been called.')
# ── Auto-derive c_star_1d from 1-D result if not provided ─────────────
if c_star_1d is None and self.delta_star_L095 is not None:
c_star_1d = self.delta_star_L095
return self._run_nsga2_for_state(
self.lib,
self.desc_array,
self.struct_map,
self.mu_labels_per_el,
getattr(self, 'embed_op_per_el', {}),
n_gen=n_gen,
pop_size=pop_size,
seed=seed,
c_star_1d=c_star_1d,
c_vec_2d_best=c_vec_2d_best,
l_tol=l_tol,
save_dir=save_dir,
)
# ── Internal: fit_iteration (used by run_loop and run_final) ────────────
def _fit_iteration(self, lib_focus, desc_array, struct_map, n_iter=None,
save=True, coarse_only=True):
"""One 'iteration': per-element auto-bw + microstate ref + coarse scan
(+ optional fine scan) + c* selection. Returns a dict with all per-element state.
coarse_only : bool
When True (default), skip the fine c_factor scan and pick c* from
the coarse peak directly — used for inner-loop iterations where
speed matters. When False, run both coarse and fine scans; the fine
scan zooms into the ±3-index coarse window around the F peak.
``run_final`` always passes False.
"""
tag = f'iter {n_iter}' if n_iter is not None else 'final'
if self.eps_noise_per_species is None:
_needs_eps = (self.granularity_kind == 'distance'
or self._backend_caps.emits_noise
or self.k_micro_ratio in ('r100', 'f_tol')
or is_rnum(self.k_micro_ratio))
if _needs_eps:
raise RuntimeError(
'eps_noise_per_species not set — pass eps_noise_per_species '
'to __init__ or call calibrate_eps_noise first.')
# count-kind scan with the N//2 reference never consumes ε_noise
logger.info(f' [{tag}] {self.cluster_backend} backend (count-kind): '
'ε_noise not required; skipping.')
self.eps_noise_per_species = {el: float('nan')
for el in self.species}
# 1. Per-element atom data
el_symbols, desc_per_el, smap_per_el = self._build_el_data(
lib_focus, desc_array, struct_map)
# 2. One current-corpus PCA embedding per element. It is reused by the
# microstate reference and every scan fit below, but never across calls:
# active-learning iterations and repeated run_final calls may carry a
# different corpus, for which a cached projection is a different metric.
embed_op_per_el = {}
f_emb_per_el = {}
t_embedded_per_el = {}
delta_auto_per_el = {}
for el in self.species:
if self._backend_caps.has_auto_bandwidth:
logger.info(f' [{tag}] auto-bw fit Judge([{el}]) …')
j_el = _fit_judge_at_delta(
[el], lib_focus, self.lib_path, desc_per_el[el], delta=0,
restart_seed=self.seed, **self._judge_backend_kwargs())
embedder = j_el.cluster.embedder
delta_auto_per_el[el] = float(j_el.cluster.granularity)
logger.info(f' [{el}] δ_auto={delta_auto_per_el[el]:.5f}')
else:
# Backends without auto-bandwidth need PCA only; δ_auto is a
# HaarWavelet/decaf concept.
logger.info(f' [{tag}] PCA embed_op fit [{el}] '
f'({self.cluster_backend} backend) …')
embedder = PCAEmbedder.fit(desc_per_el[el])
delta_auto_per_el[el] = 0.0
embed_op_per_el[el] = np.asarray(embedder.embed_op).copy()
t_embedded_per_el[el] = np.asarray(embedder.t_embedded).copy()
f_emb_per_el[el] = desc_per_el[el] @ embed_op_per_el[el]
delta_auto = float(np.mean(list(delta_auto_per_el.values())))
# 3. Per-element microstate reference
mu_per_el = {}
k_micro_per_el = {}
h_mu_per_el = {}
counts_L_per_el = {}
k_occ_L_per_el = {}
H_L_per_el = {}
for el in self.species:
logger.info(f' [{tag}] fit_microstate_reference [{el}] …')
N_el = len(desc_per_el[el])
if el in self.eps_fallback_elements:
# RDF ε calibration found no peak for this element: its ε was
# borrowed, so r100/r<NUM> against it is not meaningful — fall
# back to the conservative N//2 regardless of k_micro_ratio.
_k_explicit = max(20, N_el // 2)
_eps_force = None
elif self.k_micro_ratio == 'r100':
_k_explicit = 'r100'
_eps_force = None
elif self.k_micro_ratio == 'f_tol':
_k_explicit = 'f_tol'
_eps_force, _u_star = self.force_eps_for_element(el)
logger.info(
f' [{el}] ε_F={_eps_force:.5f} from '
f'u*={_u_star:.5f} Å, F_tol={self.eps_force_f_tol:.5f} eV/Å')
elif isinstance(self.k_micro_ratio, str) and is_rnum(self.k_micro_ratio):
_k_explicit = self.k_micro_ratio
if self.k_micro_force_active:
_eps_force, _u_star = self.force_eps_for_element(el)
logger.info(
f' [{el}] ε_F={_eps_force:.5f} from '
f'u*={_u_star:.5f} Å, F_tol={self.eps_force_f_tol:.5f} '
f'eV/Å (r{self.k_micro_ratio[1:]} force-resolution)')
else:
_eps_force = None
elif self.k_micro_ratio is None:
_k_explicit = max(20, N_el // 2)
_eps_force = None
else:
_k_explicit = max(20, N_el // self.k_micro_ratio)
_eps_force = None
k_m, mu_l, h_m = fit_microstate_reference(
f_emb_per_el[el], self.eps_noise_per_species[el],
k_micro=_k_explicit, seed=self.seed, eps_force=_eps_force)
mu_per_el[el] = mu_l
k_micro_per_el[el] = int(k_m)
h_mu_per_el[el] = float(h_m)
# minlength: see the matching comment in scan_cfactor_grid — pads
# to the requested k_micro so len(counts_L_per_el[el]) can't
# diverge from k_micro_per_el[el] (issue #82 codex-review).
counts = np.bincount(mu_l, minlength=k_micro_per_el[el])
counts_L_per_el[el] = counts
k_occ_L_per_el[el] = int((counts > 0).sum())
H_L_per_el[el] = float(_scipy_entropy(
counts[counts > 0] / counts[counts > 0].sum(), base=np.e))
logger.info(f' [{el}] k_micro={k_m} H(μ)={h_m:.4f} nats')
# f_emb_per_el (desc @ embed_op) has no consumer past the microstate
# labels above, so it still goes here.
#
# The original note at this line warned against overlapping N×D arrays
# with the scan's SHM allocations. That warning is now only partly
# honoured: t_embedded_per_el is deliberately retained past this point
# so the scan can reuse it instead of refitting PCA, so one N×D array
# per element IS live across the SHM window, alongside the SHM copy.
#
# Judged acceptable, not free: N×D is the *reduced* dimension after the
# dim=0 broken-stick cut, ~10⁻³ of the 4·N² block issue #68 deleted
# (≈21 MB at LLZO's 88k atoms, against 29 GiB). It is nonetheless
# UNMEASURED — test_memory_release's threshold is sized to catch 4·N²
# and is two orders of magnitude above N×D, so it cannot see this.
del f_emb_per_el
# Global mu_labels for backward compat (per-el with offset)
mu_labels = np.empty(len(desc_array), dtype=np.int64)
offset = 0
for el in self.species:
mask = el_symbols == el
mu_labels[mask] = mu_per_el[el] + offset
offset += k_micro_per_el[el]
k_micro = int(sum(k_micro_per_el.values()))
h_mu = float(np.mean([h_mu_per_el[el] for el in self.species]))
# 4+5. Coarse + fine c_factor scans — share one SharedMemory allocation
# so the descriptors and PCA embedding are copied once, not twice.
self._bind_scan_cache_inputs(lib_focus, desc_array, struct_map, k_micro_per_el)
if self.granularity_kind == 'count':
# k-fraction axis f: k_el = round(f · k_micro_el), linear sweep
# from ~10 clusters (smallest element) up to 1.25 · k_micro.
k_micro_min = max(1, min(k_micro_per_el.values()))
c_lo = min(10.0 / k_micro_min, 1.25)
c_hi = 1.25
logger.info(f' [{tag}] count-kind scan: f ∈ [{c_lo:.4f}, {c_hi:.2f}] '
f'(k up to 1.25·k_micro; k_micro_min={k_micro_min})')
else:
c_lo, c_hi = self._get_cfactor_range()
c_coarse = np.linspace(c_lo, c_hi, self.n_coarse)
logger.info(f' [{tag}] coarse scan: c ∈ [{c_lo:.2f}, {c_hi:.2f}] '
f'({self.n_coarse} pts, n_jobs={self.n_jobs}) …')
use_shm = (self.n_jobs > 1
and getattr(self, 'max_training_size', None) is None)
with (_shm_scope(desc_per_el, self.species,
pre_embedded_per_el=t_embedded_per_el)
if use_shm else contextlib.nullcontext((None, None))) as _shm:
if use_shm:
desc_per_el = None # SHM owns the copy; release Python-side reference
df_coarse = self._scan_cfactor(
lib_focus, desc_per_el, smap_per_el,
mu_per_el, counts_L_per_el, k_occ_L_per_el, H_L_per_el,
embed_op_per_el, c_coarse,
pre_embedded_per_el=t_embedded_per_el,
_ext_shm_meta=_shm,
)
if coarse_only:
# Coarse peak is adequate for inner-loop OOD evaluation
qualifying = (df_coarse[df_coarse['L'] >= self.l_tol]
if self.l_tol is not None else df_coarse.iloc[0:0])
best_row = (qualifying.loc[qualifying['F'].idxmax()]
if not qualifying.empty
else df_coarse.loc[df_coarse['F'].idxmax()])
best_constrained = (qualifying.loc[qualifying['F'].idxmax()]
if not qualifying.empty else None)
df_fine = df_coarse
else:
# 5. Fine c_factor scan (reuses same SHM — no second embedding pass)
logger.info(f' [{tag}] fine scan ({self.n_fine} pts) …')
df_fine, best_row, best_constrained = self._fine_scan_cfactor(
lib_focus, desc_per_el, smap_per_el,
mu_per_el, counts_L_per_el, k_occ_L_per_el, H_L_per_el,
embed_op_per_el, df_coarse,
pre_embedded_per_el=t_embedded_per_el,
_ext_shm_meta=_shm,
)
# 6. c* selection
if best_constrained is not None:
c_star = float(best_constrained['c_factor'])
sel_mode = f'L≥{self.l_tol:.2f}' + (' (coarse)' if coarse_only else '')
else:
c_star = float(best_row['c_factor'])
if self.l_tol is None:
sel_mode = ('F peak (coarse, L floor disabled)' if coarse_only
else 'F peak (L floor disabled)')
else:
sel_mode = ('F peak (coarse, unconstrained)' if coarse_only
else 'F peak (unconstrained)')
delta_star = c_star
logger.info(f' [{tag}] c*={c_star:.5f} ({sel_mode})')
return dict(
judger = None,
df_coarse = df_coarse,
df_fine = df_fine,
best_row = best_row,
best_constrained = best_constrained,
mu_labels = mu_labels,
mu_labels_per_el = mu_per_el,
k_micro = k_micro,
k_micro_per_el = k_micro_per_el,
h_mu = h_mu,
h_mu_per_el = h_mu_per_el,
embed_op_per_el = embed_op_per_el,
delta_auto = delta_auto,
delta_star = delta_star,
sel_mode = sel_mode,
)
# ── Stage 2: active-learning loop ───────────────────────────────────────
[docs]
def run_loop(self,
lib_focus_init,
desc_focus_init,
full_struct_lib_init,
sample_next_fn,
sample_next_ood_fn=None,
max_iters=None,
save=True,
skip_inner_fine=True,
inner_nsga2=False,
inner_nsga2_n_gen=5,
inner_nsga2_pop_size=None,
checkpoint_source_indices_fn=None,
checkpoint_source_fingerprints_fn=None):
"""Active-learning loop.
Parameters
----------
lib_focus_init : list[Atoms]
Initial focus-only library.
desc_focus_init : list[ndarray]
Per-structure focus-atom descriptors (parallel to lib_focus_init).
full_struct_lib_init : list[Atoms]
Initial full structures (parallel index).
sample_next_fn : callable() → tuple or None
Callback that samples the next batch given the current state.
Returns (next_lib_focus, next_desc_focus_list, next_full_structs)
or None to terminate the loop.
sample_next_ood_fn : callable(Judge) → tuple or None, optional
OOD-ranked sampler used only by loop+NSGA-II steering. Returns
(next_lib_focus, next_desc_focus_list, next_full_structs,
next_batch_ood, next_batch_ood_frac).
max_iters : int, optional
Hard cap on iterations.
save : bool
Save hifi_iter{NN}.{npz,png}, descp_iter{NN}.npy, and loop_records.json.
skip_inner_fine : bool
When True (default), each inner iteration runs only the coarse
c_factor scan and picks c* from the coarse peak. The fine scan runs
once in Stage 3 (run_final) on the fully-accumulated corpus.
inner_nsga2 : bool
When True, run a cheap NSGA-II optimisation at each iteration and
refit the next-batch classifier at its selected per-element c-vector.
checkpoint_source_indices_fn : callable() -> ndarray, optional
Exact source-library indices for the current ordered corpus. The
runner supplies this so resume can reconstruct OOD-ranked as well as
randomly sampled loops without replaying a different selection rule.
checkpoint_source_fingerprints_fn : callable() -> ndarray, optional
Content fingerprints parallel to the source indices. Both callbacks
must be supplied together so resume can detect a reordered or
replaced source library.
Sets self.records and returns it.
"""
if not self._backend_caps.has_classifier:
raise RuntimeError(
"run_loop requires a backend with a deployable classifier — "
"the loop's OOD early-stop has no "
f'{self.cluster_backend} analogue (ADR-0009).')
if self.cluster_backend == 'kmeans':
# Name the lever, not just the risk: #81 deliberately added no
# loop.max_iters config key, so a YAML-driven run has NO bound at
# all and a user told only "this may consume the full library" has
# nothing to act on. max_iters is the Python-API escape hatch.
logger.warning(
'kmeans loop uses resolution-tracking lambda_open abstention; '
'OOD may remain positive as k grows with the corpus, so the '
'existing stop policy may consume the full library (gitlab#81). '
'There is no config key for this: bound the run with '
'run_loop(max_iters=...) from the Python API, or with a '
'smaller library.')
cur_lib_focus = list(lib_focus_init)
cur_focus_desc_list = list(desc_focus_init)
cur_full_lib = list(full_struct_lib_init)
records = []
while True:
n_iter = len(records) + 1
if max_iters is not None and n_iter > max_iters:
break
DESC_cur = np.vstack(cur_focus_desc_list)
struct_map_cur = np.repeat(
np.arange(len(cur_lib_focus)),
[len(x) for x in cur_focus_desc_list],
)
logger.info(f'\n{"─" * 70}')
logger.info(f'Iteration {n_iter} | training: {len(cur_lib_focus)} structs / '
f'{DESC_cur.shape[0]} focus atoms')
logger.info('─' * 70)
res = self._fit_iteration(cur_lib_focus, DESC_cur, struct_map_cur,
n_iter=n_iter, save=save,
coarse_only=skip_inner_fine)
# Save per-iter artifacts
if save:
fr_npz = self.output_dir / f'hifi_iter{n_iter:02d}.npz'
_dc = res['df_coarse']
_df = res['df_fine']
_dcol = 'c_factor' if 'c_factor' in _dc.columns else 'delta'
_fcol = 'c_factor' if 'c_factor' in _df.columns else 'delta'
np.savez(fr_npz,
coarse_delta=_dc[_dcol].values,
coarse_L=_dc['L'].values,
coarse_R=_dc['R'].values,
coarse_F=_dc['F'].values,
coarse_nS=_dc['n_structs_S'].values,
coarse_Ng=_dc.get('N_groups', _dc['n_structs_S'] * 0).values,
fine_delta=_df[_fcol].values,
fine_L=_df['L'].values,
fine_R=_df['R'].values,
fine_F=_df['F'].values,
fine_nS=_df['n_structs_S'].values,
fine_Ng=_df.get('N_groups', _df['n_structs_S'] * 0).values,
mu_labels=res['mu_labels'],
k_micro=res['k_micro'],
h_mu=res['h_mu'],
delta_auto=res['delta_auto'],
delta_star=res['delta_star'],
sel_mode=res['sel_mode'],
eps_noise=self.eps_noise_global)
self._plot_FR_panels(
res['df_coarse'], res['df_fine'],
res['delta_auto'], res['delta_star'],
fname=str(self.output_dir / f'hifi_iter{n_iter:02d}.png'),
title=f'Iter {n_iter}',
)
np.save(self.output_dir / f'descp_iter{n_iter:02d}.npy',
DESC_cur)
np.save(self.output_dir / f'descp_iter{n_iter:02d}_struct_map.npy',
struct_map_cur)
have_indices = checkpoint_source_indices_fn is not None
have_fingerprints = checkpoint_source_fingerprints_fn is not None
if have_indices != have_fingerprints:
raise RuntimeError(
'checkpoint source indices and fingerprints must be '
'provided together.')
if have_indices:
source_idx = np.asarray(
checkpoint_source_indices_fn(), dtype=np.int64)
source_hash = np.asarray(
checkpoint_source_fingerprints_fn())
if len(source_idx) != len(cur_lib_focus):
raise RuntimeError(
'checkpoint source-index count does not match the '
f'current corpus: {len(source_idx)} != '
f'{len(cur_lib_focus)}.')
if source_hash.ndim != 1 or len(source_hash) != len(source_idx):
raise RuntimeError(
'checkpoint source fingerprints must be a '
'one-dimensional array parallel to source indices.')
np.save(
self.output_dir
/ f'descp_iter{n_iter:02d}_source_idx.npy',
source_idx)
np.save(
self.output_dir
/ f'descp_iter{n_iter:02d}_source_hash.npy',
source_hash)
c_per_el = None
extension_method = 'random'
if inner_nsga2 and len(self.species) > 1:
pop_size = (inner_nsga2_pop_size if inner_nsga2_pop_size is not None
else max(8, 4 * len(self.species)))
logger.info(f' [iter {n_iter}] inner NSGA-II: n_gen={inner_nsga2_n_gen} '
f'pop_size={pop_size}')
_, _, _, c_per_el = self._run_nsga2_for_state(
cur_lib_focus,
DESC_cur,
struct_map_cur,
res['mu_labels_per_el'],
res['embed_op_per_el'],
n_gen=inner_nsga2_n_gen,
pop_size=pop_size,
seed=self.seed + n_iter,
c_star_1d=res['delta_star'],
l_tol=self.l_tol,
save_dir=self.output_dir,
pareto_filename=f'pareto_iter{n_iter:02d}.csv',
history_filename=f'pareto_history_iter{n_iter:02d}.csv',
plot=False,
)
extension_method = 'ood'
# Refit at the selected scan value on the current corpus for
# next-batch OOD. Count-kind backends scale by this iteration's
# k_micro; distance-kind backends scale by eps_noise.
_, desc_per_el_cur, _ = self._build_el_data(
cur_lib_focus, DESC_cur, struct_map_cur)
def _iteration_scan_value_to_delta(scan_value, el):
if self.granularity_kind == 'count':
return _count_kind_k_from_fraction(
scan_value, res['k_micro_per_el'][el])
eps_el = (self.eps_noise_per_species or {}).get(
el, self.eps_noise_global)
return float(scan_value) * float(eps_el)
if c_per_el is None:
delta_per_el = {
el: _iteration_scan_value_to_delta(res['delta_star'], el)
for el in self.species
}
else:
delta_per_el = {
el: _iteration_scan_value_to_delta(c_per_el[el], el)
for el in self.species
}
best_judger = _fit_judge_at_delta(self.species, cur_lib_focus,
self.lib_path, desc_per_el_cur,
delta=delta_per_el,
restart_seed=self.seed,
**self._judge_backend_kwargs())
# Judge owns the fitted corpus if it needs it; this temporary
# per-element container has no later consumer.
del desc_per_el_cur
rep_size = len(best_judger.sub_lib_idx)
n_groups = int(best_judger.max_label)
logger.info(f' |S|={rep_size} groups={n_groups}')
# Pull next batch
sample = (sample_next_ood_fn(best_judger)
if sample_next_ood_fn is not None
else sample_next_fn())
if sample is None:
# Record this iteration without next-batch info, then stop
records.append(self._record_dict(
n_iter, cur_lib_focus, DESC_cur, res, rep_size, n_groups,
next_batch_size=0, next_batch_ood=0, next_batch_ood_frac=0.0,
extension_method=extension_method,
c_per_el=c_per_el,
))
if save:
with open(self.output_dir / 'loop_records.json', 'w') as f:
json.dump(records, f, indent=2)
# No next batch will consume this iteration's fitted/stacked
# state. Keep only the list-backed final corpus below.
del best_judger, res, DESC_cur, struct_map_cur
break
if sample_next_ood_fn is not None:
(next_lib_focus, next_desc_list, next_full_structs,
ood_next, ood_frac) = sample
else:
next_lib_focus, next_desc_list, next_full_structs = sample
DESC_next = np.vstack(next_desc_list)
if sample_next_ood_fn is None:
smap_next = np.repeat(np.arange(len(next_lib_focus)),
[len(x) for x in next_desc_list])
_, desc_next_per_el, _ = self._build_el_data(
next_lib_focus, DESC_next, smap_next)
# OOD = focus atoms falling outside every trained per-element cluster
ood_next = 0
for el in self.species:
clu = best_judger.cluster_per_el.get(el)
if clu is None or el not in desc_next_per_el:
continue
embedded_next = clu.embedder.embed(desc_next_per_el[el])
gID_distinct, _ = clu.classify(embedded_next)
ood_next += int(len(gID_distinct[-1]))
del embedded_next
ood_frac = ood_next / len(DESC_next) if len(DESC_next) else 0.0
# The labels/count above are all that survives next-batch OOD
# classification; avoid retaining its descriptor containers.
del desc_next_per_el, smap_next
logger.info(f' next batch: {len(next_lib_focus)} structs / '
f'{len(DESC_next)} focus atoms → OOD={ood_next} '
f'({100*ood_frac:.2f}%)')
# The list-backed batch is accumulated below; this stacked view is
# not part of the next iteration's state.
del DESC_next
records.append(self._record_dict(
n_iter, cur_lib_focus, DESC_cur, res, rep_size, n_groups,
next_batch_size=len(next_lib_focus),
next_batch_ood=ood_next,
next_batch_ood_frac=ood_frac,
extension_method=extension_method,
c_per_el=c_per_el,
))
if save:
with open(self.output_dir / 'loop_records.json', 'w') as f:
json.dump(records, f, indent=2)
# Classification of the next batch is this Judge's last consumer.
# Drop prior-corpus DECAF/raw-descriptor state before allocating the
# next iteration's larger embedding and scan matrices.
del best_judger
del res, DESC_cur, struct_map_cur
# DECAF fits are corpus-specific; stale for the larger next-iter dataset
self._el_fit_cache.clear()
cur_lib_focus += list(next_lib_focus)
cur_focus_desc_list += list(next_desc_list)
cur_full_lib += list(next_full_structs)
self.records = records
# Stash final state — caller can pass to run_final directly
self._loop_final_lib_focus = cur_lib_focus
self._loop_final_desc_list = cur_focus_desc_list
self._loop_final_full_lib = cur_full_lib
logger.info(f'\nLoop done in {len(records)} iterations.')
return records
@staticmethod
def _record_dict(n_iter, cur_lib_focus, DESC_cur, res,
rep_size, n_groups,
next_batch_size, next_batch_ood, next_batch_ood_frac,
extension_method='random', c_per_el=None):
bc = res['best_constrained']
if bc is not None:
L_, R_, fid_ = float(bc['L']), float(bc['R']), float(bc['F'])
else:
L_ = float(res['best_row']['L'])
R_ = float(res['best_row']['R'])
fid_ = float(res['best_row']['F'])
return dict(
iteration = n_iter,
n_training_structs = len(cur_lib_focus),
n_training_focus_atoms = int(DESC_cur.shape[0]),
delta_auto = res['delta_auto'],
delta_star = res['delta_star'],
selection_mode = res['sel_mode'],
L=L_, R=R_, F=fid_,
rep_set_size = int(rep_size),
n_groups = int(n_groups),
k_micro = res['k_micro'],
h_mu = res['h_mu'],
next_batch_size = int(next_batch_size),
next_batch_ood = int(next_batch_ood),
next_batch_ood_frac = float(next_batch_ood_frac),
extension_method = extension_method,
c_per_el = ({el: float(v) for el, v in c_per_el.items()}
if c_per_el is not None else None),
)
# ── Stage 3: final L/R scan (MANDATORY) ─────────────────────────────────
[docs]
def run_final(self,
lib_focus,
desc_array,
full_struct_lib=None,
struct_map=None,
save=True):
"""Final L/R scan on the (full) training corpus.
Stores all state needed by `compute_patient_ood` and `plot_final_fit`.
With ``save=True`` writes hifi_final.npz (plus el_fit_cache.npz when
the scan cache is non-empty), fine_scan.out and hifi_final.png.
"""
if struct_map is None:
sizes = [len(atoms) for atoms in lib_focus]
struct_map = np.repeat(np.arange(len(lib_focus)), sizes)
logger.info('=' * 70)
logger.info('Stage 3: final L/R scan')
logger.info('=' * 70)
res = self._fit_iteration(lib_focus, desc_array, struct_map,
n_iter=None, save=save,
coarse_only=False)
# Persist state for downstream methods
self.lib = lib_focus
self.desc_array = desc_array
self.struct_map = struct_map
self.full_struct_lib = full_struct_lib
self.judge_emb = res['judger']
self.delta_auto = res['delta_auto']
self.mu_labels = res['mu_labels']
self.k_micro = res['k_micro']
self.h_mu = res['h_mu']
self.df_coarse = res['df_coarse']
self.df_fine = res['df_fine']
self.mu_labels_per_el = res['mu_labels_per_el']
self.k_micro_per_el = res['k_micro_per_el']
self.h_mu_per_el = res['h_mu_per_el']
self.embed_op_per_el = res['embed_op_per_el']
self.delta_star_F = float(res['best_row']['c_factor'])
self.delta_star_L095 = (float(res['best_constrained']['c_factor'])
if res['best_constrained'] is not None
else None)
if save:
self._save_hifi_npz(out_name='hifi_final.npz')
self.save_fine_table()
self._plot_FR_panels(
self.df_coarse, self.df_fine,
self.delta_auto,
self.delta_star_L095 or self.delta_star_F,
fname=str(self.output_dir / 'hifi_final.png'),
title='Final',
)
return res
# ── Stage 4: patient OOD per fine δ (optional) ──────────────────────────
[docs]
def compute_patient_ood(self, patient_descriptors, save=True):
"""Refit Judge at each fine δ and count patient atoms in group -1.
Reuses (lib_focus, desc_array, df_fine) from `run_final`.
Updates `self.fine_pat_ood`, saves hifi_final_ood.npz,
and re-renders hifi_final.png with the OOD overlay on the ``|S|`` panel.
"""
if not self._backend_caps.has_classifier:
raise RuntimeError(
'compute_patient_ood requires a backend with a deployable '
f'classifier — {self.cluster_backend} has none (ADR-0009).')
if self.df_fine is None or self.lib is None:
raise RuntimeError('compute_patient_ood: call run_final first.')
logger.info('=' * 70)
logger.info(f'Stage 4: patient OOD per fine δ ({self.n_fine} pts, '
f'n_jobs={self.n_jobs})')
logger.info('=' * 70)
from joblib import Parallel, delayed
_scan_col = 'c_factor' if 'c_factor' in self.df_fine.columns else 'delta'
delta_arr = self.df_fine[_scan_col].values
patient_desc_per_el = _normalise_patient_descriptors(
patient_descriptors, self.species)
_, desc_per_el, _ = self._build_el_data(
self.lib, self.desc_array, self.struct_map)
def _delta_fit(scan_value):
if _scan_col == 'c_factor':
return {el: self._scan_value_to_delta(scan_value, el)
for el in self.species}
return {el: float(scan_value) for el in self.species}
results = Parallel(n_jobs=self.n_jobs, prefer='threads')(
delayed(_ood_worker)(
d, _delta_fit(d), self.species, self.lib, self.lib_path,
desc_per_el, patient_desc_per_el,
self._backend_spec, self.min_cluster_size,
self.eps_noise_per_species,
)
for d in delta_arr
)
delta_to = {r['delta']: r['pat_ood'] for r in results if r is not None}
ood = np.array([delta_to.get(float(d), -1) for d in delta_arr])
self.fine_pat_ood = ood
if save:
self._save_hifi_npz(out_name='hifi_final_ood.npz')
self._plot_FR_panels(
self.df_coarse, self.df_fine,
self.delta_auto,
self.delta_star_L095 or self.delta_star_F,
fname=str(self.output_dir / 'hifi_final.png'),
title='Final',
fine_pat_ood=ood,
)
return ood
# ── Stage 5: plot final fit at chosen δ (MANDATORY) ─────────────────────
def _apply_refined_group_ids(self, j, save):
"""Relabel the final Judge with refined/sorted global ids (issue #40).
Final-fit-only: scans and caches keep backend-native labels. When
``save``, writes ``group_id_mapping.csv`` (one row per group — including
rare-environment groups promoted from the noise slot, flagged by the
``noise_derived`` column, issue #63 — with meta_group_id and adjacency
diagnostics) and ``group_order.npz`` (per-element linkage for post-hoc
meta recuts).
"""
mapping_df = j.apply_refined_group_ids()
# Provenance column: which rows came from promoted noise reps.
if (mapping_df is not None and 'element' in mapping_df.columns
and 'old_group_id' in mapping_df.columns):
def _is_noise(row):
fit = (j.cluster_per_el or {}).get(row['element'])
nd = getattr(fit, 'noise_derived', None)
if nd is None:
return False
nd = np.asarray(nd)
old = int(row['old_group_id'])
return bool(nd[old]) if 0 <= old < len(nd) else False
mapping_df['noise_derived'] = mapping_df.apply(_is_noise, axis=1)
for el, order in j.group_order_per_el.items():
tag = (f'{order.n_meta} meta-groups (silhouette '
f'{order.meta_score:.2f})' if not order.fallback
else 'no macro structure (single meta-group)')
logger.info(f' [{el}] refined group ids: '
f'{len(order.old_of_local)} groups, {tag}')
if save:
csv_path = self.output_dir / 'group_id_mapping.csv'
mapping_df.to_csv(csv_path, index=False)
logger.info(f' Saved {csv_path}')
arrs = {}
for el, order in j.group_order_per_el.items():
arrs[f'linkage_Z_{el}'] = order.linkage_Z
arrs[f'old_of_local_{el}'] = order.old_of_local
arrs[f'meta_of_local_{el}'] = order.meta_of_local
arrs[f'offset_{el}'] = np.array(j.group_id_offset[el])
np.savez(self.output_dir / 'group_order.npz', **arrs)
return mapping_df
def _best_grid_operating_point(self, df_grid):
"""Return {el: c_factor} for the highest F row with L >= l_tol, or None.
When ``l_tol`` is None the L floor is disabled and the max-F row
(best-fidelity point) is picked.
"""
if 'L' not in df_grid.columns:
return None
ok = df_grid if self.l_tol is None else df_grid[df_grid['L'] >= self.l_tol]
if ok.empty:
return None
rec = ok.loc[ok['F'].idxmax()]
result = {el: float(rec[f'c_{el}'])
for el in self.species if f'c_{el}' in rec.index}
return result if len(result) == len(self.species) else None
_best_2d_operating_point = _best_grid_operating_point # deprecated alias
[docs]
def plot_final_fit(self,
delta_pick=None,
patient_descriptors=None,
selection_method=None,
save=True,
save_representatives=True):
"""Refit Judge(s) at `delta_pick` on the full corpus and save the rep set.
`delta_pick` can be a scalar c_factor (1-D path) or a
``{el: c_factor}`` dict (N-D path: one multi-element Judge with
per-element clustering and a joint set-cover). When ``None``, uses the
best grid operating point (``_best_grid_operating_point``) if
``self.df_grid`` yields one, else ``delta_star_L095``, else
``delta_star_F``.
Parameters
----------
delta_pick : float or dict[str, float] or None
Scalar: common c_factor for all elements (1-D path).
Dict: per-element c_factor values in c_factor units (N-D path).
selection_method : {None, 'baseline', 'mu_tiebreak', 'atom_aware'}
Dict mode accepts None, 'baseline' or 'mu_tiebreak'.
save_representatives : bool
Write backend-selected representative structures when true.
Piecing disables this because the runner writes full parents.
Returns
-------
rep_structs : list[ase.Atoms]
rep_idx : list[int]
"""
if self.lib is None:
raise RuntimeError('plot_final_fit: call run_final first.')
# ── Auto-detect operating point ───────────────────────────────────
if delta_pick is None:
_nd_df = self.df_grid
if _nd_df is not None:
pt = self._best_grid_operating_point(_nd_df)
delta_pick = pt if pt is not None else (
self.delta_star_L095 or self.delta_star_F)
else:
delta_pick = self.delta_star_L095 or self.delta_star_F
# ── Type branch ───────────────────────────────────────────────────
_dict_mode = isinstance(delta_pick, dict)
if not _dict_mode:
delta_pick = float(delta_pick)
# ── δ scaling and label ───────────────────────────────────────────
if _dict_mode:
if self.granularity_kind == 'distance' and not self.eps_noise_per_species:
raise RuntimeError(
'plot_final_fit dict mode requires eps_noise_per_species; '
'run calibrate_eps_noise or restore_state from a new-format NPZ.')
delta_bw_per_el = {
el: self._scan_value_to_delta(c, el)
for el, c in delta_pick.items()
}
self.delta_star_per_el = {el: float(c) for el, c in delta_pick.items()}
if self.granularity_kind == 'count':
label = 'grid optimum: ' + ' '.join(
f'{el} f={delta_pick[el]:.4f} k={int(delta_bw_per_el[el])}'
for el in self.species if el in delta_pick
)
else:
label = 'grid optimum: ' + ' '.join(
f'{el} c={delta_pick[el]:.4f} δ={delta_bw_per_el[el]:.5f}'
for el in self.species if el in delta_pick
)
else:
# delta_star_F/F095 store dimensionless c_factor; fit_decaf needs
# actual δ = c × ε_noise. Old-format NPZs (scan axis 'delta')
# stored raw δ directly — leave unchanged.
_scan_col = ('c_factor' if (self.df_fine is not None
and 'c_factor' in self.df_fine.columns)
else 'delta')
if _scan_col != 'c_factor':
delta_bw = delta_pick
elif self.granularity_kind == 'count':
delta_bw = float(delta_pick) # fraction axis; per-el k below
else:
_eps = self.eps_noise_global
delta_bw = delta_pick * _eps
is_constrained = (self.delta_star_L095 is not None
and abs(delta_pick - self.delta_star_L095) < 1e-9)
if self.granularity_kind == 'count':
_pt = f'f={delta_pick:.5f} (k = f·k_micro per element)'
else:
_pt = f'c={delta_pick:.5f} δ={delta_bw:.5f}'
label = (f'L≥{self.l_tol} optimum {_pt}' if is_constrained
else f'user {_pt}')
if selection_method is not None and selection_method not in (
'baseline', 'mu_tiebreak', 'atom_aware'):
raise ValueError(
f"selection_method must be None, 'baseline', 'mu_tiebreak', "
f"or 'atom_aware'; got {selection_method!r}")
if _dict_mode and selection_method not in (None, 'baseline', 'mu_tiebreak'):
raise ValueError(
"dict-mode plot_final_fit only supports selection_method=None, 'baseline', or 'mu_tiebreak'")
logger.info('=' * 70)
logger.info(f'Stage 5: plot_final_fit ({label}) '
f'[selection={selection_method or "baseline"}]')
logger.info('=' * 70)
if _dict_mode:
# ── Single multi-element Judge ────────────────────────────────
_, desc_per_el, _ = self._build_el_data(
self.lib, self.desc_array, self.struct_map)
counts_L_per_el = {}
if self.mu_labels_per_el:
for el in self.species:
if el in self.mu_labels_per_el:
# minlength: see the matching comment in
# scan_cfactor_grid (issue #82 codex-review).
counts_L_per_el[el] = np.bincount(
self.mu_labels_per_el[el],
minlength=(self.k_micro_per_el or {}).get(el, 0))
j = _fit_judge_at_delta(
self.species, self.lib, self.lib_path,
desc_array=desc_per_el,
delta=delta_bw_per_el,
mu_per_el=self.mu_labels_per_el if self.mu_labels_per_el else None,
counts_L_per_el=counts_L_per_el if counts_L_per_el else None,
use_mu_tiebreak=selection_method == 'mu_tiebreak',
restart_seed=self.seed,
**self._judge_backend_kwargs(),
)
rep_idx = sorted(j.sub_lib_idx)
n_groups = j.n_groups
for el in self.species:
if el in j.cluster_per_el:
n_el_grp = int(j.cluster_per_el[el].n_groups)
n_el_reps = sum(1 for i in rep_idx
if j.per_el_group_libs[el][i])
logger.info(f' [{el}] |groups|={n_el_grp} |S_el|={n_el_reps}')
logger.info(f' Joint set-cover: |S|={len(rep_idx)} '
f'total (el,gid) pairs={n_groups}')
ood_idx = []
n_patient_atoms = 0
ood_per_el = {}
if patient_descriptors is not None:
patient_desc_per_el = _normalise_patient_descriptors(
patient_descriptors, self.species)
n_ood, n_patient_atoms, ood_per_el = _classify_patient_ood(
j, patient_desc_per_el, self.species)
ood_idx = list(range(n_ood))
self._apply_refined_group_ids(j, save)
_label_safe = '_'.join(
f'{el}{delta_pick[el]:.4f}' for el in sorted(delta_pick))
out_xyz = self.output_dir / f'final_fit_reps_grid_{_label_safe}.xyz'
out_png_stem = self.output_dir / f'final_fit_grid_{_label_safe}'
if save and save_representatives and self.full_struct_lib is not None:
rep_struct = _build_reps_with_group_id(
rep_idx, j, self.full_struct_lib, self.species)
write(str(out_xyz), rep_struct)
logger.info(f' Saved {out_xyz}')
if save:
annotation = (f'L={j.L:.4f} R={j.R:.4f} F={j.F:.4f} |S|={len(rep_idx)}'
if j.L is not None else f'|S|={len(rep_idx)}')
for el in self.species:
if el not in j.cluster_per_el:
continue
el_png = str(out_png_stem) + f'_{el}.png'
self._plot_final_fit_panels(
cluster=j.cluster_per_el[el],
per_struct_groups=j.per_el_group_libs[el],
sub_lib_idx=j.sub_lib_idx,
annotation=annotation,
label=(f'{el} f={delta_pick[el]:.4f} k={int(delta_bw_per_el[el])}'
if self.granularity_kind == 'count'
else f'{el} c={delta_pick[el]:.4f} δ={delta_bw_per_el[el]:.5f}'),
fname=el_png,
group_order=j.group_order_per_el.get(el),
)
logger.info(f' Saved {el_png}')
else:
# ── Scalar path ───────────────────────────────────────────────
# Clear any per-element c_factors from a prior dict-mode call so
# extract_rep_surroundings falls back to the 1-D c_star.
self.delta_star_per_el = None
_, desc_per_el, _ = self._build_el_data(
self.lib, self.desc_array, self.struct_map)
if self.granularity_kind == 'count' or self.eps_noise_per_species:
delta_bw_per_el_1d = {
el: self._scan_value_to_delta(delta_pick, el)
for el in self.species
}
else:
delta_bw_per_el_1d = {el: delta_bw for el in self.species}
counts_L_per_el = {}
if self.mu_labels_per_el:
for el in self.species:
if el in self.mu_labels_per_el:
# minlength: see the matching comment in
# scan_cfactor_grid (issue #82 codex-review).
counts_L_per_el[el] = np.bincount(
self.mu_labels_per_el[el],
minlength=(self.k_micro_per_el or {}).get(el, 0))
j = _fit_judge_at_delta(
self.species, self.lib, self.lib_path,
desc_array=desc_per_el,
delta=delta_bw_per_el_1d,
mu_per_el=self.mu_labels_per_el if self.mu_labels_per_el else None,
counts_L_per_el=counts_L_per_el if counts_L_per_el else None,
use_mu_tiebreak=selection_method == 'mu_tiebreak',
restart_seed=self.seed,
**self._judge_backend_kwargs(),
)
if selection_method == 'atom_aware':
sizes = np.bincount(self.struct_map, minlength=len(self.lib))
j.sub_lib_idx = frozenset(
int(i) for i in atom_aware_greedy(j, sizes))
j.recompute_hifi_metrics()
logger.info(f' atom_aware: |S|={len(j.sub_lib_idx)}')
rep_idx = sorted(j.sub_lib_idx)
n_groups = j.n_groups
ood_idx = []
n_patient_atoms = 0
ood_per_el = {}
if patient_descriptors is not None:
patient_desc_per_el = _normalise_patient_descriptors(
patient_descriptors, self.species)
n_ood, n_patient_atoms, ood_per_el = _classify_patient_ood(
j, patient_desc_per_el, self.species)
ood_idx = list(range(n_ood))
self._apply_refined_group_ids(j, save)
out_xyz = self.output_dir / f'final_fit_reps_delta{delta_pick:.4f}.xyz'
if save and save_representatives and self.full_struct_lib is not None:
rep_struct = _build_reps_with_group_id(
rep_idx, j, self.full_struct_lib, self.species)
write(str(out_xyz), rep_struct)
logger.info(f' Saved {out_xyz}')
if save:
annotation_parts = []
if j.L is not None:
annotation_parts.append(
f'L={j.L:.4f} R={j.R:.4f} F={j.F:.4f}')
annotation_parts.append(f'|S|={len(rep_idx)}')
annotation = '\n'.join(annotation_parts)
for el in self.species:
if el not in j.cluster_per_el:
continue
out_png = self.output_dir / f'final_fit_delta{delta_pick:.4f}_{el}.png'
self._plot_final_fit_panels(
cluster=j.cluster_per_el[el],
per_struct_groups=j.per_el_group_libs[el],
sub_lib_idx=j.sub_lib_idx,
annotation=annotation,
label=f'{el} {label}',
fname=str(out_png),
group_order=j.group_order_per_el.get(el),
)
logger.info(f' Saved {out_png}')
if self.full_struct_lib is not None:
rep_structs = _build_reps_with_group_id(
rep_idx, j, self.full_struct_lib, self.species)
else:
rep_structs = [self.lib[i] for i in rep_idx]
# Expose the fitted Judge so callers (e.g. the piecing assign-back path
# in runner) can read per_el_group_libs at the chosen operating point.
self._last_final_judge = j
# Persist it too: everything else the run writes (hifi_final.npz,
# el_fit_cache.npz, representatives.xyz) is a projection of this object,
# and without the pickle it dies with the process. Both the 1-D and
# dict branches converge here, after _apply_refined_group_ids, so the
# saved fit carries the same group ids as group_id_mapping.csv.
if save and self.save_judge:
out_pkl = self.output_dir / 'judge_final.pkl'
try:
j.save(out_pkl)
logger.info(
f' Saved {out_pkl} '
f'({out_pkl.stat().st_size / 1e6:.1f} MB)')
except Exception as exc:
# A finished selection is the deliverable; an unpicklable
# object (a custom backend holding a live handle, say) must not
# discard it. PickleIO.save writes through a temporary sibling,
# so a failure here leaves no partial file and any artifact from
# a previous run is still intact.
logger.warning(f' {out_pkl} not written: {exc}')
logger.debug('')
logger.info('-' * 70)
logger.info('Final-fit summary')
logger.info('-' * 70)
logger.info(f' {label}')
logger.info(f' |S|: {len(rep_idx)}')
logger.info(f' n_groups: {n_groups}')
if patient_descriptors is not None:
frac = (100 * len(ood_idx) / n_patient_atoms
if n_patient_atoms else 0.0)
logger.info(f' patient OOD: {len(ood_idx)} / {n_patient_atoms} '
f'({frac:.2f}%)')
if ood_per_el:
logger.info(f' patient OOD per element: {ood_per_el}')
return rep_structs, rep_idx