HiFiWorkflow — L/R scan workflow

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.

class atom_hifi.training.hifi_workflow.HiFiWorkflow(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)

Bases: _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 (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).

• 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 (int) – Number of δ points in coarse / fine scan.

• 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.

force_eps_for_element(el)

Return (ε_F, u*) for one element from cached thermal ε(σ).

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

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

Columns: c_{el} for each element, L, R, F, n_structs_S, n_atoms_S.

Return type:

pd.DataFrame

run_grid_scan(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 – Merged coarse + fine grid with duplicates removed.

Return type:

pd.DataFrame

run_nsga2(n_gen=50, pop_size=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 (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.

• pop_size (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.

• 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:

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

Return type:

tuple

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

• it. (Sets self.records and returns)

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

extract_rep_surroundings(el=None, r_cut=None, n_per_group=1, save_dir=None, output_mode='sphere')

Extract coordination spheres for each DECAF group representative.

Refits a per-element Judge at c* for each focus element, then calls Judge.extract_rep_surroundings on each.

Parameters:

• el (str or None) – Focus element to process. None → all self.species.

• r_cut (float or None) – Sphere cutoff radius in Å. Forwarded to Judge.extract_rep_surroundings.

• n_per_group (int) – Representatives per group (default 1).

• save_dir (str or None) – If given, write output xyz files here.

• output_mode ({'sphere', 'full_structure'}) – Forwarded to Judge.extract_rep_surroundings. See that docstring.

Returns:

{el: {group_id: [Atoms, …]}} always (first key = element).

Return type:

dict[str, dict[int, list[Atoms]]]

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

plot_final_fit(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])