feat(stats): projection statistics (cluster-validity + faithfulness)

[jcoludar]

Summary

Adds a protspace.stats subsystem that computes per-projection quality statistics at
preparation time and bakes them into the .parquetbundle as an optional fifth part. This is the
engine half of the projection-statistics MVP (tracking issue: tsenoner/protspace_web#219;
related: #31). The protspace_web PR that consumes it follows separately — this PR lands first.

Today the pipeline (embed → project → annotate → bundle) ships coordinates with zero quality
metrics, so judging a projection is purely visual. This change answers two questions per projection:

• Cluster structure — cluster_validity: KMeans with a distance-to-chord elbow estimate of
  K, scored by silhouette, Davies–Bouldin, Calinski–Harabasz on the projection coords.
• Projection faithfulness — faithfulness: kNN-overlap and trustworthiness / continuity
  between the source embedding and the projection — i.e. how much the reduction distorted the
  neighbourhood structure.

What's in this PR

• src/protspace/stats/ — a generalized Statistic contract (each statistic declares whether it
  needs the embedding and returns one or more StatRows) + a lazy STATISTICS registry mirroring
  the existing REDUCERS pattern. sklearn imports stay function-local to preserve CLI startup.
  • stats/cluster/kmeans_elbow.py — KMeans sweep K ∈ [2, min(round(√n), 50)], elbow via
    perpendicular deviation from the first→last inertia chord (argmax index → K).
  • stats/metrics/validity.py — silhouette (seeded sample above threshold) / DB / CH at that K.
  • stats/metrics/faithfulness.py — kNN-overlap@k and trustworthiness/continuity vs the embedding.
  • stats/driver.py — compute_statistics(...) iterating registered statistics per projection,
    isolating per-statistic and per-projection failures (a bad reduction is logged and skipped,
    never sinks the report).
• Tidy long-format output — an 8-column table (space_kind, space_name, stat_family, label_kind, metric, metric_kind, value, extra_json). New statistics add rows, not columns.
• Bundle I/O (data/io/bundle.py) — an optional fifth part statistics.parquet. Layout
  core(3) + settings? + statistics?; when statistics is present without settings, a zero-byte
  settings slot keeps the fifth position unambiguous. write_bundle / read_bundle /
  extract_bundle_to_dir and replace_settings_in_bundle are all updated — the last so
  protspace style no longer silently drops a trailing stats part.
• Wiring — ReductionPipeline.run (the one stage holding embeddings and projections)
  computes stats behind prepare --stats/--no-stats; new protspace stats -i emb.h5 -p project_dir -o statistics.parquet for the discrete path; bundle -s/--statistics folds a stats parquet in.
• No new dependency — scikit-learn (KMeans, silhouette/DB/CH, manifold.trustworthiness) is
  already a core dep.

Robustness hardening (post adversarial review)

This branch was put through a multi-agent review; the confirmed findings are fixed here:

• Faithfulness no longer silently drops trustworthiness/continuity for small n. trustworthiness
  requires n_neighbors < n/2 (strict); k is now clamped to (n-1)//2 instead of n-2.
• cluster_validity scores the full projection; faithfulness scores the embedding-aligned
  subset. Previously clustering could be scored on the id-intersection subset only.
• High-dim metric is threaded through (default_metric) so faithfulness uses the run's metric
  rather than defaulting to euclidean for PCA/MDS/PaCMAP.
• Multi-embedding attribution — exact id-set cover wins immediately (no ambiguous tie-break); a
  source column maps each projection to its own embedding.
• Precomputed (n×n) similarity/distance matrices are skipped for faithfulness (not embeddings).
• kNN self-exclusion is tie-safe on coincident points (filter self, not column-0 slice).
• n_clusters reports the achieved distinct-cluster count (KMeans can collapse on coincident
  points), keeping the requested K in extra.

Tests

• tests/test_stats.py, tests/test_stats_bundle.py, tests/test_stats_cli.py — known-answer
  numeric fixtures (blob separation; faithfulness on identity vs random projections; label-permutation
  alignment), the 8-column schema, the 5-part bundle round-trip, and the protspace style round-trip.
• Full suite: 534 passed locally.

Scope (MVP) & non-goals

In: per-projection cluster_validity (unsupervised/elbow) + faithfulness, baked at prep time,
carried in the bundle. Explicitly out (non-breaking future expansions): embedding-space
cluster-validity, annotation-feature label sources, on-demand recompute, the broader
ProtSpaceExtractor pair/edge/set analyses (future typed bundle parts), and frontend rendering. The
registry + long-format table leave seams for the scalar expansions.

Data-format change: additive, backward compatible — existing 3- and 4-part bundles read and
write unchanged.

Refs tsenoner/protspace_web#219, #31

[jcoludar]

Web counterpart (consumes this, lands after a stats-bearing release): tsenoner/protspace_web#295. Tracking issue: tsenoner/protspace_web#219.

[tsenoner]

Where to put the information:

• Cluster Elbow Kx -> save as an annotation
• Cluster Silhoutte Kx -> save as annotation with x beeing variable
• kNN-overlap + trustworthiness + continuity -> metadata per dimensionality reduction method (store with the other infos)
• silhouette, Davies–Bouldin, Calinski–Harabasz -> separate parquet file (as currently)

[jcoludar]

Routing is implemented across this PR and protspace_web#295. Where each statistic lands:

• kNN-overlap, trustworthiness, continuity → projection metadata (info_json.quality), shown per-metric in the projection panel.
• silhouette, Davies–Bouldin, Calinski–Harabasz → statistics.parquet (now the only thing in that part).
• cluster elbow-K label → per-protein cluster_ annotation column (categorical, auto-styled with a Kelly palette).
• per-point silhouette → per-protein silhouette_ column (numeric ramp).

Computed columns appear in color-by automatically, grouped under a new "Statistics" section. protspace stats handles carriage: faithfulness into projections_metadata, per-protein columns into the annotations parquet (-a), and the cluster legend styles via --settings-out for bundle --settings. The prep pipeline picks all of it up. Design + phased plan: the route-projection-statistics openspec change in the web repo. Both PRs green.

Defaults chosen:

• Naming: cluster_ / silhouette_, elbow K only (no sweep), K stored in metadata, not the column name (keeps names stable if K shifts on re-run). Your review said "Kx with x variable" — if K belongs in the visible label, that's a small change.
• Coloring: clusters are colored when the column is selected, not forced as the initial color-by. Making a cluster column the default initial selection is separate frontend work (initial-view).

Deferred (not blocking):

• Re-style round-trip: running protspace style on a routed bundle doesn't yet preserve the generated cluster styles — replace_settings_in_bundle rewrites settings wholesale.
• protspace prepare has no --no-stats-annotations toggle to skip the O(n²) per-point silhouette; the prep path is gated by -a. Defaulted on.
• statistics.parquet validity scalars are carried but not surfaced as a UI table.
• The prep service still pins protspace>=0.6; needs bumping to the release that includes this once it's cut. Older engines degrade gracefully — columns ship unstyled.

Sequencing: I folded the routing directly into #61/#295 rather than landing the opaque fifth part first and stacking. Can split the per-protein annotations into a separate PR if that's easier to review.

[tsenoner]

Post-review hardening pushed (fc7a1cd)

Ran a /simplify pass + an xhigh /code-review on this branch, then researched and implemented fixes for the deferred findings. Every change was checked by an adversarial verification pass. 565 tests pass, ruff clean.

⚠️ Behavior change

• prepare --stats now defaults to False (opt-in). It previously ran the full cluster-validity + faithfulness compute on every run and injected cluster_*/silhouette_* columns + auto legend styles into every bundle. Per safe-default / least-surprise guidance, heavy + output-shape-changing work is now explicit. --stats still works when passed.

Correctness fixes

• Unbundled output dropped settings. save_output(bundled=False) wrote statistics.parquet but silently ignored settings, so auto cluster-legend styles were lost with --no-bundle. Now writes settings.parquet too.
• protspace stats multi-input collapse. Repeated same-name -i inputs (two files sharing one embedding) were deduped to the last set, scoring faithfulness on a partial point set. Now unions them via merge_same_name_sets, matching prepare.
• --settings-out without -a was a silent no-op; now fails fast.
• Non-atomic rewrites. stats rewrote the user's projections_metadata.parquet / annotations parquet in place; an interrupt could truncate them. Now temp-file + atomic rename.
• 3D truncation. _load_reductions defaulted dims=2 when metadata was absent, dropping the z coordinate; now inferred from the data.
• Continuity metric. sklearn's trustworthiness applies metric only to its first arg, so continuity always ranked the embedding by euclidean even for cosine/manhattan projections. Added a correct dual (_continuity) — bit-identical on the euclidean default path.
• Silhouette inconsistency. The headline aggregate silhouette was a sampled estimate while the per-point column was exact; they disagreed for 5k<n≤20k. The aggregate is now the exact mean of the per-point values whenever that column is computed.
• Non-reproducible subsample. Faithfulness derived its subsample seed from sorted ids (order-invariant) but selected positionally; now selects on canonical id order, so scores are row-order invariant.

Scaling & cleanup

• KMeans elbow at scale. Above 50k points the K sweep now fits a bounded subsample with MiniBatchKMeans + predict (labels for all points); full-batch below the threshold is unchanged. Removed the write-only silhouette_optimal_k cross-check and its duplicate silhouette compute.
• Collapsed duplicated faithfulness try/except into a loop; replaced a self._stats_settings side-channel with an explicit (table, settings) return; removed a redundant empty-table branch, a needless list copy, a double np.unique, and an unused param; centralized the reduction source stamp via a closure.
• Left as-is (by decision): legend colors wrap past 21 clusters — >~22 categories exceed human discriminability anyway, and the frontend already repeats colors, shows the top 10 by default, and supports hover + hide/show toggles.

Notes for review

• Numeric outputs change only for large n (>5k–50k, on the now-opt-in --stats path): continuity for non-euclidean metrics, silhouette for 5k<n≤20k, and which proteins are subsampled. All existing fixtures use n≤300, so none changed, and no large-n golden fixtures exist.
• Added regression tests: continuity↔sklearn-dual equivalence, subsample determinism + order-invariance, aggregate = per-point mean, unbundled settings, --settings-out guard.

🤖 Generated with Claude Code