Feature/mil benchmark and representation methods by Dinesh-Adhithya-H · Pull Request #154 · lueckenlab/patpy

Dinesh-Adhithya-H · 2026-05-06T10:47:52Z

No description provided.

…ions New modules ----------- src/patpy/tl/mil_models.py TorchMILWrapper adapter for torchmil (ABMIL, TransMIL, DSMIL) with per-label training, cell-importance extraction, and multi-class support. src/patpy/tl/benchmark.py MILBenchmark: cross-validated benchmark for inductive MIL models. RepresentationBenchmark: benchmark for any method that returns sample representations (unsupervised Pseudobulk / CellComposition / GroupedPseudobulk and supervised MixMIL / ABMIL used as feature extractors), fitting a LogisticRegression / Ridge probe on train and evaluating on val/test. src/patpy/pp/splits.py make_sample_splits: stratified train/val/test splitting at sample level. src/patpy/pl/mil.py plot_attention_umap, plot_attention_by_cell_type, plot_attention_celltype_heatmap: cell-level attention visualisations. Changes to existing files -------------------------- src/patpy/tl/supervised.py MixMIL: get_sample_representations, fine_tune, predict_on_adata. PULSAR / PaSCient: get_sample_representations. src/patpy/pl/__init__.py, src/patpy/pp/__init__.py, src/patpy/tl/__init__.py Export new public symbols. pyproject.toml Add torchmil and related optional dependencies. docs/notebooks/mil_methods_example.ipynb Tutorial covering MIL training, cross-validated benchmarking, patient UMAPs, representation quality, ROC curves, and attention plots for multiple labels.

n_classes in _build_tensors now reads from _label_mappings (built on full adata) instead of np.unique on the training split, preventing index-out-of- bounds when a class is absent from the current split. Also rename `col` → `lk_vals` in MILBenchmark._run_one_split to avoid shadowing the outer split-column variable.

…tcome) _build_label_mappings and MILBenchmark pre-population both skipped numeric columns. Numeric classification labels (e.g. Outcome with values 2.0-5.0) were cast directly to int, producing out-of-range class indices when target values don't start at 0. Fix: build a sorted 0-indexed class mapping for all classification labels regardless of dtype.

ABMIL does not expose intermediate bag embeddings so it cannot be used as a representation extractor. Remove it from the rep_models dict.

Use a forward hook on model.classifier to capture the bag embedding (attention-weighted aggregation of cell features) as the patient representation. Re-adds ABMIL to the RepresentationBenchmark.

review-notebook-app · 2026-05-06T10:47:57Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

VladimirShitov · 2026-05-06T11:38:48Z

We moved the tutorials to https://github.com/lueckenlab/patpy_tutorials , can you push the notebook there instead?

VladimirShitov · 2026-05-06T11:40:35Z

Also please render the notebook. It will appear in the documentation in the same way as it looks as a notebook. So all the images and cell outputs must be there

VladimirShitov

This is a very valuable contribution, thanks for implementing it! I left some comments that need to be addressed before the PR can be merged. Briefly:

The tutorial notebook should be moved to a separate repo and rendered. Also, I'd advocate for extending supervised methods example notebook instead of creating a new one.
Benchmark functionality should be deleted from this PR and moved to another one. It's definitely very interesting, but too much for this PR, and it will slow down the merging.
Some code pieces need improvement, see the comments below.

VladimirShitov · 2026-05-06T11:42:53Z

+            "raw pre-softmax scores — returning softmax-normalised weights instead.",
+            stacklevel=3,
+        )
+    return model.get_cell_importance(label=label)


This should simply be an extension of get_cell_importance method in the corresponding class. I don't see necessity in a separate function, and certainly not in the plotting module.

VladimirShitov · 2026-05-06T11:51:13Z

