feat: add abmil_benchmark CLI for precision benchmarking

[raylim]

Summary

Adds a new abmil_benchmark CLI to the Mussel package that trains a Gated-ABMIL classifier on per-slide H5 feature files and reports AUROC across multiple seeds with bootstrap 95% CIs.

This supports benchmarking how lower-precision feature storage (float32 / float16 / bfloat16) affects downstream slide-level ABMIL classification performance, complementing the existing tile-level linear_probe_benchmark.

Used by: mussel-nf PR #20 — wires this CLI into the precision-benchmarking pipeline as ABMIL_BENCHMARK_WORKFLOW

---

Design

• Follows linear_probe_benchmark conventions (Hydra config, ConfigStore, same JSON output format)
• Reads per-slide H5 files (h5["features"] shape: n_tiles × feature_dim)
• cfg.dtype (float32 / float16 / bfloat16) cast at load time to simulate information loss from reduced-precision storage; model always trains in float32 for a clean benchmark signal
• AbmilClassifier wraps mussel.models.abmil.ABMIL (n_branches=1) + nn.Linear head
• Supports an optional pre-defined split column in the labels parquet, or falls back to a random slide-level split
• bfloat16 H5 storage handled via _numpy_to_torch (|V2 opaque void detection)

Output format (results.json)

{
  "dtype": "float32",
  "target_col": "label",
  "n_slides": 200,
  "pos_rate": 0.48,
  "n_seeds": 3,
  "seeds": [42, 43, 44],
  "val":  { "auroc": { "mean": 0.82, "std": 0.03 } },
  "test": { "auroc": { "mean": 0.79, "std": 0.02, "bootstrap_ci_95": [0.71, 0.86] } }
}

Val AUROC mean will be null when the val split is single-class (too few slides). The Nextflow summary process handles this gracefully.

CLI usage

abmil_benchmark \
    features_dir=/path/to/h5s \
    labels_parquet=/path/to/labels.parquet \
    target_col=my_binary_label \
    dtype=float16 \
    n_seeds=3 \
    n_epochs=20 \
    output_summary_json=results.json

---

Files changed

File	Change
mussel/cli/abmil_benchmark.py	New CLI (~400 lines)
pyproject.toml	Add abmil_benchmark entry point
tests/mussel/cli/test_abmil_benchmark.py	New test file (9 tests)

---

Tests

All 9 tests pass:

Test	What it covers
test_load_h5_features_float32	H5 loading, float32 passthrough
test_load_h5_features_float16_cast	dtype cast at load time
test_slide_dataset_getitem	SlideDataset.__getitem__
test_collate_fn_padding	padding to N_max, mask shape
test_collate_fn_mask_correctness	mask marks padded positions
test_abmil_classifier_forward	model forward without padding
test_abmil_classifier_forward_with_padding	model forward with mask
test_split_no_leakage	slide IDs don't appear in multiple splits
test_bootstrap_ci	95% CI bounds are ordered and within [0,1]
test_end_to_end	smoke test: CPU, synthetic dataset, valid JSON output

---

Bug fixes (discovered during testing)

1. best_state initialization: initialized to None, causing a crash when the val AUROC was always undefined (single-class val split). Fixed: initialize to the model's initial weights before the training loop.
2. _sanitize_for_json inf/nan: json.dumps(float('-inf')) produces -Infinity — not valid JSON. Groovy's JsonSlurper (used in Nextflow) raises JsonInternalException. Fixed: replace inf/nan with None before serialization.

[Copilot]

Pull request overview

Adds a new abmil_benchmark command-line tool to benchmark ABMIL slide classification performance across simulated feature-storage precisions, plus supporting tests and a small precision-handling fix in aggregate_sample_features.

Changes:

• Introduces mussel.cli.abmil_benchmark (Hydra-configured) to train/evaluate an ABMIL classifier over per-slide H5 features and report AUROC across seeds with bootstrap CIs.
• Adds an abmil_benchmark console script entry point.
• Extends aggregate_sample_features to upcast float16/bfloat16 features to float32 on read, with new regression tests.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
mussel/cli/abmil_benchmark.py	New benchmark CLI: H5 feature loading, slide dataset/collate, ABMIL classifier, split/train/eval, JSON summary output.
pyproject.toml	Registers abmil_benchmark as a package entry point.
tests/mussel/cli/test_abmil_benchmark.py	Adds unit + integration smoke tests for the new CLI components.
mussel/cli/aggregate_sample_features.py	Upcasts float16/bfloat16 H5 feature arrays to float32 for downstream compatibility.
tests/mussel/cli/test_aggregate_sample_features.py	Adds tests validating float16/bfloat16 upcast behavior; minor formatting updates.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.