"""atom_hifi/training/cluster_backends.py — pluggable clusterer backends (ADR-0009).
The backend is the **clusterer only**: the PCA embedding is computed once per
element (decaf's ``get_embed(dim=0)``) and shared by every backend via
:class:`~atom_hifi.training.cluster_base.PCAEmbedder`.
:func:`fit_cluster_backend` dispatches to one of the built-in arms and returns a
:class:`~atom_hifi.training.cluster_base.ClusterFit`.
Backends
--------
* ``decaf`` (distance): mean-shift at bandwidth δ; classifies by native 1-NN
radius abstention (OOD/patient/loop support).
* ``kmeans`` (count): δ is the cluster count k; pinned ``n_init=10``. The fit
calibrates ``lambda_open`` — the DP-means cluster-opening price in
SQUARED-distance units (gitlab#80) — and ``ClusterFit.classify`` abstains via
the nearest-centroid rule (OOD/patient/loop support, gitlab#81).
* ``hdbscan`` (distance): δ is ``cluster_selection_epsilon``; identical
constructor params to decaf's own ``get_HDBSCAN`` (``min_samples`` =
``min_cluster_size``, ``allow_single_cluster=True``), but raw sklearn labels
(−1 = noise; decaf's re-classify step is skipped because its ``LabelAll``
branch is broken). Noise atoms are deduped at ε_noise and promoted to
first-class groups before :func:`fit_cluster_backend` returns.
Issue #69 moved the backend-agnostic pieces to
:mod:`atom_hifi.training.cluster_base` and the group ordering to
:mod:`atom_hifi.training.group_order`. ``_PrecomputedDecaf`` deliberately
stays here because its import path is load-bearing: ``Judge.save`` pickles
whatever a ``ClusterFit`` holds, and between issues #68 and #76 that included a
live ``_PrecomputedDecaf``. Since #76 no fit retains one, so *new* pickles do
not contain it — but the path must keep resolving or every pickle written in
that window becomes unloadable.
"""
import copy
import logging
import numpy as np
from numpy.lib import NumpyVersion
from sklearn.cluster import HDBSCAN, KMeans
import decaf
from atom_hifi.training.cluster_base import (
NOISE_FRAC_WARN,
Classifier,
ClusterBackend,
ClusterFit,
PCAEmbedder,
medoids_from_labels,
promote_trailing_slot,
)
from atom_hifi.training.lambda_open import lambda_local
# Re-exports at the historical import path (issue #69) — these moved to
# cluster_base/group_order but are imported from here by tests and downstream.
# NearestNeighborClassifier / NearestCentroidClassifier are re-exported so all
# native Classifier implementations are reachable from one module, next to the
# DecafClassifier they replace.
from atom_hifi.training.cluster_base import ( # noqa: F401
NearestCentroidClassifier,
NearestNeighborClassifier,
dedup_noise_greedy,
)
from atom_hifi.training.group_order import ( # noqa: F401
GroupOrder,
refined_group_order,
)
logger = logging.getLogger(__name__)
# GRANULARITY_KIND and BUILTIN_BACKENDS are derived from BACKEND_REGISTRY at
# the bottom of this module (issue #69) — they cannot drift from the classes.
#: NumPy 2.4.0 expired the 1.25 deprecation and now raises ``TypeError`` when an
#: ``ndim > 0`` array is converted to a scalar. sklearn's HDBSCAN trips exactly
#: that in ``_hdbscan/_tree.pyx`` ``traverse_upwards``, which assigns a
#: boolean-mask lookup — a shape-(1,) array — into a ``cdef intp_t``. That code
#: is reachable only from ``epsilon_search``, i.e. only when
#: ``cluster_selection_epsilon > 0`` — which here means any fit at non-zero δ,
#: because δ *is* that epsilon. decaf and kmeans never enter the branch, and
#: neither does hdbscan at δ=0.
#:
#: **Not every positive-δ fit fails either.** ``traverse_upwards`` runs only
#: for tree leaves below δ, so the trigger is data- and δ-dependent: on numpy
#: 2.4.6 the same 40-point blob pair survives δ=0.1 and dies at δ=0.5. Above
#: zero the guard is therefore deliberately coarser than the true condition —
#: "will this condensed tree reach that branch" cannot be answered without
#: building the tree, and a conservative refusal beats a Cython ``TypeError``
#: surfacing mid-scan.
#:
#: The lines are identical in sklearn 1.8, 1.9 and main, so this is not a
#: scikit-learn version axis: scikit-learn#33219 is open and its fix PR #34243
#: is unmerged. Delete this guard once a released sklearn carries the fix.
def _numpy_breaks_hdbscan_eps(version=None):
"""True when *version* is a numpy that raises on ndim>0 → scalar.
A function, not an inline expression, so a test can feed it version strings
it will never run under. Compared against ``'2.4.0.dev0'`` rather than
``'2.4.0'`` because NumpyVersion orders prereleases *below* the release —
``NumpyVersion('2.4.0rc1') >= '2.4.0'`` is False — and an rc carries the
expired conversion behaviour exactly as the release does.
"""
return NumpyVersion(np.__version__ if version is None else version) >= '2.4.0.dev0'
#: Module attribute, read at call time by unsupported_reason(), so tests can
#: monkeypatch either side of the guard on whatever numpy they run under.
_NUMPY_BREAKS_HDBSCAN_EPS = _numpy_breaks_hdbscan_eps()
class _PrecomputedDecaf(decaf.embed_cluster):
"""embed_cluster subclass that skips PCA and short-circuits the fit-path classify.
Usage (main process):
ec = _PrecomputedDecaf(desc_array, embed_str='PCA', cluster_str='MSC')
ec.inject_precomputed(t_embedded) # skip get_embed()
ec.get_cluster(bandwidth=delta, ...)
`inject_precomputed` is optional: with no embedding supplied the class
behaves exactly like its base, so it is constructed unconditionally and the
guard below covers every fit path (issue #68).
"""
def inject_precomputed(self, t_embedded):
"""Set the pre-computed embedding; replaces get_embed()."""
self.t_embedded = t_embedded
self.f_embedded = t_embedded
def MSC_classify(self, n_embedded, dist_in=np.array([]),
multibandwidth=False, multiscales=np.array([]), **kwargs):
"""Short-circuit decaf's fit-path self-classification (issue #68).
On our pinned path (``classify_mode='TrainingData'``, ``max_training_size=0``
so ``f_data is t_data``) ``get_MeanShift`` ends with
``self.f_gID, _ = self.MSC_classify(self.f_embedded)`` (decaf.py:1455) —
the training set classified against *itself*. Every atom is then its own
nearest neighbour at distance exactly 0 ≤ bandwidth, so the answer is the
identity ``t_gID`` with an empty distinct slot, and ``atom_hifi`` never
reads ``f_gID``. Returning it directly removes the only consumer of the
N×N distance matrix (probe: gitlab#67).
The predicate is **call-level**, not object-level. ``get_MeanShift``
passes ``self.f_embedded`` itself, so argument identity is what
distinguishes that self-classification from a genuine ``classify(novel)``
on the same fitted object, which can also arrive with an empty
``dist_in``. Since #76 no fresh fit routes here: such calls come from a
``DecafClassifier`` (explicitly supplied, or restored from a Judge pickled
between #69 and #76 and used by any OOD consumer, e.g.
``scripts/predict_ood.py``) and, in tests, from the decaf reference and
direct ``ec.classify(novel)`` calls. Keying on ``f_data is t_data`` alone
would short-circuit those too and return training labels for patient
data: silently wrong OOD, no crash.
Each remaining clause closes a way the premise could fail:
- ``f_data is t_data`` — a non-zero ``max_training_size`` makes ``t_data``
a strict subset (decaf.py:307-312), so ``f_gID`` is then *not* the
identity and must be computed.
- ``classify_mode`` and ``t_gID`` — ``inject_precomputed`` aliases the two
embeddings, and on decaf's ``ClusterCenter`` branch the first call is
``MSC_classify(self.t_embedded)`` *before* ``t_gID`` exists
(decaf.py:1441).
- empty ``dist_in`` and no ``multibandwidth`` — an explicit distance
matrix or a multiscale request is a real classification job, whatever
the other clauses say.
None of those states is reachable from ``fit_cluster_backend`` today; the
clauses are free and keep the guard honest if a caller changes.
"""
if (n_embedded is self.f_embedded
and self.f_data is self.t_data
and self.classify_mode == 'TrainingData'
and getattr(self, 't_gID', None) is not None
and dist_in.shape[0] == 0 and not multibandwidth):
# Mirror the base method's return convention: decaf.py:1452-1455
# unpacks a 2-tuple only under label_all='Both' (our pinned value).
# Preserve decaf's object semantics too: its normal classification
# creates independent lists and arrays for f_gID. Aliasing t_gID
# would let a caller mutate fitted training labels through f_gID.
def copy_gid():
return [np.array(slot, copy=True) for slot in self.t_gID]
return (copy_gid(), copy_gid()) if self.label_all == 'Both' else copy_gid()
return super().MSC_classify(n_embedded, dist_in=dist_in,
multibandwidth=multibandwidth,
multiscales=multiscales, **kwargs)
[docs]
class DecafClassifier(Classifier):
"""1-NN label transfer with radius abstention, delegated to decaf.
**No longer on any production path** (gitlab#76): fits now carry an
``abstain_radius`` and ``ClusterFit.classify`` builds a
:class:`~atom_hifi.training.cluster_base.NearestNeighborClassifier`.
Retained deliberately, as the *reference* the native port is checked
against — ``test_native_classifier_matches_decaf_exactly`` fits real decaf
and asserts the two agree slot-for-slot. Delete this class and that test
can no longer distinguish a faithful port from a plausible one. It also
keeps any explicitly-supplied ``ClusterFit.classifier`` working.
"""
def __init__(self, decaf_cluster):
self.decaf_cluster = decaf_cluster
[docs]
def classify(self, embedded, **kwargs):
return self.decaf_cluster.classify(embedded, **kwargs)
def _labels_from_t_gID(t_gID, n_atoms):
"""decaf t_gID (incl. trailing distinct slot) → flat label array.
Reproduces fit_decaf's historical atomic_tags loop exactly: the trailing
distinct/OOD slot keeps the catch-all label ``groupnum``.
"""
labels = np.zeros(n_atoms, dtype=int)
for gid, atom_list in enumerate(t_gID):
for atom in atom_list:
labels[atom] = gid
return labels
[docs]
class DecafBackend(ClusterBackend):
"""Mean-shift at bandwidth δ with 1-NN radius classification."""
name = 'decaf'
granularity_kind = 'distance'
has_classifier = True
has_auto_bandwidth = True
[docs]
def fit(self, desc_el, granularity, *, t_embedded=None, **kw):
# _PrecomputedDecaf is constructed unconditionally: it is the only place
# the fit-path short-circuit lives (issue #68), and decaf itself cannot
# be patched — it is a third-party dependency. With no embedding
# injected it behaves exactly like its base class.
precomputed = t_embedded is not None
ec = _PrecomputedDecaf(desc_el, embed_str='PCA', cluster_str='MSC')
if precomputed:
ec.inject_precomputed(t_embedded)
else:
ec.get_embed(dim=0)
ec.get_cluster(bandwidth=granularity,
bandwidth_estimate='HaarWavelet',
label_all='Both',
classify_mode='TrainingData')
labels = _labels_from_t_gID(ec.t_gID, desc_el.shape[0])
# On the SHM fast path, inject_precomputed() skips get_embed(), so
# ec.embed_op was never computed — defer it via PCAEmbedder's lazy
# desc_for_lazy path (same pattern as the kmeans/hdbscan backends
# below) instead of pulling the missing attribute eagerly.
embedder = (PCAEmbedder(np.asarray(ec.t_embedded), desc_for_lazy=desc_el)
if precomputed
else PCAEmbedder(np.asarray(ec.t_embedded), np.asarray(ec.embed_op)))
return ClusterFit(
labels=labels,
n_groups=int(ec.groupnum),
centers=np.asarray(ec.cluster_centers),
granularity=float(ec.bandwidth),
backend=self.name,
embedder=embedder,
# No classifier object: ClusterFit.classify builds a
# NearestNeighborClassifier from this radius on demand. δ IS
# decaf's abstention radius (`self.bandwidth`, decaf.py:1542), and
# `granularity` is that same float, so the answer is unchanged —
# the radius is passed explicitly rather than read off `granularity`
# because for a count-kind backend the two are different questions.
abstain_radius=float(ec.bandwidth),
)
[docs]
class KMeansBackend(ClusterBackend):
"""k-means in the shared PCA embedding; δ is the cluster count k."""
name = 'kmeans'
granularity_kind = 'count'
# gitlab#81 deliberately opens both classifier consumers. lambda_open
# tracks local partition spacing, so a kmeans loop can retain OOD atoms as
# k grows with the corpus and consume the full library. That is accepted
# existing termination behavior, warned once by run_loop rather than hidden
# behind a capability flag that contradicts ClusterFit.classify.
has_classifier = True
[docs]
def fit(self, desc_el, granularity, *, t_embedded=None, seed=42, **kw):
embedder = self.shared_embedder(desc_el, t_embedded)
X = embedder.t_embedded
n = len(X)
k = max(1, min(int(round(float(granularity))), n - 1))
km = KMeans(n_clusters=k, n_init=10, random_state=seed).fit(X)
labels = km.labels_.astype(int)
# lambda_open is calibrated from the FITTED PARTITION alone, using
# the real fitted sklearn centres — never recomputed means (the
# Stage A probe had to reconstruct means only because its cached
# artifacts never stored centres).
cal = lambda_local(X, labels,
centers=km.cluster_centers_, seed=seed)
cal['backend'] = 'kmeans'
cal['k'] = int(k)
cal['n_init'] = 10
cal['granularity'] = float(granularity)
if cal['lambda_open'] == 0.0:
# Exactly zero only — no near-zero threshold exists. Not a hard
# fail: a degenerate lambda must not break a granularity scan
# that never classifies anything; the hard fail belongs at the
# point of use, in the classifier constructor.
logger.warning(
'kmeans lambda_open is exactly 0.0 for k=%d, N=%d: every '
'cluster has occupancy <= 1, so no split can improve any '
'SSE — classification through ClusterFit.classify will '
'raise.', k, n)
# debug, not info: a granularity scan fits every δ for every element,
# so one info line per fit would bury the scan's own progress output.
# The same numbers ride on the fit as `abstain_calibration`.
logger.debug(
'kmeans k=%d: lambda_open=%.6g sqrt(lambda_open)=%.6g r100=%.6g '
'calibration %.3fs', k, cal['lambda_open'], cal['sqrt_lambda'],
cal['r100'], cal['wall_time_s'])
return ClusterFit(
labels=labels,
n_groups=k,
centers=km.cluster_centers_,
granularity=float(granularity),
backend=self.name,
embedder=embedder,
abstain_lambda=float(cal['lambda_open']),
abstain_calibration=cal,
)
[docs]
class HDBSCANBackend(ClusterBackend):
"""Sklearn density clustering with frozen hierarchy deployment (issue #97).
δ remains ``cluster_selection_epsilon`` for fitting. Prediction uses the
fitted density tree; promoted noise groups never become deployment targets.
"""
name = 'hdbscan'
granularity_kind = 'distance'
emits_noise = True
has_classifier = True
[docs]
@classmethod
def unsupported_reason(cls, granularity=None):
# Read the module attribute rather than closing over it, so tests can
# exercise both sides of the guard on whatever numpy they happen to run.
if not _NUMPY_BREAKS_HDBSCAN_EPS:
return None
# δ=0 is the one value provably safe: epsilon_search — the only caller
# of the broken traverse_upwards — runs only for a non-zero
# cluster_selection_epsilon. Refusing it too would strand the noise
# promotion and cache paths, which fit at δ=0 and work fine here.
if granularity is not None and float(granularity) <= 0.0:
return None
return (
f'the hdbscan backend needs numpy < 2.4 (found {np.__version__}): '
"on numpy >= 2.4 sklearn's HDBSCAN can raise TypeError for a "
'non-zero cluster_selection_epsilon, and δ IS that epsilon here. '
'Whether a given fit trips it depends on the data and on δ, so it '
'is refused rather than risked (scikit-learn#33219; the fix PR '
'#34243 is still unmerged). Use cluster_backend="decaf" or '
'"kmeans", or install numpy<2.4.')
[docs]
def fit(self, desc_el, granularity, *, t_embedded=None,
min_cluster_size=2, **kw):
embedder = self.shared_embedder(desc_el, t_embedded)
X = embedder.t_embedded
# identical constructor params to decaf's get_HDBSCAN
m = int(min_cluster_size)
# copy=True: SHM-served embeddings are read-only; in-place edits would abort
model = HDBSCAN(min_cluster_size=m, min_samples=m,
cluster_selection_epsilon=float(granularity),
allow_single_cluster=True, copy=True).fit(X)
labels = model.labels_.astype(int)
from atom_hifi.training.hdbscan_prediction import build_prediction_state
prediction = build_prediction_state(model, X, m)
n_groups = int(labels.max()) + 1 if (labels >= 0).any() else 0
fit = ClusterFit(
labels=labels,
n_groups=n_groups,
centers=medoids_from_labels(X, labels, n_groups),
granularity=float(granularity),
backend=self.name,
embedder=embedder,
hdbscan_prediction=prediction,
)
if fit.noise_frac > NOISE_FRAC_WARN:
logger.warning(
'hdbscan noise fraction %.2f > %.2f — softness/superionic-continuum '
'signature (cf. LLZO 0.42); density clustering may be unsuitable for '
'this sublattice (BACKEND_STRATEGY).', fit.noise_frac, NOISE_FRAC_WARN)
return fit
#: name → backend class. The single source of truth for what backends exist
#: and what they can do; ``GRANULARITY_KIND`` and ``BUILTIN_BACKENDS`` are
#: derived from it so they cannot drift.
BACKEND_REGISTRY = {b.name: b for b in (DecafBackend, KMeansBackend, HDBSCANBackend)}
[docs]
def register_backend(backend):
"""Register a custom :class:`ClusterBackend` under its ``name``.
Registration makes the backend resolvable by name in this process, which is
what ``config`` validation and ``state_io`` cache tagging need. It does NOT
reach process workers — the child interpreter re-imports the module and gets
only the built-ins — so the executable spec is additionally threaded through
the worker kwargs (``HiFiWorkflow._backend_spec``).
The backend is stored as given. In particular, an instance is not replaced
by its class: constructor state is retained and required arguments need not
be reconstructed. That registered instance is a configured prototype;
resolution deep-copies it so every fit starts from the same state.
Returns the registered name, so callers can store the tag.
"""
impl = backend if isinstance(backend, type) else type(backend)
if not issubclass(impl, ClusterBackend):
raise TypeError(
'custom cluster backend must subclass '
f'atom_hifi.training.cluster_base.ClusterBackend; got {impl!r}.')
name = getattr(backend, 'name', None)
if not isinstance(name, str) or not name:
raise ValueError(
'custom cluster backend must define a non-empty str `name` — it is '
'the cache tag, and state_io enforces equality on it across runs.')
if name in BUILTIN_BACKENDS:
raise ValueError(
f'custom cluster backend name {name!r} collides with a built-in '
f'{BUILTIN_BACKENDS}; caches tagged with it would be ambiguous.')
granularity_kind = getattr(impl, 'granularity_kind', None)
if granularity_kind not in ('count', 'distance'):
raise ValueError(
'custom cluster backend `granularity_kind` must be either '
f"'count' or 'distance'; got {granularity_kind!r}.")
existing = BACKEND_REGISTRY.get(name)
if existing is not None and existing is not backend:
raise ValueError(
f'a different backend is already registered as {name!r}; pick '
'another name rather than shadowing it.')
BACKEND_REGISTRY[name] = backend
GRANULARITY_KIND[name] = granularity_kind
return name
[docs]
def resolve_backend(backend):
"""Resolve a name, subclass, or configured prototype to a fresh instance."""
if isinstance(backend, ClusterBackend):
entry = backend
elif isinstance(backend, type) and issubclass(backend, ClusterBackend):
return backend()
else:
entry = BACKEND_REGISTRY.get(backend) if isinstance(backend, str) else None
if entry is None:
raise ValueError(
f'Unknown cluster backend {backend!r}; choose from '
f'{tuple(BACKEND_REGISTRY)} or register a ClusterBackend via '
'register_backend (Python API only — see issue #69).')
if isinstance(entry, type):
return entry()
# A configured instance carries constructor state, but one fit must not
# mutate the object reused by the next fit. Serial scans otherwise observe
# one accumulating instance while process workers each receive a serialized
# copy, making results depend on n_jobs. A fresh deep copy preserves the
# configured state and gives both paths the same per-fit lifecycle.
try:
return copy.deepcopy(entry)
except Exception as exc:
raise TypeError(
f'custom cluster backend {entry.name!r} must be deepcopyable so '
'each fit starts from the same configured state.') from exc
[docs]
def fit_cluster_backend(backend, desc_el, delta, *, t_embedded=None,
min_cluster_size=2, seed=42,
eps_noise=None):
"""Fit one clustering backend for one element; return a :class:`ClusterFit`.
The returned fit is already **promoted**: any trailing noise/distinct slot
has been deduped at ``eps_noise`` and appended as real groups (issue #63), so
every consumer sees one uniform group namespace with no ``-1`` atoms.
Parameters
----------
backend : {'decaf', 'kmeans', 'hdbscan'} or a :class:`ClusterBackend`
desc_el : (N_el, D) ndarray — raw per-element descriptors.
delta : float — granularity: bandwidth δ (decaf), cluster count k (kmeans),
cluster_selection_epsilon (hdbscan).
t_embedded : ndarray, optional — precomputed PCA embedding (SHM fast path).
min_cluster_size : int — hdbscan only.
seed : int — kmeans only (n_init=10 pinned).
eps_noise : float, optional — dedup radius for trailing-slot promotion.
Required only when the backend actually produces a trailing slot
(hdbscan); kmeans and decaf leave it empty, so promotion is a no-op.
"""
impl = resolve_backend(backend)
# The one choke point every backend fit passes through, so an unusable
# environment cannot leak past it into raw upstream tracebacks. Guarding
# only the callers is what let plot_final_fit and extract_rep_surroundings
# — both of which refit outside any scan — through in the first draft.
#
# Not sufficient on its own: the scan's per-point `except Exception`
# (_fit_cfactor_worker) would swallow this into 'zero valid fits', so the
# scan entry points check as well and report before any point is tried.
reason = impl.unsupported_reason(granularity=delta)
if reason is not None:
raise RuntimeError(reason)
fit = impl.fit(desc_el, delta, t_embedded=t_embedded,
min_cluster_size=min_cluster_size, seed=seed)
return _apply_promotion(fit, eps_noise)
def _apply_promotion(fit, eps_noise):
"""Promote the trailing slot in place and stamp the provenance fields."""
n_before = fit.n_groups
labels, n_groups, centers, noise_derived, n_trailing = promote_trailing_slot(
fit.labels, fit.n_groups, fit.centers, fit.embedder.t_embedded, eps_noise)
fit.pre_promotion_noise_frac = (
float(n_trailing) / max(len(np.asarray(fit.labels)), 1))
fit.labels = labels
fit.n_groups = n_groups
fit.centers = centers
fit.noise_derived = noise_derived
if n_trailing:
logger.info(
' [%s] %d trailing/noise atoms → %d promoted groups '
'(density groups: %d)', fit.backend, n_trailing,
n_groups - n_before, n_before)
return fit
#: backend name → 'distance' | 'count'. Derived; do not hand-maintain.
#: register_backend() extends both this and BACKEND_REGISTRY together.
GRANULARITY_KIND = {n: b.granularity_kind for n, b in BACKEND_REGISTRY.items()}
#: the three shipped backends — frozen, and NOT extended by register_backend.
#: Validation checks BACKEND_REGISTRY (which grows); this tuple stays the
#: answer to "what does atom-hifi ship", e.g. for name-collision checks.
BUILTIN_BACKENDS = tuple(BACKEND_REGISTRY)