Skip to content

GRASS-MIL 1.0: preprint method, nested CV, niche discovery, statistics - #22

Merged
LRabuzin merged 23 commits into
mainfrom
feat/preprint-method
Oct 14, 2026
Merged

GRASS-MIL 1.0: preprint method, nested CV, niche discovery, statistics#22
LRabuzin merged 23 commits into
mainfrom
feat/preprint-method

Conversation

@LRabuzin

@LRabuzin LRabuzin commented Oct 14, 2026

Copy link
Copy Markdown
Collaborator

GRASS-MIL 1.0: the preprint method, nested CV, niche discovery, statistics

Makes the public package the clean implementation of the method in
https://www.biorxiv.org/content/10.64898/2026.10.08.757722v1. Every preprint setting is a
default of a Hydra preset and stays overridable.

  • Method fidelity: bounded log edge length, GINE 4x256 (learnable eps, layer norm, max pool,
    xavier inputs), projected gated per-class attention (32/8, kaiming), logit-space pooling, one
    region per bag, region-level selection; model=grass_mil, regime=mil|mil_frozen| linear_probehead|linear_ftall|linear_ftall_reg; BGRL "A3" pretraining preset.
  • New CLIs: grass-mil-probe, grass-mil-nested-cv, grass-mil-extract, grass-mil-niches,
    grass-mil-metaclusters, grass-mil-scores, grass-mil-spatial-stats; patient-level
    aggregation in grass-mil-predict.
  • Docs: README landing page with the preprint, method correspondence, nested CV, niche
    discovery, evaluation statistics, spatial statistics; CITATION.cff; CHANGELOG.
  • Cleanup: stale configs, dummy data, shims and internal notes removed; version 1.0.0;
    hydra-core 1.3.5+.
  • Vignettes re-executed (01, 02) and a third added (03: nested CV, niches, metaclusters).
  • Validation against the original NSCLC runs recorded in docs/method_correspondence.md.

Tests: 476 passed (+1 slow end-to-end chain).

LRabuzin and others added 23 commits August 17, 2026 10:09
@LRabuzin @claude
Strip a trailing blank line that failed end-of-file-fixer
6ad7bf5 
The code-quality workflow runs pre-commit, and end-of-file-fixer rewrote
docs/interpretability_suite.md. pre-commit exits non-zero whenever a hook
modifies a file, so the job failed on a single trailing newline left by an
earlier docs edit.

Local checks missed it because they ran ruff only, which does not police
markdown end-of-file. Running `pre-commit run --all-files` reproduces it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@LRabuzin @claude
Untrack the 20 MB mibitof demo file; vignettes download it at runtime
86636b6 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 1: preprint fidelity for graph, encoder and attention; composab…
e7ed3aa 
…le model presets

- log_distance edge feature clip((ln d - 2)/3, 0, 1) in the Delaunay builder
- GIN/GINE train_eps, xavier input embeddings, kaiming attention linears
- attention defaults 32/8 with dropout, attention_space() for the pre-tanh space
- categorical unassigned_label/unassigned_index as an alternative to a reserved row
- loss weights accept lists from Hydra
- model/encoder and model/attention become real config groups; model/grass_mil preset;
  dead heads/loss/ssl groups and components.yaml removed

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 2: one region per bag, regimes with scaled learning rates, regi…
b16e4bc 
…on-level selection

- samplers: subsample_count, fixed_per_unit_seed, unit_as_bag, num_neighbors=-1 (unlimited);
  test/predict loaders score every root once
- SupervisedModule: per-component learning-rate scales (0 = frozen, encoder kept in eval),
  head pre-fit loading and freezing, inverse-frequency class weights from the training
  split, fine-tune-time edge/feature drop, gradient clipping and AMP-aware scheduler
  stepping in the hyperbatch flush, per-step cosine with T_max auto, region-level
  val/test loss, AUROC and balanced accuracy, test predictions export
