"""Saved-judge labelling and PCA-policy identities (experimental, opt-in).
Saved-judge labeller
--------------------
:func:`label_rows` is the nearest-centroid rule of a saved KMeans judge on already-embedded
rows ``z`` of any embedding dimension: Atom-HiFi's own :class:`NearestCentroidClassifier`
(squared distances, first minimum on exact ties, ``q = d1**2 / lambda_open``, OOD iff
``q > 1`` STRICTLY -> label ``-1``), then the judge's refined group order ``local_of_old`` and
element offset. Diagnostics are the decision margins ``q - 1`` (OOD iff > 0) and the
nearest-centroid gap ``d2 - d1`` (an exact tie gives 0). :func:`label_from_judge` reads those
inputs from a loaded judge. Nothing is fitted or re-embedded.
:func:`load_judge` verifies the file sha256 BEFORE unpickling and the element's stored
embedding operator against an identity the CALLER supplies (no system constants live here).
The operator is read from the stored field only: a judge without one is refused, never
re-embedded (recomputing it would be a PCA on the reference data).
PCA policies
------------
:class:`PCAPolicy` names how an embedding operator is chosen: ``fixed_k`` (top ``k``
eigenvectors) or ``native_auto`` (the broken-stick dimension). Identities are distinct strings
(``'fixed_k:21'`` vs ``'native_auto:broken_stick'``) even when the broken-stick dimension
equals ``k``. :func:`eigendecomposition` / :func:`broken_stick_dim` / :meth:`PCAPolicy.operator`
implement decaf's ``get_PCA`` rule with decaf's own arithmetic: covariance of the centred rows
``(X - mean)^T (X - mean) / N``, full ``eigh``, eigenpairs in descending order
(``argsort()[::-1]``), ``d_auto = #{k : bs_k <= lambda_k / sum(lambda)}`` with
``bs_k = (1/D) sum_{j=k..D} 1/j``, operator = leading eigenvectors, projection UNCENTRED
(``rows @ operator``), no whitening, no randomness. Calling it is the caller's decision:
importing this module computes nothing.
Importing this module does no IO and imports neither decaf, DScribe nor the Atom-HiFi
training stack; those load when a judge is loaded or rows are labelled.
"""
from __future__ import annotations
import hashlib
import json
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
__all__ = ['array_sha256', 'file_sha256', 'load_judge', 'label_rows', 'label_from_judge', 'PCAPolicy',
'PCAOperator', 'broken_stick', 'broken_stick_dim', 'eigendecomposition', 'pca_operator',
'ARRAY_SHA256_DEFINITION']
#: hash definition used for array identities (matches the stage's historical array hash)
ARRAY_SHA256_DEFINITION = 'sha256(dtype.str + json(list(shape)) + C-order bytes)'
[docs]
def array_sha256(a) -> str:
"""``sha256(dtype.str + json(list(shape)) + C-order bytes)`` of an array (0-d kept 0-d)."""
a = np.asarray(a)
h = hashlib.sha256(a.dtype.str.encode('ascii') + json.dumps(list(a.shape)).encode('ascii'))
h.update(np.ascontiguousarray(a).tobytes())
return h.hexdigest()
[docs]
def file_sha256(path, block=16 * 1024 ** 2) -> str:
h = hashlib.sha256()
with open(path, 'rb') as f:
for b in iter(lambda: f.read(block), b''):
h.update(b)
return h.hexdigest()
# ----------------------------------------------------------------------------------- saved judges
[docs]
def load_judge(path, *, expected_sha256, element, expected_n_groups=None, expected_embed_op_sha256=None,
expected_embed_op=None):
"""Load a saved :class:`atom_hifi.Judge` after verifying its identity.
``expected_sha256`` (file) is checked before unpickling. The element's stored embedding
operator must match ``expected_embed_op_sha256`` (:func:`array_sha256`) and/or
``expected_embed_op`` (dtype, shape and bytes); at least one is required.
``expected_n_groups`` optionally pins the fitted group count. Raises ``ValueError`` on
any mismatch.
"""
if expected_embed_op_sha256 is None and expected_embed_op is None:
raise TypeError('load_judge needs expected_embed_op_sha256 and/or expected_embed_op: rows embedded with a '
'different operator would be labelled silently')
path = Path(path)
got = file_sha256(path)
if got != str(expected_sha256):
raise ValueError(f'{path}: file sha256 {got} != expected {expected_sha256}; refusing to unpickle')
from atom_hifi import Judge
judge = Judge.load(path)
fit = (getattr(judge, 'cluster_per_el', None) or {}).get(element)
if fit is None:
raise ValueError(f'{path}: no fitted clustering for element {element!r}')
if expected_n_groups is not None and int(fit.n_groups) != int(expected_n_groups):
raise ValueError(f'{path}: element {element!r} has n_groups={fit.n_groups}, expected {expected_n_groups}')
op = getattr(getattr(fit, 'embedder', None), '_embed_op', None) # stored field only; the property recomputes a PCA
if op is None:
raise ValueError(f'{path}: element {element!r} has no stored embedding operator')
op = np.asarray(op)
if expected_embed_op_sha256 is not None and array_sha256(op) != str(expected_embed_op_sha256):
raise ValueError(f'{path}: element {element!r} embedding operator sha256 {array_sha256(op)} != expected '
f'{expected_embed_op_sha256}')
if expected_embed_op is not None:
ref = np.asarray(expected_embed_op)
if not (op.dtype == ref.dtype and op.shape == ref.shape and op.tobytes() == ref.tobytes()):
raise ValueError(f'{path}: element {element!r} embedding operator differs from expected_embed_op')
return judge
def _label_dtype(max_id):
for dt in (np.int16, np.int32, np.int64):
if max_id <= np.iinfo(dt).max:
return dt
raise OverflowError('group ids do not fit int64')
[docs]
def label_rows(z, centers, lambda_open, local_of_old, offset=0, *, chunk=32, diagnostics=False):
"""Nearest-centroid labels of embedded rows ``z`` (shape ``(..., d)``), ``-1`` = OOD.
``centers`` (n_groups, d) and ``lambda_open`` (squared-distance units) are the saved fit;
``local_of_old`` maps a fitted centre index to its refined local id and ``offset`` is the
element's group-id offset. ``chunk`` leading-axis entries are classified per call (the
rule is row-wise, so chunking never changes a value). Labels use the smallest of
int16/int32/int64 holding every id. With ``diagnostics`` returns
``(labels, q_margin, gap)`` where ``q_margin = q - 1`` and ``gap = d2 - d1``.
"""
from atom_hifi.training.cluster_base import NearestCentroidClassifier
z = np.asarray(z)
centers = np.asarray(centers, dtype=float)
if z.ndim < 2:
raise ValueError(f'z must have shape (..., d); got {z.shape}')
d = z.shape[-1]
if centers.ndim != 2 or centers.shape[1] != d:
raise ValueError(f'centers shape {centers.shape} does not match embedding dimension {d}')
n_groups = len(centers)
order = np.asarray(local_of_old)
if order.shape != (n_groups,) or not np.array_equal(np.sort(order), np.arange(n_groups)):
raise ValueError('local_of_old must be a permutation of range(n_groups)')
offset = int(offset)
if offset < 0:
raise ValueError(f'offset must be >= 0, got {offset}')
dtype = _label_dtype(n_groups - 1 + offset)
order = order.astype(dtype)
if int(chunk) != chunk or chunk < 1:
raise ValueError(f'chunk must be a positive int, got {chunk!r}')
clf = NearestCentroidClassifier(centers, n_groups, lambda_open)
lead = z.shape[:-1]
out = np.empty(lead, dtype)
if diagnostics:
qm = np.empty(lead, np.float64)
gap = np.empty(lead, np.float64)
for a in range(0, lead[0], int(chunk)):
b = min(a + int(chunk), lead[0])
p = clf.predict(np.asarray(z[a:b]).reshape(-1, d), return_confidence=diagnostics)
pred = p['label']
gid = np.full(len(pred), -1, dtype)
ok = pred >= 0
gid[ok] = order[pred[ok]] + offset
out[a:b] = gid.reshape((b - a,) + lead[1:])
if diagnostics:
qm[a:b] = (p['q'] - 1.0).reshape((b - a,) + lead[1:])
gap[a:b] = (p['d2'] - p['d1']).reshape((b - a,) + lead[1:])
return (out, qm, gap) if diagnostics else out
[docs]
def label_from_judge(judge, z, element, *, chunk=32, diagnostics=False):
""":func:`label_rows` with the centres, ``lambda_open``, refined order and offset of
``judge``'s ``element`` fit.
Only the case where this rule IS the saved judge's own deployed classifier is accepted: a
KMeans fit with ``abstain_lambda`` and no ``abstain_radius``, no explicit ``classifier``
override (which production ``ClusterFit.classify`` would use first) and applied refined
group ids. Anything else raises ``ValueError`` instead of labelling with a different rule.
"""
fit = (getattr(judge, 'cluster_per_el', None) or {}).get(element)
if fit is None:
raise ValueError(f'element {element!r}: the judge has no fitted clustering for this element')
backend = getattr(fit, 'backend', None)
if backend != 'kmeans':
raise ValueError(f'element {element!r}: backend {backend!r} is not supported; this labeller reproduces the '
'kmeans nearest-centroid classifier only')
if getattr(fit, 'classifier', None) is not None:
raise ValueError(f'element {element!r}: the saved fit carries an explicit classifier override, which '
'production ClusterFit.classify uses first; refusing to label with another rule')
if getattr(fit, 'abstain_radius', None) is not None:
raise ValueError(f'element {element!r}: the saved fit sets abstain_radius (a 1-NN rule); production refuses '
'a fit with both abstention fields, and this labeller implements only abstain_lambda')
lam = getattr(fit, 'abstain_lambda', None)
if lam is None:
raise ValueError(f'element {element!r}: the saved fit has no abstain_lambda (not a nearest-centroid judge)')
order = (getattr(judge, 'group_order_per_el', None) or {}).get(element)
offsets = getattr(judge, 'group_id_offset', None) or {}
if not getattr(judge, '_refined_ids_applied', False) or order is None or element not in offsets:
raise ValueError(f'element {element!r}: refined group ids are not applied on this judge (no refined order '
'or offset); labels would not be in the judge\'s published group-id space')
return label_rows(z, fit.centers, lam, order.local_of_old, offsets[element], chunk=chunk,
diagnostics=diagnostics)
# ----------------------------------------------------------------------------------- PCA policies
PCA_KINDS = ('fixed_k', 'native_auto')
[docs]
@dataclass(frozen=True)
class PCAPolicy:
"""How an embedding operator is chosen: ``fixed_k`` (with ``k``) or ``native_auto``."""
kind: str
k: int | None = None
def __post_init__(self):
if self.kind == 'fixed_k':
if isinstance(self.k, bool) or not isinstance(self.k, (int, np.integer)) or int(self.k) < 1:
raise ValueError(f"fixed_k needs an int k >= 1, got {self.k!r}")
object.__setattr__(self, 'k', int(self.k))
elif self.kind == 'native_auto':
if self.k is not None:
raise ValueError('native_auto takes no k: its dimension is an output (broken stick)')
else:
raise ValueError(f'kind must be one of {PCA_KINDS}, got {self.kind!r}')
[docs]
@classmethod
def fixed(cls, k):
return cls('fixed_k', k)
[docs]
@classmethod
def native_auto(cls):
return cls('native_auto')
@property
def identity(self) -> str:
"""Stable identity string; never equal across kinds, whatever the realized dimension."""
return f'fixed_k:{self.k}' if self.kind == 'fixed_k' else 'native_auto:broken_stick'
[docs]
@classmethod
def from_identity(cls, text):
if text == 'native_auto:broken_stick':
return cls.native_auto()
kind, _, k = str(text).partition(':')
if kind == 'fixed_k' and k.isdigit():
return cls.fixed(int(k))
raise ValueError(f'not a PCAPolicy identity: {text!r}')
[docs]
def operator(self, rows):
"""The operator this policy selects on ``rows`` (see :func:`pca_operator`)."""
return pca_operator(rows, self)
[docs]
@dataclass(frozen=True, eq=False)
class PCAOperator:
"""An embedding operator and the eigen-spectrum it was cut from.
``embed_op`` is ``(D, dim)``; ``eigenvalues`` are all ``D`` covariance eigenvalues in
descending order; ``broken_stick_dim`` is reported for both policies (for ``fixed_k`` it is
information only).
"""
policy: PCAPolicy
embed_op: np.ndarray = field(repr=False)
dim: int
eigenvalues: np.ndarray = field(repr=False)
broken_stick_dim: int
n_rows: int
@property
def identity(self):
return self.policy.identity
[docs]
def project(self, rows):
"""Uncentred projection ``rows @ embed_op`` (decaf ``embed_PCA``)."""
return np.dot(rows, self.embed_op)
[docs]
def broken_stick(segments) -> np.ndarray:
"""``bs_k = (1/D) sum_{j=k..D} 1/j`` for ``k = 1..D`` (decaf ``get_broken_stick`` arithmetic:
each suffix is its own ``np.sum``)."""
D = int(segments)
if D < 1:
raise ValueError(f'segments must be >= 1, got {segments!r}')
inv = 1 / np.arange(1, D + 1)
return np.array([np.sum(inv[i:]) for i in range(D)]) / D
[docs]
def broken_stick_dim(eigenvalues) -> int:
"""``#{k : bs_k <= lambda_k / sum(lambda)}`` over all ``D`` eigenvalues (descending order)."""
ev = np.asarray(eigenvalues, dtype=np.float64)
if ev.ndim != 1 or len(ev) == 0 or not np.all(np.isfinite(ev)):
raise ValueError('eigenvalues must be a non-empty finite 1-D array')
if np.any(ev[1:] > ev[:-1]):
raise ValueError('eigenvalues must be in descending order')
D = len(ev)
return int(np.sum(broken_stick(D)[:D] <= (ev[:D] / np.sum(ev[:D]))))
[docs]
def eigendecomposition(rows):
"""Descending eigenvalues and eigenvectors of the covariance of the centred rows."""
X = np.asarray(rows)
if X.ndim != 2 or X.shape[0] < 2 or X.shape[1] < 1:
raise ValueError(f'rows must be a 2-D array with at least two rows; got shape {X.shape}')
if not np.all(np.isfinite(X)):
raise ValueError('rows must be finite')
centred = X - np.tile(np.mean(X, axis=0), (X.shape[0], 1))
cov = np.dot(centred.T, centred) / X.shape[0]
v, F = np.linalg.eigh(cov)
order = v.argsort()[::-1]
return v[order], F[:, order]
[docs]
def pca_operator(rows, policy) -> PCAOperator:
"""Operator chosen by ``policy`` on ``rows``: top-``k`` (fixed_k) or top-``d_auto``
(native_auto) eigenvectors of one eigendecomposition."""
if not isinstance(policy, PCAPolicy):
raise TypeError('policy must be a PCAPolicy')
evals, evecs = eigendecomposition(rows)
d_auto = broken_stick_dim(evals)
dim = policy.k if policy.kind == 'fixed_k' else d_auto
if dim > len(evals):
raise ValueError(f'{policy.identity}: k={dim} exceeds the descriptor dimension {len(evals)}')
if dim < 1:
raise ValueError(f'{policy.identity}: selected dimension {dim} < 1')
return PCAOperator(policy=policy, embed_op=evecs[:, :dim], dim=int(dim), eigenvalues=evals,
broken_stick_dim=d_auto, n_rows=int(np.asarray(rows).shape[0]))