"""
atom_hifi/training/judge.py — Judge: atomic-environment clustering orchestrator.
Public API
----------
Classes
~~~~~~~
Judge
DECAF local atomic state cover fitter. Wraps DECAF clustering and the
structure-wise greedy set-cover that selects representative structures.
Key methods:
fit_decaf(desc_array=, delta=, save=, ...) — DECAF clustering + set-cover
save(path) / Judge.load(path) — pickle persistence, from PickleIO. The
raw descriptor block ``training_set`` is NOT saved (see _PICKLE_DROP);
``attach_descriptors`` puts it back when a reloaded Judge needs
``extract_rep_surroundings``.
Key attributes after fitting:
.sub_lib_idx — indices of representative structures in the library
.rep_lib — representative Atoms; set (with .rep_path) only by fit_decaf(save=True)
.max_label — number of distinct environment groups
.cluster — ClusterFit for species[0] (.embedder, .classify())
Note
~~~~
Issue #69 split this module. ``Judge`` orchestrates; the pieces it orchestrates
now live next door:
* :mod:`atom_hifi.training.selection` — set cover (``find_min_covering_subset``,
``atom_aware_greedy``, ``length_bool``)
* :mod:`atom_hifi.training.calibration` — σ_thermal, ε_noise, microstate reference
* :mod:`atom_hifi.training.cluster_base` / ``cluster_backends`` — clustering
* :mod:`atom_hifi.training.group_order` — group ordering
HiFiWorkflow and the parallel worker functions live in
atom_hifi.training.hifi_workflow.
"""
# Standard library
import logging
import sys as _sys
import types as _types
import math
import random
from copy import deepcopy
from pathlib import Path
from typing import List, Optional
# Scientific stack
logger = logging.getLogger(__name__)
import numpy as np
import pandas as pd
# ASE
from ase import Atoms
from ase.io import read, write
from ase.io.trajectory import Trajectory
# Internal
from atom_hifi.training.cluster_base import ClusterBackend, PickleIO
from atom_hifi.training.cluster_backends import (
BACKEND_REGISTRY,
fit_cluster_backend,
)
from atom_hifi.training.group_order import refined_group_order
from atom_hifi.training.selection import (
_hifi_metrics_sparse,
_one_restart,
find_min_covering_subset,
length_bool,
)
# ── Re-exports at the historical import path (issue #69) ──────────────────────
# ``judge.py`` declared these as public API and downstream code imports them
# from here; the split must not break that.
from atom_hifi.training import calibration as _calibration
from atom_hifi.training.selection import atom_aware_greedy # noqa: F401
from atom_hifi.training.calibration import ( # noqa: F401
SIGMA_THERMAL_OVERRIDE,
SIGMA_THERMAL_REF_MASS,
SIGMA_THERMAL_REF_TEMP,
SIGMA_THERMAL_REF_VALUE,
ED_SOAP_R_CUT,
fit_microstate_reference,
measure_eps_noise,
sigma_thermal,
)
# ``sigma_thermal`` now resolves its globals in ``calibration``, so a plain
# re-export would make ``judge.SIGMA_THERMAL_OVERRIDE = {...}`` — valid before
# #69 and documented expert API — a silent no-op returning the default. Item
# assignment (``judge.SIGMA_THERMAL_OVERRIDE['C'] = ...``) stays coherent
# because the dict is shared by reference, but REBINDING would fork it.
# Forward the rebind to the owning module so both spellings keep working.
_FORWARDED_TO_CALIBRATION = frozenset({
'SIGMA_THERMAL_OVERRIDE', 'SIGMA_THERMAL_REF_VALUE',
'SIGMA_THERMAL_REF_MASS', 'SIGMA_THERMAL_REF_TEMP',
})
class _JudgeModule(_types.ModuleType):
def __getattribute__(self, name):
if name in _FORWARDED_TO_CALIBRATION:
return getattr(_calibration, name)
return super().__getattribute__(name)
def __setattr__(self, name, value):
if name in _FORWARDED_TO_CALIBRATION:
setattr(_calibration, name, value)
return
super().__setattr__(name, value)
def __delattr__(self, name):
if name in _FORWARDED_TO_CALIBRATION:
delattr(_calibration, name)
if name in super().__getattribute__('__dict__'):
super().__delattr__(name)
return
super().__delattr__(name)
_judge_module = _sys.modules[__name__]
_judge_module.__class__ = _JudgeModule
# The first import assigned copied values before the module gained its proxy
# type. Remove those hidden stale entries; reads now resolve at the owner.
for _forwarded_name in _FORWARDED_TO_CALIBRATION:
_judge_module.__dict__.pop(_forwarded_name, None)
DESCRIPTORS = ['SOAP', 'ACE', 'ED_SOAP']
# ── Judge ─────────────────────────────────────────────────────────────────────
[docs]
class Judge(PickleIO):
"""DECAF local atomic state clustering and representative-set selection.
The live path is ``fit_decaf``: descriptors are supplied by the caller,
DECAF identifies local atomic states, and greedy set-cover selects
structures that cover those states.
Persistence comes from :class:`~atom_hifi.training.cluster_base.PickleIO`
and is therefore identical for every backend.
"""
lib0: List[Atoms]
#: ``training_set`` is the caller's raw descriptor block, kept only so
#: ``extract_rep_surroundings`` can re-embed it. It dwarfs everything the
#: fit actually produced (259 MB of descriptors vs a few MB of labels,
#: centres and group libraries on the mp-NiO full run), it is an input
#: rather than a result, and the run already persists it as
#: ``desc_lib.pkl`` — so it is dropped and ``attach_descriptors`` puts it
#: back on demand.
_PICKLE_DROP = ('training_set',)
def __init__(
self,
species: List[str],
lib0: Optional[List[Atoms]] = None,
lib_path: Optional[str] = None,
selection: Optional[List[str]] = None,
descriptor: str = 'SOAP',
cluster_backend: str = 'decaf',
min_cluster_size: int = 2,
**soap_params,
):
"""
Parameters
----------
species : list of str
Chemical elements to consider (e.g. ``['Pd', 'Mg', 'O']``).
lib0 : list of Atoms, optional
Training structure library supplied directly; takes precedence
over ``lib_path`` when both are given.
lib_path : str, optional
Path to a ``.xyz`` or ``.traj`` training library, read only when
``lib0`` is None. Its directory receives ``representatives.xyz``
on ``fit_decaf(save=True)`` (the current directory when
``lib_path`` is not given).
descriptor : str
Must be ``'ED_SOAP'`` for ``fit_decaf`` (asserted). The ``'SOAP'``
default is kept for compatibility, so pass
``descriptor='ED_SOAP'``.
cluster_backend : {'decaf', 'kmeans', 'hdbscan'} or ClusterBackend
Clusterer backend name, instance, or subclass (ADR-0009). ``decaf``
and ``kmeans`` provide distance-abstaining classifiers; ``hdbscan``
provides frozen density hierarchy assignment with rejection.
min_cluster_size : int
HDBSCAN minimum cluster size; ignored by other backends.
**soap_params
Accepted for compatibility with older callers; unused by the current
DECAF path.
"""
# May be a name or a live ClusterBackend (a custom plug travels as an
# object so it survives the trip into a process worker — issue #69).
is_backend_class = (
isinstance(cluster_backend, type)
and issubclass(cluster_backend, ClusterBackend)
)
if not isinstance(cluster_backend, ClusterBackend) and not is_backend_class:
if cluster_backend not in BACKEND_REGISTRY:
raise ValueError(
f'Unknown cluster backend {cluster_backend!r}; choose from '
f'{tuple(BACKEND_REGISTRY)} (register a custom ClusterBackend '
'first — Python API only, issue #69).')
self.descriptor = descriptor
self.species = species
self.path = Path(lib_path) if lib_path else Path('.')
self.delta_per_el = {} # granularity used per element in last fit_decaf call
self.cluster_backend = cluster_backend
self.min_cluster_size = int(min_cluster_size)
if lib_path is None and lib0 is None:
raise ValueError("Either lib0 or lib_path must be provided.")
elif lib0 is not None:
self.lib0 = lib0
elif lib0 is None and lib_path is not None:
self.path = Path(lib_path)
if self.path.suffix == '.xyz':
self.lib0 = read(self.path, index=':')
elif self.path.suffix == '.traj':
self.lib0 = Trajectory(self.path).atoms
else:
raise ValueError("Unsupported file format. Please provide an .xyz or .traj file.")
[docs]
def fit_decaf(self, desc_array=None, delta=None,
mu_per_el=None, counts_L_per_el=None,
t_embedded=None,
metric_func=length_bool, n_jobs_restarts=1, save=False,
cluster_cache_per_el=None, c_vals_dict=None,
use_mu_tiebreak=None, restart_seed=None,
eps_noise_per_el=None,
ace_array=None):
"""Fit DECAF clustering for all species and find the joint minimum covering subset.
Parameters
----------
desc_array : ndarray or dict[str, ndarray]
Bare ndarray → single-element mode (species[0]).
``{el: ndarray}`` dict → per-element mode; each array has shape
``(N_el_atoms, D)`` where atoms are ordered by structure index.
delta : float or dict[str, float]
Per-element backend granularity passed to ``fit_cluster_backend``:
bandwidth δ for decaf, cluster count k for kmeans,
``cluster_selection_epsilon`` for hdbscan. Scalar only in
single-element mode; dict ``{el: value}`` otherwise.
mu_per_el : dict[str, ndarray], optional
Per-atom microstate labels ``{el: (N_el_atoms,) int array}``.
When provided (with ``counts_L_per_el``), enables the cross-restart
L/R fidelity tiebreak and computes ``self.L``, ``self.R``, ``self.F``.
counts_L_per_el : dict[str, ndarray], optional
``{el: np.bincount(mu_per_el[el])}`` — precomputed microstate
distribution for the full library. Required alongside ``mu_per_el``.
use_mu_tiebreak : bool or None, optional
True uses μ coverage for greedy/restart ties. False gives pure
baseline selection while still computing L/R/F. None preserves
legacy behaviour: no greedy μ tiebreak, but keep the cross-restart
L/R tie when μ data is present.
restart_seed : int, optional
Seeds the shuffled set-cover restarts and the backend fit, which
receives ``fit_cluster_backend(seed=...)`` (42 when None); among
the built-in backends only kmeans uses it, and custom backends
receive it. With ``n_jobs_restarts == 1`` all restarts draw from
one ``random.Random(restart_seed)`` stream, or from the global
``random`` module when None. Otherwise restart ``i`` uses
``random.Random(int(restart_seed) + i)``, or ``random.Random(i)``
when None. Only with an explicit ``restart_seed`` do the two
modes share the first restart's shuffle, so the selection can
depend on ``n_jobs_restarts``.
t_embedded : ndarray or dict[str, ndarray], optional
Pre-computed PCA embedding for SharedMemory fast-path.
Bare ndarray → single-element; dict → per-element.
eps_noise_per_el : dict[str, float], optional
Per-element ε_noise; required when the backend emits a trailing
noise slot: it is the dedup δ used to promote those atoms to
first-class groups inside ``fit_cluster_backend`` (issue #63).
"""
if ace_array is not None:
import warnings
warnings.warn(
"ace_array is deprecated; use desc_array instead.",
DeprecationWarning, stacklevel=2)
desc_array = ace_array
if desc_array is None:
raise TypeError("fit_decaf() missing required argument: 'desc_array'")
assert self.descriptor in DESCRIPTORS, f"Unsupported model. Choose from {DESCRIPTORS}"
assert self.descriptor in ('ED_SOAP', 'APE'), "Please use fit method for non-DECAF models."
self.lib = self.lib0
lib_len = len(self.lib)
# ── Mode detection: wrap bare arrays into dicts ────────────────────
_multi = isinstance(desc_array, dict)
if not _multi:
el0 = self.species[0]
desc_array = {el0: desc_array}
delta = {el0: float(delta)}
t_embedded = {el0: t_embedded}
else:
if not isinstance(delta, dict):
raise ValueError('desc_array is a dict but delta is not; '
'pass delta as {el: float}')
t_embedded = t_embedded if isinstance(t_embedded, dict) else {}
has_mu_metrics = mu_per_el is not None and counts_L_per_el is not None
if use_mu_tiebreak is True and not has_mu_metrics:
raise ValueError(
'use_mu_tiebreak=True requires mu_per_el and counts_L_per_el')
greedy_mu_tiebreak = use_mu_tiebreak is True
restart_hifi_tiebreak = has_mu_metrics and use_mu_tiebreak is not False
# ── Per-element backend fitting ────────────────────────────────────
_backend = getattr(self, 'cluster_backend', 'decaf')
self.cluster_per_el = {}
self.per_el_group_libs = {el: [[] for _ in range(lib_len)]
for el in self.species}
all_el_gid_pairs = set()
merged_groups = [[] for _ in range(lib_len)]
# Per-element focus-atom counts: used by fresh fits, the joint-mu setup,
# and (as _el_sizes_per_el) by attach_descriptors / extract_rep_surroundings
el_sizes_per_el = {
el: np.array([sum(1 for a in s if a.symbol == el) for s in self.lib])
for el in self.species
}
self._el_sizes_per_el = el_sizes_per_el
for el in self.species:
delta_el = float(delta[el])
desc_el = desc_array[el]
el_sizes = el_sizes_per_el[el]
# Cache-hit path: reuse pre-computed per-element DECAF result
c_val_el = float(c_vals_dict[el]) if c_vals_dict else None
cache_key = (el, round(c_val_el, 10)) if c_val_el is not None else None
if cluster_cache_per_el and cache_key and cache_key in cluster_cache_per_el:
# entry = (pegl, all_gid[, resolved granularity]); only pegl is
# read here. In-repo scans pass only entries whose recorded
# granularity serves this fit (equal to its δ, or, for distance
# kind, a request for the cover's own cache key;
# hifi_workflow._fit_cfactor_worker_body); direct callers own
# that check.
entry = cluster_cache_per_el[cache_key]
pegl = entry[0]
self.per_el_group_libs[el] = pegl
self.cluster_per_el[el] = None
self.delta_per_el[el] = delta_el
for i_full, gids in enumerate(pegl):
for gid in gids:
if gid < 0:
continue
pair = (el, int(gid))
merged_groups[i_full].append(pair)
all_el_gid_pairs.add(pair)
logger.debug(f' cache hit: el={el} c={c_val_el:.4f} delta={delta_el:.6f}')
continue
# Fresh backend fit
t_emb_el = t_embedded.get(el)
fit = fit_cluster_backend(
_backend, desc_el, delta_el,
t_embedded=t_emb_el,
min_cluster_size=getattr(self, 'min_cluster_size', 2),
seed=restart_seed if restart_seed is not None else 42,
eps_noise=(eps_noise_per_el or {}).get(el))
logger.debug(f'Number of groups {fit.n_groups}')
logger.debug(f'Choice of granularity {fit.granularity}')
self.cluster_per_el[el] = fit
self.delta_per_el[el] = delta_el
assert el_sizes.sum() == desc_el.shape[0], (
f'desc_array[{el!r}] has {desc_el.shape[0]} rows but lib has '
f'{el_sizes.sum()} atoms of {el!r}')
atomic_tags = fit.labels
max_label_el = int(fit.n_groups)
logger.debug(f'max label: {max_label_el}')
# Noise reps were deduped and promoted to real groups inside
# fit_cluster_backend (issue #63), so there is no side-channel here:
# every atom carries a coverable (el, gid) pair and the ordinary
# set-cover picks up rare environments like any other group.
if fit.noise_derived is not None and fit.noise_derived.any():
logger.info(
f' [{el}] {int(fit.noise_derived.sum())} rare-environment '
f'groups promoted from the noise slot '
f'({max_label_el - int(fit.noise_derived.sum())} density groups)')
# Distribute to full-lib per-structure lists
N = 0
for i_full, n_el in enumerate(el_sizes):
group_ids = atomic_tags[N:N + n_el].tolist()
self.per_el_group_libs[el][i_full] = group_ids
for gid in group_ids:
if gid < 0:
continue
pair = (el, int(gid))
merged_groups[i_full].append(pair)
all_el_gid_pairs.add(pair)
N += n_el
# ── Joint group space ──────────────────────────────────────────────
self.per_structure_group_libs = merged_groups # list of (el,gid) tuples
self.n_groups = len(all_el_gid_pairs)
group_range = list(all_el_gid_pairs)
# Joint alias of n_groups (public attribute, see module docstring); read
# by HiFiWorkflow.run_loop and atom_aware_greedy's int-tag domain check
self.max_label = self.n_groups
# ── Cross-restart L/R fidelity tiebreak setup (joint microstate space) ─────
_per_struct_mu_joint = None
_per_struct_mu_sets = None
_k_joint = _H_L_joint = _N_L_joint = _k_occ_L_joint = None
_best_F = -1.0
self._hifi_metric_state = None
if has_mu_metrics:
_offset_per_el = {}
_offset = 0
for el in self.species:
_offset_per_el[el] = _offset
_offset += len(counts_L_per_el[el])
_k_joint = _offset
# Build per-structure sparse mu dicts in the joint space
_per_struct_mu_joint = [{} for _ in range(lib_len)]
for el in self.species:
mu_el = mu_per_el[el]
el_sizes = el_sizes_per_el[el]
struct_map_el = np.repeat(np.arange(lib_len), el_sizes)
offset_el = _offset_per_el[el]
for atom_i in range(len(mu_el)):
s_i = int(struct_map_el[atom_i])
joint_m = int(mu_el[atom_i]) + offset_el
d = _per_struct_mu_joint[s_i]
d[joint_m] = d.get(joint_m, 0) + 1
_per_struct_mu_sets = [
frozenset(d.keys()) for d in _per_struct_mu_joint
]
counts_L_joint = np.concatenate(
[counts_L_per_el[el] for el in self.species])
_k_occ_L_joint = int((counts_L_joint > 0).sum())
_N_L_joint = int(counts_L_joint.sum())
occ_L = counts_L_joint[counts_L_joint > 0]
p_L = occ_L / _N_L_joint
_H_L_joint = -float(np.dot(p_L, np.log(p_L)))
self._hifi_metric_state = (
_per_struct_mu_joint, _k_joint, _H_L_joint,
_N_L_joint, _k_occ_L_joint)
# ── Joint set-cover with restarts ──────────────────────────────────
tiebreak_microstates = (
_per_struct_mu_sets if greedy_mu_tiebreak else None)
_, sub_lib_idx = find_min_covering_subset(
merged_groups, group_range,
struct_microstates=tiebreak_microstates)
min_sub_len = len(sub_lib_idx)
min_sub_idx = list(sub_lib_idx)
if _per_struct_mu_joint is not None:
_best_F = _hifi_metrics_sparse(
min_sub_idx, _per_struct_mu_joint,
_k_joint, _H_L_joint, _N_L_joint, _k_occ_L_joint)
n_restarts = round(1.5 * math.log2(lib_len))
if n_jobs_restarts == 1:
rng = random.Random(restart_seed) if restart_seed is not None else None
for _ in range(n_restarts):
combined = list(zip(deepcopy(merged_groups), range(lib_len)))
if rng is None:
random.shuffle(combined)
else:
rng.shuffle(combined)
shuf_groups, shuf_index = zip(*combined)
shuf_microstates = None
if tiebreak_microstates is not None:
shuf_microstates = [
tiebreak_microstates[shuf_index[i]]
for i in range(lib_len)
]
_, shuf_idx = find_min_covering_subset(
list(shuf_groups), group_range, metric_func=metric_func,
struct_microstates=shuf_microstates)
real_idx = [shuf_index[k] for k in shuf_idx]
if len(shuf_idx) < min_sub_len:
min_sub_len = len(shuf_idx)
min_sub_idx = real_idx
if _per_struct_mu_joint is not None:
_best_F = _hifi_metrics_sparse(
min_sub_idx, _per_struct_mu_joint,
_k_joint, _H_L_joint, _N_L_joint, _k_occ_L_joint)
logger.debug(f'Updating lowest min_sub_len: {min_sub_len}')
elif (len(shuf_idx) == min_sub_len
and restart_hifi_tiebreak
and _per_struct_mu_joint is not None):
fr = _hifi_metrics_sparse(
real_idx, _per_struct_mu_joint,
_k_joint, _H_L_joint, _N_L_joint, _k_occ_L_joint)
if fr > _best_F:
min_sub_idx = real_idx
_best_F = fr
else:
from joblib import Parallel, delayed
seed0 = int(restart_seed) if restart_seed is not None else 0
results = Parallel(n_jobs=n_jobs_restarts, prefer='threads')(
delayed(_one_restart)(
merged_groups, lib_len, group_range, metric_func,
seed=seed0 + i,
struct_microstates=tiebreak_microstates)
for i in range(n_restarts)
)
for r_len, _, r_idx in results:
if r_len < min_sub_len:
min_sub_len = r_len
min_sub_idx = r_idx
if _per_struct_mu_joint is not None:
_best_F = _hifi_metrics_sparse(
min_sub_idx, _per_struct_mu_joint,
_k_joint, _H_L_joint, _N_L_joint, _k_occ_L_joint)
elif (r_len == min_sub_len
and restart_hifi_tiebreak
and _per_struct_mu_joint is not None):
fr = _hifi_metrics_sparse(
r_idx, _per_struct_mu_joint,
_k_joint, _H_L_joint, _N_L_joint, _k_occ_L_joint)
if fr > _best_F:
min_sub_idx = r_idx
_best_F = fr
logger.debug(f'Find the minimum covering subset, size: {min_sub_len}')
self.sub_lib_idx = sorted(min_sub_idx)
logger.debug(f'Number of representatives: {len(self.sub_lib_idx)}')
logger.debug(f'Number of atoms in representatives: {sum(len(self.lib[i]) for i in self.sub_lib_idx)}')
# ── Joint L / R / F (= L/R) ─────────────────────────────────────────
if self._hifi_metric_state is not None:
self.recompute_hifi_metrics()
else:
self.L = self.R = self.F = None
# Store raw descriptors for extract_rep_surroundings
self.training_set = desc_array.get(self.species[0])
if save:
rep_path = str(self.path.parent / 'representatives.xyz')
rep_lib = [self.lib[idx] for idx in self.sub_lib_idx]
self.rep_path = rep_path
self.rep_lib = rep_lib
write(rep_path, rep_lib)
@property
def cluster(self):
"""Fitted ClusterFit of species[0] (any backend).
None when that element's cover came from ``cluster_cache_per_el``.
"""
return self.cluster_per_el[self.species[0]]
[docs]
def recompute_hifi_metrics(self):
"""Recompute L/R/F for the current ``sub_lib_idx``."""
state = getattr(self, '_hifi_metric_state', None)
if state is None:
self.L = self.R = self.F = None
return
per_struct_mu, k_joint, h_L_joint, N_L_joint, k_occ_L_joint = state
counts_S = np.zeros(k_joint, dtype=np.int64)
for i in self.sub_lib_idx:
for m, v in per_struct_mu[int(i)].items():
counts_S[m] += v
occ_S = counts_S[counts_S > 0]
if len(occ_S) == 0:
self.L = 0.0
self.R = np.inf
self.F = 0.0
return
p_S = occ_S / occ_S.sum()
H_S = -float(np.dot(p_S, np.log(p_S)))
k_S = int(len(occ_S))
N_S = int(occ_S.sum())
self.L = min(H_S / h_L_joint, 1.0) if h_L_joint > 0 else 0.0
self.R = ((N_S / k_S) / (N_L_joint / k_occ_L_joint)
if k_S > 0 else np.inf)
self.F = self.L / self.R if np.isfinite(self.R) and self.R > 0 else 0.0
# save() / load() come from PickleIO, so every backend persists the same
# way and `_PICKLE_DROP` above states the footprint in one place.
[docs]
def save(self, path):
"""Save a deployment-ready Judge without deferred HDBSCAN PCA fitting.
The scan fast path may carry a lazy PCA projection. Materialize that
existing projection on explicit save, never during generic pickle/IPC
or load. Selection-only caches still do not become deployable fits.
"""
for fit in getattr(self, 'cluster_per_el', {}).values():
if fit is not None and getattr(fit, 'hdbscan_prediction', None) is not None:
_ = fit.embedder.embed_op
super().save(path)
[docs]
def attach_descriptors(self, desc_array):
"""Restore ``training_set`` on a Judge reloaded from a pickle.
``save`` drops the raw descriptor block (see ``_PICKLE_DROP``); this
puts it back so ``extract_rep_surroundings`` works again. Pass the
SAME descriptors the fit consumed — the run's ``desc_lib.pkl`` filtered
to ``species[0]``. Row count is checked against the fit's own
per-element atom counts, because silently re-embedding a different
descriptor block would produce plausible, wrong environments.
Parameters
----------
desc_array : ndarray or dict[str, ndarray]
Descriptor rows for ``species[0]``; a dict is indexed by that
element.
Returns
-------
Judge
``self``, so a load can be chained:
``Judge.load(p).attach_descriptors(d)``.
"""
el0 = self.species[0]
desc = desc_array[el0] if isinstance(desc_array, dict) else desc_array
desc = np.asarray(desc)
if desc.ndim != 2:
raise ValueError(
f'attach_descriptors: expected a 2-D (n_atoms, n_features) '
f'block for {el0!r}, got shape {desc.shape}. desc_lib.pkl holds '
'one array PER STRUCTURE — stack it and keep the rows whose '
'symbol is the focus element.')
# Row count: prefer the fit's own per-element counts, fall back to
# recounting the library, so a Judge reconstructed without
# `_el_sizes_per_el` is still checked rather than silently accepted.
sizes = getattr(self, '_el_sizes_per_el', {}).get(el0)
if sizes is None and getattr(self, 'lib', None) is not None:
sizes = [sum(1 for a in s if a.symbol == el0) for s in self.lib]
if sizes is not None:
n_expected = int(np.sum(sizes))
if desc.shape[0] != n_expected:
raise ValueError(
f'attach_descriptors: got {desc.shape[0]} rows for {el0!r} '
f'but this fit clustered {n_expected} — these are not the '
'descriptors it was fitted on.')
# Feature width: the embedder's operator is what re-embedding multiplies
# by, so a mismatch here would otherwise surface as a matmul error deep
# inside extract_rep_surroundings.
fit = (self.cluster_per_el or {}).get(el0)
embed_op = getattr(getattr(fit, 'embedder', None), '_embed_op', None)
if embed_op is not None and desc.shape[1] != embed_op.shape[0]:
raise ValueError(
f'attach_descriptors: {el0!r} descriptors have '
f'{desc.shape[1]} features but this fit embeds '
f'{embed_op.shape[0]} — these are not the descriptors it was '
'fitted on.')
self.training_set = desc
return self
[docs]
def apply_refined_group_ids(self, meta_kmax=8):
"""Relabel fitted groups with refined/sorted global ids (issue #40).
Final-fit-only, bijective: per element, orders group centers via
:func:`refined_group_order` (average linkage + optimal leaf ordering,
gap-statistic meta gate), then offsets by species order so the final
``group_id`` is globally unique. Rewrites ``per_el_group_libs`` and
``per_structure_group_libs`` in place; selection, ``|S|``, L/R/F, and
noise-derived provenance are unchanged. Never call on scan/cache fits.
Returns
-------
pandas.DataFrame
One row per group (density-supported or noise-derived): element,
old_group_id,
local_group_id, group_id, meta_group_id, bridge_meta_a,
bridge_meta_b, bridge_t, bridge_slack, population, granularity,
nearest_center_dist, prev_jump, next_jump, pc1_coord.
``bridge_t`` is the between-DOMAIN detector: each group is
projected onto the chord between its two nearest reservoir cores
(leave-one-out population-weighted meta means, ``bridge_meta_a``/
``bridge_meta_b``). ``t`` = 0/1 at the cores (domain interiors),
~0.5 midway; interior t (~0.15-0.85) = a state between two
domains — NEB images near a TS, MD basin crossings, kMC
intermediates. Cores average ~10^3 atoms so t is ε_noise-stable
(item-level top-25 % survival 8.7/10 at 1ε on ema) and, unlike a
nearest-center ratio, does not conflate isolated groups with
crossings (an isolated group projects near a core → extreme t).
``bridge_slack`` = (d1+d2-D)/D triangle excess is a SECONDARY
curvature flag (0 = on the straight A→B chord; large = off it),
not a detector on its own. NaN/-1 when n_meta < 2. Shortlist by
``|t - 0.5|``; cross-check ``nearest_center_dist`` and inspect the
environments before claiming reaction coverage.
``granularity`` is the per-element operating point actually used
(δ = c·ε_noise for distance-kind backends, k for count-kind) — the
physical yardstick for the distance columns: centers closer than
~δ cannot coexist under mean-shift, so ``next_jump/granularity``
says how many bandwidths apart adjacent sorted groups are.
"""
if getattr(self, '_refined_ids_applied', False):
raise RuntimeError(
'apply_refined_group_ids: refined ids already applied — '
'group ids are global; a second pass would corrupt them.')
if not getattr(self, 'cluster_per_el', None):
raise RuntimeError(
'apply_refined_group_ids: call fit_decaf first.')
rows = []
offset = 0
self.group_order_per_el = {}
self.group_id_offset = {}
new_of_old_per_el = {}
for el in self.species:
fit = self.cluster_per_el.get(el)
if fit is None:
continue
n_el = int(fit.n_groups)
_lab = np.asarray(fit.labels)
pops = np.bincount(_lab[(_lab >= 0) & (_lab < n_el)],
minlength=n_el)[:n_el]
order = refined_group_order(fit.centers, meta_kmax=meta_kmax,
populations=pops)
self.group_order_per_el[el] = order
self.group_id_offset[el] = offset
observed = sorted({
int(g)
for gs in self.per_el_group_libs.get(el, [])
for g in gs
if int(g) >= 0
})
extra_old = [g for g in observed if g >= n_el]
new_of_old = {
old: int(offset + loc)
for old, loc in enumerate(order.local_of_old)
}
for extra_loc, old in enumerate(extra_old):
new_of_old[old] = int(offset + n_el + extra_loc)
new_of_old_per_el[el] = new_of_old
for loc in range(n_el):
old = int(order.old_of_local[loc])
rows.append(dict(
element=el, old_group_id=old, local_group_id=loc,
group_id=offset + loc,
meta_group_id=int(order.meta_of_local[loc]),
bridge_meta_a=int(order.bridge_pair[loc][0]),
bridge_meta_b=int(order.bridge_pair[loc][1]),
bridge_t=float(order.bridge_t[loc]),
bridge_slack=float(order.bridge_slack[loc]),
population=pops[old],
granularity=float(fit.granularity),
nearest_center_dist=float(order.nearest_center_dist[loc]),
prev_jump=float(order.prev_jump[loc]),
next_jump=float(order.next_jump[loc]),
pc1_coord=float(order.pc1_coord[loc]),
))
for extra_loc, old in enumerate(extra_old):
rows.append(dict(
element=el, old_group_id=old,
local_group_id=n_el + extra_loc,
group_id=new_of_old[old],
meta_group_id=-1,
bridge_meta_a=-1, bridge_meta_b=-1,
bridge_t=np.nan, bridge_slack=np.nan,
population=sum(int(g) == old
for gs in self.per_el_group_libs[el]
for g in gs),
granularity=float(fit.granularity),
nearest_center_dist=np.nan,
prev_jump=np.nan,
next_jump=np.nan,
pc1_coord=np.nan,
))
offset += n_el + len(extra_old)
for el, new_of_old in new_of_old_per_el.items():
self.per_el_group_libs[el] = [
[int(new_of_old[int(g)]) if int(g) >= 0 else -1 for g in gs]
for gs in self.per_el_group_libs[el]
]
if getattr(self, 'per_structure_group_libs', None) is not None:
self.per_structure_group_libs = [
[(el, int(new_of_old_per_el[el][int(g)]))
if int(g) >= 0 else (el, -1)
for el, g in pairs]
for pairs in self.per_structure_group_libs
]
self._refined_ids_applied = True
mapping = pd.DataFrame(rows, columns=[
'element', 'old_group_id', 'local_group_id', 'group_id',
'meta_group_id',
'bridge_meta_a', 'bridge_meta_b', 'bridge_t', 'bridge_slack',
'population', 'granularity',
'nearest_center_dist', 'prev_jump', 'next_jump', 'pc1_coord'])
if any(getattr(fit, 'backend', None) == 'hdbscan'
for fit in self.cluster_per_el.values()):
# Deployment provenance is additive for density fits; preserve the
# established mapping schema for DECAF/KMeans and custom clusterers.
mapping['deployable'] = [
int(row['old_group_id']) < getattr(
self.cluster_per_el[row['element']], 'n_deployable_groups', 0)
for row in rows
]
return mapping
def _extract_one_surrounding(self, flat_idx, g, full_lib, r_cut,
output_mode, mu_labels, atom_to_struct,
cum_sizes):
"""Extract one representative surrounding (sphere or tagged full cell).
``g`` is the first-class group id stored in
``atoms.info['group_id']``.
"""
from ase.neighborlist import NeighborList
focus_el = self.species[0]
struct_idx = int(atom_to_struct[flat_idx])
local_focus_idx = flat_idx - int(cum_sizes[struct_idx])
parent_struct = full_lib[struct_idx]
# Map local focus index to full-atom index, skipping is_focus=0 atoms
full_syms = np.array(parent_struct.get_chemical_symbols())
if 'is_focus' in parent_struct.arrays:
_ifmask = parent_struct.arrays['is_focus'].astype(bool)
focus_positions_in_full = np.where(
(full_syms == focus_el) & _ifmask)[0]
else:
focus_positions_in_full = np.where(full_syms == focus_el)[0]
full_atom_idx = int(focus_positions_in_full[local_focus_idx])
center_pos = parent_struct.positions[full_atom_idx].copy()
mu_id = int(mu_labels[flat_idx]) if mu_labels is not None else -1
if output_mode == 'full_structure':
out_atoms = parent_struct.copy()
t = np.zeros(len(out_atoms), dtype=int)
if parent_struct.pbc.any():
cutoffs = [r_cut / 2.0] * len(parent_struct)
nl = NeighborList(cutoffs, self_interaction=False, bothways=True)
nl.update(parent_struct)
nbr_indices, _ = nl.get_neighbors(full_atom_idx)
t[list(nbr_indices)] = 2
else:
dists = parent_struct.get_distances(
full_atom_idx, np.arange(len(parent_struct)), mic=True)
nbr = (dists <= r_cut)
nbr[full_atom_idx] = False
t[nbr] = 2
# Center tag set last: takes priority over neighbour tag
# (atom can appear in its own PBC neighbour list via images)
t[full_atom_idx] = 1
out_atoms.set_tags(t)
elif parent_struct.pbc.any():
cutoffs = [r_cut / 2.0] * len(parent_struct)
nl = NeighborList(cutoffs, self_interaction=False, bothways=True)
nl.update(parent_struct)
nbr_indices, offsets = nl.get_neighbors(full_atom_idx)
cell = parent_struct.get_cell()
symbols = [focus_el] + [parent_struct[i].symbol for i in nbr_indices]
positions = [center_pos] + [
parent_struct.positions[i] + offset @ cell
for i, offset in zip(nbr_indices, offsets)
]
tags = [1] + [0] * len(nbr_indices)
out_atoms = Atoms(symbols=symbols, positions=positions, pbc=False)
out_atoms.set_tags(tags)
else:
dists = parent_struct.get_distances(
full_atom_idx, np.arange(len(parent_struct)), mic=True)
mask = dists <= r_cut
mask_indices = np.where(mask)[0]
out_atoms = parent_struct[mask].copy()
out_atoms.pbc = False
center_in_cluster = int(
np.where(mask_indices == full_atom_idx)[0][0])
tags = np.zeros(len(out_atoms), dtype=int)
tags[center_in_cluster] = 1
out_atoms.set_tags(tags)
out_atoms.info['group_id'] = g
out_atoms.info['microstate_id'] = mu_id
out_atoms.info['center_pos'] = center_pos
out_atoms.info['struct_idx'] = struct_idx
out_atoms.info['atom_idx'] = full_atom_idx
return out_atoms