- configs: regime group (mil, mil_frozen, linear_probehead, linear_ftall, linear_ftall_reg),
  data/sampler as a composable group with ego_graph_mil, scheduler/cosine_mil_step,
  callbacks mil_selection and refit, trainer/gpu_amp
- train: metrics_per_epoch.jsonl and test_predictions.csv artefacts

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 3: explicit splits, probe pre-fit and nested cross-validation
7902572 
- data/components/splits.py: ExplicitSplit (id lists or files), fold json
  reader/writer, LabelTable/load_labels, make_inner_folds (sorted ids,
  StratifiedGroupKFold, reseeds when a fold lacks a class)
- datamodule: data.split.explicit re-partitions an existing precompute
  (val may equal test for refit runs); graph_labels.patient_column ->
  data.patient_id and index entries carry patient_id / graph_y
- models/training/probe.py + grass-mil-probe: frozen-encoder ego-graph
  embeddings (fixed roots per region), StandardScaler + LogisticRegression
  folded into graph_head.net.0 (class-0 row zero)
- grass-mil-nested-cv: outer folds from json, patient-grouped inner folds,
  best epoch from metrics_per_epoch.jsonl, median refit, all-cells test,
  oof_predictions.csv, oof_patient_predictions.csv (mean), resume
- linear_* regimes use task.head_init_path: auto (resolved by nested CV)
- EpochMetricsWriter hooks on_validation_end so it sees the current epoch

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 4: BGRL A3 pretraining preset and fidelity
a6a365b 
- sampler runtime.merge_units: ego-graph roots drawn from the union of all
  regions (data/components/merge.py::merge_graph_units)
- BGRLModule: momentum_final / momentum_warmup_steps, EMA update after each
  optimizer step with the step's momentum, scheduler T_max: auto
- encoder_init_map=spacegm_bgrl remaps legacy SPACE-GM BGRL checkpoints
- presets: model/bgrl_a3, task/pretrain_bgrl_a3, data/sampler/ego_graph_pretrain,
  callbacks/pretrain_checkpoint, experiment/grass_mil_pretrain_a3
- graphs carry molecular_feature_names (identity reducer) for name-based
  feature noise; metadata.json fallback

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 5: cell cache extraction, patient-level aggregation, evaluation…
8eb9d60 
… statistics

- grass-mil-extract: per-cell cache (E fp16, instance logits, attention
  logits, cell types, positions, graph) of one or many models over all
  cells, from a nested-CV summary (held-out regions recorded) or explicit
  checkpoints; inference/model_loading.load_trained_module
- aggregation bag_scope=patient via a region->patient map: mean of core
  probabilities (mean_probability) or one merged bag with joint attention
  (attention_weighted); predict reads aggregation.patient_map
- evaluation package: bootstrap AUROC CI (units or groups), DeLong
  (Sun & Xu), within-fold z/rank standardisation, per-fold and stratified
  AUROC, operating points at within-fold quantile rules with bootstrap CIs,
  likelihood ratios, PPV/NPV at other prevalences; grass-mil-scores CLI

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 6: niche discovery on the cell cache
fc152aa 
- core/clustering: fit_niche_clustering (reduce on a sample, optional L2,
  any registry backend, k-NN assignment) and cosine_hdbscan with the
  preprint defaults; the report pipeline can use the same recipe
- interpretability/niches: relattn margins and attention space, per-niche
  table (share, regions present, composition, enrichment, empirical
  association, share AUROC, margin with bootstrap CI over regions, signed
  share, sign consistency, lift), co-localisation log2(obs/exp), figures
  (PCA/UMAP triptychs, overview, co-localisation), pipeline and report
- grass-mil-niches CLI with configs/niches.yaml (settings, no sweep)
- margin attribution plugin checks the additive identity against the bag table

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 7: cross-fold metaclusters, whole-dataset transfer, shared basis
4b63f2c 
- metaclusters/records: cluster every fold model's held-out cells with one
  fixed setting, fold-cluster records with composition
