Judge — cover + set-cover + microstate reference

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:

• atom_hifi.training.selection — set cover (find_min_covering_subset, atom_aware_greedy, length_bool)

• atom_hifi.training.calibration — σ_thermal, ε_noise, microstate reference

• atom_hifi.training.cluster_base / cluster_backends — clustering

• atom_hifi.training.group_order — group ordering

HiFiWorkflow and the parallel worker functions live in atom_hifi.training.hifi_workflow.

class atom_hifi.training.judge.Judge(species, lib0=None, lib_path=None, selection=None, descriptor='SOAP', cluster_backend='decaf', min_cluster_size=2, **soap_params)

Bases: 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 PickleIO and is therefore identical for every backend.

Parameters:

• species (List[str])

• lib0 (List[Atoms])

• lib_path (str | None)

• selection (List[str] | None)

• descriptor (str)

• cluster_backend (str)

• min_cluster_size (int)

lib0: List[Atoms]

fit_decaf(desc_array=None, delta=None, mu_per_el=None, counts_L_per_el=None, t_embedded=None, metric_func=<function 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).

property cluster

Fitted ClusterFit of species[0] (any backend).

None when that element’s cover came from cluster_cache_per_el.

recompute_hifi_metrics()

Recompute L/R/F for the current sub_lib_idx.

save(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.

attach_descriptors(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:

self, so a load can be chained: Judge.load(p).attach_descriptors(d).

Return type:

Judge

apply_refined_group_ids(meta_kmax=8)

Relabel fitted groups with refined/sorted global ids (issue #40).

Final-fit-only, bijective: per element, orders group centers via 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:

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.

Return type:

pandas.DataFrame

extract_rep_surroundings(full_lib, r_cut=None, n_per_group=1, mu_labels=None, save_dir=None, output_mode='sphere', gid_map=None, meta_map=None)

Extract coordination sphere for each DECAF group’s representative atom.

For each group g of species[0]’s fit (0 .. self.cluster.n_groups-1; empty groups skipped), finds the atom whose embedded descriptor is closest to the group centroid, then extracts output atoms centred on that atom from the corresponding full structure in full_lib.

Parameters:

• full_lib (list[Atoms]) – Full structures (all atom types), parallel-indexed to self.lib. Required because self.lib contains only focus-element atoms.

• r_cut (float or None) – Sphere cutoff radius in Å. Default 5.0.

• n_per_group (int) – Representatives per group. 1 → centroid-closest only. >1 → centroid-closest + greedy FPS in embedded space.

• mu_labels (ndarray (N_focus_atoms,) or None) – Microstate label per focus atom. Stored as atoms.info[‘microstate_id’].

• save_dir (str or None) – If given, write output xyz file to this directory.

• output_mode ({'sphere', 'full_structure'}) –

  ‘sphere’ (default): non-periodic cluster of atoms within r_cut.

  Tags: center=1, neighbours=0. File: surroundings_{el}.xyz.

  ’full_structure’: copy of the full original structure with cell/pbc

  intact. Tags: center=1, neighbours within r_cut=2, others=0. File: surroundings_full_{el}.xyz.

• gid_map (dict[int, int] or None) – Backend-native group id → refined/sorted global group id (issue #40). Applied to result keys and atoms.info[‘group_id’].

• meta_map (dict[int, int] or None) – Final (mapped) group id → meta_group_id; written to atoms.info[‘meta_group_id’] when given.

Returns:

{group_id: [Atoms, …]} — list always, even for n_per_group=1. Groups with zero atoms are omitted.

Return type:

dict[int, list[Atoms]]