Clustering primitives and persistence

The backend-agnostic pieces every clusterer shares: the PCA embedding, the deployable classifiers, the per-element fit contract, and PickleIO — the one implementation behind Judge.save / Judge.load and the judge_final.pkl a run writes.

atom_hifi/training/cluster_base.py — backend-agnostic clustering primitives.

The pieces every cluster backend shares, independent of which clusterer runs: the PCA embedding (PCAEmbedder), the deployable classifier (NearestNeighborClassifier), the per-element result contract (ClusterFit), and the trailing-noise-slot promotion machinery (issue #63). Backend implementations live in atom_hifi.training.cluster_backends; this module must not import them (the dependency runs one way, backends -> base).

Split out of cluster_backends.py in issue #69; the native classifier landed in issue #76.

atom_hifi.training.cluster_base.NOISE_FRAC_WARN = 0.3

HDBSCAN noise fraction above which the softness/superionic warning fires (cf. LLZO Li noise-frac 0.42 in the §39 three-way study).

atom_hifi.training.cluster_base.CHUNK_MAX_EL = 50000000

distance-matrix element budget per prediction chunk (a test shrinks this)

class atom_hifi.training.cluster_base.Embedder

Bases: ABC

Descriptors → the metric space the clusterer and δ both live in.

Closed on purpose — do not add a second implementation. δ = c·ε_noise, the |S| targets and the QHO σ table are all distances in decaf’s PCA metric (broken-stick dim, uncentred projection). A different embedder is a change of units, not a swap, and no artifact records which embedder produced it — caches tag the backend only. Embedder-tagging is the prerequisite for opening this up, not a detail.

abstract property t_embedded

(N, dim) embedded training rows.

abstract property embed_op

(D, dim) projection operator; embed is a matmul with it.

abstractmethod embed(desc)

Project raw descriptor rows into the embedding.

class atom_hifi.training.cluster_base.Classifier

Bases: ABC

Assigns novel embedded rows to fitted groups, abstaining when far.

Internal abstraction. The mechanism is backend-agnostic and, since issue #76, actually written that way — see NearestNeighborClassifier. Loop and patient OOD consume either supported mechanism. Concretely:

• decaf sets ClusterFit.abstain_radius -> the 1-NN-over-training-points rule (NearestNeighborClassifier).

• kmeans sets ClusterFit.abstain_lambda -> the nearest-centroid rule priced by the DP-means cluster-opening price (gitlab#80, NearestCentroidClassifier). kmeans’ granularity is still a cluster count, but the FITTED PARTITION has a resolution, and lambda_open prices it — that is what changed.

• hdbscan stores frozen density hierarchy prediction state. It uses mutual-reachability attachment, not a radius or a centroid price.

ClusterBackend.has_classifier is True for all three built-ins. Kmeans lambda_open tracks local partition spacing, so when k grows with the corpus its OOD rate can stay approximately scale-free and the loop may consume the full library. Gitlab#81 accepts that existing termination behavior; HiFiWorkflow.run_loop() warns once instead of changing the predictor or inventing an unvalidated stop threshold.

abstractmethod classify(embedded, **kwargs)

Return decaf’s (gID_distinct, gID_label_all) pair.

Both lists have one slot per selection group plus a trailing slot. gID_distinct[-1] holds rejected rows. DECAF and KMeans provide unconditional assignments in gID_label_all (empty trailing slot). HDBSCAN has no unconditional assignment: rejected rows also occupy its second list’s trailing slot. Promoted noise slots stay empty.

class atom_hifi.training.cluster_base.ClusterBackend

Bases: ABC

The clusterer, and the one public extension point.

Capabilities are class attributes, not name literals: code asks what a backend can do, never what it is called. name is the exception — it is the cache tag, and state_io enforces equality on it across runs.

name: str

cache tag — must be stable across runs (state_io enforces equality)

granularity_kind: str

‘distance’ (δ is a radius) | ‘count’ (δ is a cluster count k)

has_classifier: bool = False

fit yields a deployable classifier → OOD / patient / loop

emits_noise: bool = False

fit can leave a trailing noise slot needing ε_noise promotion (#63)

has_auto_bandwidth: bool = False

fit yields its own bandwidth estimate (decaf’s HaarWavelet δ_auto)

abstractmethod fit(desc_el, granularity, *, t_embedded=None, **kw)

Cluster one element; return an unpromoted ClusterFit.

Takes raw desc_el rather than a ready embedding because decaf builds its own embedding internally — inject_precomputed needs the decaf object itself. Backends that cluster in the shared embedding should call shared_embedder() and use its t_embedded; that is the one-liner a custom backend needs.

classmethod unsupported_reason(granularity=None)

Why this backend cannot run in this environment, or None if it can.

The capabilities above are static facts about the algorithm; this one is not. It reports a dependency fault — a broken upstream version, a missing optional package — so it has to be asked at runtime rather than read off the class.

Two questions, one method:

• granularity=None asks the general one — can this backend be relied on here at all? Scan entry points and the test skip marker use this, because a scan sweeps δ and will reach the bad values.

• passing a δ asks about that one fit, which lets a backend admit the δ values it can still serve. fit_cluster_backend uses this, so a dependency fault that only bites part of the δ range does not take the whole backend down with it.

static shared_embedder(desc_el, t_embedded=None)

The shared PCA embedding: reuse the SHM-served one, else fit it.

class atom_hifi.training.cluster_base.PickleIO

Bases: object

Pickle persistence with a declared footprint, shared by all backends.

A mixin, not a fourth abstraction: it declares no abstract method, so it stays a plain base class and leaves its subclasses’ metaclass alone.

Backend-neutral by construction: pickle serialises whatever the object holds, so one implementation here covers decaf, kmeans, hdbscan and any registered custom backend with no per-backend code — the reason it lives beside the abstractions rather than on Judge.

What it adds over a bare pickle.dump is the footprint contract. A fit artifact keeps references to arrays that are inputs, not results (raw descriptors above all), and pickling them silently turns a 10 MB object into a 300 MB file. Subclasses list those attributes in _PICKLE_DROP; they come back as None and the loader is expected to say so. Saves stamp the schema and package version so a stale file is diagnosable rather than mysterious.

The footprint applies to serialization only. __getstate__ also drives copy.copy/copy.deepcopy, so without the overrides below a caller who merely copied a fitted object would silently lose the dropped attributes on a live object no one asked to persist. Those overrides assume __dict__-based subclasses, which is what Judge is.

Serialization never computes: these objects also travel over joblib IPC, where a hook that materialised a lazy field would cost real throughput.

PICKLE_SCHEMA = 1

Bump when a load of an older file would be wrong rather than merely thin. 1 — first stamped format; _PICKLE_DROP attributes restore as None.

save(path)

Pickle this object to path (footprint per _PICKLE_DROP).

Written to a temporary sibling and moved into place with os.replace(), because open(path, 'wb') truncates first: a save that failed halfway — the very case the workflow’s save site anticipates for an unpicklable custom backend — would otherwise destroy the artifact a previous run wrote and leave a truncated file that fails at load time instead of here.

classmethod load(path)

Unpickle an object saved by save().

A file older than the schema is loaded, not rejected: the whole point of the artifact is to survive the run that wrote it, and every migration so far has been additive. A mismatch warns so the reader knows which version’s semantics they are looking at.

class atom_hifi.training.cluster_base.PCAEmbedder(t_embedded, embed_op=None, desc_for_lazy=None)

Bases: Embedder

decaf PCA embedding, computed once per element and shared by all backends.

embed is a pure matmul with embed_op (decaf’s embed_PCA), so a precomputed (t_embedded, embed_op) pair reproduces decaf’s embedding exactly.

On the SHM fast path only t_embedded is available; embed_op is then recomputed lazily from the raw descriptors on first access (scan workers never touch it, so the extra PCA is never paid in the hot loop).

classmethod fit(desc_array)

Fit decaf’s PCA on desc_array (no clustering).

property t_embedded

(N, dim) embedded training rows.

property embed_op

(D, dim) projection operator; embed is a matmul with it.

embed(desc)

Project raw descriptor rows into the embedding.

class atom_hifi.training.cluster_base.NearestNeighborClassifier(t_embedded, labels, n_groups, radius)

Bases: Classifier

1-NN label transfer over training points + radius abstention.

The native implementation ADR-0011’s abstraction table names. It replaces the delegation to decaf, and is a verbatim port of MSC_classify on the one path atom_hifi uses (classify_mode='TrainingData', label_all='Both'; decaf.py:1510-1554) — not a reimplementation of the idea. ADR-0011 left this open precisely because a loose port moves OOD counts silently, so three conventions are copied rather than improved:

• ``argsort(axis=1)[:, 0]``, not ``argmin`` (decaf.py:1508/1518). numpy’s default sort is not stable, so on an exact distance tie the two may resolve to different training atoms — and hence different groups. Whether they actually differ is numpy-version-dependent (2.2 diverges on an all-tied row, 2.4/2.5 do not), which is precisely why the call is copied rather than reasoned about: copying makes this class agree with decaf on every numpy, including future ones where the tie-break changes again. Exact ties are not hypothetical in a symmetric crystal. argmin would also be cheaper in memory; that trade is deferred to gitlab#77, which measures tie frequency first.

• ``<=``, not ``<`` for the radius test (decaf.py:1542).

• The abstention test is over the training set, not the centers, and is therefore independent of how the partition was drawn (:1540-1542). That independence is exactly why this class is backend-agnostic.

Inputs are (t_embedded, labels, n_groups, radius). Note labels is the fit’s promoted label array, where decaf reads t_gID through reformat_groupID. The two agree on the production path — MeanShift runs with cluster_all=True (decaf.py:1434) so no atom is ever -1 and the trailing slot stays empty — but they would not agree after issue-#63 promotion renumbered a backend that does emit noise. Using the promoted labels is what keeps a promoted fit self-consistent.

classify(embedded, dist_in=None, **kwargs)

Return decaf’s (gID_distinct, gID_label_all) pair.

dist_in is a precomputed (N_new, N_train) distance matrix. Both “absent” spellings are accepted: None, and decaf’s own empty-array sentinel (decaf.py:1504/1511), so a call ported straight off a decaf one does not land on a shape error deep inside argsort. A present matrix of the wrong shape is rejected, which decaf does not do — see the branch below for why that cannot move a number.

The multiscale arguments are deliberately NOT accepted, because nothing in atom_hifi passes them and quietly ignoring a multiscale request would answer a different question than the caller asked.

class atom_hifi.training.cluster_base.NearestCentroidClassifier(centers, n_groups, lambda_open)

Bases: Classifier

Nearest-centroid label transfer + DP-means opening-price abstention.

The kmeans counterpart of NearestNeighborClassifier. The assignment rule IS kmeans’ own definition of a label — argmin_j ||x - mu_j|| — which is the point: the 1-NN-over-training-points rule assigns a DIFFERENT cluster on 10-41% of the queries both rules resolve (gitlab#80 Stage A, measured on 7 real element arms across 3 systems), so a kmeans fit cannot be answered through NearestNeighborClassifier.

Abstention uses the DP-means cluster-opening price: q(x) = d1**2 / lambda_open, with q == 1 EXACTLY in distribution (STRICTLY greater means OOD). Note the units: lambda_open is in SQUARED-distance units — d1**2, not d1, is what is compared to it.

predict(embedded, return_confidence=False)

Centroid rule: assign nearest centroid, abstain where q > 1.

q(x) = d1**2 / lambda_open, with the boundary q == 1 exactly IN distribution (STRICTLY greater means OOD). d2 is the second-smallest distance (inf and confidence 1.0 when k == 1); confidence is 1 - d1/d2, 0 where d1 == d2 == 0. Distances are computed in row chunks (each chunk’s (rows, k) matrix under CHUNK_MAX_EL elements) so memory stays bounded; d1 and d2 come from the SAME distance pass.

Returns a dict with label (-1 where OOD), assigned (the raw nearest-centroid index, never -1), d1, d2, q, confidence, ood and ood_fraction. d2 and confidence are None unless return_confidence.

classify(embedded, **kwargs)

Return decaf’s (gID_distinct, gID_label_all) pair.

dist_in is deliberately NOT accepted: it is a (N_new, N_train) distance matrix over TRAINING POINTS, which is meaningless for a centroid rule — there is no training-point distance anywhere in this classifier.

class atom_hifi.training.cluster_base.ClusterFit(labels, n_groups, centers, granularity, backend, embedder, classifier=None, noise_derived=None, pre_promotion_noise_frac=None, abstain_radius=None, abstain_lambda=None, abstain_calibration=None, hdbscan_prediction=None)

Bases: object

Per-element clustering result — the backend-agnostic contract.

Fits returned by fit_cluster_backend() are self-consistent: every label is in [0, n_groups), centers and noise_derived have n_groups rows, and the trailing noise/distinct slot is empty.

Parameters:

• labels (ndarray)

• n_groups (int)

• centers (ndarray)

• granularity (float)

• backend (str)

• embedder (Embedder)

• classifier (object)

• noise_derived (ndarray)

• pre_promotion_noise_frac (float)

• abstain_radius (float)

• abstain_lambda (float)

• abstain_calibration (dict)

• hdbscan_prediction (object)

labels: ndarray

n_groups: int

centers: ndarray

granularity: float

backend: str

embedder: Embedder

classifier: object = None

Optional Classifier override. Since issue #76 no backend sets it: classify() builds a NearestNeighborClassifier from abstain_radius instead, so nothing stateful rides on the fit. It stays honoured first, which keeps an explicitly-supplied implementation — and any pickle written before #76 — working unchanged.

Renamed from decaf_cluster in issue #69 — deliberately a NEW name, not a repurposed one. A pickle written between #68 and #69 still unpickles (a dataclass restores __dict__ wholesale, so the stale decaf_cluster key simply rides along); had this field kept that name, the stale value would satisfy the is not None test and the fit would classify through a stale decaf object with no error at all. Under the new name such a fit raises, which is the honest outcome.

noise_derived: ndarray = None

pre_promotion_noise_frac: float = None

abstain_radius: float = None

Radius for the 1-NN abstention test, or None when this backend has no deployable classifier. It carries the policy, which is why it is a field and not derived from granularity: for decaf the two coincide (δ IS decaf’s radius), but kmeans’ granularity is a cluster count with no distance meaning. DecafBackend sets this, kmeans sets abstain_lambda (gitlab#80), and hdbscan stores density state instead.

abstain_lambda: float = None

DP-means cluster-opening price in SQUARED-distance units (gitlab#80), set by KMeansBackend.fit only. It is a SEPARATE field from abstain_radius and not a unit conversion of it, because the two select different classifiers — nearest-centroid vs 1-NN-over-training points — and Stage A measured those two disagreeing on 10-41% of jointly-resolved queries. ClusterFit.classify raises if both are set, rather than silently preferring one.

abstain_calibration: dict = None

Provenance dict returned by lambda_local (gitlab#80), persisted so a fit’s abstention threshold can be audited after the fact.

hdbscan_prediction: object = None

Versioned frozen density state, built from sklearn’s fitted hierarchy. Separate from promotion: its labels remain backend-local density ids.

property n_density_groups

Coverable groups produced by clustering, excluding noise promotion.

property n_deployable_groups

Number of targets supported by this saved fit’s prediction state.

classify(embedded, **kwargs)

Assign novel embedded rows to this fit’s groups, abstaining when far.

The classifier is built here rather than stored on the fit — whether NearestNeighborClassifier (from abstain_radius) or NearestCentroidClassifier (from abstain_lambda). It holds only references plus a float, so construction is O(1) and there is nothing to amortise — whereas a cached instance would put a second reference to t_embedded on an object Judge.save pickles. This is also what closes ADR-0011’s open pickle-footprint item: no live decaf object needs retaining on a fit any more.

classifier is still honoured first, so an explicitly-supplied implementation (and any pickle written before this change) keeps working.

property noise_frac

Fraction of atoms the backend put in the trailing slot.

Reads the pre-promotion value when promotion has run, so the softness signal survives noise reps becoming real groups (issue #63).

atom_hifi.training.cluster_base.dedup_noise_greedy(coords, delta)

Greedy single-pass dedup of noise atoms at radius δ (= ε_noise).

Walk the noise atoms; keep one as a DISTINCT rare-environment rep if it is farther than δ from every already-kept rep, else it is a thermal duplicate absorbed into that rep. Deterministic (natural order). O(n_noise · n_kept).

Returns:

• kept_idx (list[int]) – Indices (into coords) of the kept distinct reps.

• assign (ndarray, shape (len(coords),)) – assign[i] is the position within ``kept_idx`` of the rep that coords[i] belongs to — identity for kept reps themselves. Used to absorb thermal duplicates into their rep’s group rather than discarding them (issue #63).

atom_hifi.training.cluster_base.promote_trailing_slot(labels, n_groups, centers, t_embedded, eps_noise)

Dedup the trailing slot and append each survivor as its own group.

The backend-agnostic half of issue #63: a noise/distinct atom is a rare environment, not a non-result. Deduplicating at ε_noise and promoting each survivor to a real group makes those environments coverable by the ordinary set-cover, addressable by group id, and visible to the meta-group linkage — none of which was true while they lived in a forced_idx side-channel.

Thermal duplicates are absorbed into their kept rep’s group, so no atom is left unlabelled.

No-op when the trailing slot is empty, which is the case for kmeans (no such slot) and for decaf on the classify_mode='TrainingData' path (every atom is its own neighbour at distance 0 ≤ bandwidth).

Return type:

labels, n_groups, centers, noise_derived, n_trailing

atom_hifi.training.cluster_base.medoids_from_labels(t_embedded, labels, n_groups)

Per-cluster medoid (point nearest the cluster mean) in the embedding.

HDBSCAN sets no cluster_centers; medoids give a center where the surroundings extraction needs one. Noise atoms (label < 0) are ignored.