- metaclusters/match: Hellinger (or Aitchison) distance, constrained
  average linkage with cannot-link within fold, threshold scan (max
  replication, then fewest) or fixed, metacluster table, patient-level
  shares (each region labelled once) with Mann-Whitney/BH and AUROC
- metaclusters/transfer: every model labels every region in its own space,
  train vs held-out table, pairwise ARI, consensus patient shares
- metaclusters/basis: shared k-means basis with held-out/train/all
  attribution per niche
- grass-mil-metaclusters CLI (stage=all|cluster|match|transfer|basis);
  extract accepts explicit held-out lists per model

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 8: classical spatial statistics
42f4b42 
- spatial_stats: neighbourhood enrichment z-scores on the cell graph
  (numpy label permutations; squidpy backend optional), k-NN window
  cellular neighbourhoods (MiniBatchKMeans on a sample, per-patient CLR
  shares), contact fractions and infiltration, Ripley cross-K normalised by
  pi r^2, Mann-Whitney + BH at the patient level
- grass-mil-spatial-stats CLI with configs/spatial_stats.yaml (fine -> coarse
  group map, reference/partner groups, radii, permutations)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 9: cleanup, documentation, README, citation, vignettes
c291def 
- remove stale configs (experiment/debug), the embeddings shim, the dummy
  dataset and its skipped test, unreachable Moran helper, scripts/.env/.cursor
  residue; package __version__; version 1.0.0; hydra-core 1.3.5+; extras
  stats/plots; CI installs them
- docs: method_correspondence, nested_cv, niche_discovery, evaluation_stats,
  spatial_statistics; hyperparameter reference, workflows, capabilities,
  interpretability, dataset preparation updated; repository map regenerated;
  CITATION.cff, CHANGELOG.md; README landing page with the preprint
- units may be patches (tiles as units): split id_kind, nested_cv.id_kind,
  extract.unit_id; regime presets compose from predict/eval
- zero-validation callback filter resolves interpolations before filtering
- vignettes: 01 on the GRASS-MIL graph defaults, 02 on the pretraining
  experiment and regime=mil, new 03 (nested CV, niches, metaclusters, scores,
  spatial statistics)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 10 (part): end-to-end synthetic test, bf16-safe exports, balanc…
2d04415 
…ed synthetic folds

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Phase 10 (part): per-unit eval metrics for patch/sample bags, resume …
a10e785 
…validation, output dirs

- region-level metrics and test predictions follow task.bag_key (patch/sample
  units), validation instances grouped by their bag
- nested CV resume re-runs a refit whose predictions do not cover the test ids
- extras create the output directory before writing config/tag logs
- method correspondence: first fidelity results

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
nested CV: unit ids of patch and sample bags from the namespaced bag id
37f7df4 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Vignette 03 executed (nested CV, niches, metaclusters, scores, spatia…
acff9bf 
…l statistics)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
method correspondence: probehead and spatial-statistics fidelity results
8648b84 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
method correspondence: ftALL-reg fold-2 fidelity result
f31d500 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
model loading ignores probe pre-fit and init checkpoints (the trained…
e009a4d 
… checkpoint carries the head)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
method correspondence: niche fidelity result
9d64073 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Lint: names and unused variables flagged by the pinned pre-commit ruf…
16e5650 
…f; notebooks formatted

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Docs: getting-started walkthrough, configuration guide, cohort config…
676156d 
… template; pretraining as regime=pretrain_a3

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
Repository map regenerated; generator kept in tools/
c4d4ff5 
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin @claude
e2e test: structural checks instead of a seed-dependent ranking asser…
2cdc594 
…tion; test seed overridable

The planted niche is detected model-free by the spatial statistics step;
ranking it first with a two-epoch model failed on 3 of 5 seeds.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019CUeHkuT7ofCfUfC3Kty7A
@LRabuzin
LRabuzin merged commit 2303218 into main Oct 14, 2026
5 checks passed
@LRabuzin
LRabuzin deleted the feat/preprint-method branch October 14, 2026 09:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@LRabuzin