+    att_col = f"{label}_importance"
+    weight_label = "Mean attention (softmax)" if normalized else "Mean attention score (pre-softmax)"
+
+    frame = pd.DataFrame(


This is repeated in every function, worth a separate util function

VladimirShitov · 2026-05-06T12:06:57Z

+            # Capture pre-softmax inputs via a hook on nn.Softmax modules
+            pre_softmax: list[torch.Tensor] = []
+
+            def _make_hook(store: list):


Why is it needed on what does it do? Please leave a comment in the code

VladimirShitov · 2026-05-06T12:07:42Z

+        model.eval()
+        weights_all: list[np.ndarray] = []
+
+        if normalized:


This function looks very complicated and some blocks of code repeat in different branches. Can it be simplified?

VladimirShitov · 2026-05-06T12:08:30Z

+        captured: list[torch.Tensor] = []
+
+        def _hook(module, inp, out):
+            captured.append(inp[0].detach().cpu())


Is there a simpler way of doing it?

VladimirShitov · 2026-05-06T12:10:00Z

+        lr: float = 1e-3,
+        weight_decay: float = 1e-4,
+        device: str = "auto",
+        seed: int = 42,


Just to double check, are the defaults the same as in original repo or the ones authors recommend?

VladimirShitov · 2026-05-06T12:10:38Z

+            weight_decay=weight_decay,
+            device=device,
+            seed=seed,
+        )


I love how simple these classes look, this is great!

torchmil models expose return_rep in their forward signature; use it directly instead of a classifier-input hook. Falls back to the hook for models that don't support return_rep.

- Add TransMIL and DSMIL to make_models with max_cells_per_bag=2000 to prevent OOM on large COMBAT bags (up to ~9k cells/patient) - Use return_rep=True (local torchmil extension) for bag embeddings with fallback to classifier-input hook for stock torchmil - Add paper-ready summary table: mean ± std on test set, pivoted by label × model × metric - Add patient representation UMAPs for all methods (unsupervised + MixMIL/ABMIL) using model.plot_embedding()

- Remove module-level docstring header from pl/mil.py (reviewer request) - Replace _get_cell_importance wrapper with _store_importance that writes weights into adata.obs for sc.pl compatibility - plot_attention_umap: use sc.pl.umap per-class subset instead of raw matplotlib scatter - plot_attention_by_cell_type: use sc.pl.violin grouped by cell type - Add normalized= parameter to MixMIL.get_cell_importance for API consistency with TorchMILWrapper (MixMIL always returns softmax weights) - Fix val_size docstring in splits.py: clarifies it is a fraction of all samples, not of remaining (non-test) samples Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

VladimirShitov · 2026-05-06T14:46:30Z

Could you also add tests to make sure the code is running well? Coding agents typically do it well

LArnoldt · 2026-05-06T23:15:11Z

+        label_keys: list[str],
+        tasks: list[_PREDICTION_TASKS],
+        n_splits: int = 3,
+        split_col: str = "split",


There is in common cases no split for unsupervised representation methods, no?

LArnoldt · 2026-05-06T23:16:48Z

+        else:
+            return pd.Series(probe.predict(X), index=index, name=label)
+
+    def _compute_metrics(


Running time could be interesting to report as well. @VladimirShitov wdyt?

Yes, definitely!

VladimirShitov · 2026-06-22T11:28:18Z

@Dinesh-Adhithya-H , could you address comments in this PR and move the benchmarking code to another one?

Dinesh-Adhithya-H added 5 commits May 6, 2026 10:08

Remove ABMIL from RepresentationBenchmark (no bag embeddings)

fe7f340

ABMIL does not expose intermediate bag embeddings so it cannot be used as a representation extractor. Remove it from the rep_models dict.

Implement get_sample_representations for TorchMILWrapper (ABMIL etc.)

c416ae9

Use a forward hook on model.classifier to capture the bag embedding (attention-weighted aggregation of cell features) as the patient representation. Re-adds ABMIL to the RepresentationBenchmark.

VladimirShitov reviewed May 6, 2026

View reviewed changes

Dinesh-Adhithya-H and others added 4 commits May 6, 2026 14:15

Use return_rep=True for bag embeddings in get_sample_representations

737d573

torchmil models expose return_rep in their forward signature; use it directly instead of a classifier-input hook. Falls back to the hook for models that don't support return_rep.

Add max_cells_per_bag to TransMIL and DSMIL constructors

0f31ad5

LArnoldt reviewed May 6, 2026

View reviewed changes

Uh oh!

Conversation

Dinesh-Adhithya-H commented May 6, 2026

Uh oh!

review-notebook-app Bot commented May 6, 2026

Uh oh!

VladimirShitov commented May 6, 2026

Uh oh!

VladimirShitov commented May 6, 2026

Uh oh!

VladimirShitov left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VladimirShitov commented May 6, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VladimirShitov commented Jun 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants