From 5bb97f9f9f7b5e8fe6986265f0244e737a3e99f9 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Wed, 29 Jul 2026 14:29:24 +0200
Subject: [PATCH 01/14] sample_key to sample_key_list: save all sample keys

---
 .claude/agents/security-reviewer.md      |  13 +++
 .claude/background.md                    | 142 +++++++++++++++++++++++
 .gitignore                               |   2 +-
 CLAUDE.md                                | 109 +++++++++++++++++
 src/interscale/model/base/_base_model.py |   6 +-
 tests/test_base_model.py                 |  53 +++++++++
 6 files changed, 321 insertions(+), 4 deletions(-)
 create mode 100644 .claude/agents/security-reviewer.md
 create mode 100644 .claude/background.md
 create mode 100644 CLAUDE.md
 create mode 100644 tests/test_base_model.py

diff --git a/.claude/agents/security-reviewer.md b/.claude/agents/security-reviewer.md
new file mode 100644
index 0000000..02da9c7
--- /dev/null
+++ b/.claude/agents/security-reviewer.md
@@ -0,0 +1,13 @@
+---
+name: security-reviewer
+description: Reviews code for security vulnerabilities
+tools: Read, Grep, Glob, Bash
+model: opus
+---
+You are a senior security engineer. Review code for:
+- Injection vulnerabilities (SQL, XSS, command injection)
+- Authentication and authorization flaws
+- Secrets or credentials in code
+- Insecure data handling
+
+Provide specific line references and suggested fixes.
diff --git a/.claude/background.md b/.claude/background.md
new file mode 100644
index 0000000..882c83a
--- /dev/null
+++ b/.claude/background.md
@@ -0,0 +1,142 @@
+# Scientific background
+
+This document summarizes the biological motivation, model design, and terminology from the
+InterScale preprint so that the *purpose* behind the code — not just its structure — is clear.
+Read this when working on anything touching model architecture, loss functions, the
+interpretability/evaluation pipeline, or when the code's intent isn't obvious from names alone.
+
+> Drummer, F., Jiménez, S., Di Marco, F., Schaar, A.C., Pentimalli, T.M., Beckman, J.L., Rajewsky, N.,
+> Theis, F.J. *InterScale reveals multi-scale cellular interaction programs in spatial
+> transcriptomics.* bioRxiv (2026). https://doi.org/10.64898/2026.05.07.723456
+> (see also `docs/references.bib`, key `DrummerJimenez_2026`)
+
+## Motivation
+
+Cell–cell communication happens at multiple spatial scales simultaneously: autocrine/juxtacrine
+signaling within a few micrometers, paracrine gradients across a local neighborhood, and
+tissue-wide (or morphogen-gradient) coordination. Existing computational methods for inferring
+cell-cell communication from spatial transcriptomics typically model *one* scale — either
+adjacency-based local neighborhoods (GNN/niche methods like NCEM, NicheCompass) or generic
+long-range dependencies — but not both jointly, and many rely on ligand-receptor priors that are
+incomplete or infeasible for gene-panel-limited imaging-based platforms.
+
+InterScale's core idea: combine a graph neural network (local, short-range, adjacency-constrained)
+with a transformer (global, long-range, effectively fully-connected) in one architecture, so a
+cell's transcriptional state is decomposed into local-neighborhood and tissue-wide contributions
+that can be interpreted *separately*. This is why the codebase is organized around parallel
+"local component" / "global component" module hierarchies rather than a single monolithic model.
+
+## Architecture, mapped to code
+
+The model follows the `GraphTrans` architecture (local GNN → global transformer), with these
+building blocks:
+
+- **Local component** (`interscale.module.local_modules`) — produces `H_local`, a per-cell
+  embedding informed only by `k`-hop spatial neighbors (`k` = number of GCN/GIN layers). Four
+  variants are implemented:
+  - `GCN` — aggregates neighbor expression via the (self-loop-augmented, row-normalized) adjacency
+    matrix; fixed at 2 layers to avoid oversmoothing/oversquashing.
+  - `GIN` — Graph Isomorphism Network update rule with a learnable epsilon and an MLP.
+  - `SCVI` — wraps `scvi.nn.Encoder`; captures gene-expression structure, *not* spatial
+    neighborhood information (despite living in the "local" module tree).
+  - `Precomputed` — loads embeddings computed outside InterScale (e.g. CellCharter, BANKSY) from
+    `adata.obsm`; not trainable, so a dual-decoder setup isn't possible with this option.
+- **From local to global**: `H_local` is LayerNorm'd, then padded/truncated to a fixed sequence
+  length `S` per sliding window (`pad_batch`), and a CLS token is prepended. No positional
+  encoding is added — the GNN already encoded spatial structure, and the transformer is meant to
+  find interactions *without* distance restriction.
+- **Global component** (`interscale.module.global_modules`,
+  `TransformerNodeEncoderHook`) — a BERT-style multi-head self-attention transformer producing
+  `H_global`, i.e. tissue/window-wide context. Notably, the attention mask is the *inverse*
+  adjacency matrix (`M = 1 - A`): the transformer is explicitly steered toward attending to cells
+  that are **not** already connected in the local graph, so local and global signal don't
+  duplicate each other.
+- **Decoders** (`interscale.nn`) — separate linear (or nonlinear) decoders for `H_local` and
+  `H_global` reconstruct/predict gene expression or labels independently. Keeping the decoders
+  separate is deliberate: it lets downstream analysis attribute an effect to "local" or "global"
+  origin. The linear decoder is used for all reported experiments because it stays interpretable
+  (decoder weights → standardized gene loadings, see below); the nonlinear decoder trades that
+  away for accuracy.
+- **Masking / self-supervision** — a fraction `pct_mask_nodes` of cells per graph have expression
+  zeroed out (`tl.masking.apply_mask`); loss is computed only on masked cells. Classification uses
+  weighted cross-entropy; regression uses a scaled cosine error loss (captures direction +
+  magnitude) or Gaussian NLL.
+- **`CombinedModule` vs `DualDecoderCombinedModule`** — the latter is what gives each component
+  its own decoder (`cfg.model.decoder.dual_decoder = True`); this is required for the
+  local-vs-global attribution analyses described below.
+
+## Downstream interpretability — the actual point of the model
+
+Prediction accuracy is a *proxy* used to confirm that local/global information is real signal, not
+the end goal. The paper's real contribution is post-hoc, scale-resolved interpretation, at three
+levels (`interscale.evaluation`):
+
+- **Tissue/graph level** — the CLS token (connected to every cell) acts as a tissue-level summary.
+  Its attention scores are decomposed into "vertical"/"horizontal" (sender/receiver) components and
+  can be projected back onto the spatial slide to show which regions/cell types drive a
+  classification (e.g. condition prediction).
+- **Cell level** — raw attention weights aren't directly interpretable (attention-sink effects,
+  softmax normalization differs per sliding window, layer/head averaging can wash out signal). The
+  paper uses gradient-based **attention relevance** (averaging heads weighted by gradient
+  importance, clamped ≥0, renormalized) rather than raw attention, and focuses on the **net
+  attention flow** (incoming − outgoing, normalized per window) to get directional sender→receiver
+  interaction strength between cell types, aggregated as a "flow map" and decomposed spatially into
+  divergence-based **source/sink domains** via K-means.
+- **Gene level** — decoder weights are converted into **standardized gene loadings** (decoder
+  weight scaled by embedding-dimension-stdev / gene-stdev ratio) — this is what makes the linear
+  decoder choice matter: it's a scale-independent, cell-type-agnostic measure of "how much does
+  this embedding dimension drive this gene's expression," usable to rank informative embedding
+  dimensions and run functional enrichment separately for local vs. global programs. A
+  complementary approach compares each decoder's per-gene reconstruction rank (local rank − global
+  rank) to double-check which scale "owns" a gene.
+- **Moran's I across neighborhood size** is used throughout as an independent, model-free way to
+  quantify a gene's effective spatial length scale (fast autocorrelation decay = local; slow decay
+  = global) — used to validate that what the local/global embeddings pick up biologically matches
+  genes' actual spatial organization.
+
+## Key terminology / notation (paper → code)
+
+| Paper notation | Meaning | Code |
+|---|---|---|
+| `X ∈ R^{N,F}` | gene expression matrix (cells × genes) | expression layer registered via `AnnDataManager` |
+| `A` | adjacency / spatial connectivity matrix | built via `squidpy.gr.spatial_neighbors`, converted to edge index by `geome.transforms.AddEdgeIndex` |
+| `H_local` / `E_local` | local (neighborhood) embedding | local module output |
+| `H_global` / `E_global` | global (tissue-wide) embedding | global module (transformer) output |
+| `CLS` | classification/tissue-summary token, attends to all cells | prepended before transformer; extracted in `CombinedModel.get_model_output` |
+| `M = 1 - A` | transformer attention mask (blocks locally-connected pairs) | global module attention mask |
+| `k` / `L_GCN` | number of GCN/GIN layers = local receptive field radius | `cfg.model.local_component.parameters.*` |
+| `S` | max transformer sequence length per sliding window | `pad_batch` / dataloader |
+| `p_m` / `pct_mask_nodes` | fraction of nodes masked per graph | `cfg.dataset` / `GraphAnnDataModule` |
+| `M_net` | net attention flow matrix (directional sender→receiver) | attention relevance / net-flow analysis in `interscale.evaluation` |
+| sliding window `w`, overlap `o`, step `T = w - o` | tissue partitioning for scalability | `prepare_geome_dataset` / squidpy sliding-window util |
+
+## Datasets used in the paper (for context, not bundled with the repo)
+
+- **Molecular Cartography SHH organoids** (Legnini et al.) — optogenetically induced Sonic
+  Hedgehog signaling in neural tube organoids; this is the dataset behind `datasets._legnini` and
+  `config_files/legnini_example.yaml`. Used to show InterScale recovers spatially-local neuronal
+  differentiation programs (`GLI1`, `ISL1`, `TUBB3`) vs. broader progenitor/morphogen-regulation
+  programs (`PROM1`, `NEUROG2`, `HHIP`) from the *same* tissue, without cell-type labels.
+- **CosMx human pancreas (ND vs. T1D)** (Melton, Jiménez et al.) — used for node classification and
+  the sender-receiver attention analysis; shows disease-associated reorganization (mast cell
+  infiltration into islets in T1D) and separates local cell-state programs (endocrine/metabolic,
+  e.g. `INS`, oxidative stress) from global immune/stromal signaling programs (`C1QC`, `IL32`,
+  PI3K–AKT, focal adhesion).
+- **10X Visium human brain, AD vs. control** (Chen et al., 2022) — used for graph-level
+  classification benchmarking (local-only vs. global-only vs. combined), motivating why the global
+  component matters even for "simple" classification tasks.
+- **IMC pancreas T1D progression** (Damond et al., 2019) — robustness/generalization check.
+
+## Known limitations (relevant when extending the model)
+
+- Sliding windows cap the transformer's context length for scalability, so interactions spanning
+  window boundaries are not modeled — this is a real ceiling on the "global" scale, not just an
+  implementation detail (true long-range/endocrine signaling across whole organs is out of scope).
+- No causal/directional inference between cell groups (e.g. A → C → B chains) — attention flow is
+  correlational, not causal.
+- No ground-truth for cell-cell communication exists, so validation is via proxy classification
+  tasks and consistency with known biology (e.g. Moran's I, known marker genes), not direct
+  benchmarking of inferred interactions.
+- Raw attention weights are intentionally *not* used for interpretation (attention-sink effects,
+  window-dependent softmax normalization) — always go through the relevance/net-flow pipeline in
+  `interscale.evaluation`, not raw `attn_output_weights`.
diff --git a/.gitignore b/.gitignore
index 3751d1b..f2375b9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -25,4 +25,4 @@ src/confi_files/
 InterScale_hyperparameter_sweep/
 src/interscale/tmp/
 src/interscale/logs/
-CLAUDE.md
+CLAUDE.local.md
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 0000000..fb42c99
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,109 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project overview
+
+InterScale is a Python package for multi-scale cell interaction analysis in spatial transcriptomics data. It combines a local, graph-based component (cell-neighborhood scale) with a global, transformer-based component (whole-sample scale) into a single trainable model. It is built on top of `geome`/PyTorch Geometric, AnnData, scanpy/squidpy, `scvi-tools` (for AnnData registration/data management), and `lightning`/`pytorch_lightning` for training.
+
+For the scientific motivation, model design rationale, and terminology behind this architecture (from the InterScale preprint — Drummer, Jiménez et al., bioRxiv 2026), see [`.claude/background.md`](.claude/background.md). Read it before working on model/module architecture, loss functions, or the interpretability/evaluation pipeline — it explains *why* the local/global split and dual-decoder design exist, not just what the code does.
+
+## Common commands
+
+This project uses `hatch` as the primary project manager (also works with `uv` or `pip`; see `docs/installation.md`).
+
+```bash
+# Install dev dependencies (uv)
+uv sync --all-extras
+
+# Install dev dependencies (pip)
+pip install -e ".[dev,test,doc]"
+
+# Run the full test suite
+hatch test            # or: uv run pytest
+hatch test --all      # across the full Python version matrix
+
+# Run a single test file / test
+uv run pytest tests/test_package.py
+uv run pytest tests/test_package.py::test_import
+
+# Lint / format (ruff + biome are also run via pre-commit)
+uv run ruff check .
+uv run ruff format .
+pre-commit run --all-files
+
+# Build docs
+hatch run docs:build
+hatch run docs:open
+```
+
+Training entrypoints are run as scripts, not via a CLI command:
+
+```bash
+python src/interscale/main.py --cfg config_files/legnini_example.yaml --model_type CombinedModel
+python src/interscale/main_sweep.py --cfg <cfg.yaml> --sweep_cfg <sweep.yaml> --model_type CombinedModel --sweep_goal hyperparmeter
+```
+
+`--model_type` is one of `LocalModel`, `GlobalModel`, `CombinedModel`. `main_sweep.py` additionally requires `wandb` and a sweep config with `sweep_goal` in `{robustness, segmentation, hyperparmeter, loss}`.
+
+Ruff config lives in `pyproject.toml` (`[tool.ruff]`): line length 120, numpy docstring convention, many rules relaxed for legacy code (see `lint.ignore`). Tests are exempt from docstring rules (`D`).
+
+## Architecture
+
+### Config-driven design (yacs)
+
+Everything is driven by a single `yacs` `CfgNode` (`cfg`), built via `interscale.config.load_config(cfg_path)`:
+- `get_cfg_defaults()` assembles defaults from `config/{wandb,model,optim,dataset}_config.py`.
+- `load_config` additionally loads component-specific defaults based on `cfg.model.local_component.name` / `cfg.model.global_component.name` (from `local_component_config.py` / `global_component_config.py`) *before* merging the user's YAML file, then freezes the config.
+- User-facing YAML files (see `config_files/legnini_example.yaml`) only need to override values relevant to their run — everything else falls back to defaults.
+- The `cfg` object is threaded through nearly every layer (model, module, dataloader, geome dataset prep) rather than passed as individual kwargs.
+
+### Model hierarchy (`interscale.model`)
+
+`BaseModel` (`model/base/_base_model.py`) mimics the `scvi-tools` model API pattern:
+- Uses an `scvi.data.AnnDataManager` (registered via `_setup_anndata`/`register_fields`) to validate and access fields on the `AnnData` object (expression layer, sample/library keys, prediction obs, split key). This is why models require `ModelClass._setup_anndata(...)` to be called before instantiation — it registers the manager keyed by an AnnData UUID stored in `adata.uns`.
+- `prediction_task` is `"classification"` or `"regression"`; `prediction_level` is `"node"` or `"graph"` — these gate a lot of branching throughout the codebase (loss selection, decoder output shape, evaluation output storage).
+- `save`/`load` persist only the module `state_dict`, with filenames derived from `tl.utils.get_model_filename_prefix(cfg, local_component, global_component)`. `load` includes legacy state-dict key remapping (`global_module.module.` → `global_module.`, stripped `module.` prefixes, and `tl.utils.detect_and_remap_state_dict_keys`) to stay compatible with older/wandb-saved checkpoints.
+
+Three concrete models subclass `BaseModel`, differing in which components they instantiate and which module class they wrap:
+- `LocalModel` — local (graph) component only.
+- `GlobalModel` — global (transformer) component only.
+- `CombinedModel` — both; wraps either `CombinedModule` or `DualDecoderCombinedModule` (when `cfg.model.decoder.dual_decoder` is `True`, both local and global components get their own decoder instead of sharing one).
+
+`BaseModel._register_local_component` / `_register_global_component` look up `cfg.model.local_component.name` / `cfg.model.global_component.name` and instantiate the matching module class (currently `"GCN"` and `"self-attn-transformer"` respectively) — adding a new local/global component means adding a branch here plus a corresponding config file.
+
+### Module hierarchy (`interscale.module`)
+
+Mirrors the model hierarchy at the `pytorch_lightning`/`LightningModule` level:
+- `BaseModule` (`module/base/_base_module.py`) owns the decoder (`interscale.nn`: `LinearDecoder`, `LinearLSEDecoder`, `NonLinearDecoder`, or `None` when a wrapping combined module owns decoding), and shared node-masking logic (`tl.masking.apply_mask`) used for self-supervised/robustness training.
+- `module/local_modules/` — local (graph) encoders: `GCN`, `GIN`, `Precomputed` (precomputed embeddings), `SCVI` (SCVI-based encoder).
+- `module/global_modules/` — transformer-based global encoder (`TransformerNodeEncoderHook`) plus supporting transformer encoder/layer/utils, encoding whole-sample (long-range) context across cells.
+- `module/combined_module/` — `CombinedModule` and `DualDecoderCombinedModule` compose a local module + global module sequentially: local embeddings feed into the global (transformer) component, and predictions/attention/CLS tokens are extracted from there (see `CombinedModel.get_model_output` in `model/combined_model.py` for the full inference/evaluation flow, including how attention matrices and horizontal/vertical CLS tokens are extracted and padded to `max_seq_len`).
+
+### Data pipeline: AnnData → PyG graphs
+
+`interscale.tl.geome_utils.prepare_geome_dataset` / `prepare_a2d_dataset` bridge AnnData and PyTorch Geometric using `geome`:
+- Iterates per-sample/library (`cfg.dataset.sample_key`), builds a spatial neighbor graph per sample (`squidpy`-style `spatial_neigbors_kwargs`, converted to an edge index via `geome.transforms.AddEdgeIndex`), and yields one PyG `Data` object per sample.
+- Splits data by `cfg.dataset.split_key` (an `adata.obs` column that must contain `train`/`val`, optionally `test`) — this must exist in the AnnData before calling `prepare_geome_dataset`.
+- Handles both classification (one-hot encodes `prediction_obs`) and regression prediction tasks, and optionally attaches precomputed embeddings (`cfg.model.global_component.parameters.type_gex_embedding == "Precomputed"`) from `adata.obsm`.
+- `interscale.geome_dataloader.GraphAnnDataModule` wraps the resulting `list[Data]` splits into a `LightningDataModule`. For node-level learning it randomly masks a `pct_mask_nodes` fraction of nodes per graph (at least 1) for each dataloader construction — this is the masking scheme used for self-supervised node reconstruction/robustness experiments referenced in `config_files/legnini_example.yaml`'s `pct_mask_nodes` and `dataset.segmentation_robustness`.
+
+### Preprocessing / robustness utilities (`interscale.pp`)
+
+`pp.segmentation_noise.apply_segmentation_noise` simulates cell-segmentation errors (used when `cfg.dataset.segmentation_robustness` is set in `main.py`/`main_sweep.py`) for robustness sweeps.
+
+### Datasets (`interscale.datasets`)
+
+- `datasets/_legnini.py` provides `legnini()`, a loader for the Legnini et al. 2023 molecular cartography dataset used in the manuscript/tutorials — downloads from Zenodo and caches to `~/.cache/interscale/legnini_2023.h5ad`.
+
+
+### Evaluation (`interscale.evaluation`)
+
+Post-hoc analysis utilities operating on the `adata` produced by `save_evaluation_results`/`get_model_output`: gene loadings, gene-rank analysis, gene-set covariance, latent-space analysis, graph classification metrics, and network/attention stream visualization (`net_streams.py`).
+
+## Repository conventions
+
+- Numpy-style docstrings (see `docs/contributing.md`); many docstring lint rules are intentionally disabled in `pyproject.toml` for legacy modules, but new public APIs should still be documented in numpy style since `sphinx-autodoc-typehints`/napoleon render them for the docs site.
+- Code style is enforced by `pre-commit` (ruff-check/format, biome-format, pyproject-fmt, whitespace/merge-conflict hooks) — run `pre-commit install` once locally.
+- Tests live under `tests/`, use `pytest` with `--import-mode=importlib`; shared fixtures (e.g. a small synthetic `AnnData`) are defined in `conftest.py` at the repo root.
+- Coverage is measured over the `interscale` package only (`[tool.coverage] run.source = ["interscale"]`), excluding `test_*.py` files.
diff --git a/src/interscale/model/base/_base_model.py b/src/interscale/model/base/_base_model.py
index 8d20ffa..6c6a7fa 100644
--- a/src/interscale/model/base/_base_model.py
+++ b/src/interscale/model/base/_base_model.py
@@ -173,8 +173,8 @@ def _setup_anndata(
         """
         anndata_fields = [fields.LayerField("x", layer=layer_key)]
 
-        for i, sample_key in enumerate(sample_key_list):
-            anndata_fields.append(fields.CategoricalObsField(registry_key=f"sample_key_{i}", attr_key=sample_key))
+        for i, key in enumerate(sample_key_list):
+            anndata_fields.append(fields.CategoricalObsField(registry_key=f"sample_key_{i}", attr_key=key))
 
         if prediction_task == "classification":
             anndata_fields.append(fields.CategoricalObsField(registry_key="prediction_obs", attr_key=prediction_obs))
@@ -195,7 +195,7 @@ def _setup_anndata(
         if _SCVI_UUID_KEY not in adata.uns:
             adata.uns[_SCVI_UUID_KEY] = str(id(adata))
         cls._setup_adata_manager_store[adata.uns[_SCVI_UUID_KEY]] = manager
-        cls.sample_key = sample_key
+        cls.sample_key_list = sample_key_list
 
     # adjusted from scvi-tools
     # https://github.com/scverse/scvi-tools/blob/main/src/scvi/model/base/_base_model.py
diff --git a/tests/test_base_model.py b/tests/test_base_model.py
new file mode 100644
index 0000000..4d0983b
--- /dev/null
+++ b/tests/test_base_model.py
@@ -0,0 +1,53 @@
+import numpy as np
+import pandas as pd
+from anndata import AnnData
+
+from interscale.model.local_model import LocalModel
+
+
+def _make_adata(sample_key_columns: dict[str, list[str]], split: list[str], n_var: int = 4) -> AnnData:
+    rng = np.random.default_rng(0)
+    adata = AnnData(X=rng.integers(0, 10, size=(len(split), n_var)).astype(np.float32))
+    for key, values in sample_key_columns.items():
+        adata.obs[key] = pd.Categorical(values)
+    adata.obs["split"] = pd.Categorical(split)
+    return adata
+
+
+def test_setup_anndata_stores_full_sample_key_list():
+    """`_setup_anndata` must keep every key from `sample_key_list`, not just the last one seen in its
+    registration loop."""
+    adata = _make_adata(
+        sample_key_columns={
+            "sample_a": ["s1", "s1", "s1", "s2", "s2", "s2"],
+            "sample_b": ["fov1", "fov2", "fov1", "fov2", "fov1", "fov2"],
+        },
+        split=["train", "train", "val", "train", "val", "val"],
+    )
+    sample_key_list = ["sample_a", "sample_b"]
+
+    LocalModel._setup_anndata(
+        adata=adata,
+        layer_key=None,
+        sample_key_list=sample_key_list,
+        prediction_task="regression",
+        view_registry=False,
+    )
+
+    assert LocalModel.sample_key_list == sample_key_list
+
+
+def test_setup_anndata_with_empty_sample_key_list_does_not_raise():
+    """An empty `sample_key_list` used to raise `NameError`, since the loop variable it was read from
+    was never assigned."""
+    adata = _make_adata(sample_key_columns={}, split=["train", "train", "val", "val"])
+
+    LocalModel._setup_anndata(
+        adata=adata,
+        layer_key=None,
+        sample_key_list=[],
+        prediction_task="regression",
+        view_registry=False,
+    )
+
+    assert LocalModel.sample_key_list == []

From f54914e86ce85af41680430b2f7fd62db495b261 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Tue, 4 Aug 2026 09:17:56 +0200
Subject: [PATCH 02/14] fix masking - random reselection each training round

---
 src/interscale/geome_dataloader.py | 52 ++++++++++++++-----
 src/interscale/train/_training.py  | 12 ++++-
 src/interscale/train/_utils.py     | 13 +++++
 tests/test_geome_dataloader.py     | 80 ++++++++++++++++++++++++++++++
 4 files changed, 142 insertions(+), 15 deletions(-)
 create mode 100644 tests/test_geome_dataloader.py

diff --git a/src/interscale/geome_dataloader.py b/src/interscale/geome_dataloader.py
index dda0657..921d58c 100644
--- a/src/interscale/geome_dataloader.py
+++ b/src/interscale/geome_dataloader.py
@@ -127,8 +127,44 @@ def _smallest_data_batch_length(self, data_list: list[BaseData]):
         lengths = [data.num_nodes for data in data_list]
         return min(lengths)
 
+    def _assign_random_mask(self, data: BaseData, num_nodes_to_mask: int) -> None:
+        """Overwrites `data.mask` in place with a fresh, uniformly random selection of masked nodes.
+
+        This is the single seam where a future non-uniform sampling strategy (e.g. downweighting
+        nodes that were already masked in previous epochs) would be substituted in.
+        """
+        if data.num_nodes < num_nodes_to_mask:
+            raise ValueError("Cannot sample more nodes than available in any graph.")
+
+        mask_indices = random.sample(range(data.num_nodes), num_nodes_to_mask)
+        data.mask = torch.zeros(data.num_nodes, dtype=torch.bool)
+        data.mask[mask_indices] = True
+
+    def _num_nodes_to_mask(self, data_list: list[BaseData]) -> int:
+        """Number of nodes to mask per graph, derived from the smallest graph in `data_list`."""
+        num_nodes_to_mask = int(self._smallest_data_batch_length(data_list) * self.pct_mask_nodes)
+        return max(1, num_nodes_to_mask)  # must mask at least one node
+
+    def resample_train_mask(self) -> None:
+        """Redraws the masked node set for every training graph, in place.
+
+        `_spatial_node_loader` only ever assigns `.mask` once (when the train dataloader is first
+        built in `setup()`), so without this the same fixed subset of nodes is masked for the
+        entire training run. Call this once per epoch (see `NodeMaskResampleCallback` in
+        `interscale.train`) so every node eventually gets used as a supervision target.
+
+        Mutates the `Data` objects already held in `self.train_data` in place — no cloning, no
+        dataloader rebuild, so this adds no meaningful memory overhead beyond the `.mask` boolean
+        tensor that is already allocated.
+        """
+        if not self.setup_called:
+            return
+        num_nodes_to_mask = self._num_nodes_to_mask(self.train_data)
+        for data in self.train_data:
+            self._assign_random_mask(data, num_nodes_to_mask)
+
     def _spatial_node_loader(self, data_list: list[BaseData], shuffle: bool = False, **kwargs) -> DataListLoader:
-        """Adds a one-node mask to each Data object. TODO: load each graph multiple times with a different mask.
+        """Adds a node mask to each Data object.
 
         Args:
         ----
@@ -140,19 +176,9 @@ def _spatial_node_loader(self, data_list: list[BaseData], shuffle: bool = False,
         -------
             NeighborLoader: the node dataloader
         """
-        smallest_length = self._smallest_data_batch_length(data_list)
-        num_nodes_to_mask = int(smallest_length * self.pct_mask_nodes)
-        if num_nodes_to_mask == 0:  # must mask at least one node
-            num_nodes_to_mask = 1
-
+        num_nodes_to_mask = self._num_nodes_to_mask(data_list)
         for data in data_list:
-            if data.num_nodes < num_nodes_to_mask:
-                raise ValueError("Cannot sample more nodes than available in any graph.")
-
-            # Randomly select a ndoe to mask
-            mask_indices = random.sample(range(data.num_nodes), num_nodes_to_mask)
-            data.mask = torch.zeros(data.num_nodes, dtype=torch.bool)
-            data.mask[mask_indices] = True
+            self._assign_random_mask(data, num_nodes_to_mask)
 
         return DataLoader(
             dataset=data_list,
diff --git a/src/interscale/train/_training.py b/src/interscale/train/_training.py
index e7c8fb2..6cf5698 100644
--- a/src/interscale/train/_training.py
+++ b/src/interscale/train/_training.py
@@ -9,7 +9,7 @@
 
 from interscale.tl.utils import get_model_filename_prefix
 from interscale.train._trainingplans import TrainingPlan
-from interscale.train._utils import MetricsHistory
+from interscale.train._utils import MetricsHistory, NodeMaskResampleCallback
 
 # from interscale.model.base._trainer import TrainRunner
 
@@ -104,6 +104,7 @@ def train(
         print("Steps per epoch", steps_per_epoch)
         lr_monitor = LearningRateMonitor(logging_interval="epoch")
         self.history_ = MetricsHistory()
+        mask_resample_callback = NodeMaskResampleCallback() if self._cfg.dataset.pct_mask_nodes > 0 else None
         checkpoint_callback = None
         loss_callback = None
         performance_callback = None
@@ -185,7 +186,14 @@ def train(
         # Create list of callbacks and filter out None values
         callbacks = [
             callback
-            for callback in [lr_monitor, performance_callback, loss_callback, self.history_, checkpoint_callback]
+            for callback in [
+                lr_monitor,
+                performance_callback,
+                loss_callback,
+                self.history_,
+                mask_resample_callback,
+                checkpoint_callback,
+            ]
             if callback is not None
         ]
 
diff --git a/src/interscale/train/_utils.py b/src/interscale/train/_utils.py
index 3367a09..cbb7121 100644
--- a/src/interscale/train/_utils.py
+++ b/src/interscale/train/_utils.py
@@ -3,6 +3,19 @@
 from lightning.pytorch.callbacks import Callback
 
 
+class NodeMaskResampleCallback(Callback):
+    """Redraws the training node mask at the start of every epoch.
+
+    Without this, `GraphAnnDataModule` only ever samples one fixed set of masked nodes for the
+    whole training run (see `GraphAnnDataModule.resample_train_mask`), so nodes outside that
+    initial sample never receive direct supervision. Validation/test masks are left untouched so
+    monitored metrics (e.g. `val_loss`) stay comparable across epochs.
+    """
+
+    def on_train_epoch_start(self, trainer, pl_module):
+        trainer.datamodule.resample_train_mask()
+
+
 class MetricsHistory(Callback):
     def __init__(self):
         super().__init__()
diff --git a/tests/test_geome_dataloader.py b/tests/test_geome_dataloader.py
new file mode 100644
index 0000000..cb0ac29
--- /dev/null
+++ b/tests/test_geome_dataloader.py
@@ -0,0 +1,80 @@
+import torch
+from torch_geometric.data import Data
+
+from interscale.geome_dataloader import GraphAnnDataModule
+from interscale.train._utils import NodeMaskResampleCallback
+
+
+def _make_data(num_nodes: int, num_features: int = 3) -> Data:
+    x = torch.randn(num_nodes, num_features)
+    edge_index = torch.zeros((2, 0), dtype=torch.long)
+    return Data(x=x, edge_index=edge_index)
+
+
+def _build_datamodule(pct_mask_nodes: float = 0.5) -> GraphAnnDataModule:
+    train_data = [_make_data(20), _make_data(20), _make_data(20)]
+    val_data = [_make_data(10)]
+    test_data = [_make_data(10)]
+    dm = GraphAnnDataModule(
+        datas=[train_data, val_data, test_data],
+        batch_size=1,
+        num_workers=0,
+        pct_mask_nodes=pct_mask_nodes,
+        learning_type="node",
+    )
+    dm.setup(stage="fit")
+    dm.setup(stage="test")
+    return dm
+
+
+def test_resample_train_mask_changes_masked_nodes_without_replacing_objects():
+    """The masked node set must differ across calls, without cloning/replacing the `Data` objects."""
+    dm = _build_datamodule()
+    original_masks = [data.mask.clone() for data in dm.train_data]
+    original_ids = [id(data) for data in dm.train_data]
+    original_list_id = id(dm.train_data)
+
+    changed = False
+    for _ in range(20):
+        dm.resample_train_mask()
+        if any(not torch.equal(orig, data.mask) for orig, data in zip(original_masks, dm.train_data)):
+            changed = True
+            break
+
+    assert changed, "resample_train_mask never produced a different mask across 20 redraws"
+    assert [id(data) for data in dm.train_data] == original_ids, "Data objects must be mutated in place, not replaced"
+    assert id(dm.train_data) == original_list_id, "train_data list must not be rebuilt"
+
+
+def test_resample_train_mask_leaves_val_and_test_data_untouched():
+    dm = _build_datamodule()
+    original_val_masks = [data.mask.clone() for data in dm.val_data]
+    original_test_masks = [data.mask.clone() for data in dm.test_data]
+
+    for _ in range(5):
+        dm.resample_train_mask()
+
+    for orig, data in zip(original_val_masks, dm.val_data):
+        assert torch.equal(orig, data.mask)
+    for orig, data in zip(original_test_masks, dm.test_data):
+        assert torch.equal(orig, data.mask)
+
+
+def test_resample_train_mask_is_noop_before_setup():
+    """Calling resample_train_mask before setup() must not raise (no dataloader exists yet)."""
+    dm = GraphAnnDataModule(datas=[[_make_data(20)], [_make_data(10)]], num_workers=0)
+
+    dm.resample_train_mask()
+
+
+def test_node_mask_resample_callback_delegates_to_datamodule():
+    dm = _build_datamodule()
+    calls = []
+    dm.resample_train_mask = lambda: calls.append(1)
+
+    class _StubTrainer:
+        datamodule = dm
+
+    NodeMaskResampleCallback().on_train_epoch_start(_StubTrainer(), pl_module=None)
+
+    assert calls == [1]

From 9c55bee2ec736d8b0540fc696c190201abf83cb7 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Tue, 4 Aug 2026 21:24:48 +0200
Subject: [PATCH 03/14] fix sampling in test set

---
 10                                       |  0
 src/interscale/config/__init__.py        | 15 ++++-
 src/interscale/config/optim_config.py    |  6 ++
 src/interscale/main.py                   |  6 ++
 src/interscale/main_sweep.py             |  6 ++
 src/interscale/model/base/_base_model.py | 75 +++++++++++++++++++++---
 src/interscale/tl/geome_utils.py         | 12 ++--
 src/interscale/train/_training.py        | 63 +++++++++++++++-----
 src/interscale/train/_trainingplans.py   |  7 ++-
 9 files changed, 159 insertions(+), 31 deletions(-)
 create mode 100644 10

diff --git a/10 b/10
new file mode 100644
index 0000000..e69de29
diff --git a/src/interscale/config/__init__.py b/src/interscale/config/__init__.py
index f6f88a4..45d0bcd 100644
--- a/src/interscale/config/__init__.py
+++ b/src/interscale/config/__init__.py
@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from yacs.config import CfgNode as CN
 
 from .dataset_config import get_dataset_cfg
@@ -26,7 +28,7 @@ def load_config(cfg_path=None):
 
     Parameters
     ----------
-    cfg_path : str, optional
+    cfg_path : str or pathlib.Path, optional
         Path to the config file to load. If None, only default values are used.
 
     Returns
@@ -37,6 +39,17 @@ def load_config(cfg_path=None):
     # First get all default configs including local component defaults
     cfg = get_cfg_defaults()
 
+    # Documented as defaults-only, but the code below dereferences cfg_path
+    # unconditionally, so None used to raise AttributeError too.
+    if cfg_path is None:
+        cfg.freeze()
+        return cfg
+
+    # Callers pass a str: main.py and main_sweep.py both declare --cfg as
+    # type=str, and the docstring says str. Normalise instead of requiring
+    # every caller to wrap it.
+    cfg_path = Path(cfg_path)
+
     with cfg_path.open() as f:
         # Create a temporary config to load the model type
         temp_cfg = CN.load_cfg(f)
diff --git a/src/interscale/config/optim_config.py b/src/interscale/config/optim_config.py
index 9282fa1..8798e8c 100644
--- a/src/interscale/config/optim_config.py
+++ b/src/interscale/config/optim_config.py
@@ -23,4 +23,10 @@ def get_optim_cfg(cfg):
     cfg.optim.cross_corr = "cell"  # Currently cell is the only one that really works
     cfg.optim.n_epochs = 100
     cfg.optim.early_stopping = True
+    cfg.optim.patience = 5  # EarlyStopping patience in epochs
+    cfg.optim.min_delta = 0.0  # EarlyStopping min_delta
+    cfg.optim.min_epochs = 1  # floor on training length; set above lr_warmup to clear warm-up
+    # Metric driving EarlyStopping / ModelCheckpoint / the LR scheduler.
+    # "auto" -> val_f1_macro for classification, val_loss for regression.
+    cfg.optim.monitor = "auto"
     return cfg
diff --git a/src/interscale/main.py b/src/interscale/main.py
index cf25802..589fae8 100644
--- a/src/interscale/main.py
+++ b/src/interscale/main.py
@@ -1,4 +1,5 @@
 import argparse
+import warnings
 
 import scanpy as sc
 import squidpy as sq
@@ -9,6 +10,11 @@
 from interscale.pp import apply_segmentation_noise
 from interscale.tl import prepare_geome_dataset, remove_zero_expression_cells, set_full_reproducibility
 
+# geome calls the deprecated `sq.gr.spatial_neighbors` entrypoint; silence its
+# FutureWarnings so they don't flood the training logs.
+warnings.filterwarnings("ignore", category=FutureWarning, message=r".*spatial_neighbors.*")
+warnings.filterwarnings("ignore", category=FutureWarning, message=r".*n_neighs.*")
+
 
 def main(cfg_path, model_type):
 
diff --git a/src/interscale/main_sweep.py b/src/interscale/main_sweep.py
index cc8e5c4..63f8e7f 100644
--- a/src/interscale/main_sweep.py
+++ b/src/interscale/main_sweep.py
@@ -1,5 +1,6 @@
 import argparse
 import os
+import warnings
 
 import psutil
 import scanpy as sc
@@ -14,6 +15,11 @@
 from interscale.tl import prepare_geome_dataset
 from interscale.tl.utils import get_model_filename_prefix
 
+# geome calls the deprecated `sq.gr.spatial_neighbors` entrypoint; silence its
+# FutureWarnings so they don't flood the training logs.
+warnings.filterwarnings("ignore", category=FutureWarning, message=r".*spatial_neighbors.*")
+warnings.filterwarnings("ignore", category=FutureWarning, message=r".*n_neighs.*")
+
 
 def print_memory_usage(stage=""):
     """Print current memory usage for both CPU and GPU"""
diff --git a/src/interscale/model/base/_base_model.py b/src/interscale/model/base/_base_model.py
index 6c6a7fa..9270f71 100644
--- a/src/interscale/model/base/_base_model.py
+++ b/src/interscale/model/base/_base_model.py
@@ -14,7 +14,6 @@
     _SCVI_UUID_KEY,
 )
 from scvi.data._utils import _assign_adata_uuid, _check_if_view
-from sklearn.utils.class_weight import compute_class_weight
 from yacs.config import CfgNode as CN
 
 from interscale.module.base import GlobalModule, LocalModule
@@ -126,14 +125,74 @@ def __init__(
 
         self.class_weights = None
         if self._cfg.optim.loss == "WeightedCE":
-            self.class_weights = torch.tensor(
-                compute_class_weight(
-                    "balanced",
-                    classes=np.unique(self._adata.obs[self._cfg.dataset.prediction_obs]),
-                    y=self._adata.obs[self._cfg.dataset.prediction_obs],
+            self.class_weights = self._compute_train_class_weights()
+
+    def _training_unit_labels(self) -> pd.Series:
+        """Labels of the units the loss is computed over, restricted to the train split.
+
+        For ``prediction_level == "node"`` a unit is a cell. For ``"graph"`` a unit is one PyG
+        graph, i.e. one category of each key in ``cfg.dataset.sample_key`` -- matching how
+        :func:`interscale.tl.prepare_geome_dataset` builds and concatenates graphs.
+
+        Returns
+        -------
+        pd.Series
+            One label per training unit.
+        """
+        obs = self._adata.obs
+        pred_col = self._cfg.dataset.prediction_obs
+        split_col = self._cfg.dataset.split_key
+
+        if split_col is not None and split_col in obs:
+            train_obs = obs.loc[obs[split_col].astype(str) == "train"]
+        else:
+            logger.warning("split_key '%s' not in adata.obs -- class weights use all cells.", split_col)
+            train_obs = obs
+        if len(train_obs) == 0:
+            raise ValueError(f"No observations with {split_col} == 'train'; cannot compute class weights.")
+
+        if self.prediction_level == "node":
+            return train_obs[pred_col].astype(str)
+
+        labels = []
+        for key in self._cfg.dataset.sample_key:
+            grouped = train_obs.groupby(key, observed=True)[pred_col]
+            n_per_graph = grouped.nunique()
+            if (n_per_graph > 1).any():
+                logger.warning(
+                    "%d graphs of '%s' contain more than one '%s'; using the majority label.",
+                    int((n_per_graph > 1).sum()),
+                    key,
+                    pred_col,
                 )
-            )
-            print("WeightedCE with class weights: ", self.class_weights)
+            labels.append(grouped.agg(lambda s: s.value_counts().idxmax()).astype(str))
+        return pd.concat(labels)
+
+    def _compute_train_class_weights(self) -> torch.Tensor:
+        """Compute ``"balanced"`` class weights over the training units.
+
+        Weights are ordered like :attr:`class_labels` (i.e. ``.cat.categories``), which is also
+        the column order of the one-hot ``data.y`` produced by ``geome``, so index *i* of the
+        returned tensor lines up with logit column *i*.
+
+        Returns
+        -------
+        torch.Tensor
+            Float32 tensor of per-class weights.
+        """
+        y = self._training_unit_labels()
+        classes = [str(c) for c in self.class_labels]
+        counts = y.value_counts().reindex(classes).fillna(0.0).to_numpy(dtype=np.float64)
+        if (counts == 0).any():
+            missing = [c for c, n in zip(classes, counts, strict=True) if n == 0]
+            raise ValueError(f"Classes {missing} have no training units; WeightedCE weights undefined.")
+        # Same formula as sklearn's compute_class_weight("balanced").
+        weights = counts.sum() / (len(classes) * counts)
+        print(
+            f"WeightedCE ({self.prediction_level}-level, train split '{self._cfg.dataset.split_key}'): "
+            + ", ".join(f"{c}: n={int(n)} w={w:.4f}" for c, n, w in zip(classes, counts, weights, strict=True))
+        )
+        return torch.as_tensor(weights, dtype=torch.float32)
 
     @classmethod
     def _setup_anndata(
diff --git a/src/interscale/tl/geome_utils.py b/src/interscale/tl/geome_utils.py
index 268e1dd..b152f05 100644
--- a/src/interscale/tl/geome_utils.py
+++ b/src/interscale/tl/geome_utils.py
@@ -168,16 +168,16 @@ def prepare_geome_dataset(adata, cfg: CN):
             save_preprocessed_adata=True,
         )
 
-        pyg_train, adata_train = list(a2d(adata[adata.obs[split_key] == "train"]))
-        pyg_val, adata_val = list(a2d(adata[adata.obs[split_key] == "val"]))
+        pyg_train, _ = list(a2d(adata[adata.obs[split_key] == "train"]))
+        pyg_val, _ = list(a2d(adata[adata.obs[split_key] == "val"]))
         datas_train.extend(pyg_train)
         datas_val.extend(pyg_val)
         if "test" in np.unique(adata.obs[split_key]):
-            pyg_test, adata_test = list(a2d(adata[adata.obs[split_key] == "test"]))
+            pyg_test, _ = list(a2d(adata[adata.obs[split_key] == "test"]))
             datas_test.extend(pyg_test)
 
     if "test" in np.unique(adata.obs[split_key]):
-        datas_test, adata_test = list(a2d(adata[adata.obs[split_key] == "test"]))
-        return [datas_train, datas_val, datas_test], [adata_train, adata_val, adata_test]
+        #datas_test, adata_test = list(a2d(adata[adata.obs[split_key] == "test"]))
+        return [datas_train, datas_val, datas_test], _
 
-    return [datas_train, datas_val], [adata_train, adata_val]
+    return [datas_train, datas_val], _
diff --git a/src/interscale/train/_training.py b/src/interscale/train/_training.py
index 6cf5698..01fd583 100644
--- a/src/interscale/train/_training.py
+++ b/src/interscale/train/_training.py
@@ -1,7 +1,9 @@
 import math
+import os
 
 import lightning as L
 import lightning.pytorch as pl
+import torch
 import wandb
 from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor, ModelCheckpoint
 from lightning.pytorch.loggers import WandbLogger
@@ -23,6 +25,21 @@ class NodeMaskingTrainingPlan:
     # _data_splitter_cls = DataSplitter
     _training_plan_cls = TrainingPlan
 
+    def _resolve_monitor(self) -> tuple[str, str]:
+        """Metric driving EarlyStopping / ModelCheckpoint, and the direction to optimise it.
+
+        Returns
+        -------
+        tuple[str, str]
+            The metric name and the mode (``"min"`` or ``"max"``).
+        """
+        monitor = self._cfg.optim.monitor
+        if monitor == "auto":
+            # val_loss is a poor early-stopping criterion under class imbalance: a collapsed
+            # constant predictor is a genuine minimum of the weighted loss.
+            monitor = "val_f1_macro" if "classification" in self.prediction_task else "val_loss"
+        return monitor, ("min" if monitor.endswith("loss") else "max")
+
     # @devices_dsp.dedent -TODO: Why is this here in scvi-tools?
     def train(
         self,
@@ -30,7 +47,7 @@ def train(
         shuffle_set_split: bool = True,
         load_sparse_tensor: bool = False,
         early_stopping: bool = True,
-        patience: int = 5,
+        patience: int | None = None,
         datasplitter_kwargs: dict | None = None,
         plan_kwargs: dict | None = None,
         datamodule: L.LightningDataModule | None = None,
@@ -128,6 +145,8 @@ def train(
             patience_in_steps=steps_per_epoch,
         )
 
+        monitor, mode = self._resolve_monitor()
+
         if early_stopping:
             # TODO: why does the self.history_ stop working when using loss_callback?
             # loss_callback = EarlyStopping(
@@ -137,16 +156,15 @@ def train(
             #         verbose=False,
             #         mode="min"
             #     )
-            if "classification" in self.prediction_task:
-                performance_callback = EarlyStopping(
-                    monitor="val_loss", min_delta=0.005, patience=patience, verbose=False, mode="min"
-                )
-            elif "regression" in self.prediction_task:
-                performance_callback = EarlyStopping(
-                    monitor="val_loss", min_delta=0.005, patience=patience, verbose=False, mode="min"
-                )
-            else:
+            if not ("classification" in self.prediction_task or "regression" in self.prediction_task):
                 raise Exception("Training must be classification or regression based.")
+            performance_callback = EarlyStopping(
+                monitor=monitor,
+                min_delta=float(self._cfg.optim.min_delta),
+                patience=int(patience if patience is not None else self._cfg.optim.patience),
+                verbose=True,
+                mode=mode,
+            )
 
         if self._cfg.model.save is not None:
             run_name = get_model_filename_prefix(self._cfg, self.local_component, self.global_component)
@@ -154,9 +172,10 @@ def train(
                 checkpoint_callback = ModelCheckpoint(
                     dirpath=self._cfg.model.save,
                     filename=run_name,
-                    monitor="val_loss",
-                    mode="min",
-                )  # save model if validation accuracy increases
+                    monitor=monitor,
+                    mode=mode,
+                    save_top_k=1,
+                )  # save the best model according to `monitor`
             elif "regression" in self._cfg.dataset.prediction_task:
                 if self._cfg.optim.loss == "MSELoss":
                     checkpoint_callback = ModelCheckpoint(
@@ -213,7 +232,7 @@ def train(
             print(f"Trainable parameters: {trainable_params:,}")
 
         trainer = pl.Trainer(
-            min_epochs=1,
+            min_epochs=int(self._cfg.optim.min_epochs),
             max_epochs=int(max_epochs),
             # enable_progress_bar=True,
             callbacks=callbacks,
@@ -225,6 +244,22 @@ def train(
         )
 
         trainer.fit(training_plan, datamodule)
+
+        # EarlyStopping does not restore best weights and ModelCheckpoint only writes them to
+        # disk. Without this, the validate()/test() calls below -- and the save() further down --
+        # all report the LAST epoch, which for an early-stopped run is `patience` epochs past the
+        # best one.
+        best_ckpt = getattr(checkpoint_callback, "best_model_path", "") if checkpoint_callback else ""
+        if best_ckpt and os.path.exists(best_ckpt):
+            best_score = checkpoint_callback.best_model_score
+            print(f"Restoring best checkpoint ({monitor}={float(best_score):.4f}) from {best_ckpt}")
+            state = torch.load(best_ckpt, map_location="cpu", weights_only=False)["state_dict"]
+            missing, unexpected = training_plan.load_state_dict(state, strict=False)
+            if missing or unexpected:
+                print(f"restore: missing={missing}, unexpected={unexpected}")
+        else:
+            print("WARNING: no best checkpoint found; evaluating FINAL-epoch weights.")
+
         trainer.validate(training_plan, datamodule)
         if self.train_size + self.validation_size < 1:
             trainer.test(training_plan, datamodule)
diff --git a/src/interscale/train/_trainingplans.py b/src/interscale/train/_trainingplans.py
index 91d5067..96e3aef 100644
--- a/src/interscale/train/_trainingplans.py
+++ b/src/interscale/train/_trainingplans.py
@@ -98,7 +98,9 @@ def __init__(
         if "classification" in self.prediction_task:
             metrics = self._setup_classification_metrics(self.module.n_output)
             self.loss = self._setup_classification_loss(self.loss_type, self.class_weights)
-            self.monitor_metric = "val_f1"
+            # Must name a metric that is actually logged -- this is handed to Lightning as the
+            # LR-scheduler monitor. "val_f1" never existed; only val_f1_micro/macro/<class> do.
+            self.monitor_metric = "val_f1_macro"
         elif "regression" in self.prediction_task:
             metrics = self._setup_regression_metrics(self.module.n_output)
             self.loss = self._setup_regression_loss(self.loss_type)
@@ -121,7 +123,8 @@ def _setup_classification_loss(
         elif loss == "WeightedCE":
             assert class_weights is not None, "Class weights must be provided for WeightedCE loss."
             assert isinstance(class_weights, torch.Tensor), "class_weights must be a torch tensor"
-            return nn.CrossEntropyLoss(class_weights)
+            # .float() guards against a float64 weight buffer meeting float32 logits.
+            return nn.CrossEntropyLoss(weight=class_weights.float())
 
     def _setup_regression_loss(self, loss: Literal[REGRESSION_LOSSES]):
         """Setup loss function based on prediction task and configuration."""

From 635185da229a93a5997ed7692d31c1080ccd4d2c Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Fri, 7 Aug 2026 14:40:53 +0200
Subject: [PATCH 04/14] hyperparam sweep

---
 10                                 |   0
 src/interscale/geome_dataloader.py |  54 ++++++-----
 src/interscale/main.py             |   2 +-
 src/interscale/main_sweep.py       | 149 +++++++++++++++++++----------
 4 files changed, 132 insertions(+), 73 deletions(-)
 delete mode 100644 10

diff --git a/10 b/10
deleted file mode 100644
index e69de29..0000000
diff --git a/src/interscale/geome_dataloader.py b/src/interscale/geome_dataloader.py
index 921d58c..1897b00 100644
--- a/src/interscale/geome_dataloader.py
+++ b/src/interscale/geome_dataloader.py
@@ -73,11 +73,18 @@ def _nodewise_setup(self, stage: str | None) -> None:
         -------
             None
         """
+        # For node-level learning the masked nodes *are* the supervision/evaluation targets
+        # (see `_common_step`: y_pred/y_true are indexed by mask_idx), so every split must be
+        # masked. For graph-level learning mask_idx is discarded and the graph label is used
+        # instead, making masking pure input augmentation -- so it belongs to train only,
+        # otherwise val/test inputs are corrupted for no benefit.
+        eval_mask = self.learning_type == "node"
+
         if stage == "fit" or stage is None:
             self._train_dataloader = self._spatial_node_loader(data_list=self.train_data, shuffle=True)
-            self._val_dataloader = self._spatial_node_loader(data_list=self.val_data, shuffle=False)
+            self._val_dataloader = self._spatial_node_loader(data_list=self.val_data, shuffle=False, mask=eval_mask)
         if stage == "test" or stage is None:
-            self._test_dataloader = self._spatial_node_loader(data_list=self.test_data, shuffle=False)
+            self._test_dataloader = self._spatial_node_loader(data_list=self.test_data, shuffle=False, mask=eval_mask)
 
     def _graphwise_setup(self, stage: str | None) -> None:
         """Sets up the data loaders for graph-wise learning.
@@ -122,28 +129,22 @@ def _get_dataloader(self, dataloader):
             raise RuntimeError("setup method should be called before getting dataloaders")
         return dataloader
 
-    def _smallest_data_batch_length(self, data_list: list[BaseData]):
-        """Returns the number of nodes in the smallest graph from the list of BaseData."""
-        lengths = [data.num_nodes for data in data_list]
-        return min(lengths)
+    def _assign_random_mask(self, data: BaseData) -> None:
+        """Overwrites `data.mask` in place with a fresh Bernoulli draw at rate `pct_mask_nodes`.
 
-    def _assign_random_mask(self, data: BaseData, num_nodes_to_mask: int) -> None:
-        """Overwrites `data.mask` in place with a fresh, uniformly random selection of masked nodes.
+        Each node is masked independently with probability `pct_mask_nodes`, so the expected
+        masked fraction is the same for every graph regardless of its size. An earlier version
+        derived a single absolute node count from the smallest graph in the split and applied it
+        to every graph, which made the realised rate depend both on a graph's size and on which
+        split it landed in.
 
         This is the single seam where a future non-uniform sampling strategy (e.g. downweighting
         nodes that were already masked in previous epochs) would be substituted in.
         """
-        if data.num_nodes < num_nodes_to_mask:
-            raise ValueError("Cannot sample more nodes than available in any graph.")
-
-        mask_indices = random.sample(range(data.num_nodes), num_nodes_to_mask)
-        data.mask = torch.zeros(data.num_nodes, dtype=torch.bool)
-        data.mask[mask_indices] = True
-
-    def _num_nodes_to_mask(self, data_list: list[BaseData]) -> int:
-        """Number of nodes to mask per graph, derived from the smallest graph in `data_list`."""
-        num_nodes_to_mask = int(self._smallest_data_batch_length(data_list) * self.pct_mask_nodes)
-        return max(1, num_nodes_to_mask)  # must mask at least one node
+        mask = torch.rand(data.num_nodes) < self.pct_mask_nodes
+        if not mask.any():  # must mask at least one node
+            mask[random.randrange(data.num_nodes)] = True
+        data.mask = mask
 
     def resample_train_mask(self) -> None:
         """Redraws the masked node set for every training graph, in place.
@@ -159,26 +160,31 @@ def resample_train_mask(self) -> None:
         """
         if not self.setup_called:
             return
-        num_nodes_to_mask = self._num_nodes_to_mask(self.train_data)
         for data in self.train_data:
-            self._assign_random_mask(data, num_nodes_to_mask)
+            self._assign_random_mask(data)
 
-    def _spatial_node_loader(self, data_list: list[BaseData], shuffle: bool = False, **kwargs) -> DataListLoader:
+    def _spatial_node_loader(
+        self, data_list: list[BaseData], shuffle: bool = False, mask: bool = True, **kwargs
+    ) -> DataListLoader:
         """Adds a node mask to each Data object.
 
         Args:
         ----
         data: PyTorch geometric.Batch
         shuffle (bool, optional): whether to shuffle the data. Defaults to False.
+        mask (bool, optional): whether to mask nodes at all. Graph-level evaluation splits pass
+            False so that val/test inputs are not corrupted. Defaults to True.
         kwargs: arguments passed to the pyg.NeighborLoader
 
         Returns
         -------
             NeighborLoader: the node dataloader
         """
-        num_nodes_to_mask = self._num_nodes_to_mask(data_list)
         for data in data_list:
-            self._assign_random_mask(data, num_nodes_to_mask)
+            if mask:
+                self._assign_random_mask(data)
+            else:
+                data.mask = torch.zeros(data.num_nodes, dtype=torch.bool)
 
         return DataLoader(
             dataset=data_list,
diff --git a/src/interscale/main.py b/src/interscale/main.py
index 589fae8..2d88c8a 100644
--- a/src/interscale/main.py
+++ b/src/interscale/main.py
@@ -71,7 +71,7 @@ def main(cfg_path, model_type):
         num_workers=1,
         batch_size=int(cfg.dataset.batch_size),
         pct_mask_nodes=cfg.dataset.pct_mask_nodes,
-        learning_type="node",
+        learning_type=cfg.dataset.prediction_level,
     )
 
     model.train(max_epochs=cfg.optim.n_epochs, datamodule=dm, early_stopping=cfg.optim.early_stopping)
diff --git a/src/interscale/main_sweep.py b/src/interscale/main_sweep.py
index 63f8e7f..b9c2db8 100644
--- a/src/interscale/main_sweep.py
+++ b/src/interscale/main_sweep.py
@@ -61,7 +61,7 @@ def print_memory_debug():
         print(f"Memory debug failed: {e}")
 
 
-def main_sweep(cfg_path, model_type, sweep_goal):
+def main_sweep(cfg_path, model_type, sweep_goal, sweep_params=None):
 
     print_memory_usage("Start of main_sweep")
 
@@ -103,38 +103,61 @@ def main_sweep(cfg_path, model_type, sweep_goal):
             cfg.optim.seed = sweep_config["optim.seed"]
         elif sweep_goal == "hyperparmeter":
             print("hyperparameter sweep")
-            cfg.optim.lr = sweep_config["optim.lr"]
-            cfg.optim.lr_warmup = sweep_config["optim.lr_warmup"]
-            cfg.optim.wd = sweep_config["optim.wd"]
-            cfg.dataset.batch_size = sweep_config["dataset.batch_size"]
-            cfg.dataset.pct_mask_nodes = sweep_config["dataset.pct_mask_nodes"]
-            cfg.model.n_embed = sweep_config["model.n_embed"]
-            if model_type == "LocalModel" or model_type == "CombinedModel":
+            applied = []
+
+            def _apply(node, attr, key):
+                """Assign a swept value only if the sweep actually declares that key.
+
+                Staged sweeps vary a subset of the parameters (e.g. optimiser only), so an
+                unconditional lookup would KeyError on every key the stage omits.
+                """
+                if key in sweep_config.keys():
+                    setattr(node, attr, sweep_config[key])
+                    applied.append(key)
+
+            _apply(cfg.optim, "lr", "optim.lr")
+            _apply(cfg.optim, "lr_warmup", "optim.lr_warmup")
+            _apply(cfg.optim, "wd", "optim.wd")
+            _apply(cfg.dataset, "batch_size", "dataset.batch_size")
+            _apply(cfg.dataset, "pct_mask_nodes", "dataset.pct_mask_nodes")
+            _apply(cfg.model, "n_embed", "model.n_embed")
+
+            # Two separate `if`s, not if/elif: CombinedModel has BOTH components, and an
+            # `elif model_type == "GlobalModel" or model_type == "CombinedModel"` is
+            # unreachable for CombinedModel, so its transformer was never swept.
+            local = cfg.model.local_component.parameters
+            glob = cfg.model.global_component.parameters
+            if model_type in ("LocalModel", "CombinedModel"):
                 print("LocalModel configs")
-                cfg.model.local_component.parameters.num_layers = sweep_config[
-                    "model.local_component.parameters.num_layers"
-                ]
-                cfg.model.local_component.parameters.hidden_dim = sweep_config[
-                    "model.local_component.parameters.hidden_dim"
-                ]
-            elif model_type == "GlobalModel" or model_type == "CombinedModel":
+                _apply(local, "num_layers", "model.local_component.parameters.num_layers")
+                _apply(local, "hidden_dim", "model.local_component.parameters.hidden_dim")
+                _apply(local, "dropout_local", "model.local_component.parameters.dropout_local")
+            if model_type in ("GlobalModel", "CombinedModel"):
                 print("transformer configs")
-                cfg.model.global_component.parameters.dim_feedforward = sweep_config[
-                    "model.global_component.parameters.dim_feedforward"
-                ]
-                cfg.model.global_component.parameters.num_layers = sweep_config[
-                    "model.global_component.parameters.num_layers"
-                ]
-                cfg.model.global_component.parameters.n_heads = sweep_config[
-                    "model.global_component.parameters.n_heads"
-                ]
-                cfg.model.global_component.parameters.dropout = sweep_config[
-                    "model.global_component.parameters.dropout"
-                ]
-                # cfg.transformer.max_seq_len = sweep_run.config.transformer.max_seq_len
+                _apply(glob, "dim_feedforward", "model.global_component.parameters.dim_feedforward")
+                _apply(glob, "num_layers", "model.global_component.parameters.num_layers")
+                _apply(glob, "n_heads", "model.global_component.parameters.n_heads")
+                # The config key is `dropout_global` (see global_component_config.py); assigning
+                # to `dropout` silently created a dead key because set_new_allowed(True) is on.
+                _apply(glob, "dropout_global", "model.global_component.parameters.dropout_global")
+
+            if sweep_params is not None:
+                ignored = sorted(set(sweep_params) - set(applied))
+                if ignored:
+                    print(
+                        f"WARNING: sweep declares parameters that nothing applies, so they vary "
+                        f"between trials with no effect: {ignored}"
+                    )
+            print(f"applied sweep parameters: {sorted(applied)}")
         elif sweep_goal == "loss":
             print("loss sweep")
             cfg.optim.loss = sweep_config["optim.loss"]
+        else:
+            raise ValueError(
+                f"Unknown --sweep_goal '{sweep_goal}'. Must be one of: "
+                "robustness, segmentation, hyperparmeter, loss. "
+                "(Nothing would be swept otherwise -- every trial would train the base config.)"
+            )
         cfg.freeze()
 
     ####### PREPROCESSING #######
@@ -196,7 +219,7 @@ def main_sweep(cfg_path, model_type, sweep_goal):
         num_workers=1,
         batch_size=int(cfg.dataset.batch_size),
         pct_mask_nodes=cfg.dataset.pct_mask_nodes,
-        learning_type="node",
+        learning_type=cfg.dataset.prediction_level,
     )
     print_memory_usage("After datamodule creation")
 
@@ -216,7 +239,25 @@ def main_sweep(cfg_path, model_type, sweep_goal):
         dest="sweep_goal",
         type=str,
         required=True,
-        help="Choose sweep goal: (1) hyperparameter or (2) robustness.",
+        # Note the spelling of "hyperparmeter" -- it is what main_sweep() matches on. Without
+        # choices=, a typo silently trained the unmodified base config on every trial.
+        choices=["robustness", "segmentation", "hyperparmeter", "loss"],
+        help="Choose sweep goal: robustness, segmentation, hyperparmeter (sic) or loss.",
+    )
+    parser.add_argument(
+        "--count",
+        dest="count",
+        type=int,
+        default=30,
+        help="Number of sweep trials this agent runs before exiting. Without a bound the agent "
+        "runs until the SLURM walltime kills it.",
+    )
+    parser.add_argument(
+        "--sweep_project",
+        dest="sweep_project",
+        type=str,
+        default="InterScale_hyperparameter_sweep",
+        help="wandb project the sweep is registered under.",
     )
     parser.add_argument(
         "--prediction_task",
@@ -234,31 +275,43 @@ def main_sweep(cfg_path, model_type, sweep_goal):
 
     sweep_config = yaml_config["sweep_config"]
 
+    # "val_acc" was never a logged metric name -- the classification MetricCollection logs
+    # val_accuracy / val_f1_micro / val_f1_macro / val_f1_<class>. val_f1_macro also matches
+    # what EarlyStopping and ModelCheckpoint monitor. Only override the yaml when the flag
+    # is given, so a sweep config carrying its own metric block still works without it.
     if args.prediction_task == "classification":
-        sweep_config.update(
-            {
-                "metric": {"name": "val_acc", "goal": "maximize"},
-            }
-        )
+        sweep_config["metric"] = {"name": "val_f1_macro", "goal": "maximize"}
     elif args.prediction_task == "regression":
-        sweep_config.update(
-            {
-                "metric": {"name": "val_r2", "goal": "maximize"},  # Use 'val_r2' for regression tasks
-            }
-        )
+        sweep_config["metric"] = {"name": "val_r2", "goal": "maximize"}
 
-    if "GlobalModel" not in args.model_type or "CombinedModel" not in args.model_type:
-        transformer_keys = [key for key in sweep_config["parameters"] if key.startswith("transformer.")]
-        for key in transformer_keys:
-            del sweep_config["parameters"][key]
+    if "metric" not in sweep_config:
+        raise ValueError(
+            "Sweep config declares no `metric`, and --prediction_task was not given to supply "
+            "one. wandb would have nothing to rank trials by (and method: bayes cannot run)."
+        )
 
+    # Drop parameters for components this model_type does not have, so wandb does not sample
+    # values that nothing consumes. The previous filter looked for a `transformer.` prefix,
+    # which no key has ever used, under a condition that is true for every model_type.
+    drop_prefixes = []
+    if args.model_type not in ("LocalModel", "CombinedModel"):
+        drop_prefixes.append("model.local_component.")
+    if args.model_type not in ("GlobalModel", "CombinedModel"):
+        drop_prefixes.append("model.global_component.")
+    for key in [k for k in sweep_config["parameters"] if any(k.startswith(p) for p in drop_prefixes)]:
+        print(f"dropping sweep parameter not used by {args.model_type}: {key}")
+        del sweep_config["parameters"][key]
+
+    sweep_params = list(sweep_config["parameters"])
     print(sweep_config)
 
-    sweep_id = wandb.sweep(sweep_config, project="InterScale_hyperparameter_sweep")
+    sweep_id = wandb.sweep(sweep_config, project=args.sweep_project)
 
     def train_sweep_function():
         # Pass the sweep run object to main
-        main_sweep(args.cfg, args.model_type, args.sweep_goal)
+        main_sweep(args.cfg, args.model_type, args.sweep_goal, sweep_params=sweep_params)
 
-    # Run the sweep agent
-    wandb.agent(sweep_id, function=train_sweep_function)
+    # Without an explicit count the agent runs until the job's walltime kills it: `random`
+    # over this grid has ~1.5M combinations, so it never exhausts them on its own.
+    print(f"running {args.count} sweep trials (sweep_id={sweep_id})")
+    wandb.agent(sweep_id, function=train_sweep_function, count=args.count)

From bab7ca6e86267a54af62cf7bf2c5d39d0e252e22 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Thu, 20 Aug 2026 10:22:48 +0200
Subject: [PATCH 05/14] change config loading and estimate avg local and global
 size

---
 config_files/legnini_example.yaml     |  28 --
 src/interscale/config/__init__.py     | 172 +++++++++--
 src/interscale/config/cli.py          |  99 ++++++
 src/interscale/config/optim_config.py |   7 +-
 src/interscale/config/registry.py     | 138 +++++++++
 src/interscale/config/sweep.py        | 215 +++++++++++++
 src/interscale/main.py                |  44 ++-
 src/interscale/main_sweep.py          | 209 ++++++-------
 src/interscale/tl/__init__.py         |   3 +-
 src/interscale/tl/_preprocessing.py   |  54 ++++
 tests/test_sweep_config.py            | 426 ++++++++++++++++++++++++++
 11 files changed, 1219 insertions(+), 176 deletions(-)
 delete mode 100644 config_files/legnini_example.yaml
 create mode 100644 src/interscale/config/cli.py
 create mode 100644 src/interscale/config/registry.py
 create mode 100644 src/interscale/config/sweep.py
 create mode 100644 tests/test_sweep_config.py

diff --git a/config_files/legnini_example.yaml b/config_files/legnini_example.yaml
deleted file mode 100644
index 2bac6ab..0000000
--- a/config_files/legnini_example.yaml
+++ /dev/null
@@ -1,28 +0,0 @@
-model:
-  local_component:
-    name: GCN
-  global_component:
-    name: self-attn-transformer
-    parameters:
-      max_seq_len: 4299
-  decoder:
-    type: linear
-    dual_decoder: True
-  save: /Users/francesca.drummer/Documents/1_Projects/A3-InterScale/results/legnini23/ #Options: local, wandb, None
-optim:
-  lr: 0.001
-  loss: SmoothL1
-  seed: 44
-  accelerator: cpu
-dataset:
-  h5ad_data: /Users/francesca.drummer/Documents/1_Projects/A3-InterScale/data/legnini23.h5ad
-  name: legnini23
-  prediction_task: regression
-  prediction_level: node
-  layer_key: log1p_norm
-  sample_key: ['sample']
-  spatial_neigbors_kwargs:
-    radius: 200
-    library_key: sample
-  batch_size: 3
-  pct_mask_nodes: 0.3
diff --git a/src/interscale/config/__init__.py b/src/interscale/config/__init__.py
index 45d0bcd..5296fe3 100644
--- a/src/interscale/config/__init__.py
+++ b/src/interscale/config/__init__.py
@@ -23,13 +23,133 @@ def get_cfg_defaults():
     return cfg
 
 
-def load_config(cfg_path=None):
+def _normalise_cfg_paths(cfg_path):
+    """Coerce the ``cfg_path`` argument into a list of existing ``Path`` objects.
+
+    Accepts a single str/Path (the historical signature) or an iterable of them,
+    so callers that layer several files can pass a list.
+    """
+    if cfg_path is None:
+        return []
+    if isinstance(cfg_path, str | Path):
+        paths = [Path(cfg_path)]
+    else:
+        paths = [Path(p) for p in cfg_path]
+
+    missing = [str(p) for p in paths if not p.is_file()]
+    if missing:
+        raise FileNotFoundError("config file(s) not found: " + ", ".join(missing))
+    return paths
+
+
+def _peek_component_names(cfg_paths):
+    """Find the local/global component names declared across ``cfg_paths``.
+
+    The component *parameter* schemas (``model.local_component.parameters.*``) only exist
+    once ``get_local_component_cfg`` / ``get_global_component_cfg`` have added them, so the
+    names have to be read before any file is merged. Files are scanned in merge order and
+    the last file naming a component wins, matching what the merge itself would produce.
+
+    Uses ``getattr(..., "name", None)`` rather than attribute access: a file may set
+    ``model.global_component.parameters.max_seq_len`` without naming the component (a
+    dataset-level file layered on top of a base file that does the naming), and plain
+    attribute access raises ``AttributeError`` on the absent ``name`` key.
+    """
+    local_component_name = None
+    global_component_name = None
+
+    for path in cfg_paths:
+        with path.open() as f:
+            temp_cfg = CN.load_cfg(f)
+
+        model = getattr(temp_cfg, "model", None)
+        if model is None:
+            continue
+
+        local_component = getattr(model, "local_component", None)
+        if local_component is not None and getattr(local_component, "name", None):
+            local_component_name = local_component.name
+
+        global_component = getattr(model, "global_component", None)
+        if global_component is not None and getattr(global_component, "name", None):
+            global_component_name = global_component.name
+
+    return local_component_name, global_component_name
+
+
+def _coerce_override_values(cfg, overrides):
+    """Return ``overrides`` with ints promoted to float where the cfg default is a float.
+
+    ``merge_from_list`` rejects an int for a float key outright (yacs only tolerates a type
+    mismatch when one side is ``None``). YAML writes ``0`` as an int, so an override of
+    ``dataset.pct_mask_nodes: 0`` against the ``0.2`` default would raise even though the
+    value is perfectly valid. Promote it instead of making callers write ``0.0``.
+
+    ``overrides`` is the flat ``[key, value, key, value, ...]`` form ``merge_from_list`` takes.
+    """
+    coerced = list(overrides)
+
+    for i in range(0, len(coerced) - 1, 2):
+        key, value = coerced[i], coerced[i + 1]
+        if not isinstance(value, int) or isinstance(value, bool):
+            continue
+
+        # Walk the dotted key to the current value; a missing key is left alone so that
+        # merge_from_list raises its own (clearer) error about the unknown key.
+        node = cfg
+        try:
+            *parents, leaf = str(key).split(".")
+            for part in parents:
+                node = node[part]
+            current = node[leaf]
+        except (KeyError, TypeError):
+            continue
+
+        if isinstance(current, float):
+            coerced[i + 1] = float(value)
+
+    return coerced
+
+
+def _validate_optim(cfg):
+    """Reject configs whose training-length settings stop a run inside the LR warm-up.
+
+    ``CosineWarmupScheduler`` ramps the learning rate linearly over ``optim.lr_warmup`` epochs
+    (interval ``"epoch"``), so a run that early-stops before the ramp finishes has never trained
+    at ``optim.lr`` and reports a near-initialisation model -- typically a collapsed constant
+    predictor, which is easy to misread as a class-imbalance problem.
+
+    Raises
+    ------
+    ValueError
+        If ``optim.min_epochs`` does not exceed ``optim.lr_warmup`` while the warm-up scheduler
+        is in use.
+    """
+    if cfg.optim.lr_scheduler != "CosineWarmupScheduler":
+        return
+    if cfg.optim.min_epochs <= cfg.optim.lr_warmup:
+        raise ValueError(
+            f"optim.min_epochs ({cfg.optim.min_epochs}) must be greater than optim.lr_warmup "
+            f"({cfg.optim.lr_warmup}): with CosineWarmupScheduler the LR is still ramping until "
+            f"epoch {cfg.optim.lr_warmup}, so EarlyStopping (patience={cfg.optim.patience}) can "
+            "end the run before the model has trained at optim.lr. Raise optim.min_epochs (2x "
+            "lr_warmup is the convention here) or lower optim.lr_warmup."
+        )
+
+
+def load_config(cfg_path=None, overrides=None):
     """Loads and optionally overrides config values.
 
     Parameters
     ----------
-    cfg_path : str or pathlib.Path, optional
-        Path to the config file to load. If None, only default values are used.
+    cfg_path : str or pathlib.Path or list, optional
+        Path to the config file to load, or a list of paths merged left to right so
+        later files override earlier ones. If None, only default values are used.
+    overrides : list, optional
+        Flat ``[key, value, key, value, ...]`` overrides applied after every file, in the
+        form ``merge_from_list`` expects (e.g. ``["dataset.prediction_obs", "condition"]``).
+        Keys are dotted paths and must already exist in the config, so a typo raises
+        rather than silently doing nothing.
 
     Returns
     -------
@@ -39,39 +159,29 @@ def load_config(cfg_path=None):
     # First get all default configs including local component defaults
     cfg = get_cfg_defaults()
 
+    cfg_paths = _normalise_cfg_paths(cfg_path)
+
     # Documented as defaults-only, but the code below dereferences cfg_path
     # unconditionally, so None used to raise AttributeError too.
-    if cfg_path is None:
+    if not cfg_paths and not overrides:
+        _validate_optim(cfg)
         cfg.freeze()
         return cfg
 
-    # Callers pass a str: main.py and main_sweep.py both declare --cfg as
-    # type=str, and the docstring says str. Normalise instead of requiring
-    # every caller to wrap it.
-    cfg_path = Path(cfg_path)
-
-    with cfg_path.open() as f:
-        # Create a temporary config to load the model type
-        temp_cfg = CN.load_cfg(f)
-
-        # If model type is specified, load the corresponding local component configs
-        if hasattr(temp_cfg, "model") and hasattr(temp_cfg.model, "local_component"):
-            if temp_cfg.model.local_component.name is not None:
-                local_component_name = temp_cfg.model.local_component.name
-                if local_component_name:
-                    # Ensure local component configs are loaded before merging
-                    cfg = get_local_component_cfg(cfg, local_component_name)
-
-        # If model type is specified, load the corresponding global component configs
-        if hasattr(temp_cfg, "model") and hasattr(temp_cfg.model, "global_component"):
-            if temp_cfg.model.global_component.name is not None:
-                global_component_name = temp_cfg.model.global_component.name
-                if global_component_name:
-                    # Ensure global component configs are loaded before merging
-                    cfg = get_global_component_cfg(cfg, global_component_name)
-
-        # Now merge the full config file
-        cfg.merge_from_file(cfg_path)
+    local_component_name, global_component_name = _peek_component_names(cfg_paths)
+    if local_component_name:
+        # Ensure local component configs are loaded before merging
+        cfg = get_local_component_cfg(cfg, local_component_name)
+    if global_component_name:
+        # Ensure global component configs are loaded before merging
+        cfg = get_global_component_cfg(cfg, global_component_name)
+
+    for path in cfg_paths:
+        cfg.merge_from_file(str(path))
+
+    if overrides:
+        cfg.merge_from_list(_coerce_override_values(cfg, overrides))
 
+    _validate_optim(cfg)
     cfg.freeze()
     return cfg
diff --git a/src/interscale/config/cli.py b/src/interscale/config/cli.py
new file mode 100644
index 0000000..ed54865
--- /dev/null
+++ b/src/interscale/config/cli.py
@@ -0,0 +1,99 @@
+"""Shared ``--dataset/--task`` command-line plumbing for the training entrypoints.
+
+``main.py`` and ``main_sweep.py`` both need to turn command-line arguments into one frozen
+config. Keeping that in one place is what stops the two entrypoints drifting apart the way the
+per-dataset shell scripts did.
+"""
+
+import sys
+
+from . import load_config
+from .registry import iter_pairs, list_datasets, list_tasks, load_registry, resolve_config
+
+
+def add_config_args(parser):
+    """Add the config-selection arguments shared by every training entrypoint."""
+    parser.add_argument(
+        "--dataset",
+        dest="dataset",
+        type=str,
+        default=None,
+        help="Dataset key from config_files/registry.yaml (e.g. melton25). Use with --task.",
+    )
+    parser.add_argument(
+        "--task",
+        dest="task",
+        type=str,
+        default=None,
+        help="Task key for that dataset (e.g. graph_clas). Use with --dataset.",
+    )
+    parser.add_argument(
+        "--registry",
+        dest="registry",
+        type=str,
+        default=None,
+        help="Path to the registry (default: config_files/registry.yaml, relative to the cwd).",
+    )
+    parser.add_argument(
+        "--cfg",
+        dest="cfg",
+        type=str,
+        default=None,
+        help="A single config file, bypassing the registry. Mutually exclusive with --dataset/--task.",
+    )
+    parser.add_argument(
+        "--list",
+        dest="list_pairs",
+        action="store_true",
+        help="Print the dataset/task pairs the registry defines, then exit.",
+    )
+    return parser
+
+
+def print_registry(registry_path=None, stream=None):
+    """Print every ``(dataset, task)`` pair the registry defines."""
+    stream = sys.stdout if stream is None else stream
+    registry = load_registry(registry_path)
+
+    print("Available --dataset / --task pairs:\n", file=stream)
+    for dataset in list_datasets(registry):
+        tasks = list_tasks(registry, dataset)
+        print(f"  {dataset}", file=stream)
+        for task in tasks:
+            print(f"      --dataset {dataset} --task {task}", file=stream)
+    print("", file=stream)
+    return list(iter_pairs(registry))
+
+
+def resolve_cfg_from_args(args, parser=None):
+    """Turn parsed arguments into one frozen config.
+
+    Accepts either ``--dataset``/``--task`` (resolved through the registry) or a single
+    ``--cfg`` file, and refuses the ambiguous combination of both.
+    """
+
+    def fail(message):
+        if parser is not None:
+            parser.error(message)
+        raise SystemExit(f"error: {message}")
+
+    using_registry = args.dataset is not None or args.task is not None
+
+    if args.cfg is not None and using_registry:
+        fail("--cfg cannot be combined with --dataset/--task; the layered registry config would be ignored.")
+
+    if args.cfg is not None:
+        return load_config(args.cfg)
+
+    if not using_registry:
+        fail("give either --dataset and --task, or --cfg. Use --list to see the registered pairs.")
+
+    if args.dataset is None or args.task is None:
+        fail("--dataset and --task must be given together. Use --list to see the registered pairs.")
+
+    # KeyError from the registry carries the list of valid keys; surface it as a clean CLI
+    # error rather than a traceback several minutes into a queued job.
+    try:
+        return resolve_config(args.dataset, args.task, registry_path=args.registry)
+    except KeyError as exc:
+        fail(str(exc).strip("\"'"))
diff --git a/src/interscale/config/optim_config.py b/src/interscale/config/optim_config.py
index 8798e8c..50148f4 100644
--- a/src/interscale/config/optim_config.py
+++ b/src/interscale/config/optim_config.py
@@ -25,7 +25,12 @@ def get_optim_cfg(cfg):
     cfg.optim.early_stopping = True
     cfg.optim.patience = 5  # EarlyStopping patience in epochs
     cfg.optim.min_delta = 0.0  # EarlyStopping min_delta
-    cfg.optim.min_epochs = 1  # floor on training length; set above lr_warmup to clear warm-up
+    # Floor on training length. MUST stay above lr_warmup: CosineWarmupScheduler ramps the LR
+    # linearly over lr_warmup epochs, so a run that stops inside the ramp has only ever seen a
+    # fraction of cfg.optim.lr and is measured at (near) initialisation. The default is 2x
+    # lr_warmup, matching chen22. `_validate_optim` enforces the invariant at config-load time
+    # rather than leaving it to each dataset to remember.
+    cfg.optim.min_epochs = 40
     # Metric driving EarlyStopping / ModelCheckpoint / the LR scheduler.
     # "auto" -> val_f1_macro for classification, val_loss for regression.
     cfg.optim.monitor = "auto"
diff --git a/src/interscale/config/registry.py b/src/interscale/config/registry.py
new file mode 100644
index 0000000..712af31
--- /dev/null
+++ b/src/interscale/config/registry.py
@@ -0,0 +1,138 @@
+"""Resolve a ``(dataset, task)`` pair into the config files that describe that run.
+
+A run's configuration is layered, most general first, so that nothing is stated twice:
+
+1. ``base`` -- architecture choices shared by every run (component names, decoder type).
+2. the **dataset** file -- what the data *is*: paths, ``sample_key``, neighbour graph radius,
+   ``max_seq_len``, results dir, wandb project.
+3. the **task** file -- what is being predicted: ``prediction_task``/``prediction_level``,
+   loss, monitored metric.
+4. per-pair **overrides** -- the handful of values that belong to neither, most importantly
+   ``dataset.prediction_obs`` (chen22+graph is ``stage``, melton25+graph is ``condition``,
+   melton25+node is ``cell_type_coarse``, so the label column is a property of the pair).
+
+Dataset comes *before* task on purpose: a dataset that needs its own training length
+(chen22 trains for 400 epochs with patience 40, against the task default of 200/20) should
+win over the task-level default, because it is the more specific statement.
+
+The registry is deliberately **not** a yacs config. ``CfgNode.merge_from_file`` raises
+``KeyError`` for any key absent from the defaults schema, so a nested ``tasks:`` block
+could never live inside a config file that also gets merged into ``cfg``.
+"""
+
+from pathlib import Path
+
+import yaml
+
+from . import load_config
+
+# Relative to the current working directory. Jobs run from the repo root (scripts/run.sh
+# does `cd "$REPO"`), and the package cannot locate the repo itself: the container binds
+# src/ to /opt/interscale_src, so a path derived from __file__ would point at /opt.
+DEFAULT_REGISTRY_PATH = Path("config_files/registry.yaml")
+
+
+def load_registry(registry_path=None):
+    """Read the dataset/task registry.
+
+    Parameters
+    ----------
+    registry_path : str or pathlib.Path, optional
+        Path to ``registry.yaml``. Defaults to ``config_files/registry.yaml`` relative to
+        the current working directory.
+
+    Returns
+    -------
+    dict
+        The parsed registry, with a ``_root`` key holding the directory the registry lives
+        in, used to resolve the relative config paths it names.
+    """
+    path = Path(DEFAULT_REGISTRY_PATH if registry_path is None else registry_path)
+    if not path.is_file():
+        raise FileNotFoundError(
+            f"registry not found: {path} (cwd={Path.cwd()}). Run from the repo root or pass --registry."
+        )
+
+    with path.open() as f:
+        registry = yaml.safe_load(f) or {}
+
+    if "datasets" not in registry:
+        raise ValueError(f"{path} declares no `datasets:` block.")
+
+    registry["_root"] = path.parent
+    return registry
+
+
+def list_datasets(registry):
+    """Return the dataset keys the registry defines, sorted."""
+    return sorted(registry["datasets"])
+
+
+def list_tasks(registry, dataset):
+    """Return the task keys defined for ``dataset``, sorted."""
+    return sorted(_dataset_entry(registry, dataset).get("tasks", {}))
+
+
+def _dataset_entry(registry, dataset):
+    datasets = registry["datasets"]
+    if dataset not in datasets:
+        raise KeyError(f"unknown dataset '{dataset}'. Known datasets: {', '.join(sorted(datasets))}")
+    return datasets[dataset]
+
+
+def _task_entry(registry, dataset, task):
+    tasks = _dataset_entry(registry, dataset).get("tasks", {})
+    if task not in tasks:
+        raise KeyError(
+            f"dataset '{dataset}' defines no task '{task}'. Known tasks for {dataset}: {', '.join(sorted(tasks))}"
+        )
+    return tasks[task]
+
+
+def resolve_config_paths(dataset, task, registry_path=None, registry=None):
+    """Resolve a ``(dataset, task)`` pair into config file paths plus flat overrides.
+
+    Returns
+    -------
+    tuple[list[pathlib.Path], list]
+        The config files to merge in order, and the ``[key, value, ...]`` override list, both
+        ready to hand to :func:`interscale.config.load_config`.
+    """
+    if registry is None:
+        registry = load_registry(registry_path)
+    root = Path(registry["_root"])
+
+    dataset_entry = _dataset_entry(registry, dataset)
+    task_entry = _task_entry(registry, dataset, task)
+
+    paths = []
+    if registry.get("base"):
+        paths.append(root / registry["base"])
+    if dataset_entry.get("config"):
+        paths.append(root / dataset_entry["config"])
+    if task_entry.get("task"):
+        paths.append(root / task_entry["task"])
+
+    # Dataset-level overrides first so a task can still override them for its own pair.
+    overrides = {}
+    overrides.update(dataset_entry.get("overrides") or {})
+    overrides.update(task_entry.get("overrides") or {})
+
+    flat_overrides = []
+    for key, value in overrides.items():
+        flat_overrides.extend([key, value])
+
+    return paths, flat_overrides
+
+
+def resolve_config(dataset, task, registry_path=None, registry=None):
+    """Load the fully merged, frozen config for a ``(dataset, task)`` pair."""
+    paths, overrides = resolve_config_paths(dataset, task, registry_path=registry_path, registry=registry)
+    return load_config(paths, overrides=overrides)
+
+
+def iter_pairs(registry):
+    """Yield every ``(dataset, task)`` pair the registry defines, sorted."""
+    for dataset in list_datasets(registry):
+        for task in list_tasks(registry, dataset):
+            yield dataset, task
diff --git a/src/interscale/config/sweep.py b/src/interscale/config/sweep.py
new file mode 100644
index 0000000..7c44cb1
--- /dev/null
+++ b/src/interscale/config/sweep.py
@@ -0,0 +1,215 @@
+"""Build a wandb sweep config and apply a sampled trial onto a yacs config.
+
+Both halves used to live inline in ``main_sweep.py`` -- one in ``main_sweep()`` between an
+``h5ad`` load and a training call, the other in the ``__main__`` block -- which made the
+question "does this sweep actually vary the parameters it declares?" impossible to answer
+without launching a real run on a GPU. They are pure config transforms, so they live here
+and are covered by ``tests/test_sweep_config.py``.
+
+Application is **generic over the dotted key**: whatever ``parameters`` the sweep yaml
+declares is written to that exact path in the config, and a path that does not exist raises.
+The previous hand-maintained list of assignments produced two silent-no-op bugs that the
+config comments still record -- component keys written without the leading ``model.`` (every
+one raised ``KeyError``) and a ``dropout`` assignment against a schema whose key is
+``dropout_global``, which ``set_new_allowed(True)`` turned into a dead key that nothing read.
+Both classes are structurally impossible here: there is no allow-new-keys, and a key nothing
+can apply is an error rather than a shrug.
+"""
+
+from yacs.config import CfgNode as CN
+
+# Spelled "hyperparmeter" because that is what main_sweep.py has always matched on and what
+# existing sweep invocations pass; renaming it would silently break saved commands.
+SWEEP_GOALS = ("robustness", "segmentation", "hyperparmeter", "loss")
+
+# Metric each prediction task ranks trials by. val_f1_macro is what EarlyStopping and
+# ModelCheckpoint monitor for classification, and trainer.validate() runs after the best
+# checkpoint is restored, so the last logged value is the best epoch's score.
+TASK_METRICS = {
+    "classification": {"name": "val_f1_macro", "goal": "maximize"},
+    "regression": {"name": "val_r2", "goal": "maximize"},
+}
+
+LOCAL_COMPONENT_PREFIX = "model.local_component."
+GLOBAL_COMPONENT_PREFIX = "model.global_component."
+
+
+def unused_component_prefixes(model_type):
+    """Return the dotted key prefixes ``model_type`` has no component for.
+
+    Two independent checks, not if/elif: ``CombinedModel`` has *both* components, so an
+    ``elif`` arm mentioning it is unreachable and its transformer would never be swept.
+    """
+    prefixes = []
+    if model_type not in ("LocalModel", "CombinedModel"):
+        prefixes.append(LOCAL_COMPONENT_PREFIX)
+    if model_type not in ("GlobalModel", "CombinedModel"):
+        prefixes.append(GLOBAL_COMPONENT_PREFIX)
+    return prefixes
+
+
+def build_sweep_config(yaml_config, prediction_task=None, model_type=None):
+    """Turn a parsed sweep yaml into the dict ``wandb.sweep`` takes.
+
+    Parameters
+    ----------
+    yaml_config : dict
+        The parsed sweep yaml, containing a top-level ``sweep_config`` key.
+    prediction_task : str, optional
+        ``"classification"`` or ``"regression"``. When given, overrides the yaml's ``metric``
+        block with the metric that task actually logs. When omitted the yaml must supply its
+        own ``metric``, since wandb has nothing to rank trials by otherwise (and
+        ``method: bayes`` cannot run at all).
+    model_type : str, optional
+        When given, parameters targeting a component this model type does not have are
+        dropped, so wandb does not spend trials varying values nothing consumes.
+
+    Returns
+    -------
+    tuple[dict, list[str]]
+        The sweep config, and the sorted list of parameter names it varies.
+    """
+    if "sweep_config" not in yaml_config:
+        raise ValueError("sweep yaml declares no top-level `sweep_config:` block.")
+
+    # Copied rather than mutated in place so callers can reuse the parsed yaml.
+    sweep_config = {k: (dict(v) if isinstance(v, dict) else v) for k, v in yaml_config["sweep_config"].items()}
+
+    if prediction_task is not None:
+        if prediction_task not in TASK_METRICS:
+            raise ValueError(f"unknown prediction_task '{prediction_task}'. Must be one of: {', '.join(TASK_METRICS)}")
+        sweep_config["metric"] = dict(TASK_METRICS[prediction_task])
+
+    if "metric" not in sweep_config:
+        raise ValueError(
+            "Sweep config declares no `metric`, and no prediction_task was given to supply one. "
+            "wandb would have nothing to rank trials by (and method: bayes cannot run)."
+        )
+
+    if "parameters" not in sweep_config or not sweep_config["parameters"]:
+        raise ValueError("sweep config declares no `parameters` to vary.")
+
+    if model_type is not None:
+        drop_prefixes = unused_component_prefixes(model_type)
+        dropped = [k for k in sweep_config["parameters"] if any(k.startswith(p) for p in drop_prefixes)]
+        for key in dropped:
+            print(f"dropping sweep parameter not used by {model_type}: {key}")
+            del sweep_config["parameters"][key]
+        if not sweep_config["parameters"]:
+            raise ValueError(
+                f"every sweep parameter was dropped as unused by {model_type}; nothing left to vary."
+            )
+
+    return sweep_config, sorted(sweep_config["parameters"])
+
+
+def _set_dotted(cfg, key, value):
+    """Assign ``value`` at the dotted ``key`` in ``cfg``, requiring the path to exist."""
+    parts = str(key).split(".")
+    node = cfg
+
+    for depth, part in enumerate(parts[:-1]):
+        if not isinstance(node, CN) or part not in node:
+            raise KeyError(
+                f"sweep parameter '{key}' does not exist in the config: "
+                f"no '{'.'.join(parts[: depth + 1])}'. It would vary between trials with no effect."
+            )
+        node = node[part]
+
+    leaf = parts[-1]
+    if not isinstance(node, CN) or leaf not in node:
+        raise KeyError(
+            f"sweep parameter '{key}' does not exist in the config. "
+            f"It would vary between trials with no effect. "
+            f"Available keys under '{'.'.join(parts[:-1])}': {', '.join(sorted(node)) if isinstance(node, CN) else '<not a config node>'}"
+        )
+
+    # yacs rejects an int for a float key, and a sweep declaring `values: [0, 0.1, 0.3]`
+    # samples a genuine int for 0. Promote instead of failing mid-sweep.
+    current = node[leaf]
+    if isinstance(current, float) and isinstance(value, int) and not isinstance(value, bool):
+        value = float(value)
+
+    node[leaf] = value
+
+
+def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_params=None):
+    """Write a sampled sweep trial onto ``cfg``.
+
+    Parameters
+    ----------
+    cfg : CN
+        The base config. Defrosted, written to, and re-frozen; mutated in place and returned.
+    sweep_goal : str
+        One of :data:`SWEEP_GOALS`. Validated so a misspelling fails loudly instead of
+        skipping every branch and training the unmodified base config on every trial.
+    sweep_config : Mapping
+        The sampled trial -- ``wandb.config``, or any mapping of dotted key to value.
+    model_type : str, optional
+        When given, parameters targeting a component this model type lacks are skipped.
+    sweep_params : list, optional
+        The parameter names the sweep declares, i.e. the keys of the sweep config's
+        ``parameters`` block. This is the authoritative list of what to apply, and any
+        declared name missing from ``sweep_config`` raises. When omitted it is inferred as
+        the dotted keys of ``sweep_config``.
+
+    Returns
+    -------
+    tuple[CN, list[str]]
+        The frozen config, and the sorted list of keys actually applied.
+    """
+    if sweep_goal not in SWEEP_GOALS:
+        raise ValueError(
+            f"Unknown sweep_goal '{sweep_goal}'. Must be one of: {', '.join(SWEEP_GOALS)}. "
+            "(Nothing would be swept otherwise -- every trial would train the base config.)"
+        )
+
+    # main_sweep.py calls wandb.init(config=cfg), so wandb.config carries the whole base
+    # config alongside the sampled trial values. Only the sweep's declared parameters may be
+    # written back: iterating every key would assign the top-level `dataset` / `model` /
+    # `optim` sections over their CfgNodes with plain dicts.
+    if sweep_params is None:
+        keys = [k for k in sweep_config.keys() if "." in str(k)]
+    else:
+        keys = list(sweep_params)
+        absent = sorted(k for k in keys if k not in sweep_config)
+        if absent:
+            raise KeyError(
+                f"sweep declares parameters that the trial did not sample, so they cannot be "
+                f"applied: {absent}"
+            )
+
+    # Deliberately NOT set_new_allowed(True): allowing new keys is what let a misspelled
+    # parameter create a dead config entry that nothing ever read.
+    was_frozen = cfg.is_frozen()
+    cfg.defrost()
+
+    skipped_prefixes = unused_component_prefixes(model_type) if model_type is not None else []
+
+    applied = []
+    skipped = []
+    for key in keys:
+        if any(str(key).startswith(p) for p in skipped_prefixes):
+            skipped.append(key)
+            continue
+        _set_dotted(cfg, key, sweep_config[key])
+        applied.append(key)
+
+    if was_frozen:
+        cfg.freeze()
+
+    if skipped:
+        print(f"skipped sweep parameters for components {model_type} does not have: {sorted(skipped)}")
+
+    if sweep_params is not None:
+        # A parameter wandb varies but nothing reads makes every trial's difference invisible,
+        # which is indistinguishable from the sweep not working. Fail rather than warn.
+        unapplied = sorted(set(sweep_params) - set(applied) - set(skipped))
+        if unapplied:
+            raise KeyError(
+                f"sweep declares parameters that were not applied to the config, so they would "
+                f"vary between trials with no effect: {unapplied}"
+            )
+
+    print(f"applied sweep parameters: {sorted(applied)}")
+    return cfg, sorted(applied)
diff --git a/src/interscale/main.py b/src/interscale/main.py
index 2d88c8a..a2f6593 100644
--- a/src/interscale/main.py
+++ b/src/interscale/main.py
@@ -1,11 +1,13 @@
 import argparse
 import warnings
+from pathlib import Path
 
 import scanpy as sc
 import squidpy as sq
 
 import interscale as interscale
 from interscale.config import load_config
+from interscale.config.cli import add_config_args, print_registry, resolve_cfg_from_args
 from interscale.geome_dataloader import GraphAnnDataModule
 from interscale.pp import apply_segmentation_noise
 from interscale.tl import prepare_geome_dataset, remove_zero_expression_cells, set_full_reproducibility
@@ -16,9 +18,21 @@
 warnings.filterwarnings("ignore", category=FutureWarning, message=r".*n_neighs.*")
 
 
-def main(cfg_path, model_type):
+def main(cfg, model_type):
+    """Train a single model.
+
+    Parameters
+    ----------
+    cfg : CN or str or pathlib.Path
+        An already-resolved config, or a path to a single config file to load.
+    model_type : str
+        One of ``LocalModel``, ``GlobalModel``, ``CombinedModel``.
+    """
+    # Accepts a path as well as a cfg so the historical main(cfg_path, model_type) call still
+    # works; the CLI now resolves the config itself, since --dataset/--task layers several files.
+    if isinstance(cfg, str | Path):
+        cfg = load_config(cfg)
 
-    cfg = load_config(cfg_path)
     set_full_reproducibility(cfg.optim.seed)
     print(cfg)
     adata = sc.read_h5ad(cfg.dataset.h5ad_data)
@@ -78,16 +92,34 @@ def main(cfg_path, model_type):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="GTLongRange")
+    parser = argparse.ArgumentParser(
+        description="Train one InterScale model on a registered (dataset, task) pair.",
+        epilog=(
+            "examples:\n"
+            "  %(prog)s --dataset melton25 --task graph_clas --model_type CombinedModel\n"
+            "  %(prog)s --list\n"
+            "  %(prog)s --cfg config_files/legnini_example.yaml --model_type CombinedModel\n"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
 
-    parser.add_argument("--cfg", dest="cfg", type=str, required=True, help="The configuration file path.")
+    add_config_args(parser)
     parser.add_argument(
         "--model_type",
         dest="model_type",
         type=str,
-        required=True,
+        default=None,
+        choices=["LocalModel", "GlobalModel", "CombinedModel"],
         help="The model type: LocalModel, GlobalModel or CombinedModel.",
     )
     args = parser.parse_args()
 
-    main(args.cfg, args.model_type)
+    if args.list_pairs:
+        print_registry(args.registry)
+        raise SystemExit(0)
+
+    # Not required=True because --list must work without it.
+    if args.model_type is None:
+        parser.error("--model_type is required (LocalModel, GlobalModel or CombinedModel).")
+
+    main(resolve_cfg_from_args(args, parser=parser), args.model_type)
diff --git a/src/interscale/main_sweep.py b/src/interscale/main_sweep.py
index b9c2db8..e89fa4d 100644
--- a/src/interscale/main_sweep.py
+++ b/src/interscale/main_sweep.py
@@ -10,6 +10,8 @@
 
 import interscale as interscale
 from interscale.config import load_config
+from interscale.config.cli import add_config_args, print_registry, resolve_cfg_from_args
+from interscale.config.sweep import apply_sweep_config, build_sweep_config
 from interscale.geome_dataloader import GraphAnnDataModule
 from interscale.pp.segmentation_noise import apply_segmentation_noise
 from interscale.tl import prepare_geome_dataset
@@ -61,11 +63,31 @@ def print_memory_debug():
         print(f"Memory debug failed: {e}")
 
 
-def main_sweep(cfg_path, model_type, sweep_goal, sweep_params=None):
+def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
+    """Run one sweep trial.
+
+    Parameters
+    ----------
+    cfg_factory : callable or str or pathlib.Path
+        Called with no arguments to build a **fresh** config for this trial. It must be a
+        factory, not a config: ``wandb.agent`` invokes this function once per trial, and
+        applying a trial's values mutates the config, so a shared object would accumulate
+        every previous trial's settings. A path is accepted and wrapped for convenience.
+    model_type : str
+        One of ``LocalModel``, ``GlobalModel``, ``CombinedModel``.
+    sweep_goal : str
+        One of ``interscale.config.sweep.SWEEP_GOALS``.
+    sweep_params : list, optional
+        The parameter names the sweep declares, used to assert that every one of them was
+        actually applied to the config.
+    """
 
     print_memory_usage("Start of main_sweep")
 
-    cfg = load_config(cfg_path)
+    if callable(cfg_factory):
+        cfg = cfg_factory()
+    else:
+        cfg = load_config(cfg_factory)
 
     assert cfg.wandb.use, "Wandb is not enabled in the configuration file. Necessary for sweep."
 
@@ -88,77 +110,17 @@ def main_sweep(cfg_path, model_type, sweep_goal, sweep_params=None):
 
     # Update configuration with sweep parameters
     if sweep_config is not None:
-        cfg.set_new_allowed(True)
-        cfg.defrost()
         print("sweep config: ", sweep_config)
         print("sweep run: ", sweep_run.config)
-        if sweep_goal == "robustness":
-            print("robustness sweep")
-            cfg.dataset.pct_mask_nodes = sweep_config["dataset.pct_mask_nodes"]
-            cfg.dataset.spatial_neigbors_kwargs.radius = sweep_config["dataset.spatial_neigbors_kwargs.radius"]
-            cfg.optim.seed = sweep_config["optim.seed"]
-        elif sweep_goal == "segmentation":
-            print("segmentation sweep")
-            cfg.dataset.segmentation_robustness = sweep_config["dataset.segmentation_robustness"]
-            cfg.optim.seed = sweep_config["optim.seed"]
-        elif sweep_goal == "hyperparmeter":
-            print("hyperparameter sweep")
-            applied = []
-
-            def _apply(node, attr, key):
-                """Assign a swept value only if the sweep actually declares that key.
-
-                Staged sweeps vary a subset of the parameters (e.g. optimiser only), so an
-                unconditional lookup would KeyError on every key the stage omits.
-                """
-                if key in sweep_config.keys():
-                    setattr(node, attr, sweep_config[key])
-                    applied.append(key)
-
-            _apply(cfg.optim, "lr", "optim.lr")
-            _apply(cfg.optim, "lr_warmup", "optim.lr_warmup")
-            _apply(cfg.optim, "wd", "optim.wd")
-            _apply(cfg.dataset, "batch_size", "dataset.batch_size")
-            _apply(cfg.dataset, "pct_mask_nodes", "dataset.pct_mask_nodes")
-            _apply(cfg.model, "n_embed", "model.n_embed")
-
-            # Two separate `if`s, not if/elif: CombinedModel has BOTH components, and an
-            # `elif model_type == "GlobalModel" or model_type == "CombinedModel"` is
-            # unreachable for CombinedModel, so its transformer was never swept.
-            local = cfg.model.local_component.parameters
-            glob = cfg.model.global_component.parameters
-            if model_type in ("LocalModel", "CombinedModel"):
-                print("LocalModel configs")
-                _apply(local, "num_layers", "model.local_component.parameters.num_layers")
-                _apply(local, "hidden_dim", "model.local_component.parameters.hidden_dim")
-                _apply(local, "dropout_local", "model.local_component.parameters.dropout_local")
-            if model_type in ("GlobalModel", "CombinedModel"):
-                print("transformer configs")
-                _apply(glob, "dim_feedforward", "model.global_component.parameters.dim_feedforward")
-                _apply(glob, "num_layers", "model.global_component.parameters.num_layers")
-                _apply(glob, "n_heads", "model.global_component.parameters.n_heads")
-                # The config key is `dropout_global` (see global_component_config.py); assigning
-                # to `dropout` silently created a dead key because set_new_allowed(True) is on.
-                _apply(glob, "dropout_global", "model.global_component.parameters.dropout_global")
-
-            if sweep_params is not None:
-                ignored = sorted(set(sweep_params) - set(applied))
-                if ignored:
-                    print(
-                        f"WARNING: sweep declares parameters that nothing applies, so they vary "
-                        f"between trials with no effect: {ignored}"
-                    )
-            print(f"applied sweep parameters: {sorted(applied)}")
-        elif sweep_goal == "loss":
-            print("loss sweep")
-            cfg.optim.loss = sweep_config["optim.loss"]
-        else:
-            raise ValueError(
-                f"Unknown --sweep_goal '{sweep_goal}'. Must be one of: "
-                "robustness, segmentation, hyperparmeter, loss. "
-                "(Nothing would be swept otherwise -- every trial would train the base config.)"
-            )
-        cfg.freeze()
+        # Applies whatever dotted keys the sweep declares, and raises if any of them does not
+        # exist in the config rather than letting it vary between trials with no effect.
+        cfg, _applied = apply_sweep_config(
+            cfg,
+            sweep_goal,
+            sweep_config,
+            model_type=model_type,
+            sweep_params=sweep_params,
+        )
 
     ####### PREPROCESSING #######
     # Load adata
@@ -227,18 +189,34 @@ def _apply(node, attr, key):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="GTLongRange")
+    parser = argparse.ArgumentParser(
+        description="Run a wandb sweep for a registered (dataset, task) pair.",
+        epilog=(
+            "examples:\n"
+            "  %(prog)s --dataset melton25 --task graph_clas --model_type CombinedModel \\\n"
+            "      --sweep_cfg config_files/sweeps/hyperparameters.yaml\n"
+            "  %(prog)s --list\n"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
 
-    parser.add_argument("--cfg", dest="cfg", type=str, required=True, help="The configuration file path.")
+    add_config_args(parser)
+    parser.add_argument(
+        "--sweep_cfg", dest="sweep_cfg", type=str, required=False, help="The sweep configuration file path."
+    )
     parser.add_argument(
-        "--sweep_cfg", dest="sweep_cfg", type=str, required=True, help="The sweep configuration file path."
+        "--model_type",
+        dest="model_type",
+        type=str,
+        default=None,
+        choices=["LocalModel", "GlobalModel", "CombinedModel"],
+        help="The model type: LocalModel, GlobalModel or CombinedModel.",
     )
-    parser.add_argument("--model_type", dest="model_type", type=str, required=True)
     parser.add_argument(
         "--sweep_goal",
         dest="sweep_goal",
         type=str,
-        required=True,
+        default=None,
         # Note the spelling of "hyperparmeter" -- it is what main_sweep() matches on. Without
         # choices=, a typo silently trained the unmodified base config on every trial.
         choices=["robustness", "segmentation", "hyperparmeter", "loss"],
@@ -265,51 +243,64 @@ def _apply(node, attr, key):
         type=str,
         required=False,
         choices=["regression", "classification"],
-        help="Type of prediction task (regression or classification)",
+        help="Type of prediction task. Defaults to the resolved config's dataset.prediction_task.",
     )
     args = parser.parse_args()
 
+    if args.list_pairs:
+        print_registry(args.registry)
+        raise SystemExit(0)
+
+    # Not required=True on these, because --list must work without them.
+    for flag, value in (
+        ("--model_type", args.model_type),
+        ("--sweep_goal", args.sweep_goal),
+        ("--sweep_cfg", args.sweep_cfg),
+    ):
+        if value is None:
+            parser.error(f"{flag} is required.")
+
+    # Resolve once up front so a bad --dataset/--task fails now rather than after wandb has
+    # registered a sweep server-side. Each trial re-resolves its own fresh copy below.
+    base_cfg = resolve_cfg_from_args(args, parser=parser)
+
+    # The metric depends on the prediction task, which the resolved config already knows, so the
+    # flag is only needed to override it. The per-dataset scripts used to hardcode
+    # `--prediction_task classification`, which would rank a regression sweep by val_f1_macro.
+    prediction_task = args.prediction_task or base_cfg.dataset.prediction_task
+
     # Load both base config and sweep config from yaml
     with open(args.sweep_cfg) as f:
         yaml_config = yaml.safe_load(f)
 
-    sweep_config = yaml_config["sweep_config"]
-
-    # "val_acc" was never a logged metric name -- the classification MetricCollection logs
-    # val_accuracy / val_f1_micro / val_f1_macro / val_f1_<class>. val_f1_macro also matches
-    # what EarlyStopping and ModelCheckpoint monitor. Only override the yaml when the flag
-    # is given, so a sweep config carrying its own metric block still works without it.
-    if args.prediction_task == "classification":
-        sweep_config["metric"] = {"name": "val_f1_macro", "goal": "maximize"}
-    elif args.prediction_task == "regression":
-        sweep_config["metric"] = {"name": "val_r2", "goal": "maximize"}
-
-    if "metric" not in sweep_config:
-        raise ValueError(
-            "Sweep config declares no `metric`, and --prediction_task was not given to supply "
-            "one. wandb would have nothing to rank trials by (and method: bayes cannot run)."
-        )
-
-    # Drop parameters for components this model_type does not have, so wandb does not sample
-    # values that nothing consumes. The previous filter looked for a `transformer.` prefix,
-    # which no key has ever used, under a condition that is true for every model_type.
-    drop_prefixes = []
-    if args.model_type not in ("LocalModel", "CombinedModel"):
-        drop_prefixes.append("model.local_component.")
-    if args.model_type not in ("GlobalModel", "CombinedModel"):
-        drop_prefixes.append("model.global_component.")
-    for key in [k for k in sweep_config["parameters"] if any(k.startswith(p) for p in drop_prefixes)]:
-        print(f"dropping sweep parameter not used by {args.model_type}: {key}")
-        del sweep_config["parameters"][key]
-
-    sweep_params = list(sweep_config["parameters"])
+    sweep_config, sweep_params = build_sweep_config(
+        yaml_config,
+        prediction_task=prediction_task,
+        model_type=args.model_type,
+    )
     print(sweep_config)
 
+    # Fail before wandb.sweep() registers anything if a declared parameter does not exist in this
+    # config: the sweep would otherwise burn every trial varying a key that nothing reads.
+    apply_sweep_config(
+        base_cfg.clone(),
+        args.sweep_goal,
+        dict.fromkeys(sweep_params, None),
+        model_type=args.model_type,
+        sweep_params=sweep_params,
+    )
+
     sweep_id = wandb.sweep(sweep_config, project=args.sweep_project)
 
     def train_sweep_function():
-        # Pass the sweep run object to main
-        main_sweep(args.cfg, args.model_type, args.sweep_goal, sweep_params=sweep_params)
+        # A fresh config per trial: applying a trial's values mutates it, so a shared object
+        # would carry the previous trial's settings into this one.
+        main_sweep(
+            lambda: resolve_cfg_from_args(args, parser=parser),
+            args.model_type,
+            args.sweep_goal,
+            sweep_params=sweep_params,
+        )
 
     # Without an explicit count the agent runs until the job's walltime kills it: `random`
     # over this grid has ~1.5M combinations, so it never exhausts them on its own.
diff --git a/src/interscale/tl/__init__.py b/src/interscale/tl/__init__.py
index 5d0e539..03b4adc 100644
--- a/src/interscale/tl/__init__.py
+++ b/src/interscale/tl/__init__.py
@@ -1,4 +1,4 @@
-from ._preprocessing import remove_zero_expression_cells
+from ._preprocessing import get_average_local_and_global_size, remove_zero_expression_cells
 from .geome_utils import prepare_a2d_dataset, prepare_geome_dataset
 from .masking import apply_mask, attn_mask_diagonal, create_transformer_attention_mask_from_edges
 from .padding import pad_batch
@@ -16,4 +16,5 @@
     "create_transformer_attention_mask_from_edges",
     "attn_mask_diagonal",
     "remove_zero_expression_cells",
+    "get_average_local_and_global_size",
 ]
diff --git a/src/interscale/tl/_preprocessing.py b/src/interscale/tl/_preprocessing.py
index 97c9675..47373be 100644
--- a/src/interscale/tl/_preprocessing.py
+++ b/src/interscale/tl/_preprocessing.py
@@ -1,4 +1,6 @@
 import numpy as np
+import squidpy as sq
+from typing import Literal
 
 
 def remove_zero_expression_cells(adata):
@@ -8,3 +10,55 @@ def remove_zero_expression_cells(adata):
         nonzero_cells = np.array(adata.X.sum(axis=1) != 0).flatten()
         adata = adata[nonzero_cells].copy()
     return adata
+
+
+PIXEL_TO_UM = 0.138  # Resolve MC1 default; override via cfg if instrument differs
+
+def get_average_local_and_global_size(
+    adata,
+    cfg,
+    *,
+    coord_units: Literal["pixels", "micrometer"] = "micrometer",
+    pixel_size_um: float = PIXEL_TO_UM,
+):
+    """Average size of the local and the global window, in cells and in micrometers.
+
+    Parameters
+    ----------
+    coord_units
+        Units of ``adata.obsm['spatial']``. If ``"pixels"``, distances are
+        converted to micrometers using ``pixel_size_um``.
+    pixel_size_um
+        Micrometers per pixel. Ignored when ``coord_units='micrometer'``.
+    """
+    scale = pixel_size_um if coord_units == "pixels" else 1.0
+
+    ## Local: avg direct-neighbor count * n_layers
+    if "spatial_connectivities" not in adata.obsp:
+        kwargs = dict(cfg.dataset.spatial_neigbors_kwargs)
+        # library_key is only filled in by prepare_geome_dataset, so pick the first sample key here
+        kwargs["library_key"] = kwargs.get("library_key") or cfg.dataset.sample_key[0]
+        adata.obs[kwargs["library_key"]] = adata.obs[kwargs["library_key"]].astype("category")
+        sq.gr.spatial_neighbors(adata, **kwargs)
+    n_layers = cfg.model.local_component.parameters.num_layers
+    conn = adata.obsp["spatial_connectivities"].tocsr()
+    local_cells = conn.getnnz(axis=1).mean() * n_layers
+    # reach in micrometers: avg edge length walked over n_layers hops
+    local_dist_um = adata.obsp["spatial_distances"].data.mean() * n_layers * scale
+
+    ## Global: avg cells and avg extent per sample (or per sliding window, whichever sample_key defines)
+    groups = adata.obs.groupby(list(cfg.dataset.sample_key), observed=True)
+    coords = adata.obsm["spatial"]
+    global_cells = groups.size().mean()
+    global_dist_um = np.mean(
+        [np.linalg.norm(coords[idx].max(axis=0) - coords[idx].min(axis=0))
+         for idx in groups.indices.values()]
+    ) * scale
+
+    return {
+        "local_cells": local_cells,
+        "local_dist_um": local_dist_um,
+        "global_cells": global_cells,
+        "global_dist_um": global_dist_um,
+    }
+    
diff --git a/tests/test_sweep_config.py b/tests/test_sweep_config.py
new file mode 100644
index 0000000..9c82afe
--- /dev/null
+++ b/tests/test_sweep_config.py
@@ -0,0 +1,426 @@
+"""Tests that a sweep actually varies the hyperparameters it declares.
+
+The failure these guard against is silent: wandb samples a value, nothing in the code reads the
+key it was sampled for, every trial trains an identical model, and the sweep report looks like a
+legitimate "no effect" result. Two real instances are recorded in the config comments -- component
+keys written without the leading ``model.``, and a ``dropout`` assignment against a schema key
+actually named ``dropout_global``.
+
+The central test is :func:`test_every_declared_sweep_parameter_changes_the_config`, which walks the
+real ``config_files/sweeps/hyperparameters.yaml`` and proves each declared parameter reaches the
+config for a real registered dataset/task pair.
+"""
+
+from pathlib import Path
+
+import pytest
+import yaml
+from yacs.config import CfgNode as CN
+
+from interscale.config import load_config
+from interscale.config.registry import resolve_config
+from interscale.config.sweep import (
+    SWEEP_GOALS,
+    apply_sweep_config,
+    build_sweep_config,
+    unused_component_prefixes,
+)
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+CONFIG_DIR = REPO_ROOT / "config_files"
+REGISTRY = CONFIG_DIR / "registry.yaml"
+SWEEP_YAML = CONFIG_DIR / "sweeps" / "hyperparameters.yaml"
+
+# A pair that exercises both components, so no parameter is dropped as unused.
+REFERENCE_DATASET = "melton25"
+REFERENCE_TASK = "graph_clas"
+
+
+def get_dotted(cfg, key):
+    """Read a dotted path out of a config, raising KeyError if any segment is missing."""
+    node = cfg
+    for part in key.split("."):
+        node = node[part]
+    return node
+
+
+@pytest.fixture
+def sweep_yaml():
+    with SWEEP_YAML.open() as f:
+        return yaml.safe_load(f)
+
+
+@pytest.fixture
+def base_cfg():
+    return resolve_config(REFERENCE_DATASET, REFERENCE_TASK, registry_path=REGISTRY)
+
+
+def pick_differing_value(declared_values, current):
+    """Pick a declared sweep value that differs from the config's current value.
+
+    Using the sweep's own declared values keeps the test honest: it proves the real candidate
+    values land in the config, not just that some arbitrary sentinel can be written.
+    """
+    for value in declared_values:
+        if value != current:
+            return value
+    return None
+
+
+def test_sweep_yaml_exists():
+    assert SWEEP_YAML.is_file(), f"expected the sweep config at {SWEEP_YAML}"
+
+
+def test_every_declared_sweep_parameter_changes_the_config(sweep_yaml, base_cfg):
+    """Every parameter the sweep declares must actually reach the config.
+
+    This is the test that would have caught both historical silent-no-op bugs.
+    """
+    sweep_config, sweep_params = build_sweep_config(
+        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
+    )
+    assert sweep_params, "the sweep declares no parameters"
+
+    unverifiable = []
+
+    for key in sweep_params:
+        declared = sweep_config["parameters"][key]
+        values = declared.get("values", [declared.get("value")])
+        current = get_dotted(base_cfg, key)
+
+        target = pick_differing_value(values, current)
+        if target is None:
+            # Every declared value already equals the base config's value, so applying it
+            # cannot be observed. Reported rather than silently passed.
+            unverifiable.append(key)
+            continue
+
+        cfg = base_cfg.clone()
+        cfg, applied = apply_sweep_config(
+            cfg, "hyperparmeter", {key: target}, model_type="CombinedModel", sweep_params=[key]
+        )
+
+        assert key in applied, f"{key} was declared by the sweep but not applied"
+        assert get_dotted(cfg, key) == target, (
+            f"sweep parameter {key} did not change the config: "
+            f"expected {target!r}, config still holds {get_dotted(cfg, key)!r}"
+        )
+
+    assert not unverifiable, (
+        "these sweep parameters declare only values identical to the base config, so a trial "
+        f"varying them is indistinguishable from no trial at all: {unverifiable}"
+    )
+
+
+def test_all_sweep_parameters_applied_together(sweep_yaml, base_cfg):
+    """Applying a full sampled trial writes every parameter, not just the first."""
+    sweep_config, sweep_params = build_sweep_config(
+        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
+    )
+
+    trial = {}
+    for key in sweep_params:
+        declared = sweep_config["parameters"][key]
+        values = declared.get("values", [declared.get("value")])
+        target = pick_differing_value(values, get_dotted(base_cfg, key))
+        trial[key] = target if target is not None else values[0]
+
+    cfg, applied = apply_sweep_config(
+        base_cfg.clone(), "hyperparmeter", trial, model_type="CombinedModel", sweep_params=sweep_params
+    )
+
+    assert sorted(applied) == sorted(sweep_params)
+    for key, expected in trial.items():
+        assert get_dotted(cfg, key) == expected, f"{key} not applied"
+
+
+def test_sweep_parameters_are_all_known_to_the_config(sweep_yaml, base_cfg):
+    """No declared parameter names a config path that does not exist."""
+    _, sweep_params = build_sweep_config(
+        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
+    )
+    for key in sweep_params:
+        # Raises KeyError with the offending key if the path is absent.
+        get_dotted(base_cfg, key)
+
+
+def test_trial_does_not_leak_into_the_base_config(base_cfg):
+    """Applying a trial must not mutate a shared config.
+
+    wandb.agent runs many trials in one process. If they shared a config object, trial N would
+    inherit every earlier trial's values and the sweep results would be meaningless.
+    """
+    original = base_cfg.optim.lr
+    other = original + 0.5
+
+    applied_cfg, _ = apply_sweep_config(
+        base_cfg.clone(), "hyperparmeter", {"optim.lr": other}, sweep_params=["optim.lr"]
+    )
+
+    assert applied_cfg.optim.lr == other
+    assert base_cfg.optim.lr == original, "the trial leaked into the shared base config"
+
+
+def test_base_config_stays_frozen_after_apply(base_cfg):
+    """A frozen config comes back frozen, so later accidental writes still raise."""
+    cfg = base_cfg.clone()
+    assert cfg.is_frozen()
+    cfg, _ = apply_sweep_config(cfg, "hyperparmeter", {"optim.lr": 0.123}, sweep_params=["optim.lr"])
+    assert cfg.is_frozen()
+
+
+# --- the two historical silent-no-op bugs, as regression tests ----------------------------
+
+
+def test_component_key_without_model_prefix_raises(base_cfg):
+    """`local_component.parameters.num_layers` (no leading `model.`) must not pass silently."""
+    with pytest.raises(KeyError, match="local_component.parameters.num_layers"):
+        apply_sweep_config(
+            base_cfg.clone(),
+            "hyperparmeter",
+            {"local_component.parameters.num_layers": 3},
+            sweep_params=["local_component.parameters.num_layers"],
+        )
+
+
+def test_misspelled_dropout_key_raises(base_cfg):
+    """The transformer key is `dropout_global`; plain `dropout` used to become a dead key."""
+    with pytest.raises(KeyError, match="dropout"):
+        apply_sweep_config(
+            base_cfg.clone(),
+            "hyperparmeter",
+            {"model.global_component.parameters.dropout": 0.5},
+            sweep_params=["model.global_component.parameters.dropout"],
+        )
+
+    # ...while the correctly spelled key does land.
+    cfg, _ = apply_sweep_config(
+        base_cfg.clone(),
+        "hyperparmeter",
+        {"model.global_component.parameters.dropout_global": 0.5},
+        sweep_params=["model.global_component.parameters.dropout_global"],
+    )
+    assert cfg.model.global_component.parameters.dropout_global == 0.5
+
+
+def test_unknown_key_is_not_silently_created(base_cfg):
+    """A typo must raise rather than create a config entry nothing reads."""
+    with pytest.raises(KeyError):
+        apply_sweep_config(
+            base_cfg.clone(),
+            "hyperparmeter",
+            {"optim.learning_rate": 0.1},
+            sweep_params=["optim.learning_rate"],
+        )
+
+
+def test_declared_but_unsampled_parameter_raises(base_cfg):
+    """A parameter the sweep declares but the trial never sampled is an error, not a warning."""
+    with pytest.raises(KeyError, match="did not sample"):
+        apply_sweep_config(
+            base_cfg.clone(),
+            "hyperparmeter",
+            {"optim.lr": 0.01},
+            sweep_params=["optim.lr", "optim.wd"],
+        )
+
+
+def test_whole_config_sections_are_not_clobbered(base_cfg):
+    """wandb.config carries the whole base config; only declared parameters may be written.
+
+    main_sweep.py calls wandb.init(config=cfg), so wandb.config contains top-level `dataset` /
+    `model` / `optim` entries next to the sampled dotted keys. Writing those back would replace
+    CfgNodes with plain dicts.
+    """
+    wandb_like = {
+        "dataset": {"batch_size": 999},  # the section dump, not a sweep parameter
+        "model": {"n_embed": 999},
+        "optim.lr": 0.007,  # the actual sampled parameter
+    }
+
+    cfg, applied = apply_sweep_config(
+        base_cfg.clone(), "hyperparmeter", wandb_like, sweep_params=["optim.lr"]
+    )
+
+    assert applied == ["optim.lr"]
+    assert cfg.optim.lr == 0.007
+    assert isinstance(cfg.dataset, CN), "the dataset section was replaced by a plain dict"
+    assert cfg.dataset.batch_size == base_cfg.dataset.batch_size
+    assert cfg.model.n_embed == base_cfg.model.n_embed
+
+
+def test_inferred_params_ignore_non_dotted_keys(base_cfg):
+    """With no explicit sweep_params, only dotted keys are treated as sweep parameters."""
+    cfg, applied = apply_sweep_config(
+        base_cfg.clone(),
+        "hyperparmeter",
+        {"dataset": {"batch_size": 999}, "optim.lr": 0.003},
+    )
+    assert applied == ["optim.lr"]
+    assert isinstance(cfg.dataset, CN)
+
+
+# --- goal and model_type handling -----------------------------------------------------------
+
+
+def test_unknown_sweep_goal_raises(base_cfg):
+    """A misspelled goal must fail loudly, not train the base config on every trial."""
+    with pytest.raises(ValueError, match="Unknown sweep_goal"):
+        apply_sweep_config(base_cfg.clone(), "hyperparameter", {"optim.lr": 0.1})  # sic: correct spelling
+
+
+@pytest.mark.parametrize("goal", SWEEP_GOALS)
+def test_every_declared_goal_is_accepted(goal, base_cfg):
+    cfg, applied = apply_sweep_config(base_cfg.clone(), goal, {"optim.lr": 0.02}, sweep_params=["optim.lr"])
+    assert applied == ["optim.lr"]
+    assert cfg.optim.lr == 0.02
+
+
+def test_robustness_goal_parameters_apply(base_cfg):
+    """The robustness sweep's three keys all exist and all land."""
+    trial = {
+        "dataset.pct_mask_nodes": 0.42,
+        "dataset.spatial_neigbors_kwargs.radius": 77,
+        "optim.seed": 7,
+    }
+    cfg, applied = apply_sweep_config(
+        base_cfg.clone(), "robustness", trial, sweep_params=sorted(trial)
+    )
+    assert sorted(applied) == sorted(trial)
+    assert cfg.dataset.pct_mask_nodes == 0.42
+    assert cfg.dataset.spatial_neigbors_kwargs.radius == 77
+    assert cfg.optim.seed == 7
+
+
+def test_segmentation_goal_parameters_apply(base_cfg):
+    cfg, _ = apply_sweep_config(
+        base_cfg.clone(),
+        "segmentation",
+        {"dataset.segmentation_robustness": [0.1, 0.2]},
+        sweep_params=["dataset.segmentation_robustness"],
+    )
+    assert cfg.dataset.segmentation_robustness == [0.1, 0.2]
+
+
+def test_int_is_promoted_for_a_float_key(base_cfg):
+    """`values: [0, 0.1, 0.3]` samples a real int for 0; yacs rejects int for a float key."""
+    cfg, _ = apply_sweep_config(
+        base_cfg.clone(),
+        "hyperparmeter",
+        {"model.local_component.parameters.dropout_local": 0},
+        sweep_params=["model.local_component.parameters.dropout_local"],
+    )
+    value = cfg.model.local_component.parameters.dropout_local
+    assert value == 0
+    assert isinstance(value, float), "an int would break a later merge against this float key"
+
+
+@pytest.mark.parametrize(
+    ("model_type", "expected_prefixes"),
+    [
+        ("CombinedModel", []),
+        ("LocalModel", ["model.global_component."]),
+        ("GlobalModel", ["model.local_component."]),
+    ],
+)
+def test_unused_component_prefixes(model_type, expected_prefixes):
+    """CombinedModel must keep BOTH components -- an if/elif chain used to drop its transformer."""
+    assert unused_component_prefixes(model_type) == expected_prefixes
+
+
+@pytest.mark.parametrize(
+    ("model_type", "dropped_prefix", "kept_prefix"),
+    [
+        ("LocalModel", "model.global_component.", "model.local_component."),
+        ("GlobalModel", "model.local_component.", "model.global_component."),
+    ],
+)
+def test_single_component_model_drops_the_other_components_parameters(
+    sweep_yaml, model_type, dropped_prefix, kept_prefix
+):
+    _, sweep_params = build_sweep_config(
+        sweep_yaml, prediction_task="classification", model_type=model_type
+    )
+    assert not any(k.startswith(dropped_prefix) for k in sweep_params)
+    assert any(k.startswith(kept_prefix) for k in sweep_params), (
+        f"{model_type} should still sweep its own component's parameters"
+    )
+
+
+def test_combined_model_sweeps_both_components(sweep_yaml):
+    _, sweep_params = build_sweep_config(
+        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
+    )
+    assert any(k.startswith("model.local_component.") for k in sweep_params)
+    assert any(k.startswith("model.global_component.") for k in sweep_params)
+
+
+def test_component_parameters_are_skipped_not_applied_for_wrong_model_type(base_cfg):
+    """A global parameter reaching a LocalModel trial is skipped rather than written."""
+    cfg, applied = apply_sweep_config(
+        base_cfg.clone(),
+        "hyperparmeter",
+        {"model.global_component.parameters.n_heads": 8, "optim.lr": 0.02},
+        model_type="LocalModel",
+        sweep_params=["model.global_component.parameters.n_heads", "optim.lr"],
+    )
+    assert applied == ["optim.lr"]
+    assert cfg.model.global_component.parameters.n_heads == base_cfg.model.global_component.parameters.n_heads
+
+
+# --- metric selection ------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("prediction_task", "expected_metric"),
+    [("classification", "val_f1_macro"), ("regression", "val_r2")],
+)
+def test_metric_follows_prediction_task(sweep_yaml, prediction_task, expected_metric):
+    """A regression sweep must not be ranked by an f1 metric it never logs."""
+    sweep_config, _ = build_sweep_config(sweep_yaml, prediction_task=prediction_task)
+    assert sweep_config["metric"]["name"] == expected_metric
+    assert sweep_config["metric"]["goal"] == "maximize"
+
+
+def test_yaml_metric_is_kept_when_no_prediction_task_given(sweep_yaml):
+    sweep_config, _ = build_sweep_config(sweep_yaml)
+    assert sweep_config["metric"]["name"] == sweep_yaml["sweep_config"]["metric"]["name"]
+
+
+def test_missing_metric_without_prediction_task_raises():
+    with pytest.raises(ValueError, match="no `metric`"):
+        build_sweep_config({"sweep_config": {"method": "random", "parameters": {"optim.lr": {"values": [1]}}}})
+
+
+def test_unknown_prediction_task_raises(sweep_yaml):
+    with pytest.raises(ValueError, match="unknown prediction_task"):
+        build_sweep_config(sweep_yaml, prediction_task="clustering")
+
+
+def test_missing_sweep_config_block_raises():
+    with pytest.raises(ValueError, match="sweep_config"):
+        build_sweep_config({"parameters": {}})
+
+
+def test_empty_parameters_raises():
+    with pytest.raises(ValueError, match="no `parameters`"):
+        build_sweep_config({"sweep_config": {"metric": {"name": "val_loss", "goal": "minimize"}, "parameters": {}}})
+
+
+def test_build_sweep_config_does_not_mutate_the_parsed_yaml(sweep_yaml):
+    """Dropping component parameters must not corrupt a reusable parsed yaml."""
+    before = sorted(sweep_yaml["sweep_config"]["parameters"])
+    build_sweep_config(sweep_yaml, prediction_task="classification", model_type="LocalModel")
+    assert sorted(sweep_yaml["sweep_config"]["parameters"]) == before
+
+
+def test_sweep_parameters_apply_to_a_bare_default_config():
+    """The sweep's keys exist in the plain defaults too, not only in a dataset config."""
+    cfg = load_config(CONFIG_DIR / "base.yaml")
+    with SWEEP_YAML.open() as f:
+        sweep_config, sweep_params = build_sweep_config(
+            yaml.safe_load(f), prediction_task="classification", model_type="CombinedModel"
+        )
+    trial = {k: sweep_config["parameters"][k].get("values", [None])[0] for k in sweep_params}
+    _, applied = apply_sweep_config(cfg, "hyperparmeter", trial, model_type="CombinedModel", sweep_params=sweep_params)
+    assert sorted(applied) == sorted(sweep_params)

From e6edef6c5a667d7f28163478c94ad189e89c1a86 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Mon, 24 Aug 2026 12:29:59 +0200
Subject: [PATCH 06/14] enable plotting tools

---
 src/interscale/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/interscale/__init__.py b/src/interscale/__init__.py
index ba2b199..76dab01 100644
--- a/src/interscale/__init__.py
+++ b/src/interscale/__init__.py
@@ -1,8 +1,8 @@
 from importlib.metadata import PackageNotFoundError, version
 
-from . import config, datasets, evaluation, model, module, tl
+from . import config, datasets, evaluation, model, module, tl, pl
 
-__all__ = ["config", "datasets", "evaluation", "module", "tl", "model"]
+__all__ = ["config", "datasets", "evaluation", "module", "tl", "model", "pl"]
 
 try:
     __version__ = version("interscale")

From d3bb332110f6845123b19b2f31fbcfb48ad1c92d Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Fri, 28 Aug 2026 11:58:19 +0200
Subject: [PATCH 07/14] latent analysis - which driven, stores dim_importance
 in adata

---
 src/interscale/config/sweep.py                | 165 +++++++-
 .../evaluation/_gene_rank_analysis.py         |  13 +
 src/interscale/evaluation/_latent_analysis.py | 387 ++++++++++++++----
 src/interscale/main_sweep.py                  | 160 +++++++-
 .../module/base/_base_global_module.py        |  17 +-
 src/interscale/train/_training.py             |  40 +-
 src/interscale/train/_trainingplans.py        |  52 ++-
 tests/test_sweep_config.py                    | 238 +++++++++++
 8 files changed, 956 insertions(+), 116 deletions(-)

diff --git a/src/interscale/config/sweep.py b/src/interscale/config/sweep.py
index 7c44cb1..91ec7f3 100644
--- a/src/interscale/config/sweep.py
+++ b/src/interscale/config/sweep.py
@@ -14,6 +14,11 @@
 ``dropout_global``, which ``set_new_allowed(True)`` turned into a dead key that nothing read.
 Both classes are structurally impossible here: there is no allow-new-keys, and a key nothing
 can apply is an error rather than a shrug.
+
+**Arms** (see :func:`load_arms`) add the one thing a flat dotted-key sweep cannot express: an
+ablation axis whose levels each imply *several* config values that must move together. wandb
+searches the cartesian product of its declared parameters, so declaring the coupled keys
+separately would produce every invalid crossing of them.
 """
 
 from yacs.config import CfgNode as CN
@@ -22,6 +27,11 @@
 # existing sweep invocations pass; renaming it would silently break saved commands.
 SWEEP_GOALS = ("robustness", "segmentation", "hyperparmeter", "loss")
 
+# The one sweep parameter whose values are arm NAMES rather than config values. Reserved: a
+# config key called "arm" could never be a sweep parameter anyway, since every real one is
+# dotted.
+ARM_PARAM = "arm"
+
 # Metric each prediction task ranks trials by. val_f1_macro is what EarlyStopping and
 # ModelCheckpoint monitor for classification, and trainer.validate() runs after the best
 # checkpoint is restored, so the last logged value is the best epoch's score.
@@ -48,7 +58,7 @@ def unused_component_prefixes(model_type):
     return prefixes
 
 
-def build_sweep_config(yaml_config, prediction_task=None, model_type=None):
+def build_sweep_config(yaml_config, prediction_task=None, model_type=None, metric=None):
     """Turn a parsed sweep yaml into the dict ``wandb.sweep`` takes.
 
     Parameters
@@ -80,6 +90,13 @@ def build_sweep_config(yaml_config, prediction_task=None, model_type=None):
             raise ValueError(f"unknown prediction_task '{prediction_task}'. Must be one of: {', '.join(TASK_METRICS)}")
         sweep_config["metric"] = dict(TASK_METRICS[prediction_task])
 
+    # An explicit metric wins over the task default. The task default answers "what does this
+    # task log?", which is not always the same question as "what is this sweep searching for":
+    # a regression sweep looking for the best *correlation* must rank on val_pearson_corr, and
+    # ranking it on the task default val_r2 would select for calibration instead.
+    if metric is not None:
+        sweep_config["metric"] = dict(metric)
+
     if "metric" not in sweep_config:
         raise ValueError(
             "Sweep config declares no `metric`, and no prediction_task was given to supply one. "
@@ -100,9 +117,117 @@ def build_sweep_config(yaml_config, prediction_task=None, model_type=None):
                 f"every sweep parameter was dropped as unused by {model_type}; nothing left to vary."
             )
 
+    if ARM_PARAM in sweep_config["parameters"] and "arms" not in yaml_config:
+        raise ValueError(
+            f"sweep declares the reserved parameter '{ARM_PARAM}' but the yaml has no top-level "
+            f"`arms:` block to resolve its values against. Every trial would try to write the arm "
+            f"name to a config key called '{ARM_PARAM}', which does not exist."
+        )
+
     return sweep_config, sorted(sweep_config["parameters"])
 
 
+def load_arms(yaml_config, sweep_config=None):
+    """Read and validate a sweep yaml's optional top-level ``arms:`` block.
+
+    An *arm* is one level of an ablation axis that implies several config values at once. The
+    sliding-window ablation is the motivating case: a window size fixes the ``obs`` column the
+    windows live in, the transformer's ``max_seq_len`` (attention pads to it, so it cannot be
+    shared across arms without paying the largest arm's quadratic cost on every one) and
+    ``dataset.name`` (it reaches the checkpoint filename, so without it every arm overwrites its
+    predecessor). Those three must move together; wandb searches the cartesian product of what it
+    is given, so declaring them as three parameters would ask for 6^3 = 216 trials of which 6 are
+    meaningful.
+
+    So only the arm *name* is declared to wandb, as the reserved ``arm`` parameter, and the values
+    it stands for are looked up here. That the coupled values never round-trip through wandb is a
+    second benefit: ``dataset.sample_key`` is a list, and yacs type-checks assignments, so a value
+    wandb chose to return as a tuple or a scalar would fail mid-sweep.
+
+    Parameters
+    ----------
+    yaml_config : dict
+        The parsed sweep yaml. ``arms`` is a sibling of ``sweep_config``, not nested inside it,
+        because everything under ``sweep_config`` is sent verbatim to ``wandb.sweep``.
+    sweep_config : dict, optional
+        The built sweep config, checked against the arms so a mismatch fails before wandb
+        registers anything.
+
+    Returns
+    -------
+    dict or None
+        ``{arm_name: {dotted_key: value}}``, or None when the yaml declares no arms.
+    """
+    if "arms" not in yaml_config:
+        return None
+
+    arms = yaml_config["arms"]
+    if not isinstance(arms, dict) or not arms:
+        raise ValueError("sweep yaml's `arms:` block must be a non-empty mapping of arm name -> overrides.")
+
+    for name, overrides in arms.items():
+        if not isinstance(overrides, dict) or not overrides:
+            raise ValueError(f"arm '{name}' declares no overrides; it would be identical to every other arm.")
+        undotted = sorted(k for k in overrides if "." not in str(k))
+        if undotted:
+            raise ValueError(
+                f"arm '{name}' declares non-dotted keys {undotted}. Arm overrides are dotted config "
+                f"paths (e.g. dataset.sample_key), applied exactly as sweep parameters are."
+            )
+
+    # Every arm must set the SAME keys. An arm that omits one silently leaves it at the base
+    # config's value while its siblings override it, so the ablation would compare arms that
+    # differ in a way the yaml never states -- the same silent-no-op class this module exists to
+    # make impossible.
+    key_sets = {name: frozenset(overrides) for name, overrides in arms.items()}
+    reference_name, reference_keys = next(iter(key_sets.items()))
+    for name, keys in key_sets.items():
+        if keys != reference_keys:
+            missing = sorted(reference_keys - keys)
+            extra = sorted(keys - reference_keys)
+            raise ValueError(
+                f"arm '{name}' does not declare the same keys as arm '{reference_name}': "
+                f"missing {missing}, extra {extra}. Any key not set by every arm stays at the base "
+                f"config value for the arms that omit it, so the arms would differ in an unstated way."
+            )
+
+    if sweep_config is not None:
+        parameters = sweep_config.get("parameters", {})
+        if ARM_PARAM not in parameters:
+            raise ValueError(
+                f"sweep yaml declares an `arms:` block but no '{ARM_PARAM}' parameter to select "
+                f"between them, so no arm would ever be applied."
+            )
+
+        declared = parameters[ARM_PARAM]
+        if not isinstance(declared, dict) or "values" not in declared:
+            raise ValueError(
+                f"the '{ARM_PARAM}' parameter must declare `values:` naming the arms to run "
+                f"(got {declared!r})."
+            )
+
+        selected = list(declared["values"])
+        unknown = sorted(set(selected) - set(arms))
+        if unknown:
+            raise ValueError(f"'{ARM_PARAM}' selects arms that the `arms:` block does not define: {unknown}")
+        # Not an error -- running a subset of the defined arms is a legitimate way to re-run one
+        # of them -- but silence here is how you discover a typo only after the sweep finishes.
+        unused = sorted(set(arms) - set(selected))
+        if unused:
+            print(f"arms defined but not selected by '{ARM_PARAM}': {unused}")
+
+        # An arm key that is ALSO a declared sweep parameter would be written twice per trial with
+        # the order deciding which wins.
+        collisions = sorted(reference_keys & set(parameters))
+        if collisions:
+            raise ValueError(
+                f"keys are set both by the arms and as sweep parameters, so which value a trial "
+                f"gets would depend on application order: {collisions}"
+            )
+
+    return arms
+
+
 def _set_dotted(cfg, key, value):
     """Assign ``value`` at the dotted ``key`` in ``cfg``, requiring the path to exist."""
     parts = str(key).split(".")
@@ -133,7 +258,7 @@ def _set_dotted(cfg, key, value):
     node[leaf] = value
 
 
-def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_params=None):
+def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_params=None, arms=None):
     """Write a sampled sweep trial onto ``cfg``.
 
     Parameters
@@ -152,6 +277,9 @@ def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_par
         ``parameters`` block. This is the authoritative list of what to apply, and any
         declared name missing from ``sweep_config`` raises. When omitted it is inferred as
         the dotted keys of ``sweep_config``.
+    arms : dict, optional
+        The yaml's ``arms:`` block, from :func:`load_arms`. When given, the trial's ``arm`` value
+        is expanded into that arm's dotted overrides instead of being written to the config.
 
     Returns
     -------
@@ -179,6 +307,28 @@ def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_par
                 f"applied: {absent}"
             )
 
+    # The arm name is not a config path, so it is removed from `keys` and replaced by the dotted
+    # overrides it stands for. Done before the loop below so an arm override and a plain sweep
+    # parameter go through exactly the same _set_dotted validation.
+    arm_name = None
+    arm_overrides = {}
+    if arms is not None:
+        if ARM_PARAM not in sweep_config:
+            raise KeyError(
+                f"arms were given but the trial carries no '{ARM_PARAM}' value, so no arm can be "
+                f"applied and every trial would train the base config's windowing."
+            )
+        arm_name = sweep_config[ARM_PARAM]
+        if arm_name not in arms:
+            raise KeyError(f"trial selected arm '{arm_name}', which the `arms:` block does not define.")
+        arm_overrides = dict(arms[arm_name])
+        keys = [k for k in keys if k != ARM_PARAM]
+    elif ARM_PARAM in (sweep_params or []):
+        raise KeyError(
+            f"the sweep declares the reserved parameter '{ARM_PARAM}' but no arms were passed to "
+            f"apply_sweep_config, so it would be written to a config key that does not exist."
+        )
+
     # Deliberately NOT set_new_allowed(True): allowing new keys is what let a misspelled
     # parameter create a dead config entry that nothing ever read.
     was_frozen = cfg.is_frozen()
@@ -188,11 +338,13 @@ def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_par
 
     applied = []
     skipped = []
-    for key in keys:
+    # Arm overrides go first so a later plain parameter, if the two ever collided, would win --
+    # but load_arms rejects that collision outright, so the order is only about determinism.
+    for key, value in list(arm_overrides.items()) + [(k, sweep_config[k]) for k in keys]:
         if any(str(key).startswith(p) for p in skipped_prefixes):
             skipped.append(key)
             continue
-        _set_dotted(cfg, key, sweep_config[key])
+        _set_dotted(cfg, key, value)
         applied.append(key)
 
     if was_frozen:
@@ -204,12 +356,15 @@ def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_par
     if sweep_params is not None:
         # A parameter wandb varies but nothing reads makes every trial's difference invisible,
         # which is indistinguishable from the sweep not working. Fail rather than warn.
-        unapplied = sorted(set(sweep_params) - set(applied) - set(skipped))
+        # ARM_PARAM is excluded because it is applied as its expansion, not as itself.
+        unapplied = sorted(set(sweep_params) - set(applied) - set(skipped) - {ARM_PARAM})
         if unapplied:
             raise KeyError(
                 f"sweep declares parameters that were not applied to the config, so they would "
                 f"vary between trials with no effect: {unapplied}"
             )
 
+    if arm_name is not None:
+        print(f"applied arm '{arm_name}': {sorted(arm_overrides)}")
     print(f"applied sweep parameters: {sorted(applied)}")
     return cfg, sorted(applied)
diff --git a/src/interscale/evaluation/_gene_rank_analysis.py b/src/interscale/evaluation/_gene_rank_analysis.py
index 84ec023..be36242 100644
--- a/src/interscale/evaluation/_gene_rank_analysis.py
+++ b/src/interscale/evaluation/_gene_rank_analysis.py
@@ -22,6 +22,19 @@ def _predict_gene_r2(adata: AnnData, layers_pred: str) -> pd.DataFrame:
     # Convert y_true to a dense array
     y_true = adata.X.toarray().astype(float)
 
+    # ensure y_true is on same scale as during model training (e.g. log1p normalized) and not raw counts
+    finite_true = y_true[np.isfinite(y_true)]
+    if finite_true.size:
+        max_true = float(np.max(finite_true))
+        looks_like_counts = np.allclose(finite_true, np.round(finite_true)) and max_true > 1
+        assert max_true <= 20 and not looks_like_counts, (
+            f"adata.X looks like raw or normalized counts (max={max_true:.3g}"
+            f"{', integer-valued' if looks_like_counts else ''}), but the predictions in "
+            f"adata.layers['{layers_pred}'] are on the scale of the layer the model was trained "
+            f"on. Set adata.X to that layer (e.g. adata.X = adata.layers['log1p_norm']) before "
+            f"calling calculate_gene_ranks."
+        )
+
     # Convert predictions to NumPy arrays
     y_pred = adata.layers[layers_pred]
 
diff --git a/src/interscale/evaluation/_latent_analysis.py b/src/interscale/evaluation/_latent_analysis.py
index 51513ee..9171674 100644
--- a/src/interscale/evaluation/_latent_analysis.py
+++ b/src/interscale/evaluation/_latent_analysis.py
@@ -1,6 +1,90 @@
+import warnings
+from collections.abc import Sequence
+from typing import Literal
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import scipy.sparse as sp
+
+
+def _gene_expression_stats(X, ddof=1):
+    """Per-gene detection fraction and standard deviation, without densifying sparse input.
+
+    Returns
+    -------
+    frac : (G,) fraction of cells with X > 0
+    sd : (G,) standard deviation with the given ddof
+    """
+    n = X.shape[0]
+
+    if sp.issparse(X):
+        if X.format not in ("csr", "csc"):
+            X = X.tocsr()
+        frac = np.asarray((X > 0).sum(axis=0), dtype=float).ravel() / n
+        mean = np.asarray(X.mean(axis=0), dtype=float).ravel()
+        mean_sq = np.asarray(X.multiply(X).mean(axis=0), dtype=float).ravel()
+        var = np.maximum(mean_sq - mean**2, 0.0)
+        if ddof and n > ddof:
+            var = var * (n / (n - ddof))
+        sd = np.sqrt(var)
+    else:
+        Xd = np.asarray(X, dtype=float)
+        frac = (Xd > 0).mean(axis=0)
+        sd = Xd.std(axis=0, ddof=ddof)
+
+    return frac, sd
+
+
+def _infer_which(s_key):
+    """Infer the component ('local'/'global') from a loading key, or None if ambiguous."""
+    if s_key.startswith("_local"):
+        return "local"
+    if s_key.startswith("_global"):
+        return "global"
+    return None
+
+
+def _dim_importance_uns_key(which):
+    """Key under which `calculate_dim_importance` stores its selection in adata.uns."""
+    return f"_{which}_dim_importance"
+
+
+def _resolve_dims_from_uns(adata, uns_key, s_key):
+    """Read the dimension selection stored by `calculate_dim_importance`.
+
+    Raises with an actionable message if the record is missing, was computed without a
+    cumulative cutoff, or came from a different component than `s_key`.
+    """
+    if uns_key not in adata.uns:
+        raise KeyError(
+            f"dims=None requires a stored dimension selection, but adata.uns['{uns_key}'] "
+            f"is missing. Run calculate_dim_importance(adata, s_key='{s_key}', ...) first, "
+            f"or pass dims=[...] explicitly."
+        )
+
+    rec = adata.uns[uns_key]
+
+    if "dims_left" not in rec:
+        raise KeyError(
+            f"adata.uns['{uns_key}'] has no 'dims_left': calculate_dim_importance was run "
+            "with use_ratio=False or cumulative_cutoff=None, so no dimensions were selected. "
+            "Re-run it with use_ratio=True and a cumulative_cutoff, or pass dims=[...] explicitly."
+        )
+
+    # Guard against the mismatch this whole mechanism exists to prevent: a record built
+    # from one component being used to pick genes from the other.
+    rec_s_key = rec.get("s_key")
+    if rec_s_key is not None and str(rec_s_key) != s_key:
+        raise ValueError(
+            f"Component mismatch: adata.uns['{uns_key}'] was computed from "
+            f"s_key='{rec_s_key}', but this call resolved s_key='{s_key}'. "
+            "Check the 'which' argument, or re-run calculate_dim_importance for this component."
+        )
+
+    # Coming from storage (possibly an h5ad round-trip), so coerce rather than validate
+    # strictly -- user-supplied dims keep the strict dtype check below.
+    return np.asarray(rec["dims_left"]).ravel().astype(int)
 
 
 def _get_Z(adata, z_key):
@@ -66,12 +150,13 @@ def latent_rank_report(
 
 def get_genes_dim(
     adata,
-    dims,
+    which: Literal["global", "local"],  # required
     *,
-    which="global",  # "global" or "local"
-    n_top=20,
-    s_key=None,  # e.g. "_global_std_gene_loadings"
-    z_key=None,  # e.g. "_global_emb" (only used if residualize=True)
+    dims: Sequence[int] | None = None,  # None -> read the selection stored by calculate_dim_importance
+    n_top: int = 20,
+    s_key: str | None = None,  # e.g. "_global_std_gene_loadings"
+    z_key: str | None = None,  # e.g. "_global_emb" (only used if residualize=True)
+    uns_key: str | None = None,  # e.g. "_global_dim_importance"; None -> f"_{which}_dim_importance"
     # expression-based filtering
     X_layer=None,  # e.g. "log1p_norm"; None -> adata.X
     min_frac=0.05,  # fraction of cells with expr>0
@@ -100,39 +185,100 @@ def get_genes_dim(
 
     If plot=True, also plot a gene×dim heatmap of the returned DataFrame.
 
+    Parameters
+    ----------
+    which : {"global", "local"}
+        Which component to analyse. Required, because it selects the defaults for
+        `s_key`, `z_key` and `uns_key` — it is the only argument needed to switch
+        components.
+    dims : sequence of int, optional
+        Latent dimensions to analyse. **Leave as None (recommended)** to read the
+        selection `calculate_dim_importance` stored in adata.uns[uns_key], which keeps
+        `dims` and `which` from drifting apart. If given explicitly and a stored record
+        exists, the two are compared and a mismatch raises a warning.
+    uns_key : str, optional
+        Record written by `calculate_dim_importance`. Defaults to
+        f"_{which}_dim_importance". A record built from a different component than the
+        resolved `s_key` raises ValueError rather than silently mixing the two.
+
+    Examples
+    --------
+    >>> calculate_dim_importance(adata, "global", cumulative_cutoff=0.60)
+    >>> df = get_genes_dim(adata, "global", n_top=20)  # dims resolved from adata.uns
+
     Notes
     -----
     - Uses standardized loadings (S_std) stored in adata.varm[s_key].
     - rank_by="loading_x_sd" multiplies scores by sd(X_g) from X_layer (keeps sign).
+      Note this only cancels the per-gene standardization if X_layer is the same layer
+      that was passed to `gene_loadings()`.
+    - Genes whose score is non-finite in any requested dim are dropped before ranking.
     - enforce_specificity can help pick genes that are strong in one dim and weak in others.
+      It compares each gene's largest against its second-largest |score| across the
+      *requested* dims only, and is ignored (with a warning) for a single dim.
+    - min_frac uses (X > 0), so it is only meaningful for a non-negative layer; on a
+      centered/scaled layer roughly half of every gene's values exceed 0.
+    - Ties are broken by gene order (stable sort), so the output is reproducible.
     """
-    # --- defaults
-    if s_key is None:
-        if which == "global":
-            s_key = "_global_std_gene_loadings"
-        elif which == "local":
-            s_key = "_local_std_gene_loadings"
-        else:
-            raise ValueError("which must be 'global' or 'local' (or provide s_key explicitly)")
+    # --- defaults. `which` also names the adata.uns record, so validate it unconditionally.
+    if which not in ("local", "global"):
+        raise ValueError(f"which must be 'global' or 'local', got {which!r}")
 
+    if s_key is None:
+        s_key = f"_{which}_std_gene_loadings"
     if z_key is None:
-        if which == "global":
-            z_key = "_global_emb"
-        elif which == "local":
-            z_key = "_local_emb"
+        z_key = f"_{which}_emb"
+    if uns_key is None:
+        uns_key = _dim_importance_uns_key(which)
 
     # --- load/check S
     if s_key not in adata.varm:
         raise KeyError(f"{s_key} not found in adata.varm")
     S = np.asarray(adata.varm[s_key], dtype=float)  # (G, K)
+    if S.ndim != 2:
+        raise ValueError(f"adata.varm['{s_key}'] must be 2D (genes x dims), got shape {S.shape}")
     genes = np.asarray(adata.var_names)
-    G, K = S.shape
-
-    dims = list(dims)
-    if len(dims) == 0:
+    K = S.shape[1]
+
+    # --- resolve dims: from adata.uns by default, else cross-check what was passed
+    if dims is None:
+        dims = _resolve_dims_from_uns(adata, uns_key, s_key)
+    elif uns_key in adata.uns and "dims_left" in adata.uns[uns_key]:
+        stored = np.asarray(adata.uns[uns_key]["dims_left"]).ravel().astype(int)
+        given = np.asarray(dims).ravel()
+        # compare as sets: dims_left is importance-ordered, so a reordering is not a mismatch
+        if given.dtype != bool and set(given.tolist()) != set(stored.tolist()):
+            warnings.warn(
+                f"dims={given.tolist()} differs from the selection stored in "
+                f"adata.uns['{uns_key}']['dims_left']={stored.tolist()}. Using the dims you "
+                "passed. Omit dims to use the stored selection.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+    # --- validate dims (integer indices, in range, unique)
+    dims_arr = np.asarray(dims)
+    if dims_arr.dtype == bool:
+        raise ValueError("dims must be integer dimension indices, not a boolean mask")
+    if dims_arr.size and not np.issubdtype(dims_arr.dtype, np.integer):
+        raise ValueError(f"dims must be integer dimension indices, got dtype {dims_arr.dtype}")
+    dims_arr = dims_arr.ravel()
+    if dims_arr.size == 0:
         raise ValueError("dims must be non-empty")
-    if any((d < 0 or d >= K) for d in dims):
+    if np.any((dims_arr < 0) | (dims_arr >= K)):
         raise ValueError(f"dims must be within [0, {K - 1}]")
+    if np.unique(dims_arr).size != dims_arr.size:
+        raise ValueError("dims must not contain duplicates")
+    dims = [int(d) for d in dims_arr]
+
+    # --- validate remaining scalar arguments up front
+    if rank_by not in ("loading", "loading_x_sd"):
+        raise ValueError("rank_by must be 'loading' or 'loading_x_sd'")
+    if order_genes_by not in ("winner", "max_abs", None):
+        raise ValueError("order_genes_by must be 'winner', 'max_abs', or None")
+    n_top = int(n_top)
+    if n_top < 1:
+        raise ValueError(f"n_top must be >= 1, got {n_top}")
 
     score_full = S.copy()
 
@@ -169,52 +315,66 @@ def get_genes_dim(
 
         score_full = score_full * uniq[None, :]
 
-    # --- expression-based filters and optional sd weighting
-    X = adata.layers[X_layer] if X_layer is not None else adata.X
-    X = np.asarray(X.todense() if hasattr(X, "todense") else X, dtype=float)  # (N, G)
-
-    frac = (X > 0).mean(axis=0)
-    sd_g = X.std(axis=0, ddof=1)
+    # --- expression statistics (only if a filter or the sd weighting needs them)
+    need_frac = min_frac is not None and min_frac > 0
+    need_sd = (min_sd is not None) or (rank_by == "loading_x_sd")
 
-    mask = np.isfinite(sd_g) & (frac >= min_frac)
-    if min_sd is not None:
-        mask &= sd_g >= min_sd
-
-    valid = np.where(mask)[0]
-    if valid.size == 0:
-        raise ValueError("No genes passed filters. Relax min_frac/min_sd or check X_layer.")
+    sd_g = None
+    if need_frac or need_sd:
+        X = adata.layers[X_layer] if X_layer is not None else adata.X
+        frac, sd_g = _gene_expression_stats(X)  # sparse-aware, no densification
 
     # --- apply sd weighting to scores if requested (keeps sign)
     score_used = score_full
     if rank_by == "loading_x_sd":
         score_used = score_used * sd_g[:, None]
-    elif rank_by != "loading":
-        raise ValueError("rank_by must be 'loading' or 'loading_x_sd'")
+
+    # --- gene filters. Non-finite scores must be excluded: np.argsort sends NaN to the
+    # end of the array, so a descending sort would otherwise rank NaN genes first.
+    mask = np.isfinite(score_used[:, dims]).all(axis=1)
+    if sd_g is not None:
+        mask &= np.isfinite(sd_g)
+        if need_frac:
+            mask &= frac >= min_frac
+        if min_sd is not None:
+            mask &= sd_g >= min_sd
+
+    valid = np.where(mask)[0]
+    if valid.size == 0:
+        raise ValueError("No genes passed filters. Relax min_frac/min_sd or check X_layer.")
 
     # --- optional specificity filter across chosen dims
     valid2 = valid
-    if enforce_specificity and len(dims) >= 2:
-        A = np.abs(score_used[np.ix_(valid, dims)])  # (n_valid, n_dims)
-        maxv = A.max(axis=1)
-        second = np.partition(A, -2, axis=1)[:, -2]
-        eps = 1e-12
-
-        if specificity_mode == "ratio":
-            keep = (maxv / (second + eps)) >= specificity_min
-        elif specificity_mode == "diff":
-            keep = (maxv - second) >= specificity_min
-        else:
+    if enforce_specificity:
+        if specificity_mode not in ("ratio", "diff"):
             raise ValueError("specificity_mode must be 'ratio' or 'diff'")
 
-        valid2 = valid[keep]
-        if valid2.size == 0:
-            raise ValueError("No genes passed specificity filter. Lower specificity_min or disable it.")
+        if len(dims) < 2:
+            warnings.warn(
+                "enforce_specificity=True has no effect with a single dim; ignoring it.",
+                UserWarning,
+                stacklevel=2,
+            )
+        else:
+            A = np.abs(score_used[np.ix_(valid, dims)])  # (n_valid, n_dims)
+            maxv = A.max(axis=1)
+            second = np.partition(A, -2, axis=1)[:, -2]
+            eps = 1e-12
+
+            if specificity_mode == "ratio":
+                keep = (maxv / (second + eps)) >= specificity_min
+            else:
+                keep = (maxv - second) >= specificity_min
+
+            valid2 = valid[keep]
+            if valid2.size == 0:
+                raise ValueError("No genes passed specificity filter. Lower specificity_min or disable it.")
 
     # --- union of top genes per dim based on |score|
     selected = set()
     for d in dims:
         sc = np.abs(score_used[valid2, d])
-        top_idx = valid2[np.argsort(sc)[::-1][:n_top]]
+        top_idx = valid2[np.argsort(-sc, kind="stable")[:n_top]]
         selected.update(top_idx.tolist())
 
     selected = np.array(sorted(selected), dtype=int)
@@ -225,15 +385,15 @@ def get_genes_dim(
 
     # --- optional ordering of genes
     if order_genes_by == "winner" and df.shape[1] >= 1:
-        winner = np.argmax(np.abs(df.values), axis=1)
-        df = df.iloc[np.argsort(winner)]
+        A = np.abs(df.values)
+        winner = np.argmax(A, axis=1)
+        strength = A[np.arange(A.shape[0]), winner]
+        # primary key: winning dim; secondary: descending |score| within that dim, so
+        # each winner block reads strongest-first (np.lexsort takes the last key first)
+        df = df.iloc[np.lexsort((-strength, winner))]
     elif order_genes_by == "max_abs":
         mx = np.max(np.abs(df.values), axis=1)
-        df = df.iloc[np.argsort(-mx)]
-    elif order_genes_by is None:
-        pass
-    else:
-        raise ValueError("order_genes_by must be 'winner', 'max_abs', or None")
+        df = df.iloc[np.argsort(-mx, kind="stable")]
 
     # --- optional plot
     ax = None
@@ -274,24 +434,37 @@ def get_genes_dim(
 
 def calculate_dim_importance(
     adata,
-    s_key="_global_std_gene_loadings",
-    z_key="_global_emb",
-    mode="full",
-    use_ratio=True,
-    cumulative_cutoff=0.90,
-    spacing=2,
-    n_top=None,
+    which: Literal["global", "local"] | None = None,
+    *,
+    s_key: str | None = None,
+    z_key: str | None = None,
+    mode: Literal["full", "diag"] = "full",
+    use_ratio: bool = True,
+    cumulative_cutoff: float | None = 0.90,
+    spacing: int = 2,
+    n_top: int | None = None,
+    uns_key: str | None = None,
+    store: bool = True,
 ):
     """Calculate dimension importance scores.
 
+    Also stores the selected dimensions in ``adata.uns[uns_key]`` so that
+    `get_genes_dim` can pick them up automatically (see `store`).
+
     Parameters
     ----------
     adata : AnnData
         Annotated data matrix.
-    s_key : str
-        Key in adata.varm containing gene loadings.
-    z_key : str
-        Key in adata.obsm containing embedding.
+    which : {"global", "local"}, optional
+        Which component to score. Derives `s_key` and `z_key`, so it is normally the
+        only argument needed. Defaults to "global" when `s_key` is not given; when
+        `s_key` is given instead, the component is inferred from it. Passing a `which`
+        that contradicts `s_key` is an error.
+    s_key : str, optional
+        Key in adata.varm containing gene loadings. Defaults to
+        f"_{which}_std_gene_loadings"; only needed for non-standard keys.
+    z_key : str, optional
+        Key in adata.obsm containing embedding. Defaults to f"_{which}_emb".
     mode : str
         "full" (uses Corr(Z) off-diagonals) or "diag" (assumes dims uncorrelated).
     use_ratio : bool
@@ -302,6 +475,11 @@ def calculate_dim_importance(
         Spacing between points on x-axis.
     n_top : int, optional
         Maximum number of dimensions to include.
+    uns_key : str, optional
+        Where to store the result. Defaults to f"_{which}_dim_importance".
+    store : bool
+        If True (default), write the selection to adata.uns[uns_key] so that
+        `get_genes_dim(adata, which=which)` can resolve `dims` without being told.
 
     Returns
     -------
@@ -319,7 +497,36 @@ def calculate_dim_importance(
         - mode : mode used
         - s_key : loadings key used
         - z_key : embedding key used
+        - which : component inferred or given ("local"/"global", or None)
+        - uns_key : where the selection was stored (or None if not stored)
     """
+    # --- resolve component and keys. `which` is the normal entry point; s_key/z_key
+    # override it for non-standard keys. Never let the two disagree silently: that would
+    # file a local result under the global record (the mix-up this record exists to stop).
+    if which is not None and which not in ("local", "global"):
+        raise ValueError(f"which must be 'global' or 'local', got {which!r}")
+
+    inferred = _infer_which(s_key) if s_key is not None else None
+
+    if which is None:
+        # infer from an explicit s_key, else keep the historical "global" default
+        which = inferred if inferred is not None else ("global" if s_key is None else None)
+    elif inferred is not None and inferred != which:
+        raise ValueError(
+            f"which={which!r} contradicts s_key={s_key!r}, which looks like '{inferred}'. Pass only one of the two."
+        )
+
+    if which is None and (s_key is None or z_key is None):
+        raise ValueError(
+            f"Cannot derive the missing keys from s_key={s_key!r}. Pass which='local'/'global', "
+            "or give both s_key and z_key explicitly."
+        )
+
+    if s_key is None:
+        s_key = f"_{which}_std_gene_loadings"
+    if z_key is None:
+        z_key = f"_{which}_emb"
+
     if s_key not in adata.varm:
         raise KeyError(f"{s_key} not found in adata.varm")
     if z_key not in adata.obsm:
@@ -404,6 +611,44 @@ def calculate_dim_importance(
     y_plot = y_sorted[:K]
     dim_plot = dim_sorted[:K]
 
+    # -------------------
+    # persist the selection so get_genes_dim can resolve `dims` on its own
+    # -------------------
+
+    if uns_key is None and which is not None:
+        uns_key = _dim_importance_uns_key(which)
+
+    if store:
+        if uns_key is None:
+            warnings.warn(
+                f"Could not infer the component from s_key='{s_key}', so the dimension "
+                "selection was not stored. Pass which='local'/'global' or uns_key=... to "
+                "enable get_genes_dim(adata, ...) to resolve dims automatically.",
+                UserWarning,
+                stacklevel=2,
+            )
+        else:
+            # anndata cannot write None into .uns, so omit absent entries rather than
+            # storing None -- dims_left is None whenever no cutoff was applied.
+            record = {
+                "dims_sorted": np.asarray(dim_sorted, dtype=int),
+                "importance_sorted": np.asarray(y_sorted, dtype=float),
+                "n_dims": int(len(y_sorted)),
+                "mode": str(mode),
+                "s_key": str(s_key),
+                "z_key": str(z_key),
+                "use_ratio": bool(use_ratio),
+            }
+            if which is not None:
+                record["which"] = str(which)
+            if dims_left is not None:
+                record["dims_left"] = np.asarray(dims_left, dtype=int)
+                record["n_dims_left"] = int(n_dims_left)
+            if cumulative_cutoff is not None:
+                record["cumulative_cutoff"] = float(cumulative_cutoff)
+
+            adata.uns[uns_key] = record
+
     return {
         "y_plot": y_plot,
         "dim_plot": dim_plot,
@@ -419,4 +664,6 @@ def calculate_dim_importance(
         "z_key": z_key,
         "use_ratio": use_ratio,
         "spacing": spacing,
+        "which": which,
+        "uns_key": uns_key,
     }
diff --git a/src/interscale/main_sweep.py b/src/interscale/main_sweep.py
index e89fa4d..1113f49 100644
--- a/src/interscale/main_sweep.py
+++ b/src/interscale/main_sweep.py
@@ -1,5 +1,7 @@
 import argparse
+import gc
 import os
+import traceback
 import warnings
 
 import psutil
@@ -11,10 +13,10 @@
 import interscale as interscale
 from interscale.config import load_config
 from interscale.config.cli import add_config_args, print_registry, resolve_cfg_from_args
-from interscale.config.sweep import apply_sweep_config, build_sweep_config
+from interscale.config.sweep import ARM_PARAM, apply_sweep_config, build_sweep_config, load_arms
 from interscale.geome_dataloader import GraphAnnDataModule
 from interscale.pp.segmentation_noise import apply_segmentation_noise
-from interscale.tl import prepare_geome_dataset
+from interscale.tl import prepare_geome_dataset, remove_zero_expression_cells
 from interscale.tl.utils import get_model_filename_prefix
 
 # geome calls the deprecated `sq.gr.spatial_neighbors` entrypoint; silence its
@@ -63,7 +65,7 @@ def print_memory_debug():
         print(f"Memory debug failed: {e}")
 
 
-def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
+def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None, arms=None):
     """Run one sweep trial.
 
     Parameters
@@ -80,6 +82,9 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
     sweep_params : list, optional
         The parameter names the sweep declares, used to assert that every one of them was
         actually applied to the config.
+    arms : dict, optional
+        The sweep yaml's ``arms:`` block. Each trial's ``arm`` value expands into that arm's
+        coupled dotted overrides.
     """
 
     print_memory_usage("Start of main_sweep")
@@ -120,12 +125,35 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
             sweep_config,
             model_type=model_type,
             sweep_params=sweep_params,
+            arms=arms,
         )
 
+        # The name above was derived from the config BEFORE the trial was applied, so every run
+        # of a sweep that varies dataset.name or optim.seed -- both of which feed the prefix, and
+        # through it the checkpoint filename -- appeared under one identical name. Renaming here
+        # makes a run's name the name of the checkpoint it wrote, which is what an analysis
+        # loading models back out of the sweep has to match on.
+        file_name_prefix = get_model_filename_prefix(cfg, local_component, global_component)
+        if sweep_run.name != file_name_prefix:
+            print(f"renaming run: {sweep_run.name} -> {file_name_prefix}")
+            sweep_run.name = file_name_prefix
+        # Recorded as plain summary keys so the wandb API can be filtered on them without
+        # re-deriving anything from the nested config blob.
+        sweep_run.summary["checkpoint_prefix"] = file_name_prefix
+        sweep_run.summary["resolved_dataset_name"] = cfg.dataset.name
+        sweep_run.summary["resolved_sample_key"] = list(cfg.dataset.sample_key)
+        sweep_run.summary["resolved_seed"] = cfg.optim.seed
+        # `parameters` is only added to the schema for a model type that HAS a global component,
+        # so a LocalModel sweep must not touch it.
+        if global_component:
+            sweep_run.summary["resolved_max_seq_len"] = cfg.model.global_component.parameters.max_seq_len
+        sweep_run.summary["resolved_model_save"] = cfg.model.save
+
     ####### PREPROCESSING #######
     # Load adata
     adata = sc.read_h5ad(cfg.dataset.h5ad_data)
     print_memory_usage("After loading h5ad")
+    adata = remove_zero_expression_cells(adata)
     print(adata)
     if cfg.dataset.segmentation_robustness is not None:
         print("Applying segmentation noise...")
@@ -185,7 +213,50 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
     )
     print_memory_usage("After datamodule creation")
 
-    model.train(max_epochs=cfg.optim.n_epochs, datamodule=dm, early_stopping=cfg.optim.early_stopping)
+    # wandb.agent runs every trial of a sweep inside ONE process, so anything still holding CUDA
+    # tensors when a trial ends stays resident and the next trial starts with less GPU memory
+    # than the last. Attention memory here is multiplicative in batch x heads x layers over a
+    # 4318-long sequence, so one wide trial is enough to fill a 20 GB card -- and without this
+    # release every later trial dies of OOM even at 20 MiB allocations, which is exactly how a
+    # 60-trial sweep returned one usable result.
+    #
+    # try/finally, not a trailing statement: the trials that most need the release are precisely
+    # the ones that raise (an OOM leaves a half-built model and its activations referenced), and
+    # a cleanup placed after the call is skipped exactly then. Without this, one OOM trial
+    # poisons every trial after it and the sweep cannot be run near the memory ceiling at all.
+    # The failure is caught and re-raised as a NEW exception carrying only the message, because
+    # an exception's __traceback__ keeps every frame in the call stack alive, and those frames
+    # hold the activations that caused the OOM in the first place. wandb.agent stores the
+    # exception it catches, so the original traceback -- and through it the whole trainer.fit()
+    # stack -- outlives the trial. Deleting the locals below is then useless: measured on CosMx,
+    # GPU memory after "cleanup" went 13.0 -> 15.4 -> 16.1 -> 17.9 -> 18.1 GB across six trials
+    # and every one of them OOMed. Binding the error to a plain string lets Python's implicit
+    # `del exc` at the end of the except block drop the traceback before gc.collect() runs.
+    trial_error = None
+    try:
+        model.train(max_epochs=cfg.optim.n_epochs, datamodule=dm, early_stopping=cfg.optim.early_stopping)
+    except Exception as exc:
+        # format_exc() renders the stack to a STRING, so the full traceback survives in the log
+        # while no frame (and so no tensor) stays referenced. Keeping only str(exc) made the
+        # first CosMx OOM undiagnosable: the allocation turned out to be in the metric
+        # collection, not in attention, and nothing in the message said so.
+        trial_error = f"{type(exc).__name__}: {exc}\n{traceback.format_exc()}"
+    finally:
+        del model, dm, pyg_data_list, adata
+        gc.collect()
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                torch.cuda.reset_peak_memory_stats()
+        except ImportError:
+            pass
+        print_memory_usage("After per-trial cleanup")
+
+    # Re-raised so wandb still records the trial as failed rather than silently succeeding.
+    if trial_error is not None:
+        raise RuntimeError(f"sweep trial failed: {trial_error}")
 
 
 if __name__ == "__main__":
@@ -237,6 +308,25 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
         default="InterScale_hyperparameter_sweep",
         help="wandb project the sweep is registered under.",
     )
+    parser.add_argument(
+        "--sweep_id",
+        dest="sweep_id",
+        type=str,
+        default=None,
+        help="Join an EXISTING sweep instead of registering a new one. Two uses: several agents "
+        "working one sweep in parallel (a slurm array), and resuming a sweep whose agent hit the "
+        "walltime -- a grid sweep does not re-run trials it has already completed. The sweep's "
+        "own parameter grid is whatever was registered; --sweep_cfg is still required because it "
+        "supplies the arms and the parameter names each trial is checked against.",
+    )
+    parser.add_argument(
+        "--create_only",
+        dest="create_only",
+        action="store_true",
+        help="Register the sweep, print its id, and exit without running an agent. Use to create "
+        "a sweep on a login node and then submit a slurm array of agents against it with "
+        "--sweep_id.",
+    )
     parser.add_argument(
         "--prediction_task",
         dest="prediction_task",
@@ -245,6 +335,25 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
         choices=["regression", "classification"],
         help="Type of prediction task. Defaults to the resolved config's dataset.prediction_task.",
     )
+    parser.add_argument(
+        "--metric",
+        dest="metric",
+        type=str,
+        default=None,
+        help="Metric wandb ranks trials by, overriding the prediction task's default "
+        "(classification: val_f1_macro, regression: val_r2). Must be a metric the training plan "
+        "actually logs, e.g. val_pearson_corr for a regression sweep searching for correlation "
+        "rather than calibration. Set optim.monitor to the same metric so checkpoint selection "
+        "agrees with what the sweep ranks.",
+    )
+    parser.add_argument(
+        "--metric_goal",
+        dest="metric_goal",
+        type=str,
+        default="maximize",
+        choices=["maximize", "minimize"],
+        help="Direction for --metric (default: maximize).",
+    )
     args = parser.parse_args()
 
     if args.list_pairs:
@@ -277,20 +386,42 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None):
         yaml_config,
         prediction_task=prediction_task,
         model_type=args.model_type,
+        metric=({"name": args.metric, "goal": args.metric_goal} if args.metric else None),
     )
     print(sweep_config)
 
+    arms = load_arms(yaml_config, sweep_config)
+
     # Fail before wandb.sweep() registers anything if a declared parameter does not exist in this
     # config: the sweep would otherwise burn every trial varying a key that nothing reads.
-    apply_sweep_config(
-        base_cfg.clone(),
-        args.sweep_goal,
-        dict.fromkeys(sweep_params, None),
-        model_type=args.model_type,
-        sweep_params=sweep_params,
-    )
+    #
+    # With arms, EVERY arm is checked, not one representative: an arm's overrides are its own set
+    # of dotted keys, and a typo in the fourth arm would otherwise only surface once the first
+    # three had trained.
+    for arm_name in [None] if arms is None else list(sweep_config["parameters"][ARM_PARAM]["values"]):
+        trial = dict.fromkeys(sweep_params, None)
+        if arm_name is not None:
+            trial[ARM_PARAM] = arm_name
+        apply_sweep_config(
+            base_cfg.clone(),
+            args.sweep_goal,
+            trial,
+            model_type=args.model_type,
+            sweep_params=sweep_params,
+            arms=arms,
+        )
 
-    sweep_id = wandb.sweep(sweep_config, project=args.sweep_project)
+    if args.sweep_id:
+        sweep_id = args.sweep_id
+        print(f"joining existing sweep {args.sweep_project}/{sweep_id}")
+    else:
+        sweep_id = wandb.sweep(sweep_config, project=args.sweep_project)
+
+    if args.create_only:
+        # Printed in a grep-able form so a submit script can capture it.
+        print(f"SWEEP_ID={sweep_id}")
+        print(f"sweep url: https://wandb.ai/{wandb.Api().default_entity}/{args.sweep_project}/sweeps/{sweep_id}")
+        raise SystemExit(0)
 
     def train_sweep_function():
         # A fresh config per trial: applying a trial's values mutates it, so a shared object
@@ -300,9 +431,12 @@ def train_sweep_function():
             args.model_type,
             args.sweep_goal,
             sweep_params=sweep_params,
+            arms=arms,
         )
 
     # Without an explicit count the agent runs until the job's walltime kills it: `random`
     # over this grid has ~1.5M combinations, so it never exhausts them on its own.
     print(f"running {args.count} sweep trials (sweep_id={sweep_id})")
-    wandb.agent(sweep_id, function=train_sweep_function, count=args.count)
+    # project= is required when joining a sweep by bare id: without it the agent looks the sweep
+    # up in the default project rather than the one --sweep_project names.
+    wandb.agent(sweep_id, function=train_sweep_function, count=args.count, project=args.sweep_project)
diff --git a/src/interscale/module/base/_base_global_module.py b/src/interscale/module/base/_base_global_module.py
index 04f2950..d91b834 100644
--- a/src/interscale/module/base/_base_global_module.py
+++ b/src/interscale/module/base/_base_global_module.py
@@ -158,13 +158,24 @@ def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, p
             else torch.tensor([], device=device, dtype=torch.long)
         )
 
-        # Assertions
+        # Every adjusted index must address a valid row of y_true. This is the invariant the
+        # offset arithmetic can actually violate, so it stays.
+        #
+        # A second assertion used to sit here requiring `adjusted_mask_idx.max() >
+        # len(pad_index_nodes[0])` whenever nr_batches > 1, i.e. that masked nodes came from more
+        # than just the first graph. It was wrong twice over. Graph i's indices start at
+        # cumulative_offsets[i], so a masked node that is the *first kept node* of graph 1 gets
+        # index exactly len(pad_index_nodes[0]) -- a valid position that `>` rejected. That is
+        # reachable whenever the last batch of an epoch holds two graphs and the second is small
+        # enough for its single masked node to be node 0, which killed a legnini23 run at epoch 22.
+        # It also asserted a property that is not invariant: if a graph's masked set exceeds
+        # max_seq_len, _select_masked_nodes legitimately drops some, and that case is caught with
+        # an accurate message by the length checks in the combined modules' _common_step, where
+        # the local branch's masked-node count is compared against the global branch's.
         if len(adjusted_mask_idx) > 0:
             assert adjusted_mask_idx.max() < len(y_true), (
                 f"Mismatch: max(adjusted_mask_idx): {adjusted_mask_idx.max()}, len(y_true): {len(y_true)}"
             )
-            if nr_batches > 1:
-                assert adjusted_mask_idx.max() > len(pad_index_nodes[0]), "No masked node included from all batches"
 
         return y_true, adjusted_mask_idx
 
diff --git a/src/interscale/train/_training.py b/src/interscale/train/_training.py
index 01fd583..deb20a5 100644
--- a/src/interscale/train/_training.py
+++ b/src/interscale/train/_training.py
@@ -177,30 +177,22 @@ def train(
                     save_top_k=1,
                 )  # save the best model according to `monitor`
             elif "regression" in self._cfg.dataset.prediction_task:
-                if self._cfg.optim.loss == "MSELoss":
-                    checkpoint_callback = ModelCheckpoint(
-                        dirpath=self._cfg.model.save,
-                        filename=run_name,
-                        monitor="val_loss",
-                        mode="min",
-                    )
-                elif (
-                    self._cfg.optim.loss == "GaussianNLL"
-                    or self._cfg.optim.loss == "SmoothL1"
-                    or self._cfg.optim.loss == "BalancedPearsonCorrelationLoss"
-                    or self._cfg.optim.loss == "SCELoss"
-                    or self._cfg.optim.loss == "SCE_EntropyATT_Loss"
-                ):
-                    checkpoint_callback = ModelCheckpoint(
-                        dirpath=self._cfg.model.save,
-                        filename=run_name,
-                        monitor="val_loss",
-                        mode="min",
-                    )
-                else:
-                    raise Exception(
-                        f"Regression must be run with MSELoss, GaussianNLL, SmoothL1, BalancedPearsonCorrelationLoss or SCELoss loss. instead of {self._cfg.optim.loss}"
-                    )
+                # Every arm of the former if/elif chain built the same callback with a hardcoded
+                # monitor="val_loss", so `optim.monitor` was silently ignored for regression while
+                # classification honoured it. That is not cosmetic: train() restores the best
+                # checkpoint before trainer.validate(), so every metric reported for the run --
+                # and therefore whatever a sweep ranks on -- came from the lowest-val_loss epoch,
+                # no matter which metric the run was supposed to be selecting for.
+                #
+                # `auto` still resolves to val_loss/min for regression, so existing configs are
+                # unaffected; only a config that explicitly sets optim.monitor changes behaviour.
+                checkpoint_callback = ModelCheckpoint(
+                    dirpath=self._cfg.model.save,
+                    filename=run_name,
+                    monitor=monitor,
+                    mode=mode,
+                    save_top_k=1,
+                )
 
         # Create list of callbacks and filter out None values
         callbacks = [
diff --git a/src/interscale/train/_trainingplans.py b/src/interscale/train/_trainingplans.py
index 96e3aef..1c63498 100644
--- a/src/interscale/train/_trainingplans.py
+++ b/src/interscale/train/_trainingplans.py
@@ -13,6 +13,43 @@
 
 from .losses import BalancedPearsonCorrelationLoss, SCE_EntropyATT_Loss, SCELoss
 
+class RunningCosineSimilarity(torchmetrics.Metric):
+    """Mean per-cell cosine similarity, with state that does not grow with the dataset.
+
+    ``torchmetrics.CosineSimilarity`` is a *list-state* metric: it keeps every prediction and
+    target it is shown and concatenates them at compute time. Every other regression metric here
+    holds a few kilobytes of running sums, and this one holds ``n_cells x n_genes x 2`` floats --
+    which ``MetricCollection.forward`` then duplicates via ``_copy_state_dict`` on every step.
+
+    On legnini23 (43k cells, 88 genes) that is ~30 MB and invisible. On the CosMx pancreas
+    (387k cells, 979 genes) one epoch is ~850 MB before the copy, and it OOMed a 20 GB card
+    inside ``_regression_metrics`` on the very first trial, regardless of batch size -- the total
+    per epoch is the same however the cells are batched.
+
+    This computes the same quantity (the mean over cells of the per-cell cosine) from a running
+    sum and count, so the state is two scalars.
+    """
+
+    is_differentiable = False
+    higher_is_better = True
+    full_state_update = False
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.add_state("total", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("count", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+        """Accumulate the summed per-cell cosine similarity and the number of cells."""
+        cos = nn.functional.cosine_similarity(preds, target, dim=1)
+        self.total = self.total + cos.sum()
+        self.count = self.count + cos.numel()
+
+    def compute(self) -> torch.Tensor:
+        """Mean per-cell cosine similarity over everything seen since the last reset."""
+        return self.total / self.count
+
+
 CLASSIFICATION_LOSSES = ["CrossEntropy", "WeightedCE"]
 REGRESSION_LOSSES = [
     "MSELoss",
@@ -162,7 +199,17 @@ def _setup_regression_metrics(num_outputs: int):
                 "mse": torchmetrics.MeanSquaredError(),
                 "r2": torchmetrics.R2Score(multioutput="uniform_average"),
                 "pearson_corr": torchmetrics.PearsonCorrCoef(num_outputs=num_outputs),
-                "cosine_similarity": torchmetrics.CosineSimilarity(reduction="mean"),
+                # Pearson is invariant to any per-gene affine rescaling of the predictions, so a
+                # model whose outputs are (say) 11x too spread out still scores a high r while
+                # its R2 goes to -113. Writing predictions as k times the true sd with offset d,
+                # R2 = 2*r*k - k^2 - d^2/sigma^2, maximised at k = r -- so R2 <= r^2, and the gap
+                # between them is purely calibration. Concordance correlation folds that penalty
+                # back in, which makes it the metric to select on when both the co-variation
+                # structure AND the expression scale have to be usable.
+                "concordance_corr": torchmetrics.ConcordanceCorrCoef(num_outputs=num_outputs),
+                # Not torchmetrics.CosineSimilarity: see RunningCosineSimilarity for why its
+                # list state cannot be used on a dataset this size. Same value, O(1) memory.
+                "cosine_similarity": RunningCosineSimilarity(),
             }
         )
 
@@ -225,6 +272,9 @@ def _regression_metrics(
 
         # Take mean across pearson correlation
         metrics[f"{mode}_pearson_corr"] = torch.nanmean(metrics[f"{mode}_pearson_corr"].contiguous())
+        # Same reduction, same reason: both are per-gene vectors of length n_output, and a
+        # constant gene yields NaN rather than a number.
+        metrics[f"{mode}_concordance_corr"] = torch.nanmean(metrics[f"{mode}_concordance_corr"].contiguous())
         metrics[f"{mode}_loss"] = loss
         return loss, metrics
 
diff --git a/tests/test_sweep_config.py b/tests/test_sweep_config.py
index 9c82afe..1b0a724 100644
--- a/tests/test_sweep_config.py
+++ b/tests/test_sweep_config.py
@@ -20,9 +20,11 @@
 from interscale.config import load_config
 from interscale.config.registry import resolve_config
 from interscale.config.sweep import (
+    ARM_PARAM,
     SWEEP_GOALS,
     apply_sweep_config,
     build_sweep_config,
+    load_arms,
     unused_component_prefixes,
 )
 
@@ -424,3 +426,239 @@ def test_sweep_parameters_apply_to_a_bare_default_config():
     trial = {k: sweep_config["parameters"][k].get("values", [None])[0] for k in sweep_params}
     _, applied = apply_sweep_config(cfg, "hyperparmeter", trial, model_type="CombinedModel", sweep_params=sweep_params)
     assert sorted(applied) == sorted(sweep_params)
+
+
+# --------------------------------------------------------------------------------------------
+# Arms: one sweep parameter standing for several coupled config keys.
+#
+# The failure these guard against is the one the flat dotted-key design cannot express at all.
+# wandb searches the cartesian product of its parameters, so a window-size ablation whose three
+# implied keys were declared separately would enumerate 6^3 = 216 trials, 210 of them invalid
+# crossings (e.g. the 400-cell window column with the 3436 max_seq_len). The arms block keeps the
+# three coupled, and these tests prove the coupling survives the round trip.
+# --------------------------------------------------------------------------------------------
+
+SLIDING_WINDOW_YAML = CONFIG_DIR / "sweeps" / "sliding_window_melton25.yaml"
+
+ARM_DATASET = "melton25_sw"
+ARM_TASK = "node_reg"
+
+
+@pytest.fixture
+def arm_yaml():
+    with SLIDING_WINDOW_YAML.open() as f:
+        return yaml.safe_load(f)
+
+
+@pytest.fixture
+def arm_cfg():
+    return resolve_config(ARM_DATASET, ARM_TASK, registry_path=REGISTRY)
+
+
+def make_arm_trial(sweep_config, sweep_params, arm_name):
+    """Build the trial dict wandb would hand back for one arm, first value for everything else."""
+    trial = {k: sweep_config["parameters"][k]["values"][0] for k in sweep_params}
+    trial[ARM_PARAM] = arm_name
+    return trial
+
+
+def test_sliding_window_sweep_yaml_exists():
+    assert SLIDING_WINDOW_YAML.is_file(), f"missing sweep config: {SLIDING_WINDOW_YAML}"
+
+
+def test_load_arms_returns_none_without_an_arms_block(sweep_yaml):
+    """The flat sweeps must be entirely unaffected by the arms machinery."""
+    assert load_arms(sweep_yaml) is None
+
+
+def test_every_arm_applies_all_of_its_coupled_keys(arm_yaml, arm_cfg):
+    """Each arm's full set of dotted overrides reaches the config, for every arm."""
+    sweep_config, sweep_params = build_sweep_config(
+        arm_yaml, prediction_task="regression", model_type="CombinedModel"
+    )
+    arms = load_arms(arm_yaml, sweep_config)
+
+    for arm_name, overrides in arms.items():
+        cfg = arm_cfg.clone()
+        trial = make_arm_trial(sweep_config, sweep_params, arm_name)
+        cfg, applied = apply_sweep_config(
+            cfg, "robustness", trial, model_type="CombinedModel", sweep_params=sweep_params, arms=arms
+        )
+        for key, expected in overrides.items():
+            assert get_dotted(cfg, key) == expected, f"arm {arm_name}: {key} did not reach the config"
+            assert key in applied
+
+
+def test_arm_name_itself_is_never_written_to_the_config(arm_yaml, arm_cfg):
+    """`arm` is a selector, not a config path; writing it would need a config key called 'arm'."""
+    sweep_config, sweep_params = build_sweep_config(
+        arm_yaml, prediction_task="regression", model_type="CombinedModel"
+    )
+    arms = load_arms(arm_yaml, sweep_config)
+    trial = make_arm_trial(sweep_config, sweep_params, "w400")
+    cfg, applied = apply_sweep_config(
+        arm_cfg, "robustness", trial, model_type="CombinedModel", sweep_params=sweep_params, arms=arms
+    )
+    assert ARM_PARAM not in applied
+    assert ARM_PARAM not in cfg
+    assert ARM_PARAM not in cfg.dataset
+
+
+def test_arms_give_distinct_checkpoint_prefixes(arm_yaml, arm_cfg):
+    """dataset.name must differ per arm, or every arm overwrites the previous arm's checkpoint.
+
+    This is not hypothetical: get_model_filename_prefix keys on dataset.name, prediction task,
+    level and seed, none of which the window size touches on its own, and
+    trainer.save_checkpoint() overwrites unconditionally.
+    """
+    from interscale.tl.utils import get_model_filename_prefix
+
+    sweep_config, sweep_params = build_sweep_config(
+        arm_yaml, prediction_task="regression", model_type="CombinedModel"
+    )
+    arms = load_arms(arm_yaml, sweep_config)
+
+    prefixes = {}
+    for arm_name in arms:
+        cfg = arm_cfg.clone()
+        cfg, _ = apply_sweep_config(
+            cfg,
+            "robustness",
+            make_arm_trial(sweep_config, sweep_params, arm_name),
+            model_type="CombinedModel",
+            sweep_params=sweep_params,
+            arms=arms,
+        )
+        prefixes[arm_name] = get_model_filename_prefix(cfg, local_component=True, global_component=True)
+
+    assert len(set(prefixes.values())) == len(arms), f"arms share a checkpoint filename: {prefixes}"
+
+
+def test_max_seq_len_is_never_below_the_arms_largest_window(arm_yaml):
+    """Every arm's max_seq_len must be >= the largest window of the column it selects.
+
+    Below it, pad_batch random-subsamples the window each step and get_model_output stores an
+    attention matrix narrower than the window, which makes the downstream net-flow computation
+    fail with a shape mismatch rather than a wrong number.
+    """
+    # Largest window over ALL splits, measured on melton25_sliding_window.h5ad. Inference runs on
+    # every cell, so the train-split maximum is not the relevant bound.
+    LARGEST_WINDOW = {
+        "sliding_window_400": 89,
+        "sliding_window_800": 330,
+        "sliding_window_1200": 685,
+        "sliding_window_1600": 1218,
+        "sliding_window_2000": 1720,
+        "sliding_window_3000": 3436,
+    }
+    arms = load_arms(arm_yaml)
+    for arm_name, overrides in arms.items():
+        (column,) = overrides["dataset.sample_key"]
+        max_seq_len = overrides["model.global_component.parameters.max_seq_len"]
+        assert column in LARGEST_WINDOW, f"arm {arm_name} selects an unmeasured column {column}"
+        assert max_seq_len >= LARGEST_WINDOW[column], (
+            f"arm {arm_name}: max_seq_len {max_seq_len} < largest {column} window {LARGEST_WINDOW[column]}"
+        )
+
+
+def test_arm_with_a_missing_key_raises():
+    """An arm that omits a key its siblings set would silently keep the base config's value."""
+    yaml_config = {
+        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {ARM_PARAM: {"values": ["a", "b"]}}},
+        "arms": {
+            "a": {"dataset.name": "a", "dataset.batch_size": 8},
+            "b": {"dataset.name": "b"},
+        },
+    }
+    with pytest.raises(ValueError, match="does not declare the same keys"):
+        load_arms(yaml_config, yaml_config["sweep_config"])
+
+
+def test_arm_selecting_an_undefined_arm_raises():
+    yaml_config = {
+        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {ARM_PARAM: {"values": ["a", "typo"]}}},
+        "arms": {"a": {"dataset.name": "a"}},
+    }
+    with pytest.raises(ValueError, match="does not define"):
+        load_arms(yaml_config, yaml_config["sweep_config"])
+
+
+def test_arm_key_colliding_with_a_sweep_parameter_raises():
+    """Set twice per trial, with application order deciding the winner."""
+    yaml_config = {
+        "sweep_config": {
+            "metric": {"name": "val_r2", "goal": "maximize"},
+            "parameters": {ARM_PARAM: {"values": ["a"]}, "dataset.batch_size": {"values": [4, 8]}},
+        },
+        "arms": {"a": {"dataset.batch_size": 16}},
+    }
+    with pytest.raises(ValueError, match="both by the arms and as sweep parameters"):
+        load_arms(yaml_config, yaml_config["sweep_config"])
+
+
+def test_arms_block_without_an_arm_parameter_raises():
+    yaml_config = {
+        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {"optim.seed": {"values": [1]}}},
+        "arms": {"a": {"dataset.name": "a"}},
+    }
+    with pytest.raises(ValueError, match="no 'arm' parameter"):
+        load_arms(yaml_config, yaml_config["sweep_config"])
+
+
+def test_arm_parameter_without_an_arms_block_raises():
+    """Otherwise the arm name would be written to a config key called 'arm', which cannot exist."""
+    yaml_config = {
+        "sweep_config": {
+            "metric": {"name": "val_r2", "goal": "maximize"},
+            "parameters": {ARM_PARAM: {"values": ["a"]}},
+        }
+    }
+    with pytest.raises(ValueError, match="no top-level `arms:` block"):
+        build_sweep_config(yaml_config, prediction_task="regression", model_type="CombinedModel")
+
+
+def test_arm_overrides_must_be_dotted():
+    yaml_config = {
+        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {ARM_PARAM: {"values": ["a"]}}},
+        "arms": {"a": {"batch_size": 8}},
+    }
+    with pytest.raises(ValueError, match="non-dotted keys"):
+        load_arms(yaml_config, yaml_config["sweep_config"])
+
+
+def test_arm_with_an_unknown_dotted_key_raises(arm_cfg):
+    """Arm overrides go through the same _set_dotted validation as sweep parameters."""
+    arms = {"a": {"dataset.no_such_key": 1}}
+    with pytest.raises(KeyError):
+        apply_sweep_config(
+            arm_cfg, "robustness", {ARM_PARAM: "a"}, model_type="CombinedModel", sweep_params=[ARM_PARAM], arms=arms
+        )
+
+
+def test_reserved_arm_parameter_without_arms_passed_to_apply_raises(arm_cfg):
+    """Guards the wiring: forgetting to thread `arms` through must fail, not train the base config."""
+    with pytest.raises(KeyError, match="no arms were passed"):
+        apply_sweep_config(
+            arm_cfg, "robustness", {ARM_PARAM: "w400"}, model_type="CombinedModel", sweep_params=[ARM_PARAM]
+        )
+
+
+def test_arm_trial_does_not_leak_into_the_base_config(arm_cfg, arm_yaml):
+    """wandb.agent reuses one process per agent, so a leaked arm would poison later trials."""
+    sweep_config, sweep_params = build_sweep_config(
+        arm_yaml, prediction_task="regression", model_type="CombinedModel"
+    )
+    arms = load_arms(arm_yaml, sweep_config)
+    before = list(arm_cfg.dataset.sample_key)
+
+    clone = arm_cfg.clone()
+    apply_sweep_config(
+        clone,
+        "robustness",
+        make_arm_trial(sweep_config, sweep_params, "w3000"),
+        model_type="CombinedModel",
+        sweep_params=sweep_params,
+        arms=arms,
+    )
+    assert list(arm_cfg.dataset.sample_key) == before
+    assert list(clone.dataset.sample_key) == ["sliding_window_3000"]

From 0413000bb6c3b8e9241564febb70c8688e11849e Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Mon, 31 Aug 2026 12:01:24 +0200
Subject: [PATCH 08/14] add prefix to gene loading functions

---
 src/interscale/evaluation/_gene_loadings.py   |   1 +
 src/interscale/evaluation/_latent_analysis.py |  78 +++++++++----
 tests/test_sweep_config.py                    | 107 ++++++++++++++++++
 3 files changed, 164 insertions(+), 22 deletions(-)

diff --git a/src/interscale/evaluation/_gene_loadings.py b/src/interscale/evaluation/_gene_loadings.py
index 9baa69e..35f0f2d 100644
--- a/src/interscale/evaluation/_gene_loadings.py
+++ b/src/interscale/evaluation/_gene_loadings.py
@@ -8,6 +8,7 @@ def gene_loadings(
     adata: ad.AnnData,
     model,
     layer_key: str,
+    *,
     local_latent_key: str = "_local_emb",
     global_latent_key: str = "_global_emb",
     local_varm_key: str = "_local_std_gene_loadings",
diff --git a/src/interscale/evaluation/_latent_analysis.py b/src/interscale/evaluation/_latent_analysis.py
index 9171674..3777464 100644
--- a/src/interscale/evaluation/_latent_analysis.py
+++ b/src/interscale/evaluation/_latent_analysis.py
@@ -36,18 +36,29 @@ def _gene_expression_stats(X, ddof=1):
     return frac, sd
 
 
-def _infer_which(s_key):
-    """Infer the component ('local'/'global') from a loading key, or None if ambiguous."""
-    if s_key.startswith("_local"):
-        return "local"
-    if s_key.startswith("_global"):
-        return "global"
-    return None
+def _infer_which(s_key, prefix=""):
+    """Infer the component ('local'/'global') from a loading key, or None if ambiguous.
 
+    `prefix` is stripped first, so both "_local_std_gene_loadings" and a prefixed
+    "seed0_local_std_gene_loadings" resolve to "local".
+    """
+    body = s_key[len(prefix) :] if prefix and s_key.startswith(prefix) else s_key
+    hits = [
+        c
+        for c in ("local", "global")
+        if body.startswith(f"_{c}") or body.startswith(f"{c}_") or f"_{c}_" in body
+    ]
+    return hits[0] if len(hits) == 1 else None
+
+
+def _component_key(which, suffix, prefix=""):
+    """Compose an adata key for a component, following the f"{prefix}_{which}_..." convention."""
+    return f"{prefix}_{which}_{suffix}"
 
-def _dim_importance_uns_key(which):
+
+def _dim_importance_uns_key(which, prefix=""):
     """Key under which `calculate_dim_importance` stores its selection in adata.uns."""
-    return f"_{which}_dim_importance"
+    return _component_key(which, "dim_importance", prefix=prefix)
 
 
 def _resolve_dims_from_uns(adata, uns_key, s_key):
@@ -152,11 +163,12 @@ def get_genes_dim(
     adata,
     which: Literal["global", "local"],  # required
     *,
+    prefix: str = "",  # e.g. "seed0" -> keys "seed0_global_..."; "" -> "_global_..."
     dims: Sequence[int] | None = None,  # None -> read the selection stored by calculate_dim_importance
     n_top: int = 20,
     s_key: str | None = None,  # e.g. "_global_std_gene_loadings"
     z_key: str | None = None,  # e.g. "_global_emb" (only used if residualize=True)
-    uns_key: str | None = None,  # e.g. "_global_dim_importance"; None -> f"_{which}_dim_importance"
+    uns_key: str | None = None,  # e.g. "_global_dim_importance"; None -> f"{prefix}_{which}_dim_importance"
     # expression-based filtering
     X_layer=None,  # e.g. "log1p_norm"; None -> adata.X
     min_frac=0.05,  # fraction of cells with expr>0
@@ -191,6 +203,12 @@ def get_genes_dim(
         Which component to analyse. Required, because it selects the defaults for
         `s_key`, `z_key` and `uns_key` — it is the only argument needed to switch
         components.
+    prefix : str, default ""
+        Prefix of the stored keys, following the f"{prefix}_{which}_..." convention used
+        when the model output was saved. The default "" gives the plain "_global_emb" /
+        "_local_emb" keys; pass e.g. prefix="seed0" to analyse a run whose embeddings
+        were saved as "seed0_global_emb" / "seed0_local_emb". Ignored for any key you
+        pass explicitly.
     dims : sequence of int, optional
         Latent dimensions to analyse. **Leave as None (recommended)** to read the
         selection `calculate_dim_importance` stored in adata.uns[uns_key], which keeps
@@ -198,14 +216,18 @@ def get_genes_dim(
         exists, the two are compared and a mismatch raises a warning.
     uns_key : str, optional
         Record written by `calculate_dim_importance`. Defaults to
-        f"_{which}_dim_importance". A record built from a different component than the
-        resolved `s_key` raises ValueError rather than silently mixing the two.
+        f"{prefix}_{which}_dim_importance". A record built from a different component
+        than the resolved `s_key` raises ValueError rather than silently mixing the two.
 
     Examples
     --------
     >>> calculate_dim_importance(adata, "global", cumulative_cutoff=0.60)
     >>> df = get_genes_dim(adata, "global", n_top=20)  # dims resolved from adata.uns
 
+    >>> # a run whose embeddings were saved with prefix="seed0"
+    >>> calculate_dim_importance(adata, "global", prefix="seed0", cumulative_cutoff=0.60)
+    >>> df = get_genes_dim(adata, "global", prefix="seed0", n_top=20)
+
     Notes
     -----
     - Uses standardized loadings (S_std) stored in adata.varm[s_key].
@@ -225,11 +247,11 @@ def get_genes_dim(
         raise ValueError(f"which must be 'global' or 'local', got {which!r}")
 
     if s_key is None:
-        s_key = f"_{which}_std_gene_loadings"
+        s_key = _component_key(which, "std_gene_loadings", prefix=prefix)
     if z_key is None:
-        z_key = f"_{which}_emb"
+        z_key = _component_key(which, "emb", prefix=prefix)
     if uns_key is None:
-        uns_key = _dim_importance_uns_key(which)
+        uns_key = _dim_importance_uns_key(which, prefix=prefix)
 
     # --- load/check S
     if s_key not in adata.varm:
@@ -436,6 +458,7 @@ def calculate_dim_importance(
     adata,
     which: Literal["global", "local"] | None = None,
     *,
+    prefix: str = "",
     s_key: str | None = None,
     z_key: str | None = None,
     mode: Literal["full", "diag"] = "full",
@@ -460,11 +483,19 @@ def calculate_dim_importance(
         only argument needed. Defaults to "global" when `s_key` is not given; when
         `s_key` is given instead, the component is inferred from it. Passing a `which`
         that contradicts `s_key` is an error.
+    prefix : str, default ""
+        Prefix of the stored keys, following the f"{prefix}_{which}_..." convention used
+        when the model output was saved. The default "" gives the plain "_global_emb" /
+        "_local_emb" keys; pass e.g. prefix="seed0" to score a run whose embeddings were
+        saved as "seed0_global_emb" / "seed0_local_emb". The selection is then stored
+        under adata.uns["seed0_global_dim_importance"], so several runs can live in the
+        same AnnData without overwriting each other. Ignored for any key you pass
+        explicitly.
     s_key : str, optional
         Key in adata.varm containing gene loadings. Defaults to
-        f"_{which}_std_gene_loadings"; only needed for non-standard keys.
+        f"{prefix}_{which}_std_gene_loadings"; only needed for non-standard keys.
     z_key : str, optional
-        Key in adata.obsm containing embedding. Defaults to f"_{which}_emb".
+        Key in adata.obsm containing embedding. Defaults to f"{prefix}_{which}_emb".
     mode : str
         "full" (uses Corr(Z) off-diagonals) or "diag" (assumes dims uncorrelated).
     use_ratio : bool
@@ -476,7 +507,7 @@ def calculate_dim_importance(
     n_top : int, optional
         Maximum number of dimensions to include.
     uns_key : str, optional
-        Where to store the result. Defaults to f"_{which}_dim_importance".
+        Where to store the result. Defaults to f"{prefix}_{which}_dim_importance".
     store : bool
         If True (default), write the selection to adata.uns[uns_key] so that
         `get_genes_dim(adata, which=which)` can resolve `dims` without being told.
@@ -498,6 +529,7 @@ def calculate_dim_importance(
         - s_key : loadings key used
         - z_key : embedding key used
         - which : component inferred or given ("local"/"global", or None)
+        - prefix : key prefix used
         - uns_key : where the selection was stored (or None if not stored)
     """
     # --- resolve component and keys. `which` is the normal entry point; s_key/z_key
@@ -506,7 +538,7 @@ def calculate_dim_importance(
     if which is not None and which not in ("local", "global"):
         raise ValueError(f"which must be 'global' or 'local', got {which!r}")
 
-    inferred = _infer_which(s_key) if s_key is not None else None
+    inferred = _infer_which(s_key, prefix=prefix) if s_key is not None else None
 
     if which is None:
         # infer from an explicit s_key, else keep the historical "global" default
@@ -523,9 +555,9 @@ def calculate_dim_importance(
         )
 
     if s_key is None:
-        s_key = f"_{which}_std_gene_loadings"
+        s_key = _component_key(which, "std_gene_loadings", prefix=prefix)
     if z_key is None:
-        z_key = f"_{which}_emb"
+        z_key = _component_key(which, "emb", prefix=prefix)
 
     if s_key not in adata.varm:
         raise KeyError(f"{s_key} not found in adata.varm")
@@ -616,7 +648,7 @@ def calculate_dim_importance(
     # -------------------
 
     if uns_key is None and which is not None:
-        uns_key = _dim_importance_uns_key(which)
+        uns_key = _dim_importance_uns_key(which, prefix=prefix)
 
     if store:
         if uns_key is None:
@@ -637,6 +669,7 @@ def calculate_dim_importance(
                 "mode": str(mode),
                 "s_key": str(s_key),
                 "z_key": str(z_key),
+                "prefix": str(prefix),
                 "use_ratio": bool(use_ratio),
             }
             if which is not None:
@@ -665,5 +698,6 @@ def calculate_dim_importance(
         "use_ratio": use_ratio,
         "spacing": spacing,
         "which": which,
+        "prefix": prefix,
         "uns_key": uns_key,
     }
diff --git a/tests/test_sweep_config.py b/tests/test_sweep_config.py
index 1b0a724..7504424 100644
--- a/tests/test_sweep_config.py
+++ b/tests/test_sweep_config.py
@@ -662,3 +662,110 @@ def test_arm_trial_does_not_leak_into_the_base_config(arm_cfg, arm_yaml):
     )
     assert list(arm_cfg.dataset.sample_key) == before
     assert list(clone.dataset.sample_key) == ["sliding_window_3000"]
+
+
+# --------------------------------------------------------------------------------------------
+# Every arm-bearing sweep yaml in the repo, not just the one this file was written against.
+#
+# Discovered rather than listed: a new ablation adds a yaml and inherits these checks, which is
+# the point. Each is resolved against the (dataset, task) pair its own header names, so the test
+# proves the arms apply to the config they will actually be run with.
+# --------------------------------------------------------------------------------------------
+
+# sweep yaml -> the registered pair it is written for. Kept explicit because the yaml does not
+# name its dataset: --dataset/--task are passed on the command line.
+ARM_SWEEP_PAIRS = {
+    "sliding_window_melton25.yaml": ("melton25_sw", "node_reg"),
+    "overlap_ladder_legnini.yaml": ("legnini23_overlap", "node_reg"),
+}
+
+
+def arm_sweep_yamls():
+    """Every yaml under config_files/sweeps that declares an `arms:` block."""
+    found = []
+    for path in sorted((CONFIG_DIR / "sweeps").glob("*.yaml")):
+        with path.open() as f:
+            if "arms" in (yaml.safe_load(f) or {}):
+                found.append(path)
+    return found
+
+
+def test_every_arm_sweep_yaml_is_registered_in_this_test():
+    """A new arm sweep must be added to ARM_SWEEP_PAIRS, or it goes untested."""
+    undeclared = [p.name for p in arm_sweep_yamls() if p.name not in ARM_SWEEP_PAIRS]
+    assert not undeclared, (
+        f"arm sweep yaml(s) with no (dataset, task) declared in ARM_SWEEP_PAIRS: {undeclared}. "
+        f"Add them so the checks below cover them."
+    )
+
+
+@pytest.mark.parametrize("yaml_name", sorted(ARM_SWEEP_PAIRS))
+def test_arm_sweep_applies_to_its_registered_pair(yaml_name):
+    """Every arm of every arm sweep resolves and applies against its own dataset/task pair.
+
+    This is the check that would have caught a step of the overlap ladder naming an obs column that
+    the registry's dataset file does not point at, or a max_seq_len key that moved.
+    """
+    dataset, task = ARM_SWEEP_PAIRS[yaml_name]
+    with (CONFIG_DIR / "sweeps" / yaml_name).open() as f:
+        yaml_config = yaml.safe_load(f)
+
+    base = resolve_config(dataset, task, registry_path=REGISTRY)
+    sweep_config, sweep_params = build_sweep_config(
+        yaml_config, prediction_task=base.dataset.prediction_task, model_type="CombinedModel"
+    )
+    arms = load_arms(yaml_config, sweep_config)
+    assert arms, f"{yaml_name} declares no arms"
+
+    for arm_name in sweep_config["parameters"][ARM_PARAM]["values"]:
+        trial = {k: v["values"][0] for k, v in sweep_config["parameters"].items()}
+        trial[ARM_PARAM] = arm_name
+        cfg, applied = apply_sweep_config(
+            base.clone(), "robustness", trial, model_type="CombinedModel",
+            sweep_params=sweep_params, arms=arms,
+        )
+        for key, expected in arms[arm_name].items():
+            assert get_dotted(cfg, key) == expected, f"{yaml_name} {arm_name}: {key} not applied"
+        # A sample_key that is empty would train on nothing; one that is a bare string would be
+        # iterated character by character by prepare_geome_dataset.
+        assert isinstance(cfg.dataset.sample_key, list) and cfg.dataset.sample_key, (
+            f"{yaml_name} {arm_name}: dataset.sample_key must be a non-empty list, "
+            f"got {cfg.dataset.sample_key!r}"
+        )
+
+
+@pytest.mark.parametrize("yaml_name", sorted(ARM_SWEEP_PAIRS))
+def test_arm_sweep_gives_every_trial_its_own_checkpoint(yaml_name):
+    """No two trials of an arm sweep may write the same checkpoint filename.
+
+    get_model_filename_prefix keys on dataset.name, task, level and seed. An arm that forgets to
+    set dataset.name silently overwrites its predecessor's checkpoint, and the sweep then reports
+    metrics for models that no longer exist on disk.
+    """
+    from interscale.tl.utils import get_model_filename_prefix
+
+    dataset, task = ARM_SWEEP_PAIRS[yaml_name]
+    with (CONFIG_DIR / "sweeps" / yaml_name).open() as f:
+        yaml_config = yaml.safe_load(f)
+    base = resolve_config(dataset, task, registry_path=REGISTRY)
+    sweep_config, sweep_params = build_sweep_config(
+        yaml_config, prediction_task=base.dataset.prediction_task, model_type="CombinedModel"
+    )
+    arms = load_arms(yaml_config, sweep_config)
+
+    seeds = sweep_config["parameters"].get("optim.seed", {}).get("values", [None])
+    prefixes = {}
+    for arm_name in sweep_config["parameters"][ARM_PARAM]["values"]:
+        for seed in seeds:
+            trial = {k: v["values"][0] for k, v in sweep_config["parameters"].items()}
+            trial[ARM_PARAM] = arm_name
+            if seed is not None:
+                trial["optim.seed"] = seed
+            cfg, _ = apply_sweep_config(
+                base.clone(), "robustness", trial, model_type="CombinedModel",
+                sweep_params=sweep_params, arms=arms,
+            )
+            prefixes[(arm_name, seed)] = get_model_filename_prefix(cfg, True, True)
+
+    collisions = {v for v in prefixes.values() if list(prefixes.values()).count(v) > 1}
+    assert not collisions, f"{yaml_name}: trials sharing a checkpoint filename: {collisions}"

From ace214394abe9e8dedf6d981ed8879409560a714 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Mon, 31 Aug 2026 15:24:34 +0200
Subject: [PATCH 09/14] implement gene wise masking + pct_mask_nodes to
 mask_percentage

---
 CHANGELOG.md                                  |  39 ++
 CLAUDE.md                                     |  22 +-
 src/interscale/config/__init__.py             |  31 +-
 src/interscale/config/dataset_config.py       |  16 +-
 src/interscale/geome_dataloader.py            |  47 ++-
 src/interscale/main.py                        |   3 +-
 src/interscale/main_sweep.py                  |   3 +-
 src/interscale/model/base/_base_model.py      |   6 +-
 src/interscale/model/combined_model.py        |   6 +-
 src/interscale/model/global_model.py          |   3 +-
 src/interscale/model/local_model.py           |   3 +-
 .../module/base/_base_global_module.py        |  92 ++++-
 .../module/base/_base_local_module.py         |  14 +-
 src/interscale/module/base/_base_module.py    |  44 +-
 .../module/combined_module/combined_module.py |  21 +-
 .../dual_decoder_combined_module.py           |  49 ++-
 src/interscale/tl/__init__.py                 |  16 +-
 src/interscale/tl/masking.py                  | 234 ++++++++++-
 src/interscale/train/_training.py             |   2 +-
 src/interscale/train/_trainingplans.py        | 152 +++++--
 tests/test_gene_masking.py                    | 390 ++++++++++++++++++
 tests/test_geome_dataloader.py                |   4 +-
 tests/test_global_pca_persistence.py          | 118 ++++++
 tests/test_sweep_config.py                    |  12 +-
 24 files changed, 1208 insertions(+), 119 deletions(-)
 create mode 100644 tests/test_gene_masking.py
 create mode 100644 tests/test_global_pca_persistence.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 537ed6e..287e41d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,45 @@ and this project adheres to [Semantic Versioning][].
 [keep a changelog]: https://keepachangelog.com/en/1.1.0/
 [semantic versioning]: https://semver.org/spec/v2.0.0.html
 
+## [Unreleased]
+
+### Added
+
+- Gene-wise (per-entry) masking for the node-level reconstruction task, alongside the existing
+  whole-cell masking, selected by `dataset.mask_strategy` (`"node"` | `"gene"`). Under `"gene"`
+  a Bernoulli subset of `(cell, gene)` entries is blanked in every cell and the loss and every
+  regression metric are restricted to those entries. Registered as the `node_reg_genemask` task
+  for `legnini23` and `melton25`, with `config_files/sweeps/mask_granularity_legnini.yaml`
+  running the granularity/rate ablation.
+
+### Changed
+
+- `tl.masking.MASK_VALUE` is now **-1**, was 0. Masked positions must be distinguishable from
+  real measurements: ~62% of legnini23's `log1p_norm` entries are already exactly 0 and 648 of
+  its cells are all-zero, so a zero fill made the corruption invisible under gene masking.
+  Measured at gene rate 0.25 over 3 seeds, fill 0 scored `val_concordance_corr` 0.0696 ±0.0113
+  against -1's 0.0799 ±0.0046 — better on every seed, with 2.5x less run-to-run spread. A
+  configurable fill (`zero`/`sentinel`/`learned`, including a trainable per-gene `[MASK]` token)
+  was trialled and removed: the learnable token tied with the fixed -1 to three decimals, so the
+  option carried no information. **This also changes the cell-masking path**, whose published
+  runs used 0.
+- `dataset.pct_mask_nodes` and `dataset.pct_mask_genes` are replaced by a single
+  `dataset.mask_percentage` (and `GraphAnnDataModule`/`BaseModule` take one `mask_percentage`
+  argument). `mask_strategy` already says whether a masked unit is a cell or an entry, so a
+  second rate was always the inert half of a pair — and setting the wrong one silently produced
+  a run with no masking. `GraphAnnDataModule.mask_rate` and `BaseModule.mask_rate` are gone with
+  it. **Every config using `pct_mask_nodes` must be renamed**, including saved wandb sweeps.
+- `tl.masking.apply_mask` takes `mask_strategy` explicitly instead of inferring it from whether
+  the batch carries a `gene_mask`, and `GlobalModule._process_batch_for_metrics` gates on the
+  strategy the same way. A stale `gene_mask` can no longer hijack a cell-masking run, so
+  `GraphAnnDataModule._clear_gene_mask` is removed.
+- `tl.masking.apply_mask` and `BaseModule._common_step_masking` now return a third value, the
+  entry mask; every `_common_step` returns a sixth value carrying it through to the training plan.
+  `GlobalModule._process_batch_for_metrics` returns it as a third value.
+  `DualDecoderCombinedModule.compute_separate_losses` takes it as an optional fifth argument.
+- `LocalModule._common_step` returned a 4-tuple where the training plan unpacked 5; it now
+  matches the 6-tuple contract of the other modules.
+
 ## [0.0.1]
 
 initial release
diff --git a/CLAUDE.md b/CLAUDE.md
index fb42c99..2ad10e3 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -75,7 +75,7 @@ Three concrete models subclass `BaseModel`, differing in which components they i
 ### Module hierarchy (`interscale.module`)
 
 Mirrors the model hierarchy at the `pytorch_lightning`/`LightningModule` level:
-- `BaseModule` (`module/base/_base_module.py`) owns the decoder (`interscale.nn`: `LinearDecoder`, `LinearLSEDecoder`, `NonLinearDecoder`, or `None` when a wrapping combined module owns decoding), and shared node-masking logic (`tl.masking.apply_mask`) used for self-supervised/robustness training.
+- `BaseModule` (`module/base/_base_module.py`) owns the decoder (`interscale.nn`: `LinearDecoder`, `LinearLSEDecoder`, `NonLinearDecoder`, or `None` when a wrapping combined module owns decoding), and shared masking logic (`tl.masking.apply_mask`, via `_common_step_masking`) used for self-supervised/robustness training. `_common_step_masking` returns `(batch_masked, mask_idx, entry_mask)`; every `_common_step` returns `(local_emb, global_emb, y_pred, y_true, attn, entry_mask)`, and `entry_mask` is what restricts the loss and metrics under gene masking.
 - `module/local_modules/` — local (graph) encoders: `GCN`, `GIN`, `Precomputed` (precomputed embeddings), `SCVI` (SCVI-based encoder).
 - `module/global_modules/` — transformer-based global encoder (`TransformerNodeEncoderHook`) plus supporting transformer encoder/layer/utils, encoding whole-sample (long-range) context across cells.
 - `module/combined_module/` — `CombinedModule` and `DualDecoderCombinedModule` compose a local module + global module sequentially: local embeddings feed into the global (transformer) component, and predictions/attention/CLS tokens are extracted from there (see `CombinedModel.get_model_output` in `model/combined_model.py` for the full inference/evaluation flow, including how attention matrices and horizontal/vertical CLS tokens are extracted and padded to `max_seq_len`).
@@ -86,7 +86,25 @@ Mirrors the model hierarchy at the `pytorch_lightning`/`LightningModule` level:
 - Iterates per-sample/library (`cfg.dataset.sample_key`), builds a spatial neighbor graph per sample (`squidpy`-style `spatial_neigbors_kwargs`, converted to an edge index via `geome.transforms.AddEdgeIndex`), and yields one PyG `Data` object per sample.
 - Splits data by `cfg.dataset.split_key` (an `adata.obs` column that must contain `train`/`val`, optionally `test`) — this must exist in the AnnData before calling `prepare_geome_dataset`.
 - Handles both classification (one-hot encodes `prediction_obs`) and regression prediction tasks, and optionally attaches precomputed embeddings (`cfg.model.global_component.parameters.type_gex_embedding == "Precomputed"`) from `adata.obsm`.
-- `interscale.geome_dataloader.GraphAnnDataModule` wraps the resulting `list[Data]` splits into a `LightningDataModule`. For node-level learning it randomly masks a `pct_mask_nodes` fraction of nodes per graph (at least 1) for each dataloader construction — this is the masking scheme used for self-supervised node reconstruction/robustness experiments referenced in `config_files/legnini_example.yaml`'s `pct_mask_nodes` and `dataset.segmentation_robustness`.
+- `interscale.geome_dataloader.GraphAnnDataModule` wraps the resulting `list[Data]` splits into a `LightningDataModule`. For node-level learning it corrupts the reconstruction input for each dataloader construction (redrawn every epoch by `NodeMaskResampleCallback`) — this is the masking scheme used for self-supervised node reconstruction/robustness experiments referenced in `config_files/legnini_example.yaml`'s `pct_mask_nodes` and `dataset.segmentation_robustness`.
+
+### Masking (`tl/masking.py`)
+
+Two config axes, both under `cfg.dataset`. Read the module docstring of `tl/masking.py` before changing any of them.
+
+| key | values | meaning |
+| --- | --- | --- |
+| `mask_strategy` | `node` (default) / `gene` | Granularity. `node` blanks whole cells (`data.mask` `[N]`) and scores all G genes of them. `gene` blanks individual `(cell, gene)` entries (`data.gene_mask` `[N, G]`, every cell a target) and scores **only those entries**. |
+| `mask_percentage` | float | The Bernoulli rate, per cell under `node` and per entry under `gene` — one key, because the strategy already says what a unit is. Not comparable across strategies at equal values. `_validate_masking` rejects a regression config with a rate outside (0, 1]. |
+
+Masked positions are always filled with `tl.masking.MASK_VALUE`, which is **-1** and is not
+configurable. Two things are easy to get wrong here:
+
+- **The strategy is passed explicitly, never inferred.** `apply_mask` and `_process_batch_for_metrics` both take/consult `mask_strategy` rather than checking whether a `gene_mask` attribute exists, so a stale mask on a reused `Data` object cannot turn a cell-masking run into a gene-masking one. That is why the dataloader has no clear-the-attribute helper.
+- **The loss must be restricted to the masked entries under `gene`.** That is what `entry_mask` threads through `_common_step` → `TrainingPlan._regression_metrics` → `tl.masking.masked_loss`, and metrics go through `train._trainingplans.masked_regression_metrics` (same metric *names*, computed over masked entries only, verified against torchmetrics in `tests/test_gene_masking.py`). Without it the unmasked entries — which were handed to the model as input — make the objective the identity map. **Metric names are shared but the entry sets differ, so a `val_r2` from a `node` run and one from a `gene` run are not comparable as absolute numbers.**
+- **`MASK_VALUE` must stay outside the layer's range; do not set it back to 0.** ~62% of legnini23's `log1p_norm` entries are already exactly 0 (and 648 cells are all-zero), so a zero fill makes a masked position indistinguishable from a real measurement — under `gene` the corruption becomes invisible and the model cannot tell which entries it is being asked for. Cell masking was less exposed because an all-zero *row* is still mostly a pattern. Measured at gene rate 0.25 over 3 seeds: fill 0 gave `val_concordance_corr` 0.0696 ±0.0113, fill -1 gave 0.0799 ±0.0046 — better on every seed and with 2.5× less run-to-run spread, below the cell-masking baseline's own CV. A learnable per-gene `[MASK]` token (GraphMAE's design) was tried and tied with the fixed -1 to three decimals — it drifts only ~0.02–0.035 from any init — so it was removed rather than kept as an option.
+
+The registered ablation pairs are `<dataset>/node_reg` (cell) vs `<dataset>/node_reg_genemask` (gene) — note each carries its own `dataset.name` so the two arms do not share a checkpoint prefix. `config_files/sweeps/mask_granularity_legnini.yaml` runs the granularity/rate grid.
 
 ### Preprocessing / robustness utilities (`interscale.pp`)
 
diff --git a/src/interscale/config/__init__.py b/src/interscale/config/__init__.py
index 5296fe3..4c16fac 100644
--- a/src/interscale/config/__init__.py
+++ b/src/interscale/config/__init__.py
@@ -82,7 +82,7 @@ def _coerce_override_values(cfg, overrides):
 
     ``merge_from_list`` rejects an int for a float key outright (yacs only tolerates a type
     mismatch when one side is ``None``). YAML writes ``0`` as an int, so an override of
-    ``dataset.pct_mask_nodes: 0`` against the ``0.2`` default would raise even though the
+    ``dataset.mask_percentage: 0`` against the ``0.2`` default would raise even though the
     value is perfectly valid. Promote it instead of making callers write ``0.0``.
 
     ``overrides`` is the flat ``[key, value, key, value, ...]`` form ``merge_from_list`` takes.
@@ -111,6 +111,33 @@ def _coerce_override_values(cfg, overrides):
     return coerced
 
 
+def _validate_masking(cfg):
+    """Reject masking settings that would silently train on an uncorrupted input.
+
+    A reconstruction task with ``mask_percentage`` at 0 has the identity map as its solution, and
+    every metric looks excellent, so it fails loudly instead.
+
+    Raises
+    ------
+    ValueError
+        If ``mask_strategy`` is unknown, or ``mask_percentage`` is not a rate in (0, 1] for a
+        regression task.
+    """
+    from interscale.tl.masking import MASK_STRATEGIES
+
+    if cfg.dataset.mask_strategy not in MASK_STRATEGIES:
+        raise ValueError(f"dataset.mask_strategy must be one of {MASK_STRATEGIES}, got {cfg.dataset.mask_strategy!r}.")
+
+    if "regression" not in cfg.dataset.prediction_task:
+        return
+
+    if not 0 < cfg.dataset.mask_percentage <= 1:
+        raise ValueError(
+            f"dataset.mask_percentage must be in (0, 1] for a regression task, got "
+            f"{cfg.dataset.mask_percentage}. With no masking a reconstruction target is its own input."
+        )
+
+
 def _validate_optim(cfg):
     """Reject configs whose training-length settings stop a run inside the LR warm-up.
 
@@ -165,6 +192,7 @@ def load_config(cfg_path=None, overrides=None):
     # unconditionally, so None used to raise AttributeError too.
     if not cfg_paths and not overrides:
         _validate_optim(cfg)
+        _validate_masking(cfg)
         cfg.freeze()
         return cfg
 
@@ -183,5 +211,6 @@ def load_config(cfg_path=None, overrides=None):
         cfg.merge_from_list(_coerce_override_values(cfg, overrides))
 
     _validate_optim(cfg)
+    _validate_masking(cfg)
     cfg.freeze()
     return cfg
diff --git a/src/interscale/config/dataset_config.py b/src/interscale/config/dataset_config.py
index 8944289..4b91144 100644
--- a/src/interscale/config/dataset_config.py
+++ b/src/interscale/config/dataset_config.py
@@ -9,7 +9,9 @@ def get_dataset_cfg(cfg):
     sample_key: list of keys in adata.obs to split the data into PyG Data objects (e.i. sliding_window, FOV, sample etc)
     num_features: number of gene expressions (added in prepare_geome_function)
     num_features: number of classes in prediction_obs (added in prepare_geome_function)
-    pct_mask_nodes: percentage of single nodes to mask during training in a graph
+    mask_strategy: granularity of the reconstruction corruption, "node" or "gene"
+    mask_percentage: Bernoulli masking probability -- per cell under mask_strategy "node",
+        per (cell, gene) entry under "gene"
     """
     cfg.dataset = CN()
 
@@ -30,7 +32,17 @@ def get_dataset_cfg(cfg):
     cfg.dataset.num_features = -1
     cfg.dataset.num_classes = -1
 
-    cfg.dataset.pct_mask_nodes = 0.2
+    # Reconstruction corruption. "node" blanks whole cells and scores all G genes of them;
+    # "gene" blanks individual (cell, gene) entries in every cell and scores those entries only.
+    # See interscale.tl.masking for why the two objectives behave so differently -- under "node"
+    # the target cell contributes nothing about itself, so the population mean is already a
+    # strong solution. "node" stays the default so every existing config is unchanged.
+    cfg.dataset.mask_strategy = "node"
+    # One rate, not one per strategy: mask_strategy already says what a unit is, so a second key
+    # would only ever be the inert half of the pair -- and setting the wrong one silently gives a
+    # run with no masking. Note the two strategies are not comparable at equal values: a per-entry
+    # rate is a different quantity from a per-cell one (MAE/GraphMAE use 0.25-0.75 for features).
+    cfg.dataset.mask_percentage = 0.2
 
     # Segmentation robustness parameters
     cfg.dataset.segmentation_robustness = None  # [node_fraction, overflow_fraction] or None
diff --git a/src/interscale/geome_dataloader.py b/src/interscale/geome_dataloader.py
index 1897b00..7446b09 100644
--- a/src/interscale/geome_dataloader.py
+++ b/src/interscale/geome_dataloader.py
@@ -12,10 +12,10 @@
 VALID_SPLIT = {"node", "graph"}
 
 # TODO: Fix dataloader
-import random
-
 import torch
 
+from interscale.tl.masking import MASK_STRATEGIES, sample_gene_mask, sample_node_mask
+
 
 class GraphAnnDataModule(pl.LightningDataModule):
     """Lightning DataModule for graph data."""
@@ -25,7 +25,8 @@ def __init__(
         datas: Sequence[Sequence[Data]] | None = None,
         batch_size: int = 1,
         num_workers: int = 1,
-        pct_mask_nodes: float = 0.5,
+        mask_percentage: float = 0.5,
+        mask_strategy: Literal["node", "gene"] = "node",
         learning_type: Literal["node", "graph"] = "node",
     ):
         """Manages loading and sampling schemes before loading to GPU.
@@ -36,6 +37,13 @@ def __init__(
             List of train, val (and test) data to be loaded. Defaults to None.
         batch_size (int, optional): The batch size. Defaults to 1.
         num_workers (int, optional): The number of workers. Defaults to 1.
+        mask_percentage (float, optional): Bernoulli masking probability, per cell under
+            `mask_strategy="node"` and per (cell, gene) entry under `"gene"`. One argument rather
+            than one per strategy: the strategy already says what a unit is. Defaults to 0.5.
+        mask_strategy (Literal["node", "gene"], optional): Granularity of the corruption.
+            "node" blanks whole cells (loss over all G genes of the masked cells); "gene" blanks
+            individual (cell, gene) entries in every cell (loss over those entries only). See
+            `interscale.tl.masking` for why the two behave so differently. Defaults to "node".
         learning_type (Literal["node", "graph"], optional): The type of learning to be performed.
             If "graph" is selected, `batch_size` means the number of graphs and `datas` is expected to be a list of Data.
             If "node" is selected, `batch_size` means the number of nodes and `datas` is expected to be a list of Data objects
@@ -59,7 +67,10 @@ def __init__(
         if learning_type not in VALID_SPLIT:
             raise ValueError("Learning type must be one of %r." % VALID_SPLIT)
         self.learning_type = learning_type
-        self.pct_mask_nodes = pct_mask_nodes
+        if mask_strategy not in MASK_STRATEGIES:
+            raise ValueError("mask_strategy must be one of %r." % (MASK_STRATEGIES,))
+        self.mask_strategy = mask_strategy
+        self.mask_percentage = mask_percentage
         self.first_time = True
 
     def _nodewise_setup(self, stage: str | None) -> None:
@@ -130,21 +141,25 @@ def _get_dataloader(self, dataloader):
         return dataloader
 
     def _assign_random_mask(self, data: BaseData) -> None:
-        """Overwrites `data.mask` in place with a fresh Bernoulli draw at rate `pct_mask_nodes`.
+        """Overwrites the mask attributes of `data` in place with a fresh Bernoulli draw.
 
-        Each node is masked independently with probability `pct_mask_nodes`, so the expected
-        masked fraction is the same for every graph regardless of its size. An earlier version
-        derived a single absolute node count from the smallest graph in the split and applied it
-        to every graph, which made the realised rate depend both on a graph's size and on which
-        split it landed in.
+        Under `mask_strategy="node"` only `data.mask` `[N]` is written: each cell is masked
+        independently with probability `mask_percentage`.
 
-        This is the single seam where a future non-uniform sampling strategy (e.g. downweighting
-        nodes that were already masked in previous epochs) would be substituted in.
+        Under `mask_strategy="gene"` a second attribute `data.gene_mask` `[N, G]` is written,
+        with each (cell, gene) entry drawn independently at `mask_percentage`. `data.mask` is then
+        the row-wise OR of it -- i.e. "this cell is a supervision target" -- which is what the
+        rest of the pipeline (padding, `_process_batch_for_metrics`) keys on. 
+
+        `gene_mask` is a node-level attribute of shape `[num_nodes, ...]`, so PyG collates it by
+        concatenating along dim 0 exactly like `x` -- no custom `__cat_dim__` needed. 
         """
-        mask = torch.rand(data.num_nodes) < self.pct_mask_nodes
-        if not mask.any():  # must mask at least one node
-            mask[random.randrange(data.num_nodes)] = True
-        data.mask = mask
+        if self.mask_strategy == "gene":
+            gene_mask = sample_gene_mask(data.num_nodes, data.x.shape[1], self.mask_percentage)
+            data.gene_mask = gene_mask
+            data.mask = gene_mask.any(dim=1)
+        else:
+            data.mask = sample_node_mask(data.num_nodes, self.mask_percentage)
 
     def resample_train_mask(self) -> None:
         """Redraws the masked node set for every training graph, in place.
diff --git a/src/interscale/main.py b/src/interscale/main.py
index a2f6593..930ed92 100644
--- a/src/interscale/main.py
+++ b/src/interscale/main.py
@@ -84,7 +84,8 @@ def main(cfg, model_type):
         datas=pyg_data_list,
         num_workers=1,
         batch_size=int(cfg.dataset.batch_size),
-        pct_mask_nodes=cfg.dataset.pct_mask_nodes,
+        mask_percentage=cfg.dataset.mask_percentage,
+        mask_strategy=cfg.dataset.mask_strategy,
         learning_type=cfg.dataset.prediction_level,
     )
 
diff --git a/src/interscale/main_sweep.py b/src/interscale/main_sweep.py
index 1113f49..af3b1ce 100644
--- a/src/interscale/main_sweep.py
+++ b/src/interscale/main_sweep.py
@@ -208,7 +208,8 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None, arms=None
         datas=pyg_data_list,
         num_workers=1,
         batch_size=int(cfg.dataset.batch_size),
-        pct_mask_nodes=cfg.dataset.pct_mask_nodes,
+        mask_percentage=cfg.dataset.mask_percentage,
+        mask_strategy=cfg.dataset.mask_strategy,
         learning_type=cfg.dataset.prediction_level,
     )
     print_memory_usage("After datamodule creation")
diff --git a/src/interscale/model/base/_base_model.py b/src/interscale/model/base/_base_model.py
index 9270f71..b6c5fa1 100644
--- a/src/interscale/model/base/_base_model.py
+++ b/src/interscale/model/base/_base_model.py
@@ -458,7 +458,8 @@ def _register_local_component(self) -> LocalModule:
                 n_embed=self.n_embed,
                 decoder_type=self._cfg.model.decoder.type,
                 dropout_decoder=self._cfg.model.decoder.dropout_decoder,
-                pct_mask_nodes=self._cfg.dataset.pct_mask_nodes,
+                mask_percentage=self._cfg.dataset.mask_percentage,
+                mask_strategy=self._cfg.dataset.mask_strategy,
                 n_layers=self._cfg.model.local_component.parameters.num_layers,
                 hidden_dim=self._cfg.model.local_component.parameters.hidden_dim,
                 dropout_local=self._cfg.model.local_component.parameters.dropout_local,
@@ -487,7 +488,8 @@ def _register_global_component(self) -> GlobalModule:
                 n_embed=self.n_embed,
                 decoder_type=self._cfg.model.decoder.type,
                 dropout_decoder=self._cfg.model.decoder.dropout_decoder,
-                pct_mask_nodes=self._cfg.dataset.pct_mask_nodes,
+                mask_percentage=self._cfg.dataset.mask_percentage,
+                mask_strategy=self._cfg.dataset.mask_strategy,
                 max_seq_len=self._cfg.model.global_component.parameters.max_seq_len,
                 n_heads=self._cfg.model.global_component.parameters.n_heads,
                 dropout_global=self._cfg.model.global_component.parameters.dropout_global,
diff --git a/src/interscale/model/combined_model.py b/src/interscale/model/combined_model.py
index 6256813..477d3e5 100644
--- a/src/interscale/model/combined_model.py
+++ b/src/interscale/model/combined_model.py
@@ -36,7 +36,8 @@ def __init__(
                 decoder_type=None,  # Container doesn't need its own decoder, only submodules do
                 dropout_decoder=self._cfg.model.decoder.dropout_decoder,
                 decoder_hidden_dims=self._cfg.model.decoder.hidden_dims,
-                pct_mask_nodes=self._cfg.dataset.pct_mask_nodes,
+                mask_percentage=self._cfg.dataset.mask_percentage,
+                mask_strategy=self._cfg.dataset.mask_strategy,
             )
         else:
             self.module = CombinedModule(
@@ -47,7 +48,8 @@ def __init__(
                 decoder_type=None,  # Container doesn't need its own decoder, only global module does
                 dropout_decoder=self._cfg.model.decoder.dropout_decoder,
                 decoder_hidden_dims=self._cfg.model.decoder.hidden_dims,
-                pct_mask_nodes=self._cfg.dataset.pct_mask_nodes,
+                mask_percentage=self._cfg.dataset.mask_percentage,
+                mask_strategy=self._cfg.dataset.mask_strategy,
             )
 
         self._model_summary_string = self._model_summary_string + self.module.get_model_summary()
diff --git a/src/interscale/model/global_model.py b/src/interscale/model/global_model.py
index 06fde7e..1b0e556 100644
--- a/src/interscale/model/global_model.py
+++ b/src/interscale/model/global_model.py
@@ -38,7 +38,8 @@ def __init__(
             decoder_type=self._cfg.model.decoder.type,
             dropout_decoder=self._cfg.model.decoder.dropout_decoder,
             decoder_hidden_dims=self._cfg.model.decoder.hidden_dims,
-            pct_mask_nodes=self._cfg.dataset.pct_mask_nodes,
+            mask_percentage=self._cfg.dataset.mask_percentage,
+            mask_strategy=self._cfg.dataset.mask_strategy,
             type_gex_embedding=self._cfg.model.global_component.parameters.type_gex_embedding,
         )
 
diff --git a/src/interscale/model/local_model.py b/src/interscale/model/local_model.py
index 7ab46ce..fee3957 100644
--- a/src/interscale/model/local_model.py
+++ b/src/interscale/model/local_model.py
@@ -33,7 +33,8 @@ def __init__(
             decoder_type=self._cfg.model.decoder.type,
             dropout_decoder=self._cfg.model.decoder.dropout_decoder,
             decoder_hidden_dims=self._cfg.model.decoder.hidden_dims,
-            pct_mask_nodes=self._cfg.dataset.pct_mask_nodes,
+            mask_percentage=self._cfg.dataset.mask_percentage,
+            mask_strategy=self._cfg.dataset.mask_strategy,
         )
 
     @torch.inference_mode()
diff --git a/src/interscale/module/base/_base_global_module.py b/src/interscale/module/base/_base_global_module.py
index d91b834..eb9632b 100644
--- a/src/interscale/module/base/_base_global_module.py
+++ b/src/interscale/module/base/_base_global_module.py
@@ -1,6 +1,7 @@
 from abc import abstractmethod
 from typing import Literal
 
+import numpy as np
 import torch
 from sklearn.decomposition import NMF, PCA
 
@@ -17,6 +18,24 @@ def __init__(self, **base_module_kwargs):
 
         if self.type_gex_embedding == "PCA":
             self.pca = PCA(n_components=self.n_embed)
+            # A fitted sklearn estimator is not part of `module.state_dict()`, and
+            # `BaseModel.save` persists nothing else -- so a GlobalModel reloaded for inference
+            # would arrive with an UNFITTED pca and silently refit it on the first evaluation
+            # batch. The transformer's weights were learned on the basis fitted to the first
+            # TRAINING batch; a basis refitted elsewhere differs by rotation and by component
+            # sign, so the reloaded model would decode a different space than it was trained on
+            # and produce attention that means nothing. These buffers carry the fit through the
+            # checkpoint.
+            #
+            # Registered unconditionally for the PCA branch (not lazily on first fit) because a
+            # buffer that does not exist at construction time cannot receive a value from
+            # `load_state_dict`, which is exactly when it is needed. Older checkpoints simply
+            # have no entry for them; `BaseModel.load` uses strict=False, so they stay zeroed
+            # and `pca_fitted_` stays False, reproducing the previous refit-on-load behaviour
+            # rather than failing.
+            self.register_buffer("pca_mean_", torch.zeros(self.n_input))
+            self.register_buffer("pca_components_", torch.zeros(self.n_embed, self.n_input))
+            self.register_buffer("pca_fitted_", torch.zeros(1, dtype=torch.bool))
         elif self.type_gex_embedding == "NMF":
             self.nmf = NMF(n_components=self.n_embed, init="random", random_state=0)
         elif self.type_gex_embedding == "Precomputed":
@@ -51,12 +70,23 @@ def create_gex_embedding(self, embeddings: torch.Tensor, type: Literal["PCA", "N
             Size: [N, E]
         """
         if type == "PCA":
-            # Fit PCA only once (on first batch), then use transform for subsequent batches
-            # This avoids expensive refitting on every batch during training
-            if not hasattr(self.pca, "components_"):
-                return self.pca.fit_transform(embeddings)
-            else:
-                return self.pca.transform(embeddings)
+            # Fit PCA only once (on the first batch that arrives with no fit available), then
+            # project every later batch through the stored basis. Two sources of a fit, in
+            # order: the buffers restored from a checkpoint, then this process's own first
+            # batch. Checking the buffers FIRST is what makes a reloaded model reproduce the
+            # basis it was trained on instead of refitting on evaluation data.
+            # `_common_step` hands this a numpy array but `GlobalModel.get_model_output` hands
+            # it `batch.x` straight off the batch, which is a (possibly CUDA) tensor.
+            if isinstance(embeddings, torch.Tensor):
+                embeddings = embeddings.detach().cpu().numpy()
+            if not bool(self.pca_fitted_):
+                self.pca.fit(embeddings)
+                self.pca_mean_.copy_(torch.as_tensor(self.pca.mean_, dtype=self.pca_mean_.dtype))
+                self.pca_components_.copy_(
+                    torch.as_tensor(self.pca.components_, dtype=self.pca_components_.dtype)
+                )
+                self.pca_fitted_.fill_(True)
+            return self._pca_transform(embeddings)
         elif type == "NMF":
             if not hasattr(self.nmf, "components_"):
                 return self.nmf.fit_transform(embeddings)
@@ -65,6 +95,19 @@ def create_gex_embedding(self, embeddings: torch.Tensor, type: Literal["PCA", "N
         else:
             raise ValueError(f"Invalid embedding type: {type}")
 
+    def _pca_transform(self, embeddings):
+        """Project onto the stored PCA basis: ``(X - mean_) @ components_.T``.
+
+        Written out rather than delegated to ``self.pca.transform`` so that the projection
+        depends only on the two buffers, which are the only part of the fit that survives a
+        checkpoint round trip. ``sklearn``'s own ``transform`` would additionally require the
+        estimator's private fitted attributes to be present, which after a reload they are not.
+        Equivalent to it for ``whiten=False``, which is the default this module constructs.
+        """
+        mean = self.pca_mean_.detach().cpu().numpy()
+        components = self.pca_components_.detach().cpu().numpy()
+        return (np.asarray(embeddings, dtype=np.float64) - mean) @ components.T
+
     def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, pad_index_nodes, mask_idx_tensor):
         """Process batch to extract y_true and adjusted_mask_idx for metrics calculation.
 
@@ -90,6 +133,10 @@ def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, p
             Ground truth values
         adjusted_mask_idx: torch.Tensor [N_masked nodes]
             Adjusted indices for masked nodes
+        entry_mask: torch.Tensor [N_included_nodes, F] | None
+            The batch's `gene_mask` gathered and reordered exactly like `y_true`, so that
+            `entry_mask[adjusted_mask_idx]` lines up entry-for-entry with the scored predictions.
+            `None` whenever the batch carries no gene mask (cell masking, or classification).
         """
         assert prediction_level == "node", "Node specific retrieval only necessary for node-level prediction."
 
@@ -108,6 +155,14 @@ def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, p
         adjusted_mask_idx_list = []
         y_true_list = []
 
+        # Gathered in lockstep with y_true below. Gated on the configured strategy, not merely on
+        # the attribute being present, so a stale `gene_mask` on a reused Data object cannot turn
+        # a cell-masking run into a gene-masking one. Only regression has entries to mask; a
+        # classification target is a label per cell, not a gene vector.
+        use_gene_mask = self.mask_strategy == "gene" and "regression" in prediction_task
+        gene_mask = getattr(batch, "gene_mask", None) if use_gene_mask else None
+        entry_mask_list = [] if gene_mask is not None else None
+
         for i in range(nr_batches):
             batch_start = batch_starts[i].item()
             batch_end = batch_ends[i].item()
@@ -123,6 +178,8 @@ def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, p
                     y_true_list.append(batch.y[mask][pad_index_nodes[i]])
                 elif "regression" in prediction_task:
                     y_true_list.append(batch.x[mask][pad_index_nodes[i]])
+                if entry_mask_list is not None:
+                    entry_mask_list.append(gene_mask[mask][pad_index_nodes[i]])
                 continue
 
             # Create pad_indices tensor once
@@ -149,9 +206,12 @@ def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, p
                 y_true_list.append(batch.x[mask][pad_index_nodes[i]])
             else:
                 raise Exception("Choose a valid prediction task (classification or regression).")
+            if entry_mask_list is not None:
+                entry_mask_list.append(gene_mask[mask][pad_index_nodes[i]])
 
         # Concatenate results
         y_true = torch.cat(y_true_list, dim=0)
+        entry_mask = torch.cat(entry_mask_list, dim=0) if entry_mask_list else None
         adjusted_mask_idx = (
             torch.cat(adjusted_mask_idx_list, dim=0)
             if adjusted_mask_idx_list
@@ -177,7 +237,12 @@ def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, p
                 f"Mismatch: max(adjusted_mask_idx): {adjusted_mask_idx.max()}, len(y_true): {len(y_true)}"
             )
 
-        return y_true, adjusted_mask_idx
+        if entry_mask is not None:
+            assert entry_mask.shape == y_true.shape, (
+                f"Mismatch: entry_mask.shape: {tuple(entry_mask.shape)}, y_true.shape: {tuple(y_true.shape)}"
+            )
+
+        return y_true, adjusted_mask_idx, entry_mask
 
     # def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, pad_index_nodes, mask_idx_tensor):
     #     """Process batch to extract y_true and adjusted_mask_idx for metrics calculation.
@@ -290,9 +355,13 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
             Size: [N, C] (classification) or [N, F] (regression) with SEQ_LEN_MASK for padding nodes.
         y_true: torch.Tensor
             Size: [N, C] (classification) or [N, F] (regression) with SEQ_LEN_MASK for padding nodes.
+        attn_matrix: torch.Tensor
+            Stacked per-layer attention weights.
+        entry_mask: torch.Tensor | None
+            Size: [N, F] under gene masking, marking the entries the loss is scored on.
         """
         # Mask nodes  - before GEX embedding because otherwise embedding contains information about masked nodes
-        batch_masked, mask_idx = self._common_step_masking(batch)
+        batch_masked, mask_idx, _ = self._common_step_masking(batch)
         if hasattr(batch_masked, "embeddings"):
             embedding = batch_masked.embeddings
         else:
@@ -317,18 +386,21 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
 
         if prediction_task == "classification" and prediction_level == "graph":
             y_true = batch.y[batch.ptr[:-1]]
+            entry_mask = None
         else:
-            y_true, adjusted_mask_idx = self._process_batch_for_metrics(
+            y_true, adjusted_mask_idx, entry_mask = self._process_batch_for_metrics(
                 batch, prediction_task, prediction_level, pad_index_nodes, mask_idx
             )
             y_pred = y_pred[adjusted_mask_idx]
             y_true = y_true[adjusted_mask_idx]
+            if entry_mask is not None:
+                entry_mask = entry_mask[adjusted_mask_idx]
 
         assert len(y_pred) == len(y_true), "y_pred and y_true are not consistent"
         assert not torch.any(torch.isnan(y_pred)), "y_pred contains NaN values"
         assert not torch.any(torch.isnan(y_true)), "y_true contains NaN values"
 
-        return None, global_embedding, y_pred, y_true, attn_matrix
+        return None, global_embedding, y_pred, y_true, attn_matrix, entry_mask
 
     def get_global_embeddings(self, x, edge_index):
         return self.forward(x, edge_index)
diff --git a/src/interscale/module/base/_base_local_module.py b/src/interscale/module/base/_base_local_module.py
index 3cdf3d0..c756d40 100644
--- a/src/interscale/module/base/_base_local_module.py
+++ b/src/interscale/module/base/_base_local_module.py
@@ -40,9 +40,13 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
             Size: [B, C] (classification) or [B, F] (regression)
         y_true: torch.Tensor
             Size: [B, ] (classification) or [B, F] (regression)
+        attn: None
+            This module has no attention; returned for a uniform `_common_step` contract.
+        entry_mask: torch.Tensor | None
+            Size: [B, F] under gene masking, marking the entries the loss is scored on.
         """
         # Mask nodes
-        batch_masked, mask_idx = self._common_step_masking(batch)
+        batch_masked, mask_idx, entry_mask = self._common_step_masking(batch)
 
         local_embedding = self.forward(batch_masked.x, batch_masked.edge_index)
         y_pred = self.decoder.forward(local_embedding)
@@ -60,12 +64,16 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
         if "classification" in prediction_task:
             y_true = batch.y[mask_idx]  # batch without mask because constant otherwise
             assert y_true.shape == y_pred.shape
-            return local_embedding, None, y_pred, y_true
+            # Class labels are not gene entries, so there is nothing for an entry mask to select.
+            return local_embedding, None, y_pred, y_true, None, None
 
         if "regression" in prediction_task:
             y_true = batch.x[mask_idx]  # batch without mask because constant otherwise
             assert y_true.shape == y_pred.shape
-            return local_embedding, None, y_pred, y_true
+            if entry_mask is not None:
+                entry_mask = entry_mask[mask_idx]
+                assert entry_mask.shape == y_pred.shape
+            return local_embedding, None, y_pred, y_true, None, entry_mask
 
         assert False, "Prediction task not supported"
 
diff --git a/src/interscale/module/base/_base_module.py b/src/interscale/module/base/_base_module.py
index d40c953..a79baff 100644
--- a/src/interscale/module/base/_base_module.py
+++ b/src/interscale/module/base/_base_module.py
@@ -5,7 +5,7 @@
 import torch
 
 from interscale.nn import LinearDecoder, LinearLSEDecoder, NonLinearDecoder
-from interscale.tl.masking import apply_mask
+from interscale.tl.masking import MASK_STRATEGIES, apply_mask
 
 
 class BaseModule(L.LightningModule, ABC):
@@ -20,7 +20,8 @@ def __init__(
         dropout_decoder: float = 0.2,
         decoder_hidden_dims: list[int] = [128, 128],
         dual_decoder: bool = False,
-        pct_mask_nodes: float = 0.0,
+        mask_percentage: float = 0.0,
+        mask_strategy: Literal["node", "gene"] = "node",
         type_gex_embedding: Literal["PCA", "NMF", "scvi"] | None = None,
     ):
         """
@@ -41,8 +42,11 @@ def __init__(
             Hidden dimensions for the decoder only if decoder_type is "nonlinear".
         dual_decoder: bool
             If True, use dual decoder for combined module. Both local and global decoders are used.
-        pct_mask_nodes: float
-            percentage of nodes to mask.
+        mask_percentage: float
+            Bernoulli masking probability -- per cell under ``mask_strategy="node"``, per
+            (cell, gene) entry under ``"gene"``.
+        mask_strategy: Literal["node", "gene"]
+            Granularity of the reconstruction corruption; see ``interscale.tl.masking``.
         type_gex_embedding: Literal["PCA", "NMF","scvi"] | None
             Type of GEX embedding to use.
         """
@@ -57,12 +61,16 @@ def __init__(
         self.decoder_type = decoder_type
         self.decoder_hidden_dims = decoder_hidden_dims
         self.dual_decoder = dual_decoder
-        self.pct_mask_nodes = pct_mask_nodes
+        self.mask_percentage = mask_percentage
+        if mask_strategy not in MASK_STRATEGIES:
+            raise ValueError(f"mask_strategy must be one of {MASK_STRATEGIES}, got {mask_strategy!r}.")
+        self.mask_strategy = mask_strategy
         self.type_gex_embedding = type_gex_embedding
-        if self.pct_mask_nodes > 0:
-            self.masked_nodes = True
-        else:
-            self.masked_nodes = False
+        # `masked_nodes` means "the reconstruction input is corrupted at all", not "whole cells
+        # are blanked". Both strategies set it: under gene masking every cell is a supervision
+        # target, so `pad_batch`'s keep_indices is all-True and it degenerates to the plain
+        # subsampling branch, which is the correct behaviour.
+        self.masked_nodes = self.mask_percentage > 0
 
         # Define components
         self.local_component = None
@@ -85,7 +93,7 @@ def __init__(
             raise ValueError(f"Decoder {self.decoder_type} not found.")
 
     def _common_step_masking(self, batch):
-        """Mask nodes in the batch.
+        """Corrupt the reconstruction input of the batch.
 
         Parameters
         ----------
@@ -95,16 +103,22 @@ def _common_step_masking(self, batch):
         Returns
         -------
         batch_masked: Batch
-            Batch of data with masked nodes having value MASK_VALUE.
+            Batch of data with the masked entries set to MASK_VALUE.
         mask_idx: torch.Tensor
-            Indices of masked nodes. Size: [N_masked_nodes, ]
+            Indices of the cells that are supervision targets. Size: [N_masked_nodes, ]
+        entry_mask: torch.Tensor | None
+            ``[N, G]`` boolean over the full node ordering when ``mask_strategy == "gene"``,
+            marking the entries the loss must be restricted to; ``None`` under cell masking,
+            where the whole row of every target cell is scored. Callers must subset it with the
+            same row indices they apply to ``y_true``.
         """
-        if self.pct_mask_nodes > 0:
-            batch_masked, mask_idx = apply_mask(batch)
+        if self.masked_nodes:
+            batch_masked, mask_idx, entry_mask = apply_mask(batch, self.mask_strategy)
         else:
             mask_idx = torch.arange(batch.x.shape[0], device=batch.x.device)
             batch_masked = batch
-        return batch_masked, mask_idx
+            entry_mask = None
+        return batch_masked, mask_idx, entry_mask
 
     @abstractmethod
     def _common_step(self, batch):
diff --git a/src/interscale/module/combined_module/combined_module.py b/src/interscale/module/combined_module/combined_module.py
index 5a39031..6529566 100644
--- a/src/interscale/module/combined_module/combined_module.py
+++ b/src/interscale/module/combined_module/combined_module.py
@@ -26,7 +26,8 @@ def __init__(self, cfg: CN, **base_module_kwargs):
             decoder_type=None,  # don't need decoder for local module
             dropout_decoder=0,
             decoder_hidden_dims=[],
-            pct_mask_nodes=self.pct_mask_nodes,
+            mask_percentage=self.mask_percentage,
+            mask_strategy=self.mask_strategy,
         )
         self.global_module = GlobalModule.from_config(
             cfg,
@@ -34,7 +35,8 @@ def __init__(self, cfg: CN, **base_module_kwargs):
             n_output=self.n_output,
             n_embed=self.n_embed,
             decoder_type=cfg.model.decoder.type,
-            pct_mask_nodes=self.pct_mask_nodes,
+            mask_percentage=self.mask_percentage,
+            mask_strategy=self.mask_strategy,
         )
 
     def predict_local(self, local_embedding):
@@ -63,8 +65,12 @@ def forward(self, batch_masked):
         return local_embedding, global_embedding, src_padding_mask, pad_index_nodes, attention_mask, attn_matrix
 
     def _common_step(self, batch, prediction_task, prediction_level: Literal["node", "graph"]):
-        """Shared step between train, val and test."""
-        batch_masked, mask_idx = self._common_step_masking(batch)
+        """Shared step between train, val and test.
+
+        The trailing `entry_mask` is `None` under cell masking and `[N_masked, F]` under gene
+        masking, where it marks the entries the loss must be restricted to.
+        """
+        batch_masked, mask_idx, _ = self._common_step_masking(batch)
 
         local_embedding, global_embedding, src_padding_mask, pad_index_nodes, attention_mask, attn_matrix = (
             self.forward(batch_masked)
@@ -73,18 +79,21 @@ def _common_step(self, batch, prediction_task, prediction_level: Literal["node",
 
         if prediction_task == "classification" and prediction_level == "graph":
             y_true = batch.y[batch.ptr[:-1]]
+            entry_mask = None
         else:
-            y_true, adjusted_mask_idx = self.global_module._process_batch_for_metrics(
+            y_true, adjusted_mask_idx, entry_mask = self.global_module._process_batch_for_metrics(
                 batch, prediction_task, prediction_level, pad_index_nodes, mask_idx
             )
             y_pred = y_pred[adjusted_mask_idx]
             y_true = y_true[adjusted_mask_idx]
+            if entry_mask is not None:
+                entry_mask = entry_mask[adjusted_mask_idx]
 
         assert len(y_pred) == len(y_true), "y_pred and y_true are not consistent"
         assert not torch.any(torch.isnan(y_pred)), "y_pred contains NaN values"
         assert not torch.any(torch.isnan(y_true)), "y_true contains NaN values"
 
-        return local_embedding, global_embedding, y_pred, y_true, attn_matrix
+        return local_embedding, global_embedding, y_pred, y_true, attn_matrix, entry_mask
 
     def get_model_summary(self) -> str:
         """Returns a string containing the model's parameters summary.
diff --git a/src/interscale/module/combined_module/dual_decoder_combined_module.py b/src/interscale/module/combined_module/dual_decoder_combined_module.py
index a477c4d..7f9e00e 100644
--- a/src/interscale/module/combined_module/dual_decoder_combined_module.py
+++ b/src/interscale/module/combined_module/dual_decoder_combined_module.py
@@ -4,6 +4,7 @@
 from yacs.config import CfgNode as CN
 
 from interscale.module.base import BaseModule, GlobalModule, LocalModule
+from interscale.tl.masking import masked_loss
 
 
 class DualDecoderCombinedModule(BaseModule):
@@ -38,7 +39,8 @@ def __init__(self, cfg: CN, **base_module_kwargs):
             decoder_type=cfg.model.decoder.type,
             dropout_decoder=cfg.model.decoder.dropout_decoder,
             decoder_hidden_dims=cfg.model.decoder.hidden_dims,
-            pct_mask_nodes=self.pct_mask_nodes,
+            mask_percentage=self.mask_percentage,
+            mask_strategy=self.mask_strategy,
         )
         # Global module with decoder
         self.global_module = GlobalModule.from_config(
@@ -49,7 +51,8 @@ def __init__(self, cfg: CN, **base_module_kwargs):
             decoder_type=cfg.model.decoder.type,
             dropout_decoder=cfg.model.decoder.dropout_decoder,
             decoder_hidden_dims=cfg.model.decoder.hidden_dims,
-            pct_mask_nodes=self.pct_mask_nodes,
+            mask_percentage=self.mask_percentage,
+            mask_strategy=self.mask_strategy,
         )
 
         # Store split point for separating concatenated predictions
@@ -128,7 +131,7 @@ def _common_step(self, batch, prediction_task, prediction_level: Literal["node",
         Returns predictions and ground truth for both local and global decoders
         on masked tokens, which can be combined in the loss function.
         """
-        batch_masked, mask_idx = self._common_step_masking(batch)
+        batch_masked, mask_idx, node_entry_mask = self._common_step_masking(batch)
 
         local_embedding, global_embedding, src_padding_mask, pad_index_nodes, attention_mask, attn = self.forward(
             batch_masked
@@ -153,12 +156,14 @@ def _common_step(self, batch, prediction_task, prediction_level: Literal["node",
             # Store metadata for graph level (only global predictions)
             self._n_masked_nodes = None
             self._is_graph_level = True
+            entry_mask_combined = None
         elif prediction_level == "node":
             # For node-level predictions, get ground truth for masked nodes
-            y_true, adjusted_mask_idx = self.global_module._process_batch_for_metrics(
+            y_true, adjusted_mask_idx, entry_mask = self.global_module._process_batch_for_metrics(
                 batch, prediction_task, prediction_level, pad_index_nodes, mask_idx
             )
             y_true_masked = y_true[adjusted_mask_idx]
+            entry_mask_masked = entry_mask[adjusted_mask_idx] if entry_mask is not None else None
 
             # Filter global predictions to masked nodes (same indices as y_true)
             y_pred_global_masked = y_pred_global[adjusted_mask_idx]
@@ -185,6 +190,21 @@ def _common_step(self, batch, prediction_task, prediction_level: Literal["node",
             y_pred_combined = torch.cat([y_pred_local, y_pred_global_masked], dim=0)
             y_true_combined = torch.cat([y_true_masked, y_true_masked], dim=0)
 
+            # The entry mask has to follow the same stacking. The local branch is indexed by
+            # `mask_idx` (the batch's own node order) while the global branch goes through
+            # `adjusted_mask_idx` (the padded, per-graph-subsampled order), so the two halves are
+            # not the same rows in general -- build each half from its own indexing rather than
+            # duplicating one of them.
+            if node_entry_mask is None:
+                entry_mask_combined = None
+            else:
+                entry_mask_local = node_entry_mask[mask_idx]
+                assert entry_mask_local.shape == y_pred_local.shape, (
+                    f"Mismatch: entry_mask_local.shape: {tuple(entry_mask_local.shape)}, "
+                    f"y_pred_local.shape: {tuple(y_pred_local.shape)}"
+                )
+                entry_mask_combined = torch.cat([entry_mask_local, entry_mask_masked], dim=0)
+
         else:
             raise ValueError(f"Invalid prediction level: {prediction_level}")
 
@@ -192,7 +212,7 @@ def _common_step(self, batch, prediction_task, prediction_level: Literal["node",
         assert not torch.any(torch.isnan(y_pred_combined)), "y_pred contains NaN values"
         assert not torch.any(torch.isnan(y_true_combined)), "y_true contains NaN values"
 
-        return local_embedding, global_embedding, y_pred_combined, y_true_combined, attn
+        return local_embedding, global_embedding, y_pred_combined, y_true_combined, attn, entry_mask_combined
 
     def get_separate_predictions(self, y_pred_combined, y_true_combined):
         """Get separate predictions and ground truth for local and global decoders.
@@ -236,6 +256,7 @@ def compute_separate_losses(
         loss_type: Literal["GaussianNLL", "MSELoss", "CrossEntropy", "WeightedCE"],
         y_pred_combined: torch.Tensor,
         y_true_combined: torch.Tensor,
+        entry_mask_combined: torch.Tensor | None = None,
     ):
         """Compute separate losses for local and global predictions.
 
@@ -252,6 +273,11 @@ def compute_separate_losses(
             ``[2*N_masked, C]``; for graph-level: ``[B, C]`` where ``B`` is batch size.
         y_true_combined
             Combined ground truth from ``_common_step``. Same shape as ``y_pred_combined``.
+        entry_mask_combined
+            Optional ``[2*N_masked, F]`` boolean from ``_common_step``. When given, each half's
+            loss is computed over that half's masked entries only -- which is the whole point of
+            gene masking, since the unmasked entries were handed to the model as input and
+            reconstructing them is the identity.
 
         Returns
         -------
@@ -274,14 +300,13 @@ def compute_separate_losses(
             y_pred_global = y_pred_combined[self._n_masked_nodes :]
             y_true_local = y_true_combined[: self._n_masked_nodes]
             y_true_global = y_true_combined[self._n_masked_nodes :]
-            if loss_type == "GaussianNLL":
-                sd_local = torch.std(y_true_local, dim=1, keepdim=True)
-                sd_global = torch.std(y_true_global, dim=1, keepdim=True)
-                local_loss = loss_fn(y_pred_local, y_true_local, sd_local)
-                global_loss = loss_fn(y_pred_global, y_true_global, sd_global)
+            if entry_mask_combined is None:
+                mask_local = mask_global = None
             else:
-                local_loss = loss_fn(y_pred_local, y_true_local)
-                global_loss = loss_fn(y_pred_global, y_true_global)
+                mask_local = entry_mask_combined[: self._n_masked_nodes]
+                mask_global = entry_mask_combined[self._n_masked_nodes :]
+            local_loss = masked_loss(loss_fn, loss_type, y_pred_local, y_true_local, mask_local)
+            global_loss = masked_loss(loss_fn, loss_type, y_pred_global, y_true_global, mask_global)
             losses["local_loss"] = local_loss
             losses["global_loss"] = global_loss
 
diff --git a/src/interscale/tl/__init__.py b/src/interscale/tl/__init__.py
index 03b4adc..9e7c93a 100644
--- a/src/interscale/tl/__init__.py
+++ b/src/interscale/tl/__init__.py
@@ -1,6 +1,15 @@
 from ._preprocessing import get_average_local_and_global_size, remove_zero_expression_cells
 from .geome_utils import prepare_a2d_dataset, prepare_geome_dataset
-from .masking import apply_mask, attn_mask_diagonal, create_transformer_attention_mask_from_edges
+from .masking import (
+    MASK_STRATEGIES,
+    apply_mask,
+    attn_mask_diagonal,
+    create_transformer_attention_mask_from_edges,
+    masked_loss,
+    masked_row_std,
+    sample_gene_mask,
+    sample_node_mask,
+)
 from .padding import pad_batch
 from .self_attn_relevance import SelfAttentionRelevance
 from .utils import check_and_update_cfg, set_full_reproducibility
@@ -13,6 +22,11 @@
     "set_full_reproducibility",
     "SelfAttentionRelevance",
     "apply_mask",
+    "masked_loss",
+    "masked_row_std",
+    "sample_node_mask",
+    "sample_gene_mask",
+    "MASK_STRATEGIES",
     "create_transformer_attention_mask_from_edges",
     "attn_mask_diagonal",
     "remove_zero_expression_cells",
diff --git a/src/interscale/tl/masking.py b/src/interscale/tl/masking.py
index 1e4163f..50ebd37 100644
--- a/src/interscale/tl/masking.py
+++ b/src/interscale/tl/masking.py
@@ -1,21 +1,148 @@
+"""Input corruption for the masked-reconstruction objective.
+
+Two granularities are available, selected by ``mask_strategy``:
+
+``"node"`` -- cell masking (the original behaviour)
+    A Bernoulli subset of *cells* has its entire expression vector replaced by ``MASK_VALUE``,
+    and the loss is evaluated on all G genes of those cells. A masked cell carries no
+    information about itself, so the only thing the model can condition on is its neighbourhood,
+    and E[x_i | neighbours of i] is close to the population mean. A near-constant predictor is
+    therefore a strong solution to this objective, which is what makes it look like the model
+    "learns the mean instead of reconstructing".
+
+``"gene"`` -- per-entry masking (GraphMAE / MAE style)
+    Every cell keeps most of its expression vector; a Bernoulli subset of *entries* ``(cell,
+    gene)`` is replaced by ``MASK_VALUE``, drawn independently per cell. The loss is evaluated
+    on those entries only (see :func:`masked_loss` below). The model now has the
+    cell's remaining genes to condition on, so within-cell co-expression -- not just the
+    population mean -- is available and rewarded.
+
+Note on GraphMAE (arXiv:2205.10803): GraphMAE itself masks whole nodes, exactly as ``"node"``
+does here. What it changes to avoid the trivial solution is the *criterion* (scaled cosine
+error, already available as ``interscale.train.losses.SCELoss``), a learnable ``[MASK]`` token
+rather than a zero vector, re-mask decoding, and a GNN decoder. Per-entry masking is the
+orthogonal knob this module adds; the two are independent and can be ablated together.
+
+THE FILL VALUE IS -1, NOT 0, AND IS NOT CONFIGURABLE. 62% of legnini23's ``log1p_norm`` entries
+are exactly 0, and 648 of its cells are all-zero outright, so a zero fill makes a masked position
+indistinguishable from a real measurement -- under gene masking the corruption becomes invisible
+and the model cannot identify the entries it is being asked to reconstruct. ``MASK_VALUE = -1``
+is outside the range of any log1p-normalised layer, so a masked position is unambiguous.
+
+This was measured, not assumed (legnini23, gene masking at rate 0.25, 3 seeds):
+
+    fill 0     val_concordance 0.0696 +/- 0.0113   (CV 16.2%)
+    fill -1    val_concordance 0.0799 +/- 0.0046   (CV  5.8%)
+
+-1 wins on every seed and cuts the run-to-run spread by ~2.5x, to below the cell-masking
+baseline's own CV of 10.6%. A learnable per-gene [MASK] token (GraphMAE's design) was also tried
+and scored identically to the fixed -1 to three decimals: the token drifts only ~0.02-0.035 over
+100 epochs from any initialisation, so the fill behaves as a constant, not as something worth
+learning. It was removed rather than kept as an option.
+"""
+
 import torch
 from torch_geometric.data import Batch
 
-MASK_VALUE = 0
+# -1, not 0: outside the range of any log1p-normalised expression layer (which is >= 0), so a
+# masked position can never be confused with a real measurement. See the module docstring for the
+# measurement that settled this. Changing it back to 0 silently un-does that result.
+MASK_VALUE = -1.0
+
+MASK_STRATEGIES = ("node", "gene")
+
+
+def sample_node_mask(num_nodes: int, pct: float, generator: torch.Generator | None = None) -> torch.Tensor:
+    """Draw a per-cell mask: each cell is masked independently with probability ``pct``.
+
+    Parameters
+    ----------
+    num_nodes
+        Number of cells in the graph.
+    pct
+        Per-cell masking probability.
+    generator
+        Optional RNG, for reproducible draws.
 
+    Returns
+    -------
+    torch.Tensor
+        Boolean tensor of shape ``[num_nodes]``. At least one cell is always masked, otherwise
+        the graph contributes no supervision at all.
+    """
+    mask = torch.rand(num_nodes, generator=generator) < pct
+    if not mask.any():
+        mask[torch.randint(num_nodes, (1,), generator=generator)] = True
+    return mask
+
+
+def sample_gene_mask(
+    num_nodes: int, num_genes: int, pct: float, generator: torch.Generator | None = None
+) -> torch.Tensor:
+    """Draw a per-entry mask: each ``(cell, gene)`` entry is masked independently with prob ``pct``.
+
+    The draw is independent per cell, so different cells lose different genes. That is
+    deliberate -- a mask shared across all cells of a graph would let the model learn a fixed
+    "these G_masked genes are always missing" shortcut, and would make each step's supervision
+    a single gene subset rather than |cells| different ones.
+
+    Parameters
+    ----------
+    num_nodes
+        Number of cells in the graph.
+    num_genes
+        Number of genes (columns of ``data.x``).
+    pct
+        Per-entry masking probability.
+    generator
+        Optional RNG, for reproducible draws.
+
+    Returns
+    -------
+    torch.Tensor
+        Boolean tensor of shape ``[num_nodes, num_genes]``. Every row has at least one masked
+        entry, so every cell contributes to the loss and the per-cell cosine metric is defined
+        for all of them.
+    """
+    mask = torch.rand(num_nodes, num_genes, generator=generator) < pct
 
-def apply_mask(batched_data: Batch):
-    """Mask nodes from PyG object in .mask attribute.
+    # Rows that came up all-False would silently drop out of the loss and make per-cell metrics
+    # undefined; give each of them exactly one masked gene.
+    empty_rows = ~mask.any(dim=1)
+    if empty_rows.any():
+        fill = torch.randint(num_genes, (int(empty_rows.sum()),), generator=generator)
+        mask[empty_rows, fill] = True
+    return mask
+
+
+def apply_mask(batched_data: Batch, mask_strategy: str = "node"):
+    """Corrupt ``batched_data.x`` at the granularity named by ``mask_strategy``.
+
+    Under ``"gene"`` the batch's ``gene_mask`` ``[N, G]`` selects the entries to overwrite, and it
+    is handed back so the loss can be restricted to them. Under ``"node"`` every gene of every
+    cell selected by ``.mask`` is overwritten and the returned entry mask is ``None``, meaning
+    "score the full rows".
+
+    The strategy is a parameter rather than being inferred from whether a ``gene_mask`` attribute
+    happens to be present. Sniffing the attribute made a stale ``gene_mask`` -- left on a ``Data``
+    object reused across strategies -- silently override the configured strategy, which the
+    dataloader then had to defend against by deleting the attribute.
 
     Args:
-        batched_data (Batch): _description_
+        batched_data (Batch): batch carrying ``.mask`` ``[N]``, plus ``.gene_mask`` ``[N, G]`` when
+            ``mask_strategy == "gene"``.
+        mask_strategy: one of :data:`MASK_STRATEGIES`.
 
     Returns
     -------
         batched_data_w_mask (Batch):
-            Batch only containing nodes that were not masked
+            Copy of the batch with the masked entries set to ``MASK_VALUE``.
         mask_idx (torch.Tensor):
-            Indices of masked nodes
+            Indices of the cells that carry at least one masked entry -- i.e. the rows on which
+            predictions are scored.
+        entry_mask (torch.Tensor | None):
+            ``[N, G]`` boolean over the *full* node ordering under gene masking, ``None`` under
+            cell masking. Callers must subset it with the same indices they use for ``y_true``.
 
     Example:
         Data object:
@@ -24,18 +151,101 @@ def apply_mask(batched_data: Batch):
         mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool)
         data = Data(x=x, edge_index=edge_index, mask=mask)
         ----
-        mask_idx = torch.tensor([1, 3])
-        masked_values = torch.tensor([[0., 0.], [3., 4.], [0., 0.], [7., 8.]])
+        mask_idx = torch.tensor([0, 2])
+        masked_values = torch.tensor([[-1., -1.], [3., 4.], [-1., -1.], [7., 8.]])
     """
     assert batched_data.mask is not None, "Mask is not set in the batch."
+    assert mask_strategy in MASK_STRATEGIES, f"mask_strategy must be one of {MASK_STRATEGIES}, got {mask_strategy!r}."
 
-    mask = batched_data.mask
-    mask_idx = torch.where(mask == 1)[0]  # TODO into 2D array [B, N_batched_nodes]
+    gene_mask = getattr(batched_data, "gene_mask", None) if mask_strategy == "gene" else None
     masked_values = batched_data.x.clone()
-    masked_values[mask] = MASK_VALUE
+
+    if gene_mask is None:
+        assert mask_strategy == "node", "mask_strategy='gene' but the batch carries no gene_mask."
+        mask = batched_data.mask
+        mask_idx = torch.where(mask == 1)[0]  # TODO into 2D array [B, N_batched_nodes]
+        masked_values[mask] = MASK_VALUE
+        entry_mask = None
+    else:
+        gene_mask = gene_mask.bool()
+        assert gene_mask.shape == batched_data.x.shape, (
+            f"Mismatch: gene_mask.shape: {tuple(gene_mask.shape)}, x.shape: {tuple(batched_data.x.shape)}"
+        )
+        masked_values[gene_mask] = MASK_VALUE
+        mask_idx = torch.where(gene_mask.any(dim=1))[0]
+        entry_mask = gene_mask
+
     batched_data_w_mask = batched_data.clone()
     batched_data_w_mask.x = masked_values
-    return batched_data_w_mask, mask_idx
+    return batched_data_w_mask, mask_idx, entry_mask
+
+
+# Losses whose value depends on the *arrangement* of a row, not just on the individual entries:
+# they normalise or centre along dim=-1. Selecting entries out of them would change what a "row"
+# is, so those get the masked entries zeroed in both tensors instead -- which restricts every sum,
+# dot product and norm involved to the masked coordinates, leaving the row structure intact.
+_ROW_STRUCTURED_LOSSES = ("SCELoss", "SCE_EntropyATT_Loss", "BalancedPearsonCorrelationLoss")
+
+
+def masked_row_std(y: torch.Tensor, entry_mask: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    """Per-row standard deviation of ``y`` over the masked entries only, shape ``[N, 1]``.
+
+    Note: this is the *spread* argument the existing ``GaussianNLL`` branch passes, and it is
+    deliberately std, not variance. ``nn.GaussianNLLLoss`` documents its third argument as a
+    variance, and the pre-existing unmasked branch has always passed ``torch.std(...)`` there --
+    a real bug, but one the published cell-masking runs were trained with. Matching it keeps the
+    two masking arms comparable; fix both call sites together, never just this one, or the
+    ablation stops being an ablation.
+    """
+    m = entry_mask.to(y.dtype)
+    n = m.sum(dim=1, keepdim=True).clamp(min=1)
+    mean = (y * m).sum(dim=1, keepdim=True) / n
+    var = (((y * m) ** 2).sum(dim=1, keepdim=True) / n - mean**2).clamp(min=0)
+    return var.sqrt().clamp(min=eps)
+
+
+def masked_loss(loss_fn, loss_type: str, y_pred: torch.Tensor, y_true: torch.Tensor, entry_mask=None):
+    """Evaluate a reconstruction loss on the masked entries only.
+
+    Under cell masking (``entry_mask is None``) every entry of every scored row was blanked, so
+    this is just ``loss_fn(y_pred, y_true)`` and the behaviour is unchanged. Under gene masking
+    most entries of a scored row were *given to the model as input*; including them would let
+    the identity map dominate the objective and would make the reported loss incomparable to the
+    cell-masking arm.
+
+    Parameters
+    ----------
+    loss_fn
+        The configured criterion, e.g. ``nn.SmoothL1Loss()``.
+    loss_type
+        Its name as it appears in ``optim.loss``; selects how the restriction is applied.
+    y_pred, y_true
+        ``[N, G]`` predictions and targets for the scored rows.
+    entry_mask
+        ``[N, G]`` boolean, or ``None``.
+
+    Returns
+    -------
+    torch.Tensor
+        Scalar loss.
+    """
+    if entry_mask is None:
+        if loss_type == "GaussianNLL":
+            return loss_fn(y_pred, y_true, torch.std(y_true, dim=1, keepdim=True))
+        return loss_fn(y_pred, y_true)
+
+    if loss_type in _ROW_STRUCTURED_LOSSES:
+        m = entry_mask.to(y_pred.dtype)
+        return loss_fn(y_pred * m, y_true * m)
+
+    if loss_type == "GaussianNLL":
+        # std, matching the unmasked branch above -- see masked_row_std for why.
+        spread = masked_row_std(y_true, entry_mask).expand_as(y_true)
+        return loss_fn(y_pred[entry_mask], y_true[entry_mask], spread[entry_mask])
+
+    # Element-wise criteria (MSELoss, SmoothL1, ...) reduce over whatever they are given, so
+    # handing them the selected entries as a flat vector is exactly a mean over masked entries.
+    return loss_fn(y_pred[entry_mask], y_true[entry_mask])
 
 
 def create_transformer_attention_mask_from_edges(
diff --git a/src/interscale/train/_training.py b/src/interscale/train/_training.py
index deb20a5..176777e 100644
--- a/src/interscale/train/_training.py
+++ b/src/interscale/train/_training.py
@@ -121,7 +121,7 @@ def train(
         print("Steps per epoch", steps_per_epoch)
         lr_monitor = LearningRateMonitor(logging_interval="epoch")
         self.history_ = MetricsHistory()
-        mask_resample_callback = NodeMaskResampleCallback() if self._cfg.dataset.pct_mask_nodes > 0 else None
+        mask_resample_callback = NodeMaskResampleCallback() if self._cfg.dataset.mask_percentage > 0 else None
         checkpoint_callback = None
         loss_callback = None
         performance_callback = None
diff --git a/src/interscale/train/_trainingplans.py b/src/interscale/train/_trainingplans.py
index 1c63498..479e585 100644
--- a/src/interscale/train/_trainingplans.py
+++ b/src/interscale/train/_trainingplans.py
@@ -10,6 +10,7 @@
 
 from interscale.module.base._base_module import BaseModule
 from interscale.nn import CosineWarmupScheduler
+from interscale.tl.masking import masked_loss
 
 from .losses import BalancedPearsonCorrelationLoss, SCE_EntropyATT_Loss, SCELoss
 
@@ -50,6 +51,82 @@ def compute(self) -> torch.Tensor:
         return self.total / self.count
 
 
+def masked_regression_metrics(
+    y_pred: torch.Tensor, y_true: torch.Tensor, entry_mask: torch.Tensor, eps: float = 1e-8
+) -> dict[str, torch.Tensor]:
+    """The regression metrics of ``_setup_regression_metrics``, restricted to the masked entries.
+
+    Under gene masking the scored rows are mostly entries the model was *given*. Feeding the full
+    rows to the ``MetricCollection`` would score the identity map on those and inflate every
+    number -- including ``val_r2``, which drives early stopping and checkpoint selection. There
+    is no way to express "these entries only" to a torchmetrics per-output metric (the surviving
+    entries are ragged across genes), so the same quantities are computed here from masked sums.
+
+    Every entry that is not masked is multiplied by zero before any sum is taken, and zeros
+    contribute nothing to a sum, so each moment below is exactly the moment over the masked
+    entries -- no approximation.
+
+    Parameters
+    ----------
+    y_pred, y_true
+        ``[N, G]`` predictions and targets for the scored rows.
+    entry_mask
+        ``[N, G]`` boolean marking the masked entries.
+    eps
+        Guard for degenerate (zero-variance) genes.
+
+    Returns
+    -------
+    dict
+        Unprefixed metric names mapped to scalar tensors, matching the keys that
+        ``_setup_regression_metrics`` produces: ``mse``, ``r2`` (per-gene, uniform average),
+        ``pearson_corr``, ``concordance_corr``, ``cosine_similarity`` (per cell).
+    """
+    m = entry_mask.to(y_pred.dtype)
+    p_ = y_pred * m
+    t_ = y_true * m
+
+    n_gene = m.sum(dim=0)  # [G] masked cells per gene
+    n_total = m.sum()
+
+    mse = ((p_ - t_) ** 2).sum() / n_total.clamp(min=1)
+
+    # Per-gene first and second moments over that gene's masked cells.
+    ng = n_gene.clamp(min=1)
+    mean_p = p_.sum(dim=0) / ng
+    mean_t = t_.sum(dim=0) / ng
+    var_p = (p_**2).sum(dim=0) / ng - mean_p**2
+    var_t = (t_**2).sum(dim=0) / ng - mean_t**2
+    cov = (p_ * t_).sum(dim=0) / ng - mean_p * mean_t
+
+    # A gene with fewer than two masked cells, or with no spread in either vector, has no
+    # correlation defined; NaN it out and let nanmean skip it, exactly as the unmasked path
+    # already does for constant genes.
+    nan = torch.tensor(float("nan"), device=y_pred.device, dtype=y_pred.dtype)
+    usable = (n_gene >= 2) & (var_t > eps)
+
+    pearson = torch.where(usable & (var_p > eps), cov / (var_p.clamp(min=eps) * var_t.clamp(min=eps)).sqrt(), nan)
+    concordance = torch.where(usable, 2 * cov / (var_p + var_t + (mean_p - mean_t) ** 2 + eps), nan)
+
+    # R2 per gene, then uniform average -- the same reduction torchmetrics'
+    # R2Score(multioutput="uniform_average") applies.
+    ss_res = ((p_ - t_) ** 2).sum(dim=0)
+    ss_tot = var_t * ng
+    r2 = torch.where(usable, 1 - ss_res / ss_tot.clamp(min=eps), nan)
+
+    # Per-cell cosine over that cell's masked genes: the zeroed entries drop out of both the dot
+    # product and the two norms, so this is the cosine on the masked coordinates.
+    cosine = nn.functional.cosine_similarity(p_, t_, dim=1)
+
+    return {
+        "mse": mse,
+        "r2": torch.nanmean(r2),
+        "pearson_corr": torch.nanmean(pearson),
+        "concordance_corr": torch.nanmean(concordance),
+        "cosine_similarity": cosine.mean(),
+    }
+
+
 CLASSIFICATION_LOSSES = ["CrossEntropy", "WeightedCE"]
 REGRESSION_LOSSES = [
     "MSELoss",
@@ -241,6 +318,7 @@ def _regression_metrics(
         metrics: MetricCollection,
         mask_idx: torch.Tensor | None = None,
         attn: torch.Tensor | None = None,
+        entry_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
         """Calculate regression metrics.
 
@@ -259,22 +337,37 @@ def _regression_metrics(
             The mask indices to apply to the metrics.
         attn : torch.Tensor | None
             The attention weights to apply to the metrics.
+        entry_mask : torch.Tensor | None
+            [N, G] boolean marking the entries that were actually masked. Set under
+            ``mask_strategy="gene"``; ``None`` under cell masking, where the whole row of every
+            scored cell was blanked and there is nothing to restrict. When given, BOTH the loss
+            and the metrics are computed over those entries only -- see ``masked_regression_metrics``.
         """
-        if self.loss_type == "GaussianNLL":
-            sd = torch.std(y_true, dim=1, keepdim=True)
-            loss = self.loss(y_pred, y_true, sd)
-        elif self.loss_type == "SCE_EntropyATT_Loss":
-            loss = self.loss(y_pred, y_true, attn)
+        if self.loss_type == "SCE_EntropyATT_Loss":
+            # Takes attention as a third argument, so it cannot go through masked_loss. Zeroing
+            # the unmasked entries restricts its row-wise cosine to the masked coordinates,
+            # which is what the row-structured branch of masked_loss does too.
+            if entry_mask is None:
+                loss = self.loss(y_pred, y_true, attn)
+            else:
+                m = entry_mask.to(y_pred.dtype)
+                loss = self.loss(y_pred * m, y_true * m, attn)
         else:
-            loss = self.loss(y_pred, y_true)
-
-        metrics = metrics(y_pred, y_true)
+            loss = masked_loss(self.loss, self.loss_type, y_pred, y_true, entry_mask)
+
+        if entry_mask is None:
+            metrics = metrics(y_pred, y_true)
+            # Take mean across pearson correlation
+            metrics[f"{mode}_pearson_corr"] = torch.nanmean(metrics[f"{mode}_pearson_corr"].contiguous())
+            # Same reduction, same reason: both are per-gene vectors of length n_output, and a
+            # constant gene yields NaN rather than a number.
+            metrics[f"{mode}_concordance_corr"] = torch.nanmean(metrics[f"{mode}_concordance_corr"].contiguous())
+        else:
+            # Same metric names, so `optim.monitor`, the sweep `--metric` flag and every existing
+            # wandb panel keep working across both strategies -- what changes is only which
+            # entries they are computed over.
+            metrics = {f"{mode}_{k}": v for k, v in masked_regression_metrics(y_pred, y_true, entry_mask).items()}
 
-        # Take mean across pearson correlation
-        metrics[f"{mode}_pearson_corr"] = torch.nanmean(metrics[f"{mode}_pearson_corr"].contiguous())
-        # Same reduction, same reason: both are per-gene vectors of length n_output, and a
-        # constant gene yields NaN rather than a number.
-        metrics[f"{mode}_concordance_corr"] = torch.nanmean(metrics[f"{mode}_concordance_corr"].contiguous())
         metrics[f"{mode}_loss"] = loss
         return loss, metrics
 
@@ -293,6 +386,7 @@ def _compute_and_log_metrics(
         mode: str,
         metrics: MetricCollection,
         attn: torch.Tensor | None,
+        entry_mask: torch.Tensor | None = None,
     ):
         """Helper method to log metrics for training, validation, or test steps.
 
@@ -315,7 +409,7 @@ def _compute_and_log_metrics(
             metrics.pop(f"{mode}_f1_per_class")
 
         elif "regression" in self.prediction_task:
-            loss, metrics = self._regression_metrics(y_pred, y_true, mode, metrics, attn=attn)
+            loss, metrics = self._regression_metrics(y_pred, y_true, mode, metrics, attn=attn, entry_mask=entry_mask)
 
         # Set sync_dist=True only for test mode
         sync_dist = mode == "test"
@@ -330,13 +424,15 @@ def training_step(self, batch):
         -------
             loss: torch.nn.Module
         """
-        local_embedding, global_embedding, y_pred, y_true, attn = self.module._common_step(
+        local_embedding, global_embedding, y_pred, y_true, attn, entry_mask = self.module._common_step(
             batch, self.prediction_task, self.prediction_level
         )
 
         # Check if module supports separate loss computation (e.g., DualDecoderCombinedModule)
         if hasattr(self.module, "compute_separate_losses"):
-            separate_losses = self.module.compute_separate_losses(self.loss, self.loss_type, y_pred, y_true)
+            separate_losses = self.module.compute_separate_losses(
+                self.loss, self.loss_type, y_pred, y_true, entry_mask
+            )
 
             # Log separate losses (on_step=False, on_epoch=True to match existing pattern)
             if separate_losses.get("local_loss") is not None:
@@ -368,7 +464,7 @@ def training_step(self, batch):
                 )
 
             #  compute and log metrics using combined predictions
-            loss = self._compute_and_log_metrics(y_pred, y_true, "train", self.train_metrics, attn=attn)
+            loss = self._compute_and_log_metrics(y_pred, y_true, "train", self.train_metrics, attn=attn, entry_mask=entry_mask)
 
             if separate_losses.get("kl_loss") is not None:
                 kl_loss = separate_losses["kl_loss"]
@@ -393,18 +489,20 @@ def training_step(self, batch):
             assert not torch.isnan(loss), "loss is NaN"
             return loss
         else:
-            return self._compute_and_log_metrics(y_pred, y_true, "train", self.train_metrics, attn=attn)
+            return self._compute_and_log_metrics(y_pred, y_true, "train", self.train_metrics, attn=attn, entry_mask=entry_mask)
         # return self._compute_and_log_metrics(y_pred, y_true, 'train', self.train_metrics, attn=attn)
 
     def validation_step(self, batch):
         """Validation step for the model."""
-        local_embedding, global_embedding, y_pred, y_true, attn = self.module._common_step(
+        local_embedding, global_embedding, y_pred, y_true, attn, entry_mask = self.module._common_step(
             batch, self.prediction_task, self.prediction_level
         )
 
         # Check if module supports separate loss computation (e.g., DualDecoderCombinedModule)
         if hasattr(self.module, "compute_separate_losses"):
-            separate_losses = self.module.compute_separate_losses(self.loss, self.loss_type, y_pred, y_true)
+            separate_losses = self.module.compute_separate_losses(
+                self.loss, self.loss_type, y_pred, y_true, entry_mask
+            )
 
             # Log separate losses (on_step=False, on_epoch=True to match existing pattern)
             if separate_losses.get("local_loss") is not None:
@@ -427,7 +525,7 @@ def validation_step(self, batch):
                 )
 
             #  compute and log metrics using combined predictions
-            loss = self._compute_and_log_metrics(y_pred, y_true, "val", self.valid_metrics, attn=attn)
+            loss = self._compute_and_log_metrics(y_pred, y_true, "val", self.valid_metrics, attn=attn, entry_mask=entry_mask)
 
             if separate_losses.get("kl_loss") is not None:
                 kl_loss = separate_losses["kl_loss"]
@@ -452,18 +550,20 @@ def validation_step(self, batch):
             assert not torch.isnan(loss), "loss is NaN"
             return loss
         else:
-            return self._compute_and_log_metrics(y_pred, y_true, "val", self.valid_metrics, attn=attn)
+            return self._compute_and_log_metrics(y_pred, y_true, "val", self.valid_metrics, attn=attn, entry_mask=entry_mask)
 
         # return self._compute_and_log_metrics(y_pred, y_true, 'val', self.valid_metrics, attn=attn)
 
     def test_step(self, batch):
         """Test step for the model."""
-        local_embedding, global_embedding, y_pred, y_true, attn = self.module._common_step(
+        local_embedding, global_embedding, y_pred, y_true, attn, entry_mask = self.module._common_step(
             batch, self.prediction_task, self.prediction_level
         )
         # Check if module supports separate loss computation (e.g., DualDecoderCombinedModule)
         if hasattr(self.module, "compute_separate_losses"):
-            separate_losses = self.module.compute_separate_losses(self.loss, self.loss_type, y_pred, y_true)
+            separate_losses = self.module.compute_separate_losses(
+                self.loss, self.loss_type, y_pred, y_true, entry_mask
+            )
 
             # Log separate losses (on_step=False, on_epoch=True to match existing pattern, sync_dist=True for test)
             if separate_losses.get("local_loss") is not None:
@@ -495,7 +595,7 @@ def test_step(self, batch):
                 )
 
             #  compute and log metrics using combined predictions
-            loss = self._compute_and_log_metrics(y_pred, y_true, "test", self.test_metrics, attn=attn)
+            loss = self._compute_and_log_metrics(y_pred, y_true, "test", self.test_metrics, attn=attn, entry_mask=entry_mask)
 
             if separate_losses.get("kl_loss") is not None:
                 kl_loss = separate_losses["kl_loss"]
@@ -520,7 +620,7 @@ def test_step(self, batch):
             assert not torch.isnan(loss), "loss is NaN"
             return loss
         else:
-            return self._compute_and_log_metrics(y_pred, y_true, "test", self.test_metrics, attn=attn)
+            return self._compute_and_log_metrics(y_pred, y_true, "test", self.test_metrics, attn=attn, entry_mask=entry_mask)
         # return self._compute_and_log_metrics(y_pred, y_true, 'test', self.test_metrics,attn=attn)
 
     def configure_optimizers(self):
diff --git a/tests/test_gene_masking.py b/tests/test_gene_masking.py
new file mode 100644
index 0000000..480b63f
--- /dev/null
+++ b/tests/test_gene_masking.py
@@ -0,0 +1,390 @@
+"""Gene-wise (per-entry) masking: sampling, corruption, and the masked loss/metrics.
+
+The contract these lock down is the one that makes the ablation meaningful: under
+``mask_strategy="gene"`` the model must be scored ONLY on the entries it did not see. If the
+restriction leaks, the unmasked entries -- which were handed to the model as input -- turn the
+objective into the identity map and every number becomes incomparable to the cell-masking arm.
+"""
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+import torchmetrics
+from torch_geometric.data import Data
+
+from interscale.geome_dataloader import GraphAnnDataModule
+from interscale.tl.masking import (
+    MASK_VALUE,
+    apply_mask,
+    masked_loss,
+    sample_gene_mask,
+    sample_node_mask,
+)
+from interscale.train._trainingplans import RunningCosineSimilarity, masked_regression_metrics
+
+
+def _make_data(num_nodes: int, num_features: int = 6) -> Data:
+    x = torch.randn(num_nodes, num_features)
+    edge_index = torch.zeros((2, 0), dtype=torch.long)
+    return Data(x=x, edge_index=edge_index)
+
+
+def _build_datamodule(mask_strategy: str = "gene", mask_percentage: float = 0.4) -> GraphAnnDataModule:
+    dm = GraphAnnDataModule(
+        datas=[[_make_data(20), _make_data(20)], [_make_data(10)], [_make_data(10)]],
+        batch_size=1,
+        num_workers=0,
+        mask_percentage=mask_percentage,
+        mask_strategy=mask_strategy,
+        learning_type="node",
+    )
+    dm.setup(stage="fit")
+    dm.setup(stage="test")
+    return dm
+
+
+# --------------------------------------------------------------------------- mask sampling
+
+
+def test_sample_gene_mask_has_no_empty_rows():
+    """A cell with nothing masked would contribute no loss and an undefined per-cell cosine."""
+    mask = sample_gene_mask(500, 8, pct=0.01)
+
+    assert mask.shape == (500, 8)
+    assert mask.any(dim=1).all()
+
+
+def test_sample_gene_mask_hits_the_requested_rate():
+    mask = sample_gene_mask(2000, 50, pct=0.3)
+
+    assert mask.float().mean().item() == pytest.approx(0.3, abs=0.02)
+
+
+def test_sample_gene_mask_differs_across_cells():
+    """A gene subset shared by every cell would be a fixed shortcut rather than supervision."""
+    mask = sample_gene_mask(200, 40, pct=0.5)
+
+    unique_rows = {tuple(row.tolist()) for row in mask}
+    assert len(unique_rows) > 100
+
+
+def test_sample_node_mask_masks_at_least_one_cell():
+    assert sample_node_mask(50, pct=0.0).sum() == 1
+
+
+# ------------------------------------------------------------------------------- datamodule
+
+
+def test_gene_strategy_writes_gene_mask_and_derives_node_mask():
+    dm = _build_datamodule("gene")
+
+    for data in dm.train_data:
+        assert data.gene_mask.shape == data.x.shape
+        # Every cell is a supervision target under gene masking; selectivity moved to the entries.
+        assert torch.equal(data.mask, data.gene_mask.any(dim=1))
+        assert data.mask.all()
+
+
+def test_node_strategy_leaves_no_gene_mask():
+    dm = _build_datamodule("node")
+
+    for data in dm.train_data:
+        assert "gene_mask" not in data
+
+
+def test_a_stale_gene_mask_cannot_hijack_the_node_strategy():
+    """The strategy is a parameter, so a leftover attribute is simply ignored.
+
+    This is why the dataloader has no `_clear_gene_mask`: correctness comes from telling
+    `apply_mask` (and `_process_batch_for_metrics`) which strategy is configured, not from
+    scrubbing the attribute they used to sniff.
+    """
+    data = _make_data(12)
+    data.gene_mask = sample_gene_mask(12, data.x.shape[1], pct=0.5)  # stale, from a gene run
+    data.mask = sample_node_mask(12, pct=0.5)
+
+    out, mask_idx, entry_mask = apply_mask(data, "node")
+
+    assert entry_mask is None
+    assert torch.equal(mask_idx, torch.where(data.mask)[0])
+    assert (out.x[data.mask] == MASK_VALUE).all()
+    assert torch.equal(out.x[~data.mask], data.x[~data.mask])
+
+
+def test_gene_strategy_without_a_gene_mask_fails_loudly():
+    data = _make_data(8)
+    data.mask = sample_node_mask(8, pct=0.5)
+
+    with pytest.raises(AssertionError, match="no gene_mask"):
+        apply_mask(data, "gene")
+
+
+def test_resample_redraws_the_gene_mask_in_place():
+    dm = _build_datamodule("gene")
+    original = [data.gene_mask.clone() for data in dm.train_data]
+    ids = [id(data) for data in dm.train_data]
+
+    dm.resample_train_mask()
+
+    assert any(not torch.equal(o, d.gene_mask) for o, d in zip(original, dm.train_data))
+    assert [id(d) for d in dm.train_data] == ids
+
+
+def test_unmasked_split_gets_no_gene_mask():
+    """Graph-level eval splits pass mask=False and must not be corrupted."""
+    dm = GraphAnnDataModule(
+        datas=[[_make_data(20)], [_make_data(10)], [_make_data(10)]],
+        batch_size=1,
+        num_workers=0,
+        mask_strategy="gene",
+        learning_type="graph",
+    )
+    dm.setup(stage="fit")
+
+    for data in dm.val_data:
+        assert "gene_mask" not in data
+        assert not data.mask.any()
+
+
+def test_invalid_strategy_is_rejected():
+    with pytest.raises(ValueError, match="mask_strategy"):
+        GraphAnnDataModule(datas=[[_make_data(4)], [_make_data(4)]], mask_strategy="cell")
+
+
+# ----------------------------------------------------------------------------- apply_mask
+
+
+def test_apply_mask_gene_path_blanks_only_the_masked_entries():
+    data = _make_data(12)
+    data.gene_mask = sample_gene_mask(12, data.x.shape[1], pct=0.4)
+    data.mask = data.gene_mask.any(dim=1)
+
+    out, mask_idx, entry_mask = apply_mask(data, "gene")
+
+    assert torch.equal(entry_mask, data.gene_mask)
+    assert (out.x[data.gene_mask] == MASK_VALUE).all()
+    # Everything else survives untouched -- that surviving context is the whole point.
+    assert torch.equal(out.x[~data.gene_mask], data.x[~data.gene_mask])
+    assert torch.equal(mask_idx, torch.arange(12))
+    assert not torch.equal(out.x, data.x)
+
+
+def test_apply_mask_node_path_is_unchanged():
+    data = _make_data(12)
+    data.mask = sample_node_mask(12, pct=0.5)
+
+    out, mask_idx, entry_mask = apply_mask(data, "node")
+
+    assert entry_mask is None
+    assert torch.equal(mask_idx, torch.where(data.mask)[0])
+    assert (out.x[data.mask] == MASK_VALUE).all()
+    assert torch.equal(out.x[~data.mask], data.x[~data.mask])
+
+
+# ----------------------------------------------------------------------------- masked loss
+
+
+@pytest.mark.parametrize("loss_type,loss_fn", [("SmoothL1", nn.SmoothL1Loss()), ("MSELoss", nn.MSELoss())])
+def test_masked_loss_scores_only_the_masked_entries(loss_type, loss_fn):
+    torch.manual_seed(0)
+    y_true = torch.randn(64, 10)
+    y_pred = torch.randn(64, 10)
+    entry_mask = sample_gene_mask(64, 10, pct=0.3)
+
+    got = masked_loss(loss_fn, loss_type, y_pred, y_true, entry_mask)
+
+    assert got == pytest.approx(loss_fn(y_pred[entry_mask], y_true[entry_mask]).item())
+
+
+def test_masked_loss_ignores_error_outside_the_mask():
+    """The decisive property: a wrong answer on an entry the model was GIVEN must cost nothing."""
+    y_true = torch.zeros(4, 5)
+    y_pred = torch.zeros(4, 5)
+    entry_mask = torch.zeros(4, 5, dtype=torch.bool)
+    entry_mask[:, 0] = True
+
+    y_pred[:, 3] = 100.0  # unmasked entry, wildly wrong
+
+    assert masked_loss(nn.MSELoss(), "MSELoss", y_pred, y_true, entry_mask).item() == 0.0
+
+
+def test_masked_loss_without_entry_mask_is_the_plain_loss():
+    y_true, y_pred = torch.randn(8, 5), torch.randn(8, 5)
+    fn = nn.SmoothL1Loss()
+
+    assert masked_loss(fn, "SmoothL1", y_pred, y_true, None) == pytest.approx(fn(y_pred, y_true).item())
+
+
+def test_masked_loss_row_structured_restricts_to_masked_coordinates():
+    """SCE normalises along the row, so entries are zeroed rather than selected."""
+    from interscale.train.losses import SCELoss
+
+    torch.manual_seed(0)
+    y_true = torch.randn(16, 7)
+    y_pred = torch.randn(16, 7)
+    entry_mask = sample_gene_mask(16, 7, pct=0.5)
+    m = entry_mask.float()
+
+    got = masked_loss(SCELoss(), "SCELoss", y_pred, y_true, entry_mask)
+
+    assert got == pytest.approx(SCELoss()(y_pred * m, y_true * m).item())
+
+
+def test_masked_loss_gaussian_nll_uses_the_masked_row_spread():
+    torch.manual_seed(0)
+    y_true = torch.randn(16, 7)
+    y_pred = torch.randn(16, 7)
+    entry_mask = sample_gene_mask(16, 7, pct=0.5)
+
+    got = masked_loss(nn.GaussianNLLLoss(), "GaussianNLL", y_pred, y_true, entry_mask)
+
+    assert torch.isfinite(got)
+
+
+def test_masked_row_std_matches_torch_std_when_nothing_is_held_out():
+    """The masked spread must agree with the unmasked branch's torch.std on an all-True mask.
+
+    That agreement is what keeps a GaussianNLL cell-masking run and a gene-masking run on the
+    same footing; if the two branches computed different quantities the arms would not be
+    comparable. (Both pass std where nn.GaussianNLLLoss documents a variance -- a pre-existing
+    bug, preserved deliberately; see masked_row_std.)
+    """
+    from interscale.tl.masking import masked_row_std
+
+    torch.manual_seed(0)
+    y = torch.randn(32, 9)
+    full = torch.ones(32, 9, dtype=torch.bool)
+
+    # torch.std is Bessel-corrected, masked_row_std is not, so compare the population form.
+    expected = y.std(dim=1, keepdim=True, unbiased=False)
+    assert torch.allclose(masked_row_std(y, full), expected, atol=1e-5)
+
+
+# --------------------------------------------------------------------------- masked metrics
+
+
+def _reference_metrics(y_pred, y_true, m):
+    """Brute-force per-gene / per-cell loop over the masked entries only."""
+    per_gene_r2, per_gene_r, per_gene_ccc = [], [], []
+    for g in range(y_true.shape[1]):
+        sel = m[:, g]
+        p, t = y_pred[sel, g], y_true[sel, g]
+        if sel.sum() < 2 or t.var(unbiased=False) < 1e-8:
+            continue
+        per_gene_r2.append(torchmetrics.R2Score()(p, t).item())
+        per_gene_r.append(torchmetrics.PearsonCorrCoef()(p, t).item())
+        per_gene_ccc.append(torchmetrics.ConcordanceCorrCoef()(p, t).item())
+    per_cell_cos = [
+        torch.nn.functional.cosine_similarity(y_pred[i, m[i]], y_true[i, m[i]], dim=0).item()
+        for i in range(y_true.shape[0])
+    ]
+    return {
+        "mse": ((y_pred[m] - y_true[m]) ** 2).mean().item(),
+        "r2": float(np.mean(per_gene_r2)),
+        "pearson_corr": float(np.mean(per_gene_r)),
+        "concordance_corr": float(np.mean(per_gene_ccc)),
+        "cosine_similarity": float(np.mean(per_cell_cos)),
+    }
+
+
+def test_masked_metrics_match_a_brute_force_loop_over_masked_entries():
+    torch.manual_seed(0)
+    y_true = torch.randn(200, 12)
+    y_pred = 0.6 * y_true + 0.5 * torch.randn(200, 12)
+    m = sample_gene_mask(200, 12, pct=0.4)
+
+    got = masked_regression_metrics(y_pred, y_true, m)
+    ref = _reference_metrics(y_pred, y_true, m)
+
+    for key, expected in ref.items():
+        assert got[key].item() == pytest.approx(expected, abs=1e-4), key
+
+
+def test_masked_metrics_reduce_to_torchmetrics_when_nothing_is_held_out():
+    """An all-True mask is the cell-masking case, so the two paths must agree there."""
+    torch.manual_seed(0)
+    y_true = torch.randn(200, 12)
+    y_pred = 0.6 * y_true + 0.5 * torch.randn(200, 12)
+    full = torch.ones(200, 12, dtype=torch.bool)
+
+    got = masked_regression_metrics(y_pred, y_true, full)
+
+    assert got["mse"].item() == pytest.approx(torchmetrics.MeanSquaredError()(y_pred, y_true).item(), abs=1e-5)
+    assert got["r2"].item() == pytest.approx(
+        torchmetrics.R2Score(multioutput="uniform_average")(y_pred, y_true).item(), abs=1e-4
+    )
+    assert got["pearson_corr"].item() == pytest.approx(
+        torch.nanmean(torchmetrics.PearsonCorrCoef(num_outputs=12)(y_pred, y_true)).item(), abs=1e-4
+    )
+    assert got["cosine_similarity"].item() == pytest.approx(RunningCosineSimilarity()(y_pred, y_true).item(), abs=1e-5)
+
+
+def test_masked_metrics_are_blind_to_predictions_outside_the_mask():
+    """The inflation this guards against: scoring entries the model was handed as input."""
+    torch.manual_seed(0)
+    y_true = torch.randn(80, 10)
+    y_pred = 0.5 * y_true + 0.3 * torch.randn(80, 10)
+    m = sample_gene_mask(80, 10, pct=0.3)
+
+    before = masked_regression_metrics(y_pred, y_true, m)
+    corrupted = y_pred.clone()
+    corrupted[~m] = 1e3
+    after = masked_regression_metrics(corrupted, y_true, m)
+
+    for key in before:
+        assert after[key].item() == pytest.approx(before[key].item(), abs=1e-5), key
+
+
+def test_masked_metrics_skip_genes_with_too_few_masked_cells():
+    """A gene masked in one cell has no correlation defined; it must not poison the mean."""
+    torch.manual_seed(0)
+    y_true = torch.randn(50, 4)
+    y_pred = torch.randn(50, 4)
+    m = torch.zeros(50, 4, dtype=torch.bool)
+    m[:, 0] = True  # well-populated
+    m[0, 1] = True  # single cell -> undefined
+
+    got = masked_regression_metrics(y_pred, y_true, m)
+
+    for key, value in got.items():
+        assert torch.isfinite(value), key
+
+
+# ------------------------------------------------------------------------- the fill value
+#
+# MASK_VALUE is -1 and is no longer configurable. On a log1p layer that is ~62% exact zeros (and
+# with 648 all-zero cells), a 0 fill makes a masked position indistinguishable from a real
+# measurement; under gene masking the corruption becomes invisible entirely. Measured at gene
+# rate 0.25 over 3 seeds, 0 scored 0.0696 +/- 0.0113 against -1's 0.0799 +/- 0.0046.
+
+
+def test_mask_value_is_outside_the_range_of_a_log1p_layer():
+    """The one property the fill value has to have. A log1p layer is >= 0 by construction."""
+    assert MASK_VALUE < 0
+
+
+def test_masked_positions_are_distinguishable_from_real_zeros():
+    """The regression this guards: with a 0 fill, `x == MASK_VALUE` would also select real zeros."""
+    data = _make_data(200, num_features=6)
+    # ~60% exact zeros, as in the real layer.
+    data.x = (torch.rand(200, 6) > 0.6).float() * torch.rand(200, 6)
+    data.gene_mask = sample_gene_mask(200, 6, pct=0.3)
+    data.mask = data.gene_mask.any(dim=1)
+
+    out, _, entry_mask = apply_mask(data, "gene")
+
+    # Every masked position, and ONLY the masked positions, carry the fill value.
+    assert torch.equal(out.x == MASK_VALUE, entry_mask)
+
+
+def test_all_zero_cells_stay_distinguishable_from_masked_cells():
+    """legnini23 has 648 cells whose expression is entirely zero, so an all-zero row is real data."""
+    data = _make_data(10, num_features=6)
+    data.x = torch.zeros(10, 6)
+    data.mask = sample_node_mask(10, pct=0.5)
+
+    out, _, _ = apply_mask(data, "node")
+
+    assert (out.x[data.mask] == MASK_VALUE).all()
+    assert (out.x[~data.mask] == 0).all()
diff --git a/tests/test_geome_dataloader.py b/tests/test_geome_dataloader.py
index cb0ac29..26f5268 100644
--- a/tests/test_geome_dataloader.py
+++ b/tests/test_geome_dataloader.py
@@ -11,7 +11,7 @@ def _make_data(num_nodes: int, num_features: int = 3) -> Data:
     return Data(x=x, edge_index=edge_index)
 
 
-def _build_datamodule(pct_mask_nodes: float = 0.5) -> GraphAnnDataModule:
+def _build_datamodule(mask_percentage: float = 0.5) -> GraphAnnDataModule:
     train_data = [_make_data(20), _make_data(20), _make_data(20)]
     val_data = [_make_data(10)]
     test_data = [_make_data(10)]
@@ -19,7 +19,7 @@ def _build_datamodule(pct_mask_nodes: float = 0.5) -> GraphAnnDataModule:
         datas=[train_data, val_data, test_data],
         batch_size=1,
         num_workers=0,
-        pct_mask_nodes=pct_mask_nodes,
+        mask_percentage=mask_percentage,
         learning_type="node",
     )
     dm.setup(stage="fit")
diff --git a/tests/test_global_pca_persistence.py b/tests/test_global_pca_persistence.py
new file mode 100644
index 0000000..0207015
--- /dev/null
+++ b/tests/test_global_pca_persistence.py
@@ -0,0 +1,118 @@
+"""The PCA front-end of a GlobalModel must survive a checkpoint round trip.
+
+A `GlobalModel` has no local component, so `type_gex_embedding="PCA"` supplies the transformer's
+input: a `sklearn` PCA is fitted to the first batch and every later batch is projected through it.
+`BaseModel.save` persists only `module.state_dict()`, and a `sklearn` estimator is not part of
+one -- so before `pca_mean_`/`pca_components_`/`pca_fitted_` were registered as buffers, a
+reloaded model arrived unfitted and refit the PCA on the first *evaluation* batch. The transformer
+was trained on the basis fitted to the first training batch; a basis refitted on other data
+differs by rotation and by component sign, so the reloaded model decoded a different space than it
+was trained on, silently and without an error.
+
+That matters for any analysis that loads a global-only checkpoint back -- the component ablation
+reads attention and CLS tokens out of exactly such a reload.
+"""
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from interscale.module.global_modules import TransformerNodeEncoderHook  # noqa: E402
+
+N_INPUT, N_EMBED = 120, 16
+
+
+def build_module():
+    return TransformerNodeEncoderHook(
+        max_seq_len=64,
+        n_heads=2,
+        dropout_global=0.0,
+        act_func="relu",
+        num_layers=1,
+        dim_feedforward=32,
+        long_range_attention=False,
+        n_input=N_INPUT,
+        n_output=N_INPUT,
+        n_embed=N_EMBED,
+        decoder_type="linear",
+        dropout_decoder=0.0,
+        decoder_hidden_dims=[32],
+        mask_percentage=0.1,
+        mask_strategy="node",
+        type_gex_embedding="PCA",
+    )
+
+
+@pytest.fixture
+def batches():
+    rng = np.random.default_rng(0)
+    return (
+        rng.normal(size=(200, N_INPUT)).astype(np.float32),  # "training" batch, fits the PCA
+        rng.normal(size=(90, N_INPUT)).astype(np.float32),  # "evaluation" batch
+    )
+
+
+def test_fit_is_recorded_in_the_state_dict(batches):
+    train_batch, _ = batches
+    module = build_module()
+    assert not bool(module.pca_fitted_)
+
+    module.create_gex_embedding(train_batch, type="PCA")
+
+    assert bool(module.pca_fitted_)
+    state = module.state_dict()
+    assert {"pca_mean_", "pca_components_", "pca_fitted_"} <= set(state)
+    assert state["pca_components_"].shape == (N_EMBED, N_INPUT)
+
+
+def test_reload_projects_through_the_trained_basis(batches):
+    """The whole point: a reloaded module must not refit on the batch it is handed."""
+    train_batch, eval_batch = batches
+    module = build_module()
+    module.create_gex_embedding(train_batch, type="PCA")
+    reference = module.create_gex_embedding(eval_batch, type="PCA")
+
+    reloaded = build_module()
+    _, unexpected = reloaded.load_state_dict(module.state_dict(), strict=False)
+    assert not unexpected
+
+    np.testing.assert_allclose(reloaded.create_gex_embedding(eval_batch, type="PCA"), reference, atol=1e-6)
+
+
+def test_projection_matches_sklearn(batches):
+    """Equivalent to `PCA.transform` for whiten=False, which is what this module constructs.
+
+    Compared against the module's own fitted estimator rather than a freshly fitted one: at
+    n_components well below min(n_samples, n_features), sklearn's `svd_solver="auto"` selects the
+    randomized solver, whose `random_state` is None, so two fits of the same matrix do not give
+    the same basis.
+    """
+    train_batch, eval_batch = batches
+    module = build_module()
+    module.create_gex_embedding(train_batch, type="PCA")
+
+    np.testing.assert_allclose(
+        module.create_gex_embedding(eval_batch, type="PCA"),
+        module.pca.transform(eval_batch),
+        atol=1e-5,
+    )
+
+
+def test_checkpoint_without_pca_buffers_still_loads(batches):
+    """Backwards compatibility: checkpoints written before the buffers existed have no entry for
+    them. `BaseModel.load` uses strict=False, so those keep the old refit-on-load behaviour rather
+    than failing to load at all.
+    """
+    train_batch, eval_batch = batches
+    module = build_module()
+    module.create_gex_embedding(train_batch, type="PCA")
+    reference = module.create_gex_embedding(eval_batch, type="PCA")
+
+    legacy_state = {k: v for k, v in module.state_dict().items() if not k.startswith("pca_")}
+    legacy = build_module()
+    legacy.load_state_dict(legacy_state, strict=False)
+
+    assert not bool(legacy.pca_fitted_)
+    refit = legacy.create_gex_embedding(eval_batch, type="PCA")
+    assert not np.allclose(refit, reference, atol=1e-6)
diff --git a/tests/test_sweep_config.py b/tests/test_sweep_config.py
index 7504424..c6ec510 100644
--- a/tests/test_sweep_config.py
+++ b/tests/test_sweep_config.py
@@ -281,7 +281,7 @@ def test_every_declared_goal_is_accepted(goal, base_cfg):
 def test_robustness_goal_parameters_apply(base_cfg):
     """The robustness sweep's three keys all exist and all land."""
     trial = {
-        "dataset.pct_mask_nodes": 0.42,
+        "dataset.mask_percentage": 0.42,
         "dataset.spatial_neigbors_kwargs.radius": 77,
         "optim.seed": 7,
     }
@@ -289,7 +289,7 @@ def test_robustness_goal_parameters_apply(base_cfg):
         base_cfg.clone(), "robustness", trial, sweep_params=sorted(trial)
     )
     assert sorted(applied) == sorted(trial)
-    assert cfg.dataset.pct_mask_nodes == 0.42
+    assert cfg.dataset.mask_percentage == 0.42
     assert cfg.dataset.spatial_neigbors_kwargs.radius == 77
     assert cfg.optim.seed == 7
 
@@ -677,6 +677,14 @@ def test_arm_trial_does_not_leak_into_the_base_config(arm_cfg, arm_yaml):
 ARM_SWEEP_PAIRS = {
     "sliding_window_melton25.yaml": ("melton25_sw", "node_reg"),
     "overlap_ladder_legnini.yaml": ("legnini23_overlap", "node_reg"),
+    # Masking-granularity ablation. Resolved against the CELL-masking pair: the arms set
+    # mask_strategy/mask_token themselves, so starting from node_reg proves each arm overrides
+    # the baseline rather than relying on the genemask task file to have set it.
+    "mask_granularity_legnini.yaml": ("legnini23", "node_reg"),
+    # Rate/length follow-up. Same base pair, same reason.
+    "mask_rate_and_length_legnini.yaml": ("legnini23", "node_reg"),
+    # Cell-masking rate ladder, the counterpart to the gene ladder in mask_granularity.
+    "mask_rate_cell_ladder_legnini.yaml": ("legnini23", "node_reg"),
 }
 
 

From a568e88b0a85acabdf0ba958fac9c6955729abe3 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Tue, 8 Sep 2026 07:39:14 +0200
Subject: [PATCH 10/14] anchor GEX plots

---
 CHANGELOG.md                                  |  10 +
 docs/api/tools.md                             |  36 ++
 src/interscale/evaluation/net_streams.py      | 104 +++-
 src/interscale/geome_dataloader.py            |   4 +-
 src/interscale/model/base/_base_model.py      |  41 +-
 src/interscale/model/combined_model.py        |   6 +-
 .../module/base/_base_global_module.py        | 181 ++----
 .../module/base/_base_local_module.py         |  24 +-
 .../module/combined_module/combined_module.py |  18 +-
 .../dual_decoder_combined_module.py           | 106 ++--
 src/interscale/pl/__init__.py                 |  16 +
 src/interscale/pl/anchor_plots.py             | 561 ++++++++++++++++++
 src/interscale/tl/__init__.py                 |  16 +
 src/interscale/tl/anchors.py                  | 426 +++++++++++++
 src/interscale/tl/masking.py                  |  23 +-
 src/interscale/train/_trainingplans.py        |  38 +-
 tests/test_anchors.py                         | 392 ++++++++++++
 tests/test_gene_masking.py                    |  52 ++
 18 files changed, 1808 insertions(+), 246 deletions(-)
 create mode 100644 src/interscale/pl/anchor_plots.py
 create mode 100644 src/interscale/tl/anchors.py
 create mode 100644 tests/test_anchors.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 287e41d..46b6051 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,16 @@ and this project adheres to [Semantic Versioning][].
 
 ### Added
 
+- Anchor-point / distance-zone analysis of a latent dimension (`interscale.tl.anchors`,
+  plotted by `interscale.pl.anchor_plots`): `find_anchor_cells` localises the foci a dimension
+  is reading (per-sample tail quantile, then a spatial-coherence filter),
+  `anchor_signed_distance` gives every cell a signed distance to the anchor border,
+  `anchor_zones` cuts that into `anchor core` / `near` / `far` at the local component's reach
+  (`local_reach_um` = `num_layers` x `radius`), and `profile_by_distance` /
+  `anchor_enrichment` profile expression and composition against it. This is the evidence a
+  *global* dimension carries a long-range interaction rather than a neighbourhood effect:
+  every distance plot marks the local reach, past which no cell's neighbourhood extends.
+
 - Gene-wise (per-entry) masking for the node-level reconstruction task, alongside the existing
   whole-cell masking, selected by `dataset.mask_strategy` (`"node"` | `"gene"`). Under `"gene"`
   a Bernoulli subset of `(cell, gene)` entries is blanked in every cell and the loss and every
diff --git a/docs/api/tools.md b/docs/api/tools.md
index 9a612bd..24aae29 100644
--- a/docs/api/tools.md
+++ b/docs/api/tools.md
@@ -45,3 +45,39 @@ Downstream InterScale's output can be used for gene, cell and tissue level analy
     gene_loadings
     calculate_gene_ranks
 ```
+
+## Anchors and distance zones
+
+Where a latent dimension anchors on a slide, and how expression behaves at range from there.
+The distance axis is the point: the local component reaches at most
+{func}`~interscale.tl.local_reach_um` micrometres, so structure beyond that abscissa cannot
+have come from it.
+
+```{eval-rst}
+.. currentmodule:: interscale.tl
+
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    local_reach_um
+    find_anchor_cells
+    anchor_signed_distance
+    anchor_zones
+    profile_by_distance
+    anchor_enrichment
+```
+
+```{eval-rst}
+.. currentmodule:: interscale.pl
+
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    anchor_map
+    signed_distance_map
+    zone_map
+    gene_maps
+    distance_profile
+```
diff --git a/src/interscale/evaluation/net_streams.py b/src/interscale/evaluation/net_streams.py
index c47bc8e..eb33730 100644
--- a/src/interscale/evaluation/net_streams.py
+++ b/src/interscale/evaluation/net_streams.py
@@ -694,7 +694,18 @@ def compute_hierarchical_net_flow(
     return final_results
 
 
-def plot_global_directionality(mean_df, std_df, only_positive=True, title="Net Flow", figsize=(8, 6)):
+def plot_global_directionality(
+    mean_df,
+    std_df,
+    only_positive=True,
+    title="Net Flow",
+    figsize=(8, 6),
+    ax=None,
+    vmin=None,
+    vmax=None,
+    show=None,
+    order=None,
+):
     """
     Visualize aggregated net flow between cell types using a dot plot.
 
@@ -714,23 +725,63 @@ def plot_global_directionality(mean_df, std_df, only_positive=True, title="Net F
     title : str, optional
         Title of the plot.
     figsize : tuple, optional
-        Dimensions of the figure.
-    """
-    # 1. Melt the data
-    plot_data = mean_df.reset_index().melt(id_vars="index")
-    plot_data.columns = ["Sender", "Receiver", "Flow"]
+        Dimensions of the figure. Ignored when ``ax`` is given.
+    ax : matplotlib.axes.Axes, optional
+        Draw into an existing axes instead of a new figure. Needed to place two conditions --
+        or two model variants of an ablation -- side by side for comparison; without it every
+        call opened its own figure and the panels could only be stacked as separate outputs.
+    vmin, vmax : float, optional
+        Colour limits. Pass the same pair to every panel of a multi-panel figure: seaborn
+        otherwise normalises each panel to its own range, so equal colours in two panels would
+        stand for different flow magnitudes.
+    show : bool, optional
+        Whether to call ``plt.show()``. Defaults to True when ``ax`` is None (the previous
+        behaviour, which suits a single standalone plot) and False otherwise, since a composed
+        figure must not be shown until every panel is drawn.
+    order : list, optional
+        Category order for both axes. Defaults to alphabetical, which is right for cell-type
+        names but wrong for any label set that is already ordered -- the SHH rings sort to
+        ``SHH, far, medium, near`` alphabetically, which interleaves near and far and destroys
+        the distance axis the rings exist to represent.
 
-    # 2. Define and Enforce the same order for both axes
-    categories = sorted(mean_df.index.unique())
+    Returns
+    -------
+    matplotlib.axes.Axes
+        The axes drawn into, so a caller composing a figure can adjust it further.
+    """
+    # 1. Melt the data.
+    #
+    # The axes are renamed rather than relying on `reset_index()` producing a column literally
+    # called "index", which it only does when the index has no name. A `mean_df` built by
+    # `DataFrame.pivot` carries the pivot's index/column names, and melting it by "index" raised
+    # KeyError.
+    mean_df = mean_df.rename_axis(index="Sender", columns="Receiver")
+    std_df = std_df.rename_axis(index="Sender", columns="Receiver")
+
+    plot_data = mean_df.reset_index().melt(id_vars="Sender", var_name="Receiver", value_name="Flow")
+
+    # 2. Calculate Consistency.
+    #
+    # Merged on the (Sender, Receiver) pair, not zipped by position. `melt` emits column-major
+    # (every sender of receiver 1, then every sender of receiver 2, ...) while `.values.flatten()`
+    # is row-major, so pairing the two positionally attached each flow to the standard deviation
+    # of its TRANSPOSED entry -- invisible for a symmetric matrix, and wrong for every net flow,
+    # which is antisymmetric by construction.
+    std_long = std_df.reset_index().melt(id_vars="Sender", var_name="Receiver", value_name="Std")
+    plot_data = plot_data.merge(std_long, on=["Sender", "Receiver"], how="left")
+    plot_data["Consistency"] = 1 / (plot_data["Std"] + 1e-9)
+
+    # 3. Define and Enforce the same order for both axes
+    categories = list(order) if order is not None else sorted(mean_df.index.unique())
+    if order is not None:
+        missing = set(mean_df.index.unique()) - set(categories)
+        if missing:
+            raise ValueError(f"`order` does not cover every category in mean_df: {sorted(missing)}")
 
     # Convert to Categorical with a fixed list of categories
     plot_data["Sender"] = pd.Categorical(plot_data["Sender"], categories=categories)
     plot_data["Receiver"] = pd.Categorical(plot_data["Receiver"], categories=categories)
 
-    # 3. Calculate Consistency
-    std_flat = std_df.values.flatten()
-    plot_data["Consistency"] = 1 / (std_flat + 1e-9)
-
     # Filter only positive flows
     if only_positive:
         plot_data = plot_data[plot_data["Flow"] > 0]
@@ -738,7 +789,10 @@ def plot_global_directionality(mean_df, std_df, only_positive=True, title="Net F
         plot_data = plot_data[plot_data["Flow"].abs() > 0]
 
     # 4. Plotting
-    plt.figure(figsize=figsize)
+    if show is None:
+        show = ax is None
+    if ax is None:
+        _, ax = plt.subplots(figsize=figsize)
 
     # Now Seaborn will use the categorical order automatically
     sns.scatterplot(
@@ -748,16 +802,26 @@ def plot_global_directionality(mean_df, std_df, only_positive=True, title="Net F
         size="Consistency",
         hue="Flow",
         palette="YlOrRd" if only_positive else "coolwarm",
+        hue_norm=None if vmin is None and vmax is None else (vmin, vmax),
         sizes=(20, 500),
+        ax=ax,
     )
 
-    # Force axes to show all categories in the right order
-    plt.xticks(ticks=range(len(categories)), labels=categories, rotation=45)
-    plt.yticks(ticks=range(len(categories)), labels=categories)
+    # Force axes to show all categories in the right order.
+    #
+    # set_xticks BEFORE set_xticklabels, and both on the axes rather than through pyplot: with an
+    # empty `plot_data` (every flow filtered out, which `only_positive` can do for a whole panel)
+    # seaborn draws no categorical axis at all, and labelling ticks that do not exist raises.
+    ax.set_xticks(range(len(categories)))
+    ax.set_xticklabels(categories, rotation=45, ha="right")
+    ax.set_yticks(range(len(categories)))
+    ax.set_yticklabels(categories)
 
     # Move legend outside
-    plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0)
+    ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0)
 
-    plt.title(title)
-    plt.tight_layout()
-    plt.show()
+    ax.set_title(title)
+    if show:
+        plt.tight_layout()
+        plt.show()
+    return ax
diff --git a/src/interscale/geome_dataloader.py b/src/interscale/geome_dataloader.py
index 7446b09..a52ab36 100644
--- a/src/interscale/geome_dataloader.py
+++ b/src/interscale/geome_dataloader.py
@@ -149,10 +149,10 @@ def _assign_random_mask(self, data: BaseData) -> None:
         Under `mask_strategy="gene"` a second attribute `data.gene_mask` `[N, G]` is written,
         with each (cell, gene) entry drawn independently at `mask_percentage`. `data.mask` is then
         the row-wise OR of it -- i.e. "this cell is a supervision target" -- which is what the
-        rest of the pipeline (padding, `_process_batch_for_metrics`) keys on. 
+        rest of the pipeline (padding, `_process_batch_for_metrics`) keys on.
 
         `gene_mask` is a node-level attribute of shape `[num_nodes, ...]`, so PyG collates it by
-        concatenating along dim 0 exactly like `x` -- no custom `__cat_dim__` needed. 
+        concatenating along dim 0 exactly like `x` -- no custom `__cat_dim__` needed.
         """
         if self.mask_strategy == "gene":
             gene_mask = sample_gene_mask(data.num_nodes, data.x.shape[1], self.mask_percentage)
diff --git a/src/interscale/model/base/_base_model.py b/src/interscale/model/base/_base_model.py
index b6c5fa1..5552df2 100644
--- a/src/interscale/model/base/_base_model.py
+++ b/src/interscale/model/base/_base_model.py
@@ -36,6 +36,11 @@ class _SAVE_KEYS_NT(NamedTuple):
 
 SAVE_KEYS = _SAVE_KEYS_NT()
 
+# State entries that may legitimately be absent from an older checkpoint. The PCA front-end's
+# buffers were added after some checkpoints were written; `BaseModel.load` warns and carries on
+# for these, and raises for anything else. See tests/test_global_pca_persistence.py.
+_OPTIONAL_STATE_PREFIXES = ("pca_",)
+
 
 # adjusted from scvi-tools
 # https://github.com/scverse/scvi-tools/blob/main/src/scvi/model/base/_base_model.py
@@ -606,6 +611,7 @@ def load(
         postfix: str | None = None,
         wandb_save: bool = False,
         enable_remapping: bool = True,
+        allow_partial_load: bool = False,
     ):
         """Load a saved model.
 
@@ -625,6 +631,10 @@ def load(
             Whether this is a global component model.
         wandb_save
             Whether this was saved via wandb.
+        allow_partial_load
+            If False (default), a checkpoint whose keys do not match the model this cfg builds
+            raises, instead of silently leaving the unmatched tensors randomly initialised. Set
+            True only when a partly initialised model is deliberate.
         enable_remapping
             Whether to enable automatic state dict key remapping.
 
@@ -703,10 +713,33 @@ def load(
         # Load the state dict
         missing_keys, unexpected_keys = model.module.load_state_dict(state_dict, strict=False)
 
-        if missing_keys:
-            print(f"Warning: Missing keys when loading state dict: {missing_keys}")
-        if unexpected_keys:
-            print(f"Warning: Unexpected keys when loading state dict: {unexpected_keys}")
+        # strict=False is load_state_dict's "fill in what you can and say nothing" mode. It is
+        # needed for the legacy carve-out below, but on its own it turns an architecture mismatch
+        # -- a config that does not describe the checkpoint -- into a model whose unmatched
+        # tensors keep their random initialisation. That model runs, produces plausible-looking
+        # embeddings and gene loadings, and is wrong. Printed warnings do not survive a notebook
+        # with hundreds of lines of output, so anything but the documented benign case raises.
+        benign_missing = [k for k in missing_keys if k.rsplit(".", 1)[-1].startswith(_OPTIONAL_STATE_PREFIXES)]
+        hard_missing = [k for k in missing_keys if k not in benign_missing]
+
+        if benign_missing:
+            print(
+                f"Warning: checkpoint predates these buffers, loading without them: {benign_missing}. "
+                f"A PCA front-end will refit on the first batch it sees rather than reusing the "
+                f"basis it was trained with (see tests/test_global_pca_persistence.py)."
+            )
+
+        if (hard_missing or unexpected_keys) and not allow_partial_load:
+            raise RuntimeError(
+                f"Checkpoint does not match the model this cfg builds, so the mismatched tensors "
+                f"would keep their random initialisation and the model would be silently wrong.\n"
+                f"  missing from the checkpoint (left random): {hard_missing}\n"
+                f"  present in the checkpoint but not in the model: {unexpected_keys}\n"
+                f"  checkpoint: {model_save_path}\n"
+                f"Check that dataset/task, n_embed, decoder type and dual_decoder in cfg match the "
+                f"run that wrote it. Pass allow_partial_load=True only if a partly initialised "
+                f"model is genuinely what you want."
+            )
 
         model.is_trained_ = True
 
diff --git a/src/interscale/model/combined_model.py b/src/interscale/model/combined_model.py
index 477d3e5..d5827ec 100644
--- a/src/interscale/model/combined_model.py
+++ b/src/interscale/model/combined_model.py
@@ -1,6 +1,5 @@
 import numpy as np
 import pandas as pd
-import torch
 from anndata import AnnData
 from yacs.config import CfgNode as CN
 
@@ -114,8 +113,9 @@ def get_model_output(self, adata: AnnData | None = None, prefix: str = ""):
             ## Local model output
             local_embedding = self.module.local_module.forward(batch.x, batch.edge_index)
             if self._cfg.model.decoder.dual_decoder:
-                mask_idx = torch.arange(local_embedding.size(0), device=local_embedding.device)
-                y_pred_local = self.module.predict_local(local_embedding, mask_idx)
+                # node_idx=None: inference decodes every cell, in the batch's own order. This
+                # used to pass an explicit arange, which was the same identity gather.
+                y_pred_local = self.module.predict_local(local_embedding)
 
             sample_mask_local = local_embeddings_df.index.isin(batch.obs_names.numpy().astype(int).astype(str))
             local_embeddings_df.loc[sample_mask_local] = local_embedding.detach().cpu().numpy()
diff --git a/src/interscale/module/base/_base_global_module.py b/src/interscale/module/base/_base_global_module.py
index eb9632b..7045331 100644
--- a/src/interscale/module/base/_base_global_module.py
+++ b/src/interscale/module/base/_base_global_module.py
@@ -82,9 +82,7 @@ def create_gex_embedding(self, embeddings: torch.Tensor, type: Literal["PCA", "N
             if not bool(self.pca_fitted_):
                 self.pca.fit(embeddings)
                 self.pca_mean_.copy_(torch.as_tensor(self.pca.mean_, dtype=self.pca_mean_.dtype))
-                self.pca_components_.copy_(
-                    torch.as_tensor(self.pca.components_, dtype=self.pca_components_.dtype)
-                )
+                self.pca_components_.copy_(torch.as_tensor(self.pca.components_, dtype=self.pca_components_.dtype))
                 self.pca_fitted_.fill_(True)
             return self._pca_transform(embeddings)
         elif type == "NMF":
@@ -108,11 +106,12 @@ def _pca_transform(self, embeddings):
         components = self.pca_components_.detach().cpu().numpy()
         return (np.asarray(embeddings, dtype=np.float64) - mean) @ components.T
 
-    def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, pad_index_nodes, mask_idx_tensor):
-        """Process batch to extract y_true and adjusted_mask_idx for metrics calculation.
+    def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, pad_index_nodes):
+        """Gather the ground truth for every cell the transformer actually produced, in its order.
 
-        mask_idx = torch.tensor([0, 2, 3, 7, 8])
-        pad_index_nodes = [[0, 1, 2, 3], [0, 1], [0, 1, 2, 3]]
+        ``pad_batch`` emits one row per kept cell, per graph, in the order named by
+        ``pad_index_nodes`` -- which is a per-graph, graph-local index. This turns that into a
+        single batch-global index so predictions, truth and masks can all be lined up by it.
 
         Parameters
         ----------
@@ -123,126 +122,70 @@ def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, p
         prediction_level: str
             Level of prediction ('node' or 'graph')
         pad_index_nodes: List[List[int]]
-            List of padded node indices: [B, S] or [B,N] if number of nodes in graph are smaller than max_seq_len (S)
-        mask_idx_tensor: torch.Tensor
-            Indices of masked nodes of shape [N_masked_nodes] with range [0, N_nodes-1]
+            Per-graph indices of the cells the padding kept: ``[B, S]``, or ``[B, N]`` when the
+            graph is smaller than max_seq_len.
 
         Returns
         -------
-        y_true: torch.Tensor [N_included_nodes, C] (classification) or [N_included_nodes, F] (regression)
-            Ground truth values
-        adjusted_mask_idx: torch.Tensor [N_masked nodes]
-            Adjusted indices for masked nodes
-        entry_mask: torch.Tensor [N_included_nodes, F] | None
-            The batch's `gene_mask` gathered and reordered exactly like `y_true`, so that
-            `entry_mask[adjusted_mask_idx]` lines up entry-for-entry with the scored predictions.
-            `None` whenever the batch carries no gene mask (cell masking, or classification).
+        y_true: torch.Tensor [N_kept, C] (classification) or [N_kept, F] (regression)
+            Ground truth for the kept cells, in transformer-output order.
+        padded_node_idx: torch.Tensor [N_kept]
+            The batch-global node index of each kept row. Index anything that lives in the
+            batch's own node order (a local-decoder output, ``batch.mask``, ``batch.gene_mask``)
+            with this to bring it into the same order.
+        entry_mask: torch.Tensor [N_kept, F] | None
+            Which entries the *loss* is scored on, in the same order. Under gene masking this is
+            the batch's ``gene_mask``; under cell masking it is the per-cell mask broadcast over
+            all F genes, so a masked cell is all-True and an unmasked cell all-False -- which is
+            what lets ``masked_loss`` recognise a whole-row mask and keep the row structure that
+            row-normalising criteria need. ``None`` for classification, where the target is a
+            label per cell and there are no entries to mask.
+
+        Notes
+        -----
+        This replaces an earlier version that returned ``adjusted_mask_idx`` -- the positions of
+        the masked cells within the kept rows -- which every caller then used to subset y_pred and
+        y_true down to masked cells only. Metrics are now computed over ALL kept cells and only
+        the loss is restricted, so the subsetting is gone and with it the index arithmetic that
+        matched ``mask_idx`` against ``pad_indices`` graph by graph.
         """
         assert prediction_level == "node", "Node specific retrieval only necessary for node-level prediction."
 
-        nr_batches = batch.batch[-1] + 1
+        nr_batches = int(batch.batch[-1]) + 1
         device = batch.x.device
 
-        # Pre-compute batch boundaries (O(B×N) total)
         batch_sizes = torch.tensor([batch.batch.eq(i).sum().item() for i in range(nr_batches)], device=device)
         batch_starts = torch.cat([torch.tensor([0], device=device), batch_sizes.cumsum(0)[:-1]])
-        batch_ends = batch_starts + batch_sizes
-
-        # Pre-compute cumulative offsets for adjusted indices
-        pad_lengths = torch.tensor([len(pad) for pad in pad_index_nodes], device=device)
-        cumulative_offsets = torch.cat([torch.tensor([0], device=device), pad_lengths.cumsum(0)[:-1]])
-
-        adjusted_mask_idx_list = []
-        y_true_list = []
-
-        # Gathered in lockstep with y_true below. Gated on the configured strategy, not merely on
-        # the attribute being present, so a stale `gene_mask` on a reused Data object cannot turn
-        # a cell-masking run into a gene-masking one. Only regression has entries to mask; a
-        # classification target is a label per cell, not a gene vector.
-        use_gene_mask = self.mask_strategy == "gene" and "regression" in prediction_task
-        gene_mask = getattr(batch, "gene_mask", None) if use_gene_mask else None
-        entry_mask_list = [] if gene_mask is not None else None
-
-        for i in range(nr_batches):
-            batch_start = batch_starts[i].item()
-            batch_end = batch_ends[i].item()
-
-            # Find masked indices in this batch range (vectorized)
-            mask_in_batch = (mask_idx_tensor >= batch_start) & (mask_idx_tensor < batch_end)
-            batch_mask_idx = mask_idx_tensor[mask_in_batch]
-
-            if len(batch_mask_idx) == 0:
-                # Extract y_true even if no masked nodes
-                mask = batch.batch.eq(i)
-                if "classification" in prediction_task:
-                    y_true_list.append(batch.y[mask][pad_index_nodes[i]])
-                elif "regression" in prediction_task:
-                    y_true_list.append(batch.x[mask][pad_index_nodes[i]])
-                if entry_mask_list is not None:
-                    entry_mask_list.append(gene_mask[mask][pad_index_nodes[i]])
-                continue
-
-            # Create pad_indices tensor once
-            pad_indices = torch.tensor(pad_index_nodes[i], device=device) + batch_start
-
-            # Vectorized intersection and position finding
-            # Use broadcasting: [M, 1] == [1, P] creates [M, P] boolean matrix
-            matches = batch_mask_idx.unsqueeze(1) == pad_indices.unsqueeze(0)  # [M, P]
-            is_in_pad = matches.any(dim=1)  # [M] - which masked nodes are in pad_indices
-
-            if is_in_pad.any():
-                # Get positions of matches in pad_indices (first occurrence)
-                positions_in_pad = matches.long().argmax(dim=1)[is_in_pad]  # [M_valid]
-
-                # Adjust indices with cumulative offset
-                adjusted_indices = positions_in_pad + cumulative_offsets[i]
-                adjusted_mask_idx_list.append(adjusted_indices)
-
-            # Extract y_true for included nodes
-            mask = batch.batch.eq(i)
-            if "classification" in prediction_task:
-                y_true_list.append(batch.y[mask][pad_index_nodes[i]])
-            elif "regression" in prediction_task:
-                y_true_list.append(batch.x[mask][pad_index_nodes[i]])
-            else:
-                raise Exception("Choose a valid prediction task (classification or regression).")
-            if entry_mask_list is not None:
-                entry_mask_list.append(gene_mask[mask][pad_index_nodes[i]])
-
-        # Concatenate results
-        y_true = torch.cat(y_true_list, dim=0)
-        entry_mask = torch.cat(entry_mask_list, dim=0) if entry_mask_list else None
-        adjusted_mask_idx = (
-            torch.cat(adjusted_mask_idx_list, dim=0)
-            if adjusted_mask_idx_list
-            else torch.tensor([], device=device, dtype=torch.long)
+
+        # Graph-local kept indices -> batch-global node indices, concatenated in output order.
+        padded_node_idx = torch.cat(
+            [
+                torch.tensor(pad_index_nodes[i], device=device, dtype=torch.long) + batch_starts[i]
+                for i in range(nr_batches)
+            ]
         )
 
-        # Every adjusted index must address a valid row of y_true. This is the invariant the
-        # offset arithmetic can actually violate, so it stays.
-        #
-        # A second assertion used to sit here requiring `adjusted_mask_idx.max() >
-        # len(pad_index_nodes[0])` whenever nr_batches > 1, i.e. that masked nodes came from more
-        # than just the first graph. It was wrong twice over. Graph i's indices start at
-        # cumulative_offsets[i], so a masked node that is the *first kept node* of graph 1 gets
-        # index exactly len(pad_index_nodes[0]) -- a valid position that `>` rejected. That is
-        # reachable whenever the last batch of an epoch holds two graphs and the second is small
-        # enough for its single masked node to be node 0, which killed a legnini23 run at epoch 22.
-        # It also asserted a property that is not invariant: if a graph's masked set exceeds
-        # max_seq_len, _select_masked_nodes legitimately drops some, and that case is caught with
-        # an accurate message by the length checks in the combined modules' _common_step, where
-        # the local branch's masked-node count is compared against the global branch's.
-        if len(adjusted_mask_idx) > 0:
-            assert adjusted_mask_idx.max() < len(y_true), (
-                f"Mismatch: max(adjusted_mask_idx): {adjusted_mask_idx.max()}, len(y_true): {len(y_true)}"
-            )
+        if "classification" in prediction_task:
+            y_true = batch.y[padded_node_idx]
+            entry_mask = None
+        elif "regression" in prediction_task:
+            y_true = batch.x[padded_node_idx]
+            use_gene_mask = self.mask_strategy == "gene"
+            gene_mask = getattr(batch, "gene_mask", None) if use_gene_mask else None
+            if gene_mask is not None:
+                entry_mask = gene_mask[padded_node_idx].bool()
+            else:
+                # Whole-cell masking expressed at entry level, so one mask covers both strategies.
+                entry_mask = batch.mask[padded_node_idx].bool().unsqueeze(1).expand_as(y_true)
+        else:
+            raise Exception("Choose a valid prediction task (classification or regression).")
 
         if entry_mask is not None:
             assert entry_mask.shape == y_true.shape, (
                 f"Mismatch: entry_mask.shape: {tuple(entry_mask.shape)}, y_true.shape: {tuple(y_true.shape)}"
             )
 
-        return y_true, adjusted_mask_idx, entry_mask
+        return y_true, padded_node_idx, entry_mask
 
     # def _process_batch_for_metrics(self, batch, prediction_task, prediction_level, pad_index_nodes, mask_idx_tensor):
     #     """Process batch to extract y_true and adjusted_mask_idx for metrics calculation.
@@ -358,10 +301,12 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
         attn_matrix: torch.Tensor
             Stacked per-layer attention weights.
         entry_mask: torch.Tensor | None
-            Size: [N, F] under gene masking, marking the entries the loss is scored on.
+            Size: [N, F], marking the entries the LOSS is scored on. y_pred/y_true cover every
+            cell the transformer produced, not just the masked ones -- see
+            `_process_batch_for_metrics`.
         """
         # Mask nodes  - before GEX embedding because otherwise embedding contains information about masked nodes
-        batch_masked, mask_idx, _ = self._common_step_masking(batch)
+        batch_masked, _, _ = self._common_step_masking(batch)
         if hasattr(batch_masked, "embeddings"):
             embedding = batch_masked.embeddings
         else:
@@ -388,13 +333,15 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
             y_true = batch.y[batch.ptr[:-1]]
             entry_mask = None
         else:
-            y_true, adjusted_mask_idx, entry_mask = self._process_batch_for_metrics(
-                batch, prediction_task, prediction_level, pad_index_nodes, mask_idx
+            y_true, padded_node_idx, entry_mask = self._process_batch_for_metrics(
+                batch, prediction_task, prediction_level, pad_index_nodes
             )
-            y_pred = y_pred[adjusted_mask_idx]
-            y_true = y_true[adjusted_mask_idx]
-            if entry_mask is not None:
-                entry_mask = entry_mask[adjusted_mask_idx]
+            if entry_mask is None:
+                # Node-level classification: the masked cells ARE the supervision targets, so
+                # both loss and metrics stay restricted to them. Only reconstruction moved to
+                # scoring every cell.
+                keep = batch.mask[padded_node_idx].bool()
+                y_pred, y_true = y_pred[keep], y_true[keep]
 
         assert len(y_pred) == len(y_true), "y_pred and y_true are not consistent"
         assert not torch.any(torch.isnan(y_pred)), "y_pred contains NaN values"
diff --git a/src/interscale/module/base/_base_local_module.py b/src/interscale/module/base/_base_local_module.py
index c756d40..a927902 100644
--- a/src/interscale/module/base/_base_local_module.py
+++ b/src/interscale/module/base/_base_local_module.py
@@ -37,13 +37,14 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
         global_embedding: torch.Tensor
             Size: [N, E]
         y_pred: torch.Tensor
-            Size: [B, C] (classification) or [B, F] (regression)
+            Size: [B, C] (classification, masked cells) or [N, F] (regression, ALL cells)
         y_true: torch.Tensor
-            Size: [B, ] (classification) or [B, F] (regression)
+            Size: [B, ] (classification, masked cells) or [N, F] (regression, ALL cells)
         attn: None
             This module has no attention; returned for a uniform `_common_step` contract.
         entry_mask: torch.Tensor | None
-            Size: [B, F] under gene masking, marking the entries the loss is scored on.
+            Size: [N, F] for regression, marking the entries the LOSS is scored on; the metrics
+            use every cell. None for classification.
         """
         # Mask nodes
         batch_masked, mask_idx, entry_mask = self._common_step_masking(batch)
@@ -59,20 +60,25 @@ def _common_step(self, batch, prediction_task: str, prediction_level: Literal["n
         )
         assert y_pred.isnan().sum() == 0, "y_pred contains NaN values"
 
-        y_pred = y_pred[mask_idx]
-
         if "classification" in prediction_task:
+            # The masked cells ARE the supervision targets here, so both loss and metrics stay
+            # restricted to them. Only reconstruction moved to scoring every cell.
+            y_pred = y_pred[mask_idx]
             y_true = batch.y[mask_idx]  # batch without mask because constant otherwise
             assert y_true.shape == y_pred.shape
             # Class labels are not gene entries, so there is nothing for an entry mask to select.
             return local_embedding, None, y_pred, y_true, None, None
 
         if "regression" in prediction_task:
-            y_true = batch.x[mask_idx]  # batch without mask because constant otherwise
+            # Every cell is scored. `entry_mask` says which entries the LOSS uses: the gene mask
+            # under gene masking, the per-cell mask broadcast over all genes under cell masking
+            # (masked_loss recognises the whole-row form and subsets rows, so the loss is
+            # unchanged from when this returned masked cells only).
+            y_true = batch.x  # batch without mask because constant otherwise
             assert y_true.shape == y_pred.shape
-            if entry_mask is not None:
-                entry_mask = entry_mask[mask_idx]
-                assert entry_mask.shape == y_pred.shape
+            if entry_mask is None:
+                entry_mask = batch.mask.bool().unsqueeze(1).expand_as(y_true)
+            assert entry_mask.shape == y_pred.shape
             return local_embedding, None, y_pred, y_true, None, entry_mask
 
         assert False, "Prediction task not supported"
diff --git a/src/interscale/module/combined_module/combined_module.py b/src/interscale/module/combined_module/combined_module.py
index 6529566..5fef961 100644
--- a/src/interscale/module/combined_module/combined_module.py
+++ b/src/interscale/module/combined_module/combined_module.py
@@ -67,10 +67,11 @@ def forward(self, batch_masked):
     def _common_step(self, batch, prediction_task, prediction_level: Literal["node", "graph"]):
         """Shared step between train, val and test.
 
-        The trailing `entry_mask` is `None` under cell masking and `[N_masked, F]` under gene
-        masking, where it marks the entries the loss must be restricted to.
+        y_pred/y_true cover every cell the transformer produced. The trailing `entry_mask`
+        `[N, F]` marks the entries the LOSS is restricted to -- the gene mask under gene masking,
+        the per-cell mask broadcast over all genes under cell masking.
         """
-        batch_masked, mask_idx, _ = self._common_step_masking(batch)
+        batch_masked, _, _ = self._common_step_masking(batch)
 
         local_embedding, global_embedding, src_padding_mask, pad_index_nodes, attention_mask, attn_matrix = (
             self.forward(batch_masked)
@@ -81,13 +82,12 @@ def _common_step(self, batch, prediction_task, prediction_level: Literal["node",
             y_true = batch.y[batch.ptr[:-1]]
             entry_mask = None
         else:
-            y_true, adjusted_mask_idx, entry_mask = self.global_module._process_batch_for_metrics(
-                batch, prediction_task, prediction_level, pad_index_nodes, mask_idx
+            y_true, padded_node_idx, entry_mask = self.global_module._process_batch_for_metrics(
+                batch, prediction_task, prediction_level, pad_index_nodes
             )
-            y_pred = y_pred[adjusted_mask_idx]
-            y_true = y_true[adjusted_mask_idx]
-            if entry_mask is not None:
-                entry_mask = entry_mask[adjusted_mask_idx]
+            if entry_mask is None:  # node-level classification -- see GlobalModule._common_step
+                keep = batch.mask[padded_node_idx].bool()
+                y_pred, y_true = y_pred[keep], y_true[keep]
 
         assert len(y_pred) == len(y_true), "y_pred and y_true are not consistent"
         assert not torch.any(torch.isnan(y_pred)), "y_pred contains NaN values"
diff --git a/src/interscale/module/combined_module/dual_decoder_combined_module.py b/src/interscale/module/combined_module/dual_decoder_combined_module.py
index 7f9e00e..a8973ec 100644
--- a/src/interscale/module/combined_module/dual_decoder_combined_module.py
+++ b/src/interscale/module/combined_module/dual_decoder_combined_module.py
@@ -61,26 +61,33 @@ def __init__(self, cfg: CN, **base_module_kwargs):
         self._n_masked_nodes = None
         self._is_graph_level = False
 
-    def predict_local(self, local_embedding, mask_idx):
-        """Predict with the local decoder on masked nodes.
+    def predict_local(self, local_embedding, node_idx=None):
+        """Predict with the local decoder, optionally reordered/subset to ``node_idx``.
 
         Parameters
         ----------
         local_embedding: torch.Tensor
             Size: [N, E]
-        mask_idx: torch.Tensor
-            Indices of masked nodes. Size: [N_masked_nodes, ]
+        node_idx: torch.Tensor | None
+            Batch-global node indices to gather, in the order wanted. Pass
+            `_process_batch_for_metrics`'s ``padded_node_idx`` to line the local predictions up
+            with the global branch's output. ``None`` returns all nodes in batch order.
 
         Returns
         -------
         y_pred_local: torch.Tensor
-            Size: [N_masked_nodes, C] or [N_masked_nodes, F]
+            Size: [len(node_idx), C] or [len(node_idx), F]
+
+        Notes
+        -----
+        This used to take ``mask_idx`` -- the masked nodes in the *batch's* order -- while the
+        global branch was subset by ``adjusted_mask_idx``, the masked nodes in the *padded* order.
+        The two halves were then concatenated and compared row against row, which is only correct
+        when the padding keeps every cell in its original order. Gathering both branches by
+        ``padded_node_idx`` makes the pairing correct by construction.
         """
-        # Predict on all nodes
         y_pred_all = self.local_module.decoder.forward(local_embedding)
-        # Filter to masked nodes
-        y_pred_local = y_pred_all[mask_idx]
-        return y_pred_local
+        return y_pred_all if node_idx is None else y_pred_all[node_idx]
 
     def predict_global(self, global_embedding, src_padding_mask, prediction_level):
         """Predict with the global decoder.
@@ -128,83 +135,58 @@ def forward(self, batch_masked):
     def _common_step(self, batch, prediction_task, prediction_level: Literal["node", "graph"]):
         """Shared step between train, val and test.
 
-        Returns predictions and ground truth for both local and global decoders
-        on masked tokens, which can be combined in the loss function.
+        Returns predictions and ground truth for both local and global decoders over EVERY cell
+        the transformer produced -- not only the masked ones. `entry_mask_combined` marks the
+        entries the loss is restricted to; the metrics use everything.
         """
-        batch_masked, mask_idx, node_entry_mask = self._common_step_masking(batch)
+        batch_masked, _, _ = self._common_step_masking(batch)
 
         local_embedding, global_embedding, src_padding_mask, pad_index_nodes, attention_mask, attn = self.forward(
             batch_masked
         )
 
-        # Predict from local embedding on masked nodes
-        y_pred_local = self.predict_local(local_embedding, mask_idx)
-
-        # Predict from global embedding
+        # Predict from global embedding (one row per cell the padding kept)
         y_pred_global = self.predict_global(global_embedding, src_padding_mask, prediction_level)
 
-        # Get ground truth for masked nodes
         if prediction_task == "classification" and prediction_level == "graph":
-            # For graph-level classification, we need to handle this differently
-            # since we have one prediction per graph
+            # One prediction per graph, so there is no local/global pairing to do.
             y_true = batch.y[batch.ptr[:-1]]
-            # For graph level, we can't easily combine local and global
-            # So we'll use global predictions only for graph level
             y_pred_combined = y_pred_global
             y_true_combined = y_true
-
-            # Store metadata for graph level (only global predictions)
             self._n_masked_nodes = None
             self._is_graph_level = True
             entry_mask_combined = None
         elif prediction_level == "node":
-            # For node-level predictions, get ground truth for masked nodes
-            y_true, adjusted_mask_idx, entry_mask = self.global_module._process_batch_for_metrics(
-                batch, prediction_task, prediction_level, pad_index_nodes, mask_idx
+            y_true, padded_node_idx, entry_mask = self.global_module._process_batch_for_metrics(
+                batch, prediction_task, prediction_level, pad_index_nodes
             )
-            y_true_masked = y_true[adjusted_mask_idx]
-            entry_mask_masked = entry_mask[adjusted_mask_idx] if entry_mask is not None else None
 
-            # Filter global predictions to masked nodes (same indices as y_true)
-            y_pred_global_masked = y_pred_global[adjusted_mask_idx]
+            # Both branches gathered by the SAME index, so row i of each is the same cell.
+            y_pred_local = self.predict_local(local_embedding, padded_node_idx)
 
-            # Store split point: first half is local, second half is global
-            n_masked = len(y_pred_local)
-            self._n_masked_nodes = n_masked
-            self._is_graph_level = False
+            if entry_mask is None:
+                # Node-level classification: the masked cells are the supervision targets, so
+                # both halves stay restricted to them. Only reconstruction scores every cell.
+                keep = batch.mask[padded_node_idx].bool()
+                y_pred_local, y_pred_global, y_true = y_pred_local[keep], y_pred_global[keep], y_true[keep]
 
-            # Combine local and global predictions
-            # Both should have the same number of masked nodes
-            assert len(y_pred_local) == len(y_pred_global_masked), (
-                f"Local and global predictions have different lengths: {len(y_pred_local)} vs {len(y_pred_global_masked)}"
+            assert y_pred_local.shape == y_pred_global.shape, (
+                f"Local and global predictions differ in shape: "
+                f"{tuple(y_pred_local.shape)} vs {tuple(y_pred_global.shape)}"
             )
-            assert len(y_true_masked) == len(y_pred_local), (
-                f"Ground truth and local predictions have different lengths: {len(y_true_masked)} vs {len(y_pred_local)}"
-            )
-            assert len(y_true_masked) == len(y_pred_global_masked), (
-                f"Ground truth and global predictions have different lengths: {len(y_true_masked)} vs {len(y_pred_global_masked)}"
+            assert len(y_true) == len(y_pred_local), (
+                f"Ground truth and predictions differ in length: {len(y_true)} vs {len(y_pred_local)}"
             )
 
-            # Concatenate predictions: [N_masked, C] + [N_masked, C] -> [2*N_masked, C]
-            # This allows the loss function to compute loss on both predictions
-            y_pred_combined = torch.cat([y_pred_local, y_pred_global_masked], dim=0)
-            y_true_combined = torch.cat([y_true_masked, y_true_masked], dim=0)
-
-            # The entry mask has to follow the same stacking. The local branch is indexed by
-            # `mask_idx` (the batch's own node order) while the global branch goes through
-            # `adjusted_mask_idx` (the padded, per-graph-subsampled order), so the two halves are
-            # not the same rows in general -- build each half from its own indexing rather than
-            # duplicating one of them.
-            if node_entry_mask is None:
-                entry_mask_combined = None
-            else:
-                entry_mask_local = node_entry_mask[mask_idx]
-                assert entry_mask_local.shape == y_pred_local.shape, (
-                    f"Mismatch: entry_mask_local.shape: {tuple(entry_mask_local.shape)}, "
-                    f"y_pred_local.shape: {tuple(y_pred_local.shape)}"
-                )
-                entry_mask_combined = torch.cat([entry_mask_local, entry_mask_masked], dim=0)
+            self._n_masked_nodes = len(y_pred_local)
+            self._is_graph_level = False
 
+            # [N, C] + [N, C] -> [2N, C], so the loss can score both decoders.
+            y_pred_combined = torch.cat([y_pred_local, y_pred_global], dim=0)
+            y_true_combined = torch.cat([y_true, y_true], dim=0)
+            # The entry mask is per-cell, and both halves are now the same cells in the same
+            # order, so it is simply stacked twice.
+            entry_mask_combined = None if entry_mask is None else torch.cat([entry_mask, entry_mask], dim=0)
         else:
             raise ValueError(f"Invalid prediction level: {prediction_level}")
 
diff --git a/src/interscale/pl/__init__.py b/src/interscale/pl/__init__.py
index d61b00a..a2f1413 100644
--- a/src/interscale/pl/__init__.py
+++ b/src/interscale/pl/__init__.py
@@ -1,3 +1,12 @@
+from .anchor_plots import (
+    ANCHOR_COLORS,
+    ZONE_COLORS,
+    anchor_map,
+    distance_profile,
+    gene_maps,
+    signed_distance_map,
+    zone_map,
+)
 from .config import Plotting, settings
 from .gene_level_plots import dim_importance_elbow, gene_ranks, latent_correlation
 
@@ -7,4 +16,11 @@
     "latent_correlation",
     "dim_importance_elbow",
     "gene_ranks",
+    "ZONE_COLORS",
+    "ANCHOR_COLORS",
+    "anchor_map",
+    "signed_distance_map",
+    "zone_map",
+    "gene_maps",
+    "distance_profile",
 ]
diff --git a/src/interscale/pl/anchor_plots.py b/src/interscale/pl/anchor_plots.py
new file mode 100644
index 0000000..4dd576c
--- /dev/null
+++ b/src/interscale/pl/anchor_plots.py
@@ -0,0 +1,561 @@
+"""Figures for the anchor / distance-zone analysis of :mod:`interscale.tl.anchors`.
+
+The panels here are the visual argument that a *global* InterScale dimension carries a
+long-range interaction rather than a neighbourhood effect:
+
+    anchor, focus = find_anchor_cells(adata, "43_global_emb", 13, key_added="anchor13")
+    dist  = anchor_signed_distance(adata, anchor)
+    zone  = anchor_zones(adata, dist, near_max=local_reach_um(cfg))
+
+    anchor_map(adata, anchor, samples=slides)            # where the dimension anchors
+    signed_distance_map(adata, dist, samples=slides)     # distance to the anchor border
+    zone_map(adata, zone, samples=slides)                # core / near / far
+    gene_maps(adata, genes, samples=slides)              # the dimension's genes, per slide
+    distance_profile(prof, local_reach=local_reach_um(cfg))   # those genes vs distance
+
+Every distance plot marks ``local_reach``, the furthest a cell's own representation can
+travel through the local component (:func:`interscale.tl.local_reach_um`). That line is the
+point of the figure: structure to the right of it is out of the local component's reach for
+*every* cell, so it cannot be a neighbourhood effect.
+
+All functions are plain matplotlib on ``adata.obsm[spatial_key]`` rather than
+``squidpy.pl.spatial_scatter``, because the anchor/zone colourings are categorical overlays
+with a fixed order and the distance colouring is diverging around a hard zero (the border) --
+both need norms and z-order that the generic scatter does not expose.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from matplotlib.colors import TwoSlopeNorm
+from matplotlib.lines import Line2D
+from scipy.sparse import issparse
+
+from interscale.tl.anchors import ZONE_CORE, ZONE_FAR, ZONE_NEAR, ZONE_ORDER
+
+__all__ = [
+    "ZONE_COLORS",
+    "ANCHOR_COLORS",
+    "NA_COLOR",
+    "anchor_map",
+    "signed_distance_map",
+    "zone_map",
+    "gene_maps",
+    "distance_profile",
+]
+
+# Core and near take the dataset palette's "closest to the SHH source" end (`#005F73` /
+# `#88C8B2` in `figures/config.yml`), so a zone figure reads next to the ring figures instead
+# of inventing a second distance colour language. `far` is deliberately NOT the palette's
+# `#EE9B00`: it is the majority of every slide, and a saturated majority reads as the subject
+# of the panel. It gets the neutral tissue tan of the fibrosis-core panels this mirrors, which
+# also makes it the same colour as `outside` in :func:`anchor_map`.
+ZONE_COLORS = {ZONE_CORE: "#005F73", ZONE_NEAR: "#88C8B2", ZONE_FAR: "#D9D3C4"}
+
+# Cells with no zone at all -- a sample where the dimension found no focus, so the signed
+# distance is nan for every cell of it. Zoned and unzoned cells never share a panel (a sample
+# is either wholly anchored or wholly nan), so this only has to be distinguishable from the
+# scale, not from the tan.
+NA_COLOR = "#F0F0F0"
+
+# Binary anchor overlay: everything that is not a focus stays the same neutral tissue tan.
+ANCHOR_COLORS = {False: ZONE_COLORS[ZONE_FAR], True: ZONE_COLORS[ZONE_CORE]}
+
+
+# --------------------------------------------------------------------------------------------
+# shared plumbing
+# --------------------------------------------------------------------------------------------
+
+
+def _vector(adata, x, what):
+    """Accept either an ``adata.obs`` column name or an array over cells."""
+    if isinstance(x, str):
+        if x not in adata.obs:
+            raise KeyError(f"{x!r} not found in adata.obs")
+        v = adata.obs[x]
+    else:
+        v = x
+    v = v.to_numpy() if hasattr(v, "to_numpy") else np.asarray(v)
+    if v.shape[0] != adata.n_obs:
+        raise ValueError(f"{what} has length {v.shape[0]}, expected {adata.n_obs}")
+    return v
+
+
+def _slides(adata, sample_key, samples):
+    """[(label, row indices)] for the requested samples, in the order requested."""
+    grp = adata.obs[sample_key].astype(str).to_numpy()
+    if samples is None:
+        labels = list(map(str, pd.unique(grp)))
+    else:
+        labels = [str(s) for s in np.atleast_1d(samples)]
+        missing = [s for s in labels if s not in set(grp)]
+        if missing:
+            raise KeyError(f"adata.obs['{sample_key}'] has no group(s) {missing}")
+    return [(s, np.where(grp == s)[0]) for s in labels]
+
+
+def _grid(n, ncols, panel_size, axes=None):
+    """(fig, flat list of axes) sized so every panel is `panel_size`."""
+    if axes is not None:
+        axes = np.atleast_1d(np.asarray(axes, dtype=object)).ravel().tolist()
+        if len(axes) < n:
+            raise ValueError(f"need {n} axes, got {len(axes)}")
+        return axes[0].get_figure(), axes[:n]
+    ncols = n if ncols is None else min(int(ncols), n)
+    nrows = int(np.ceil(n / ncols))
+    fig, ax = plt.subplots(nrows, ncols, figsize=(panel_size[0] * ncols, panel_size[1] * nrows), squeeze=False)
+    flat = ax.ravel().tolist()
+    for a in flat[n:]:
+        a.set_axis_off()
+    return fig, flat[:n]
+
+
+def _bare(ax, title=None):
+    ax.set_aspect("equal")
+    ax.set_xticks([])
+    ax.set_yticks([])
+    for side in ax.spines.values():
+        side.set_visible(False)
+    if title is not None:
+        ax.set_title(title, fontsize=10)
+
+
+def _expression(adata, rows, gene, layer):
+    """One gene's values for `rows`, densified for those rows only."""
+    j = int(np.where(adata.var_names == gene)[0][0])
+    X = adata.layers[layer] if layer is not None else adata.X
+    col = X[rows, j]
+    if issparse(col):
+        col = col.toarray()
+    return np.asarray(col, dtype=float).ravel()
+
+
+# --------------------------------------------------------------------------------------------
+# spatial panels
+# --------------------------------------------------------------------------------------------
+
+
+def anchor_map(
+    adata,
+    anchor,
+    *,
+    samples: Sequence[str] | None = None,
+    sample_key: str = "sample",
+    spatial_key: str = "spatial",
+    focus=None,
+    colors: dict | None = None,
+    labels: tuple[str, str] = ("outside", "anchor core"),
+    size: float = 6.0,
+    ncols: int | None = None,
+    panel_size: tuple[float, float] = (3.0, 3.0),
+    axes=None,
+    legend: bool = True,
+    titles: Sequence[str] | None = None,
+    show: bool = True,
+):
+    """Where a dimension anchors: one panel per sample, anchor foci against the rest.
+
+    Parameters
+    ----------
+    anchor
+        Boolean over cells, or an ``adata.obs`` column name -- what
+        :func:`interscale.tl.find_anchor_cells` returned.
+    focus
+        Optional focus ids (its second return value). When given, each focus is numbered at
+        its centroid, which is what makes "there are three anchor points" checkable rather
+        than asserted.
+    titles
+        Panel titles, defaulting to the sample name plus its anchor-cell count.
+
+    Returns
+    -------
+    (fig, axes)
+    """
+    anchor = _vector(adata, anchor, "anchor").astype(bool)
+    focus = None if focus is None else _vector(adata, focus, "focus").astype(int)
+    colors = {**ANCHOR_COLORS, **(colors or {})}
+    xy = np.asarray(adata.obsm[spatial_key], dtype=float)
+
+    panels = _slides(adata, sample_key, samples)
+    fig, axs = _grid(len(panels), ncols, panel_size, axes=axes)
+
+    for k, ((label, rows), ax) in enumerate(zip(panels, axs, strict=True)):
+        inside = rows[anchor[rows]]
+        outside = rows[~anchor[rows]]
+        # Anchors last so a focus is never hidden under the tissue it sits in.
+        ax.scatter(xy[outside, 0], xy[outside, 1], s=size, c=colors[False], linewidths=0)
+        ax.scatter(xy[inside, 0], xy[inside, 1], s=size, c=colors[True], linewidths=0)
+        if focus is not None:
+            for f in np.unique(focus[inside]):
+                cxy = xy[inside[focus[inside] == f]].mean(axis=0)
+                ax.annotate(
+                    str(f),
+                    cxy,
+                    ha="center",
+                    va="center",
+                    fontsize=8,
+                    color="white",
+                    fontweight="bold",
+                )
+        default = f"{label}\n{inside.size} anchor cells"
+        _bare(ax, default if titles is None else titles[k])
+
+    if legend:
+        axs[-1].legend(
+            handles=[
+                Line2D([], [], marker="o", ls="", color=colors[False], label=labels[0]),
+                Line2D([], [], marker="o", ls="", color=colors[True], label=labels[1]),
+            ],
+            loc="center left",
+            bbox_to_anchor=(1.0, 0.5),
+            frameon=False,
+            fontsize=9,
+        )
+    fig.tight_layout()
+    if show:
+        plt.show()
+    return fig, axs
+
+
+def signed_distance_map(
+    adata,
+    dist,
+    *,
+    samples: Sequence[str] | None = None,
+    sample_key: str = "sample",
+    spatial_key: str = "spatial",
+    cmap: str = "RdBu",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    size: float = 6.0,
+    ncols: int | None = None,
+    panel_size: tuple[float, float] = (3.4, 3.0),
+    axes=None,
+    colorbar: bool = True,
+    titles: Sequence[str] | None = None,
+    show: bool = True,
+):
+    """Signed distance to the anchor border: red inside a focus, blue away from it.
+
+    The norm is a :class:`~matplotlib.colors.TwoSlopeNorm` centred on 0 -- the border -- and
+    *not* symmetric, because a focus is a few hundred um deep while the tissue outside it runs
+    to several thousand: a symmetric range would flatten the core to a single shade. The
+    colour scale is shared across panels so two slides are comparable.
+
+    Returns
+    -------
+    (fig, axes)
+    """
+    dist = _vector(adata, dist, "dist").astype(float)
+    xy = np.asarray(adata.obsm[spatial_key], dtype=float)
+
+    panels = _slides(adata, sample_key, samples)
+    shown = np.concatenate([rows for _, rows in panels])
+    d = dist[shown]
+    if not np.isfinite(d).any():
+        raise ValueError("no finite signed distances in the requested samples")
+    lo = float(np.nanmin(d)) if vmin is None else float(vmin)
+    hi = float(np.nanmax(d)) if vmax is None else float(vmax)
+    # TwoSlopeNorm insists on vmin < vcenter < vmax; a sample that is entirely outside every
+    # focus has no negative side, so nudge rather than raise.
+    lo = min(lo, -1e-6)
+    hi = max(hi, 1e-6)
+    norm = TwoSlopeNorm(vcenter=0.0, vmin=lo, vmax=hi)
+
+    fig, axs = _grid(len(panels), ncols, panel_size, axes=axes)
+    sm = None
+    for k, ((label, rows), ax) in enumerate(zip(panels, axs, strict=True)):
+        # A sample with no focus is all-nan, and scatter would simply not draw it -- an empty
+        # panel reads as a plotting failure rather than as the result it is, so the tissue is
+        # drawn in the na colour first.
+        na = rows[~np.isfinite(dist[rows])]
+        if na.size:
+            ax.scatter(xy[na, 0], xy[na, 1], s=size, c=NA_COLOR, linewidths=0)
+        # Nearest-first so the anchor neighbourhood is not buried under distant cells.
+        keep = rows[np.isfinite(dist[rows])]
+        order = keep[np.argsort(-dist[keep])]
+        if order.size:
+            sm = ax.scatter(xy[order, 0], xy[order, 1], s=size, c=dist[order], cmap=cmap, norm=norm, linewidths=0)
+        _bare(ax, label if titles is None else titles[k])
+
+    if colorbar and sm is not None:
+        cb = fig.colorbar(sm, ax=axs, fraction=0.03, pad=0.02)
+        cb.set_label("signed distance to anchor border [um]", fontsize=9)
+    else:
+        fig.tight_layout()
+    if show:
+        plt.show()
+    return fig, axs
+
+
+def zone_map(
+    adata,
+    zone,
+    *,
+    samples: Sequence[str] | None = None,
+    sample_key: str = "sample",
+    spatial_key: str = "spatial",
+    colors: dict | None = None,
+    size: float = 6.0,
+    ncols: int | None = None,
+    panel_size: tuple[float, float] = (3.0, 3.0),
+    axes=None,
+    legend: bool = True,
+    titles: Sequence[str] | None = None,
+    show: bool = True,
+):
+    """``anchor core`` / ``near`` / ``far`` bands, one panel per sample.
+
+    Zones are drawn in :data:`interscale.tl.anchors.ZONE_ORDER`, so ``core`` ends up on top
+    of ``far`` and the legend order is the distance order regardless of how many cells each
+    band holds. Cells with an undefined zone -- a sample where the dimension found no focus --
+    are drawn in :data:`NA_COLOR`.
+
+    Returns
+    -------
+    (fig, axes)
+    """
+    zone = _vector(adata, zone, "zone").astype(object)
+    colors = {**ZONE_COLORS, **(colors or {})}
+    xy = np.asarray(adata.obsm[spatial_key], dtype=float)
+
+    panels = _slides(adata, sample_key, samples)
+    fig, axs = _grid(len(panels), ncols, panel_size, axes=axes)
+
+    for k, ((label, rows), ax) in enumerate(zip(panels, axs, strict=True)):
+        unassigned = rows[pd.isna(zone[rows])]
+        if unassigned.size:
+            ax.scatter(xy[unassigned, 0], xy[unassigned, 1], s=size, c=NA_COLOR, linewidths=0)
+        for z in reversed(ZONE_ORDER):  # far first, core last
+            sel = rows[zone[rows] == z]
+            if sel.size:
+                ax.scatter(xy[sel, 0], xy[sel, 1], s=size, c=colors[z], linewidths=0)
+        _bare(ax, label if titles is None else titles[k])
+
+    if legend:
+        axs[-1].legend(
+            handles=[Line2D([], [], marker="o", ls="", color=colors[z], label=z) for z in ZONE_ORDER],
+            loc="center left",
+            bbox_to_anchor=(1.0, 0.5),
+            frameon=False,
+            fontsize=9,
+        )
+    fig.tight_layout()
+    if show:
+        plt.show()
+    return fig, axs
+
+
+def gene_maps(
+    adata,
+    genes: Sequence[str],
+    *,
+    samples: Sequence[str],
+    layer: str | None = "log1p_norm",
+    sample_key: str = "sample",
+    spatial_key: str = "spatial",
+    anchor=None,
+    cmap: str = "viridis",
+    vmax_quantile: float = 0.99,
+    size: float = 6.0,
+    panel_size: tuple[float, float] = (2.8, 2.6),
+    anchor_edge: str = "#D00000",
+    show: bool = True,
+):
+    """Expression of a dimension's top genes, genes down the rows and slides across.
+
+    One colour scale per **gene**, taken over the slides shown (``vmax_quantile`` of the
+    pooled non-zero-inclusive values), so a row compares slides. Per gene and not global
+    because the genes are on very different absolute levels and a shared scale would render
+    the low ones blank.
+
+    Parameters
+    ----------
+    anchor
+        Optional boolean/obs key. Anchor cells get a coloured ring, which is what ties the
+        expression pattern back to the foci the dimension was anchored on.
+
+    Returns
+    -------
+    (fig, axes) with ``axes`` shaped ``(len(genes), len(samples))``.
+    """
+    genes = [str(g) for g in np.atleast_1d(genes)]
+    missing = [g for g in genes if g not in set(map(str, adata.var_names))]
+    if missing:
+        raise KeyError(f"genes not in adata.var_names: {missing}")
+    anchor = None if anchor is None else _vector(adata, anchor, "anchor").astype(bool)
+    xy = np.asarray(adata.obsm[spatial_key], dtype=float)
+
+    panels = _slides(adata, sample_key, samples)
+    fig, axs = plt.subplots(
+        len(genes),
+        len(panels),
+        figsize=(panel_size[0] * len(panels) + 0.9, panel_size[1] * len(genes)),
+        squeeze=False,
+    )
+
+    for i, gene in enumerate(genes):
+        vals = {label: _expression(adata, rows, gene, layer) for label, rows in panels}
+        pooled = np.concatenate(list(vals.values()))
+        hi = float(np.quantile(pooled, vmax_quantile))
+        if hi <= 0:  # gene undetected on these slides -- keep the panel honest, not blank
+            hi = float(pooled.max()) or 1.0
+        sm = None
+        for j, (label, rows) in enumerate(panels):
+            ax = axs[i, j]
+            v = vals[label]
+            order = np.argsort(v)  # brightest on top
+            sm = ax.scatter(
+                xy[rows[order], 0],
+                xy[rows[order], 1],
+                s=size,
+                c=v[order],
+                cmap=cmap,
+                vmin=0.0,
+                vmax=hi,
+                linewidths=0,
+            )
+            if anchor is not None:
+                a = rows[anchor[rows]]
+                if a.size:
+                    ax.scatter(
+                        xy[a, 0],
+                        xy[a, 1],
+                        s=size * 2.6,
+                        facecolors="none",
+                        edgecolors=anchor_edge,
+                        linewidths=0.35,
+                    )
+            _bare(ax, label if i == 0 else None)
+            if j == 0:
+                ax.set_ylabel(gene, fontsize=10, rotation=0, ha="right", va="center", labelpad=8)
+        fig.colorbar(sm, ax=axs[i, :].tolist(), fraction=0.025, pad=0.02).set_label(layer or "X", fontsize=8)
+
+    if show:
+        plt.show()
+    return fig, axs
+
+
+# --------------------------------------------------------------------------------------------
+# distance profiles
+# --------------------------------------------------------------------------------------------
+
+
+def distance_profile(
+    prof: pd.DataFrame,
+    *,
+    genes: Sequence[str] | None = None,
+    local_reach: float | None = None,
+    colors: Sequence[str] | dict | None = None,
+    band: bool = True,
+    one_panel_per_gene: bool = False,
+    ncols: int | None = 3,
+    panel_size: tuple[float, float] = (3.2, 2.5),
+    ax=None,
+    ylabel: str = "mean expression",
+    title: str | None = None,
+    show: bool = True,
+):
+    """Binned expression against signed distance to the anchor border.
+
+    Takes what :func:`interscale.tl.profile_by_distance` returns (``gene, center, mean, sem,
+    n``). Three reference marks make the axis readable:
+
+    * ``x = 0`` -- the anchor border. Left of it is inside a focus.
+    * the shaded strip left of 0 -- the anchor core itself.
+    * ``x = local_reach`` -- :func:`interscale.tl.local_reach_um`. **Anything the profile does
+      to the right of this line is beyond the local component's reach**, so it cannot be a
+      neighbourhood effect; that is the whole reason the profile is plotted against distance.
+
+    Parameters
+    ----------
+    one_panel_per_gene
+        One axes per gene (shared x). Use it when the genes differ by more than a factor of
+        two in level, where a single panel hides the weaker ones' shape.
+
+    Returns
+    -------
+    (fig, axes)
+    """
+    required = {"gene", "center", "mean", "sem"}
+    if not required.issubset(prof.columns):
+        raise ValueError(f"prof is missing columns {sorted(required - set(prof.columns))}")
+    if prof.empty:
+        raise ValueError("prof is empty -- every distance bin fell below min_cells")
+
+    genes = list(prof["gene"].unique()) if genes is None else [str(g) for g in genes]
+    missing = [g for g in genes if g not in set(prof["gene"])]
+    if missing:
+        raise ValueError(f"genes not present in prof: {missing}")
+
+    if isinstance(colors, dict):
+        color_of = colors
+    else:
+        cyc = list(colors) if colors is not None else list(plt.rcParams["axes.prop_cycle"].by_key()["color"])
+        color_of = {g: cyc[i % len(cyc)] for i, g in enumerate(genes)}
+
+    if one_panel_per_gene:
+        fig, axs = _grid(len(genes), ncols, panel_size)
+        panels = [(g, axs[i]) for i, g in enumerate(genes)]
+    elif ax is not None:
+        fig, axs = ax.get_figure(), [ax]
+        panels = [(g, ax) for g in genes]
+    else:
+        fig, axs = plt.subplots(figsize=(panel_size[0] * 1.6, panel_size[1] * 1.3))
+        axs = [axs]
+        panels = [(g, axs[0]) for g in genes]
+
+    xmin = float(prof["center"].min())
+    xmax = float(prof["center"].max())
+
+    for gene, a in panels:
+        sub = prof[prof["gene"] == gene].sort_values("center")
+        a.plot(sub["center"], sub["mean"], "-o", ms=3, lw=1.4, color=color_of[gene], label=gene)
+        if band:
+            a.fill_between(
+                sub["center"],
+                sub["mean"] - sub["sem"],
+                sub["mean"] + sub["sem"],
+                color=color_of[gene],
+                alpha=0.18,
+                linewidth=0,
+            )
+
+    for a in dict.fromkeys(a for _, a in panels):
+        if xmin < 0:
+            a.axvspan(xmin, 0.0, color=ZONE_COLORS[ZONE_CORE], alpha=0.10, linewidth=0)
+        a.axvline(0.0, color="0.35", lw=1.0)
+        if local_reach is not None and xmin < local_reach < xmax:
+            a.axvline(local_reach, color="0.35", lw=1.0, ls="--")
+            # Rotated and inside the axes: a horizontal label overruns the panel as soon as
+            # the genes get one panel each, which is the layout this figure is usually in.
+            a.annotate(
+                f"local reach {local_reach:g} um",
+                (local_reach, 0.98),
+                xycoords=("data", "axes fraction"),
+                xytext=(4, 0),
+                textcoords="offset points",
+                rotation=90,
+                ha="left",
+                va="top",
+                fontsize=7,
+                color="0.35",
+            )
+        a.set_xlabel("signed distance to anchor border [um]", fontsize=9)
+        a.set_ylabel(ylabel, fontsize=9)
+        a.spines[["top", "right"]].set_visible(False)
+        if one_panel_per_gene:
+            a.set_title(next(g for g, aa in panels if aa is a), fontsize=10)
+        else:
+            a.legend(frameon=False, fontsize=8)
+
+    if title:
+        fig.suptitle(title, fontsize=11)
+    fig.tight_layout()
+    if show:
+        plt.show()
+    return fig, axs
diff --git a/src/interscale/tl/__init__.py b/src/interscale/tl/__init__.py
index 9e7c93a..7cb9310 100644
--- a/src/interscale/tl/__init__.py
+++ b/src/interscale/tl/__init__.py
@@ -1,4 +1,13 @@
 from ._preprocessing import get_average_local_and_global_size, remove_zero_expression_cells
+from .anchors import (
+    ZONE_ORDER,
+    anchor_enrichment,
+    anchor_signed_distance,
+    anchor_zones,
+    find_anchor_cells,
+    local_reach_um,
+    profile_by_distance,
+)
 from .geome_utils import prepare_a2d_dataset, prepare_geome_dataset
 from .masking import (
     MASK_STRATEGIES,
@@ -31,4 +40,11 @@
     "attn_mask_diagonal",
     "remove_zero_expression_cells",
     "get_average_local_and_global_size",
+    "find_anchor_cells",
+    "anchor_signed_distance",
+    "anchor_zones",
+    "profile_by_distance",
+    "anchor_enrichment",
+    "local_reach_um",
+    "ZONE_ORDER",
 ]
diff --git a/src/interscale/tl/anchors.py b/src/interscale/tl/anchors.py
new file mode 100644
index 0000000..5963a5b
--- /dev/null
+++ b/src/interscale/tl/anchors.py
@@ -0,0 +1,426 @@
+"""Anchor points of a latent dimension, and expression profiles at range from them.
+
+Motivation
+----------
+A *global* InterScale dimension is free to draw on the whole sample, so a claim that it
+carries a long-range interaction needs the interaction shown at a distance. These functions
+localise the tissue that drives a dimension (its **anchor foci**), measure every other cell's
+distance to those foci, and profile expression against that distance:
+
+    anchor, cluster = find_anchor_cells(adata, "43_global_emb", 13)   # a handful of foci
+    d = anchor_signed_distance(adata, anchor)                         # <0 inside, >0 outside
+    zone = anchor_zones(adata, d, near_max=local_reach_um(cfg))       # core / near / far
+    prof = profile_by_distance(adata, ["PTCH1", "NKX6.1"], d)         # binned mean +/- SEM
+
+The construction mirrors the fibrosis-core panels of the multi-organ supplement: a
+thresholded score defines cores, cores get a signed distance to their border, and covariates
+are plotted against that signed distance. The difference is that the core here is not an
+annotation -- it is read off the model.
+
+Why "signed distance to the border" and not "distance to the nearest anchor cell": the latter
+is zero for every cell of a focus regardless of how deep inside it sits, which collapses the
+whole core into one bin. Negative depth keeps the core resolvable, and 0 is the border in both
+directions, so a profile is continuous across it.
+
+Why the distance axis is the point
+----------------------------------
+The local component sees `num_layers` hops of a radius-`radius` graph, i.e. at most
+`num_layers * radius` micrometers (:func:`local_reach_um`). Structure in a profile *beyond*
+that abscissa cannot have come from the local component, which is precisely what makes it
+evidence of a long-range interaction rather than a neighbourhood effect. Every plot in
+``interscale.pl.anchor_plots`` therefore marks that reach.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import numpy as np
+import pandas as pd
+from scipy.sparse import csr_matrix, issparse
+from scipy.sparse.csgraph import connected_components
+from scipy.spatial import cKDTree
+
+__all__ = [
+    "local_reach_um",
+    "find_anchor_cells",
+    "anchor_signed_distance",
+    "anchor_zones",
+    "profile_by_distance",
+    "anchor_enrichment",
+]
+
+# Zone names. Module constants because the plotting side has to order and colour them, and a
+# typo in one place would silently produce an unordered categorical.
+ZONE_CORE = "anchor core"
+ZONE_NEAR = "near"
+ZONE_FAR = "far"
+ZONE_ORDER = (ZONE_CORE, ZONE_NEAR, ZONE_FAR)
+
+
+def local_reach_um(cfg) -> float:
+    """Furthest a cell's own representation can travel through the local component, in um.
+
+    ``num_layers`` message-passing hops over a radius-``radius`` graph, so ``num_layers *
+    radius`` is a hard bound -- the *typical* walk is shorter (see
+    :func:`interscale.tl.get_average_local_and_global_size`, which measures it from the
+    realised edge lengths). The bound is the useful number here: anything past it is out of
+    the local component's reach for every cell, not just the average one.
+    """
+    radius = cfg.dataset.spatial_neigbors_kwargs.radius
+    if radius is None:
+        raise ValueError(
+            "cfg.dataset.spatial_neigbors_kwargs.radius is None, so the local component's "
+            "reach is a neighbour count rather than a distance and cannot be expressed in um."
+        )
+    return float(radius) * int(cfg.model.local_component.parameters.num_layers)
+
+
+def _grouping(adata, sample_key, samples):
+    """(group labels, groups to visit). Restricting to `samples` is not cosmetic: the
+    threshold is per group, so visiting a group means thresholding inside it.
+    """
+    grp = adata.obs[sample_key].astype(str).to_numpy()
+    if samples is None:
+        return grp, np.unique(grp)
+    todo = np.asarray([str(s) for s in np.atleast_1d(samples)])
+    missing = set(todo) - set(np.unique(grp))
+    if missing:
+        raise KeyError(f"adata.obs['{sample_key}'] has no group(s) {sorted(missing)}")
+    return grp, todo
+
+
+def find_anchor_cells(
+    adata,
+    emb_key: str,
+    dim: int,
+    *,
+    sample_key: str = "sample",
+    spatial_key: str = "spatial",
+    quantile: float = 0.95,
+    sign: str = "auto",
+    link_radius: float = 400.0,
+    min_cluster_size: int = 8,
+    samples: Sequence[str] | None = None,
+    key_added: str | None = None,
+):
+    """Cells driving one latent dimension, kept only where they form a spatial focus.
+
+    Two filters, in this order:
+
+    1. **Tail of the dimension**, thresholded at ``quantile`` *within each sample*. Per
+       sample, not globally: the samples are separate graphs with their own global embedding,
+       and one sample's offset would otherwise take every anchor in the dataset.
+    2. **Spatial coherence** -- single-linkage clustering of the tail cells at
+       ``link_radius``, discarding components smaller than ``min_cluster_size``. This is what
+       turns "the top 5% of cells" into "a handful of anchor points": scattered high-score
+       cells carry no location to measure a distance from, and dropping them is the difference
+       between a few foci and hundreds of singletons.
+
+    Parameters
+    ----------
+    emb_key
+        ``adata.obsm`` key, e.g. ``f"{seed}_global_emb"``.
+    dim
+        Column of that embedding. Normally one of the dimensions
+        :func:`interscale.evaluation.calculate_dim_importance` kept.
+    quantile
+        Per-sample score quantile the tail starts at. 0.95 is a deliberate compromise: at 0.99
+        a ~1800-cell sample offers only ~18 tail cells, too few to survive
+        ``min_cluster_size`` unless they sit on top of each other, so most dimensions come
+        back with no focus at all.
+    sign
+        ``"+"`` / ``"-"`` to take the high or low tail; ``"auto"`` (default) takes whichever
+        tail reaches further from the median. A dimension's sign is arbitrary -- it flips with
+        the decoder weight column -- so hardcoding ``"+"`` would silently analyse the wrong
+        end for half the dimensions.
+    link_radius
+        Single-linkage distance for grouping tail cells into foci. Default 400 um =
+        :func:`local_reach_um` for the legnini23 configs (2 hops x 200 um), i.e. cells close
+        enough that the local component could have linked them count as one focus.
+    min_cluster_size
+        Foci smaller than this are dropped entirely (not marked as anchors).
+    samples
+        Restrict to these ``sample_key`` groups. Everything else gets ``anchor=False`` and
+        ``cluster=-1``; the thresholds are unaffected, being per sample already.
+    key_added
+        If given, write ``adata.obs[key_added]`` (bool) and
+        ``adata.obs[f"{key_added}_focus"]`` (int, -1 = not an anchor).
+
+    Returns
+    -------
+    anchor : (n_obs,) bool
+    cluster : (n_obs,) int
+        Focus id, unique across samples; ``-1`` where ``anchor`` is False.
+
+    Notes
+    -----
+    ``adata.uns[f"{key_added}_info"]`` records the dimension, the sign taken and the
+    per-sample thresholds, so a figure can state what it thresholded rather than the caller
+    having to remember.
+    """
+    if not 0.0 < quantile < 1.0:
+        raise ValueError(f"quantile must be in (0, 1), got {quantile}")
+    if sign not in ("auto", "+", "-"):
+        raise ValueError(f"sign must be 'auto', '+' or '-', got {sign!r}")
+
+    if emb_key not in adata.obsm:
+        raise KeyError(f"{emb_key} not found in adata.obsm")
+    Z = np.asarray(adata.obsm[emb_key], dtype=float)
+    if Z.ndim == 1:
+        Z = Z[:, None]
+    dim = int(dim)
+    if not -Z.shape[1] <= dim < Z.shape[1]:
+        raise ValueError(f"dim {dim} out of range for {emb_key} with {Z.shape[1]} dimensions")
+
+    xy = np.asarray(adata.obsm[spatial_key], dtype=float)
+    grp, todo = _grouping(adata, sample_key, samples)
+
+    s = Z[:, dim]
+    if sign == "auto":
+        med = np.nanmedian(s)
+        lo, hi = np.nanquantile(s, [1.0 - quantile, quantile])
+        sgn = 1.0 if abs(hi - med) >= abs(med - lo) else -1.0
+    else:
+        sgn = 1.0 if sign == "+" else -1.0
+    score = sgn * s
+
+    anchor = np.zeros(adata.n_obs, dtype=bool)
+    cluster = np.full(adata.n_obs, -1, dtype=int)
+    thresholds, next_id = {}, 0
+
+    for g in todo:
+        rows = np.where(grp == g)[0]
+        thr = float(np.nanquantile(score[rows], quantile))
+        thresholds[str(g)] = thr
+        cand = rows[score[rows] >= thr]
+        if cand.size == 0:
+            continue
+
+        # Single linkage via connected components of the <=link_radius graph. Built as a
+        # sparse distance matrix so this stays linear-ish in the number of close pairs
+        # instead of materialising cand x cand.
+        tree = cKDTree(xy[cand])
+        pairs = tree.sparse_distance_matrix(tree, link_radius, output_type="coo_matrix")
+        n_comp, labels = connected_components(
+            csr_matrix((np.ones_like(pairs.data, dtype=bool), (pairs.row, pairs.col)), shape=(cand.size, cand.size)),
+            directed=False,
+        )
+        for c in range(n_comp):
+            idx = cand[labels == c]
+            if idx.size >= min_cluster_size:
+                anchor[idx] = True
+                cluster[idx] = next_id
+                next_id += 1
+
+    if key_added is not None:
+        adata.obs[key_added] = anchor
+        adata.obs[f"{key_added}_focus"] = cluster
+        adata.uns[f"{key_added}_info"] = {
+            "emb_key": str(emb_key),
+            "dim": dim,
+            "sign": sgn,
+            "quantile": float(quantile),
+            "link_radius": float(link_radius),
+            "min_cluster_size": int(min_cluster_size),
+            "n_foci": int(next_id),
+            "thresholds": thresholds,
+        }
+    return anchor, cluster
+
+
+def anchor_signed_distance(
+    adata,
+    anchor,
+    *,
+    sample_key: str = "sample",
+    spatial_key: str = "spatial",
+    samples: Sequence[str] | None = None,
+    key_added: str | None = None,
+):
+    """Signed distance to the anchor border: negative inside a focus, positive outside.
+
+    Per sample, and per sample only -- distances are computed within a ``sample_key`` group
+    and never across groups, so a cell is never measured against an anchor on another slide.
+
+    A cell outside every focus gets ``+`` its distance to the nearest anchor cell; a cell
+    inside one gets ``-`` its distance to the nearest non-anchor cell, i.e. how deep in the
+    focus it sits. Both are zero at the border, so the two branches join up.
+
+    Returns
+    -------
+    (n_obs,) float, ``nan`` for samples that were skipped or that hold no anchor (a Ctrl
+    sample where the dimension found nothing is a *result*, so it is not an error) and for
+    samples that are entirely anchor (no border to measure against).
+    """
+    anchor = np.asarray(anchor, dtype=bool)
+    if anchor.shape != (adata.n_obs,):
+        raise ValueError(f"anchor must have shape ({adata.n_obs},), got {anchor.shape}")
+
+    xy = np.asarray(adata.obsm[spatial_key], dtype=float)
+    grp, todo = _grouping(adata, sample_key, samples)
+
+    out = np.full(adata.n_obs, np.nan)
+    for g in todo:
+        rows = np.where(grp == g)[0]
+        inside, outside = rows[anchor[rows]], rows[~anchor[rows]]
+        if inside.size == 0 or outside.size == 0:
+            continue
+        out[outside] = cKDTree(xy[inside]).query(xy[outside])[0]
+        out[inside] = -cKDTree(xy[outside]).query(xy[inside])[0]
+
+    if key_added is not None:
+        adata.obs[key_added] = out
+    return out
+
+
+def anchor_zones(adata, dist, *, near_max: float = 400.0, key_added: str | None = None):
+    """Split a signed distance into ``anchor core`` / ``near`` / ``far``.
+
+    ``near_max`` is the local component's reach (:func:`local_reach_um`), which makes the cut
+    mean something: ``near`` is the band a local model could have explained on its own, and
+    ``far`` is the band only the global component can reach. Returns an ordered categorical
+    (``nan`` distances stay unassigned) so plots and groupbys keep core -> near -> far order.
+    """
+    dist = np.asarray(dist, dtype=float)
+    zone = np.full(adata.n_obs, None, dtype=object)
+    ok = np.isfinite(dist)
+    zone[ok & (dist <= 0)] = ZONE_CORE
+    zone[ok & (dist > 0) & (dist <= near_max)] = ZONE_NEAR
+    zone[ok & (dist > near_max)] = ZONE_FAR
+    out = pd.Categorical(zone, categories=list(ZONE_ORDER), ordered=True)
+    if key_added is not None:
+        adata.obs[key_added] = out
+    return out
+
+
+def profile_by_distance(
+    adata,
+    genes: Sequence[str],
+    dist,
+    *,
+    layer: str | None = "log1p_norm",
+    mask=None,
+    n_bins: int = 14,
+    min_cells: int = 15,
+    binning: str = "quantile",
+):
+    """Mean expression per distance bin, one row per (gene, bin).
+
+    Parameters
+    ----------
+    layer
+        ``adata.layers`` key holding the expression to profile, or ``None`` for ``.X``. Use
+        the layer the decoders were trained on, so the profile is on the same scale as the
+        loadings that selected `genes`.
+    mask
+        Boolean over ``adata.n_obs``, normally one slide. Combined with ``isfinite(dist)``.
+    binning
+        ``"quantile"`` (default) puts a comparable number of cells in each bin, which keeps
+        the error bars comparable along the axis; ``"equal"`` uses equal-width bins, which is
+        the honest choice if the *shape* against distance is being read off rather than
+        compared bin to bin. Bins under ``min_cells`` are dropped either way.
+
+    Returns
+    -------
+    DataFrame with ``gene, bin, lo, hi, center, n, mean, sem``. ``center`` is the mean
+    distance of the cells in the bin, not the bin midpoint, so a point sits where its cells
+    actually are.
+    """
+    if binning not in ("quantile", "equal"):
+        raise ValueError(f"binning must be 'quantile' or 'equal', got {binning!r}")
+
+    genes = list(genes)
+    missing = [g for g in genes if g not in adata.var_names]
+    if missing:
+        raise KeyError(f"genes not in adata.var_names: {missing}")
+
+    dist = np.asarray(dist, dtype=float)
+    ok = np.isfinite(dist)
+    if mask is not None:
+        ok &= np.asarray(mask, dtype=bool)
+    if ok.sum() == 0:
+        return pd.DataFrame(columns=["gene", "bin", "lo", "hi", "center", "n", "mean", "sem"])
+
+    sub = adata[ok, genes]
+    X = sub.layers[layer] if layer is not None else sub.X
+    X = np.asarray(X.todense() if issparse(X) else X, dtype=float)
+    d = dist[ok]
+
+    if binning == "quantile":
+        edges = np.unique(np.quantile(d, np.linspace(0.0, 1.0, n_bins + 1)))
+    else:
+        edges = np.linspace(d.min(), d.max(), n_bins + 1)
+    if edges.size < 2:
+        return pd.DataFrame(columns=["gene", "bin", "lo", "hi", "center", "n", "mean", "sem"])
+    # digitize on the interior edges only, then clip, so the extremes land in the end bins
+    # rather than in out-of-range bin 0 / bin n+1.
+    which = np.clip(np.digitize(d, edges[1:-1], right=True), 0, edges.size - 2)
+
+    rows = []
+    for b in range(edges.size - 1):
+        sel = which == b
+        n = int(sel.sum())
+        if n < min_cells:
+            continue
+        vals = X[sel]
+        # ddof=1 needs n >= 2; min_cells is well above that, but sem is meaningless at n=1.
+        sem = vals.std(axis=0, ddof=1) / np.sqrt(n)
+        for j, gene in enumerate(genes):
+            rows.append(
+                {
+                    "gene": gene,
+                    "bin": b,
+                    "lo": float(edges[b]),
+                    "hi": float(edges[b + 1]),
+                    "center": float(d[sel].mean()),
+                    "n": n,
+                    "mean": float(vals[:, j].mean()),
+                    "sem": float(sem[j]),
+                }
+            )
+    return pd.DataFrame(rows)
+
+
+def anchor_enrichment(adata, anchor, obs_key, *, sample_key="sample", samples=None):
+    """What the anchor foci are made of, per sample: composition inside vs in the rest.
+
+    Answers "did this dimension anchor on anything nameable?" -- e.g. whether a dimension's
+    foci coincide with the SHH-source graft. ``log2_enrichment`` is inside-vs-rest on the
+    fraction of cells; it is ``inf`` for a category absent outside the foci and ``-inf`` for
+    one absent inside, which is the honest reading of a category that is exclusive to one
+    side.
+
+    Returns a DataFrame with ``sample, category, n_in, n_out, frac_in, frac_out,
+    log2_enrichment``, empty for samples with no anchor.
+    """
+    anchor = np.asarray(anchor, dtype=bool)
+    grp, todo = _grouping(adata, sample_key, samples)
+    labels = adata.obs[obs_key].astype(str).to_numpy()
+    cats = pd.unique(adata.obs[obs_key].astype(str))
+
+    rows = []
+    for g in todo:
+        m = grp == g
+        inside, outside = m & anchor, m & ~anchor
+        if inside.sum() == 0:
+            continue
+        for c in cats:
+            n_in = int((labels[inside] == c).sum())
+            n_out = int((labels[outside] == c).sum())
+            f_in = n_in / inside.sum()
+            f_out = n_out / outside.sum() if outside.sum() else np.nan
+            with np.errstate(divide="ignore", invalid="ignore"):
+                lfc = np.log2(f_in / f_out) if f_in or f_out else np.nan
+            rows.append(
+                {
+                    "sample": str(g),
+                    "category": c,
+                    "n_in": n_in,
+                    "n_out": n_out,
+                    "frac_in": f_in,
+                    "frac_out": f_out,
+                    "log2_enrichment": float(lfc),
+                }
+            )
+    return pd.DataFrame(rows)
diff --git a/src/interscale/tl/masking.py b/src/interscale/tl/masking.py
index 50ebd37..711496d 100644
--- a/src/interscale/tl/masking.py
+++ b/src/interscale/tl/masking.py
@@ -204,6 +204,13 @@ def masked_row_std(y: torch.Tensor, entry_mask: torch.Tensor, eps: float = 1e-8)
     return var.sqrt().clamp(min=eps)
 
 
+def _plain_loss(loss_fn, loss_type: str, y_pred: torch.Tensor, y_true: torch.Tensor):
+    """``loss_fn`` over everything it is given, with GaussianNLL's third argument supplied."""
+    if loss_type == "GaussianNLL":
+        return loss_fn(y_pred, y_true, torch.std(y_true, dim=1, keepdim=True))
+    return loss_fn(y_pred, y_true)
+
+
 def masked_loss(loss_fn, loss_type: str, y_pred: torch.Tensor, y_true: torch.Tensor, entry_mask=None):
     """Evaluate a reconstruction loss on the masked entries only.
 
@@ -222,7 +229,8 @@ def masked_loss(loss_fn, loss_type: str, y_pred: torch.Tensor, y_true: torch.Ten
     y_pred, y_true
         ``[N, G]`` predictions and targets for the scored rows.
     entry_mask
-        ``[N, G]`` boolean, or ``None``.
+        ``[N, G]`` boolean, or ``None``. Under cell masking this is the per-cell mask broadcast
+        over all genes, so it is all-True or all-False per row and the row branch below fires.
 
     Returns
     -------
@@ -230,9 +238,16 @@ def masked_loss(loss_fn, loss_type: str, y_pred: torch.Tensor, y_true: torch.Ten
         Scalar loss.
     """
     if entry_mask is None:
-        if loss_type == "GaussianNLL":
-            return loss_fn(y_pred, y_true, torch.std(y_true, dim=1, keepdim=True))
-        return loss_fn(y_pred, y_true)
+        return _plain_loss(loss_fn, loss_type, y_pred, y_true)
+
+    # Whole-row mask, i.e. cell masking expressed at entry level: every row is either entirely
+    # in or entirely out. Subset the ROWS rather than the entries, which keeps each row intact --
+    # a row-normalising criterion (SCE's cosine, Pearson) is only meaningful on a whole row, and
+    # zeroing the excluded rows instead would feed it rows of zeros that dilute the mean.
+    # This branch reproduces the pre-all-cells behaviour of cell masking exactly.
+    rows_in = entry_mask.all(dim=1)
+    if bool((rows_in | (~entry_mask).all(dim=1)).all()):
+        return _plain_loss(loss_fn, loss_type, y_pred[rows_in], y_true[rows_in])
 
     if loss_type in _ROW_STRUCTURED_LOSSES:
         m = entry_mask.to(y_pred.dtype)
diff --git a/src/interscale/train/_trainingplans.py b/src/interscale/train/_trainingplans.py
index 479e585..9ffdddd 100644
--- a/src/interscale/train/_trainingplans.py
+++ b/src/interscale/train/_trainingplans.py
@@ -338,10 +338,14 @@ def _regression_metrics(
         attn : torch.Tensor | None
             The attention weights to apply to the metrics.
         entry_mask : torch.Tensor | None
-            [N, G] boolean marking the entries that were actually masked. Set under
-            ``mask_strategy="gene"``; ``None`` under cell masking, where the whole row of every
-            scored cell was blanked and there is nothing to restrict. When given, BOTH the loss
-            and the metrics are computed over those entries only -- see ``masked_regression_metrics``.
+            [N, G] boolean marking the entries that were actually masked.
+
+            The LOSS is restricted to those entries -- that is the self-supervised objective, and
+            training on entries the model was handed as input would make it the identity map.
+            The METRICS are computed over every cell in ``y_pred``, masked or not, because a
+            reconstruction is wanted for the whole tissue and not only for the hidden part.
+            The masked-only versions are logged alongside under a ``masked_`` prefix; they are
+            the honest held-out score and the two differ a lot, so read the prefix.
         """
         if self.loss_type == "SCE_EntropyATT_Loss":
             # Takes attention as a third argument, so it cannot go through masked_loss. Zeroing
@@ -355,18 +359,20 @@ def _regression_metrics(
         else:
             loss = masked_loss(self.loss, self.loss_type, y_pred, y_true, entry_mask)
 
-        if entry_mask is None:
-            metrics = metrics(y_pred, y_true)
-            # Take mean across pearson correlation
-            metrics[f"{mode}_pearson_corr"] = torch.nanmean(metrics[f"{mode}_pearson_corr"].contiguous())
-            # Same reduction, same reason: both are per-gene vectors of length n_output, and a
-            # constant gene yields NaN rather than a number.
-            metrics[f"{mode}_concordance_corr"] = torch.nanmean(metrics[f"{mode}_concordance_corr"].contiguous())
-        else:
-            # Same metric names, so `optim.monitor`, the sweep `--metric` flag and every existing
-            # wandb panel keep working across both strategies -- what changes is only which
-            # entries they are computed over.
-            metrics = {f"{mode}_{k}": v for k, v in masked_regression_metrics(y_pred, y_true, entry_mask).items()}
+        # Primary metrics: every cell the model produced, masked or not.
+        metrics = metrics(y_pred, y_true)
+        # Take mean across pearson correlation
+        metrics[f"{mode}_pearson_corr"] = torch.nanmean(metrics[f"{mode}_pearson_corr"].contiguous())
+        # Same reduction, same reason: both are per-gene vectors of length n_output, and a
+        # constant gene yields NaN rather than a number.
+        metrics[f"{mode}_concordance_corr"] = torch.nanmean(metrics[f"{mode}_concordance_corr"].contiguous())
+
+        # Held-out companions on the masked entries only. Kept because the all-cell numbers above
+        # include entries the model was given as input, which it can partly copy -- so they are
+        # the reconstruction score for the tissue, not evidence the model generalises. These are.
+        if entry_mask is not None:
+            for name, value in masked_regression_metrics(y_pred, y_true, entry_mask).items():
+                metrics[f"{mode}_masked_{name}"] = value
 
         metrics[f"{mode}_loss"] = loss
         return loss, metrics
diff --git a/tests/test_anchors.py b/tests/test_anchors.py
new file mode 100644
index 0000000..cfc6fee
--- /dev/null
+++ b/tests/test_anchors.py
@@ -0,0 +1,392 @@
+"""Tests for `interscale.tl.anchors` and the figures in `interscale.pl.anchor_plots`.
+
+The fixture is a two-slide synthetic stand-in for legnini23: `slideA` carries two compact
+foci of one latent dimension plus a gene that decays slowly away from them (long range), a
+gene that decays fast (local), and a flat gene; `slideB` carries none of it. That asymmetry
+is what the assertions are about -- a construction that finds anchors on a slide with no
+structure, or that measures a distance across slides, would pass a single-slide test.
+"""
+
+from types import SimpleNamespace
+
+import anndata as ad
+import matplotlib
+import numpy as np
+import pandas as pd
+import pytest
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from interscale.pl import (
+    anchor_map,
+    distance_profile,
+    gene_maps,
+    signed_distance_map,
+    zone_map,
+)
+from interscale.tl import (
+    anchor_enrichment,
+    anchor_signed_distance,
+    anchor_zones,
+    find_anchor_cells,
+    local_reach_um,
+    profile_by_distance,
+)
+from interscale.tl.anchors import ZONE_CORE, ZONE_FAR, ZONE_NEAR
+
+FOCI = np.array([[1500.0, 1500.0], [4500.0, 4200.0]])
+GENES = ["LONG", "SHORT", "FLAT"]
+
+
+@pytest.fixture(scope="module")
+def anchor_adata():
+    rng = np.random.default_rng(0)
+    n_per = 1200
+    xy = np.vstack([rng.uniform(0, 6000, size=(n_per, 2)) for _ in range(2)])
+    sample = np.array(["slideA"] * n_per + ["slideB"] * n_per)
+    on_a = sample == "slideA"
+    n = len(xy)
+
+    # Distance to the nearest planted focus, the ground truth every assertion is against.
+    d = np.min(np.linalg.norm(xy[:, None, :] - FOCI[None], axis=2), axis=1)
+
+    Z = rng.normal(0, 1, size=(n, 4))
+    # Dimension 1 peaks at the foci on slideA only, and on its NEGATIVE pole -- a latent
+    # dimension's sign is arbitrary, so sign="auto" has to find this.
+    Z[:, 1] = np.where(on_a, -4.0 * np.exp(-((d / 600.0) ** 2)), 0.0) + rng.normal(0, 0.3, n)
+    # Dimension 2 is offset on slideB, which a global (rather than per-sample) threshold
+    # would let take every anchor in the dataset.
+    Z[:, 2] = np.where(on_a, 0.0, 6.0) + rng.normal(0, 0.3, n)
+
+    X = np.empty((n, 3))
+    X[:, 0] = np.where(on_a, 3.0 * np.exp(-d / 2500.0), 0.2)  # decays past the local reach
+    X[:, 1] = np.where(on_a, 2.0 * np.exp(-d / 900.0), 0.1)  # decays within it
+    X[:, 2] = 1.0  # flat
+    X = np.clip(X + rng.normal(0, 0.1, X.shape), 0, None)
+
+    adata = ad.AnnData(X=X.astype(np.float32))
+    adata.var_names = GENES
+    adata.layers["log1p_norm"] = adata.X.copy()
+    adata.obs["sample"] = pd.Categorical(sample)
+    adata.obs["truth"] = pd.Categorical(np.where(d < 800, "focus", "other"))
+    adata.obsm["spatial"] = xy
+    adata.obsm["40_global_emb"] = Z
+    adata.uns["d_to_focus"] = d
+    return adata
+
+
+@pytest.fixture(scope="module")
+def anchored(anchor_adata):
+    """(anchor, focus, dist, zone) for dimension 1 -- the planted long-range dimension."""
+    anchor, focus = find_anchor_cells(anchor_adata, "40_global_emb", 1, min_cluster_size=8)
+    dist = anchor_signed_distance(anchor_adata, anchor)
+    zone = anchor_zones(anchor_adata, dist, near_max=400.0)
+    return anchor, focus, dist, zone
+
+
+# ------------------------------------------------------------------------------------------
+# local_reach_um
+# ------------------------------------------------------------------------------------------
+
+
+def _cfg(radius, num_layers=2):
+    return SimpleNamespace(
+        dataset=SimpleNamespace(spatial_neigbors_kwargs=SimpleNamespace(radius=radius)),
+        model=SimpleNamespace(local_component=SimpleNamespace(parameters=SimpleNamespace(num_layers=num_layers))),
+    )
+
+
+def test_local_reach_is_hops_times_radius():
+    assert local_reach_um(_cfg(200, 2)) == 400.0
+
+
+def test_local_reach_raises_for_knn_graph():
+    # radius None means a neighbour-count graph, which has no reach in um to report.
+    with pytest.raises(ValueError, match="cannot be expressed in um"):
+        local_reach_um(_cfg(None))
+
+
+# ------------------------------------------------------------------------------------------
+# find_anchor_cells
+# ------------------------------------------------------------------------------------------
+
+
+def test_finds_the_planted_foci_on_the_right_slide(anchor_adata, anchored):
+    anchor, focus, _, _ = anchored
+    sample = anchor_adata.obs["sample"].to_numpy()
+    assert anchor.sum() > 0
+    assert anchor[sample == "slideB"].sum() == 0, "anchored on a slide with no structure"
+    assert len(set(focus[anchor])) == len(FOCI)
+    # anchors sit at the foci, not scattered over the slide
+    d = anchor_adata.uns["d_to_focus"]
+    assert d[anchor].mean() < 0.5 * d[sample == "slideA"].mean()
+
+
+def test_auto_sign_takes_the_informative_pole(anchor_adata):
+    anchor, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, min_cluster_size=8, key_added="a")
+    assert anchor_adata.uns["a_info"]["sign"] == -1.0
+    # forcing the wrong pole must not reproduce the foci
+    wrong, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, sign="+", min_cluster_size=8)
+    assert not (anchor & wrong).any()
+
+
+def test_threshold_is_per_sample(anchor_adata):
+    # Dimension 2 is +6 on slideB and 0 on slideA. A dataset-wide quantile would put every
+    # anchor on slideB; a per-sample one splits them.
+    anchor, _ = find_anchor_cells(anchor_adata, "40_global_emb", 2, min_cluster_size=1)
+    sample = anchor_adata.obs["sample"].to_numpy()
+    assert anchor[sample == "slideA"].sum() > 0
+    assert anchor[sample == "slideB"].sum() > 0
+
+
+def test_tighter_quantile_gives_fewer_anchors(anchor_adata):
+    loose, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, quantile=0.95, min_cluster_size=4)
+    tight, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, quantile=0.99, min_cluster_size=4)
+    assert 0 < tight.sum() < loose.sum()
+    assert tight.sum() < 0.02 * anchor_adata.n_obs, "a threshold this tight must leave few anchors"
+
+
+def test_min_cluster_size_discards_scattered_cells(anchor_adata):
+    kept, _ = find_anchor_cells(anchor_adata, "40_global_emb", 3, min_cluster_size=1)
+    dropped, _ = find_anchor_cells(anchor_adata, "40_global_emb", 3, min_cluster_size=25)
+    # dimension 3 is pure noise: its tail has no spatial focus to survive the size filter
+    assert kept.sum() > 0
+    assert dropped.sum() == 0
+
+
+def test_link_radius_merges_foci(anchor_adata):
+    # The two planted foci are ~4000 um apart, so linking at 5000 makes them one component.
+    # Restricted to slideA: at that link radius the whole of a structureless slide links up
+    # too (see test_link_radius_too_large_manufactures_a_focus), which would count as a
+    # second focus here for a reason that has nothing to do with merging.
+    kw = {"min_cluster_size": 8, "samples": ["slideA"]}
+    _, near = find_anchor_cells(anchor_adata, "40_global_emb", 1, link_radius=400.0, **kw)
+    _, far = find_anchor_cells(anchor_adata, "40_global_emb", 1, link_radius=5000.0, **kw)
+    assert len(set(near[near >= 0])) == 2
+    assert len(set(far[far >= 0])) == 1
+
+
+def test_link_radius_too_large_manufactures_a_focus(anchor_adata):
+    # The spatial-coherence filter IS link_radius: raised past the slide's own extent, every
+    # tail cell links to every other and the tail of a dimension with no structure comes back
+    # as one big "focus". slideB has nothing planted in dimension 1, so it must stay empty at
+    # the default and only appear when the radius is absurd -- which is why the default is
+    # tied to the local component's reach rather than picked for how many foci it yields.
+    sample = anchor_adata.obs["sample"].to_numpy()
+    default, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, min_cluster_size=8)
+    absurd, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, link_radius=20_000.0, min_cluster_size=8)
+    assert default[sample == "slideB"].sum() == 0
+    assert absurd[sample == "slideB"].sum() > 0
+
+
+def test_key_added_records_what_was_thresholded(anchor_adata):
+    find_anchor_cells(anchor_adata, "40_global_emb", 1, min_cluster_size=8, key_added="anchor1")
+    info = anchor_adata.uns["anchor1_info"]
+    assert anchor_adata.obs["anchor1"].dtype == bool
+    assert set(anchor_adata.obs.loc[~anchor_adata.obs["anchor1"], "anchor1_focus"]) == {-1}
+    assert info["dim"] == 1 and info["n_foci"] == 2
+    assert set(info["thresholds"]) == {"slideA", "slideB"}
+
+
+def test_samples_restricts_without_moving_the_threshold(anchor_adata):
+    both, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, min_cluster_size=8)
+    one, _ = find_anchor_cells(anchor_adata, "40_global_emb", 1, min_cluster_size=8, samples=["slideA"])
+    assert np.array_equal(both, one)  # slideB had no anchor to lose
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "match"),
+    [
+        ({"quantile": 1.0}, "quantile must be"),
+        ({"sign": "up"}, "sign must be"),
+        ({"dim": 99}, "out of range"),
+        ({"emb_key": "nope"}, "not found in adata.obsm"),
+    ],
+)
+def test_find_anchor_cells_rejects_bad_arguments(anchor_adata, kwargs, match):
+    call = {"emb_key": "40_global_emb", "dim": 1, **kwargs}
+    with pytest.raises((ValueError, KeyError), match=match):
+        find_anchor_cells(anchor_adata, **call)
+
+
+# ------------------------------------------------------------------------------------------
+# anchor_signed_distance / anchor_zones
+# ------------------------------------------------------------------------------------------
+
+
+def test_signed_distance_is_negative_inside_and_nan_without_anchors(anchor_adata, anchored):
+    anchor, _, dist, _ = anchored
+    sample = anchor_adata.obs["sample"].to_numpy()
+    assert (dist[anchor] <= 0).all()
+    outside = np.isfinite(dist) & ~anchor
+    assert (dist[outside] > 0).all()
+    # slideB holds no focus, so there is nothing to measure a distance from -- a result, not
+    # an error, and it must not silently borrow slideA's anchors.
+    assert np.isnan(dist[sample == "slideB"]).all()
+
+
+def test_signed_distance_never_crosses_samples(anchor_adata, anchored):
+    anchor, _, dist, _ = anchored
+    xy = anchor_adata.obsm["spatial"]
+    sample = anchor_adata.obs["sample"].to_numpy()
+    a_rows = np.where((sample == "slideA") & ~anchor)[0]
+    inside = np.where(anchor)[0]
+    within = np.min(np.linalg.norm(xy[a_rows][:, None] - xy[inside][None], axis=2), axis=1)
+    assert np.allclose(dist[a_rows], within)
+
+
+def test_signed_distance_rejects_wrong_length(anchor_adata):
+    with pytest.raises(ValueError, match="must have shape"):
+        anchor_signed_distance(anchor_adata, np.zeros(7, dtype=bool))
+
+
+def test_zones_split_at_zero_and_at_near_max(anchor_adata):
+    dist = np.full(anchor_adata.n_obs, np.nan)
+    dist[:5] = [-10.0, 0.0, 1.0, 400.0, 400.1]
+    zone = anchor_zones(anchor_adata, dist, near_max=400.0)
+    assert list(zone[:5]) == [ZONE_CORE, ZONE_CORE, ZONE_NEAR, ZONE_NEAR, ZONE_FAR]
+    assert pd.isna(zone[5])
+    assert zone.ordered and list(zone.categories) == [ZONE_CORE, ZONE_NEAR, ZONE_FAR]
+
+
+# ------------------------------------------------------------------------------------------
+# profile_by_distance
+# ------------------------------------------------------------------------------------------
+
+
+def test_profile_separates_long_from_short_range(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    prof = profile_by_distance(anchor_adata, GENES, dist, n_bins=12, min_cells=15)
+    piv = prof.pivot(index="center", columns="gene", values="mean").sort_index()
+    beyond = piv.index > 400.0  # past the local component's reach
+
+    # SHORT is spent by the time it leaves the local reach; LONG still falls beyond it.
+    long_fall = piv.loc[beyond, "LONG"].iloc[0] - piv.loc[beyond, "LONG"].iloc[-1]
+    short_fall = piv.loc[beyond, "SHORT"].iloc[0] - piv.loc[beyond, "SHORT"].iloc[-1]
+    assert long_fall > short_fall
+    assert piv["FLAT"].std() < 0.1 * piv["LONG"].std()
+    assert (prof["sem"] >= 0).all() and (prof["n"] >= 15).all()
+
+
+def test_profile_center_is_where_the_cells_are(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    prof = profile_by_distance(anchor_adata, ["LONG"], dist, n_bins=8)
+    assert ((prof["center"] >= prof["lo"]) & (prof["center"] <= prof["hi"])).all()
+
+
+def test_profile_mask_and_binning(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    sample = anchor_adata.obs["sample"].to_numpy()
+    q = profile_by_distance(anchor_adata, ["LONG"], dist, mask=sample == "slideA", n_bins=10)
+    e = profile_by_distance(anchor_adata, ["LONG"], dist, mask=sample == "slideA", n_bins=10, binning="equal")
+    # quantile bins hold a comparable number of cells; equal-width ones do not
+    assert q["n"].std() < e["n"].std()
+    # slideB is all-nan, so masking it out cannot change the profile
+    both = profile_by_distance(anchor_adata, ["LONG"], dist, n_bins=10)
+    assert np.allclose(q["mean"], both["mean"])
+
+
+def test_profile_returns_empty_frame_when_nothing_is_selected(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    out = profile_by_distance(anchor_adata, ["LONG"], dist, mask=np.zeros(anchor_adata.n_obs, bool))
+    assert out.empty and "mean" in out.columns
+
+
+def test_profile_rejects_unknown_genes_and_binning(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    with pytest.raises(KeyError, match="not in adata.var_names"):
+        profile_by_distance(anchor_adata, ["NOPE"], dist)
+    with pytest.raises(ValueError, match="binning must be"):
+        profile_by_distance(anchor_adata, ["LONG"], dist, binning="log")
+
+
+# ------------------------------------------------------------------------------------------
+# anchor_enrichment
+# ------------------------------------------------------------------------------------------
+
+
+def test_enrichment_finds_what_the_foci_are_made_of(anchor_adata, anchored):
+    anchor, _, _, _ = anchored
+    enr = anchor_enrichment(anchor_adata, anchor, "truth").set_index(["sample", "category"])
+    assert list(enr.index.get_level_values("sample").unique()) == ["slideA"]  # slideB has none
+    assert enr.loc[("slideA", "focus"), "log2_enrichment"] > 2
+    # a category with no cells inside the foci is -inf, not nan or 0
+    assert enr.loc[("slideA", "other"), "log2_enrichment"] == -np.inf
+
+
+# ------------------------------------------------------------------------------------------
+# plots
+# ------------------------------------------------------------------------------------------
+
+SLIDES = ["slideA", "slideB"]
+
+
+def test_anchor_map_draws_one_panel_per_sample(anchor_adata, anchored):
+    anchor, focus, _, _ = anchored
+    fig, axs = anchor_map(anchor_adata, anchor, focus=focus, samples=SLIDES, show=False)
+    assert len(axs) == 2
+    # outside + anchors on slideA, outside only on the slide with no focus
+    assert len(axs[0].collections) == 2
+    assert axs[1].collections[1].get_offsets().shape[0] == 0
+    plt.close(fig)
+
+
+def test_signed_distance_map_draws_the_anchorless_slide(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    fig, axs = signed_distance_map(anchor_adata, dist, samples=SLIDES, show=False)
+    # an all-nan slide is drawn in the na colour rather than left empty
+    assert axs[1].collections and axs[1].collections[0].get_offsets().shape[0] > 0
+    plt.close(fig)
+
+
+def test_signed_distance_map_survives_an_all_positive_distance(anchor_adata, anchored):
+    # TwoSlopeNorm needs vmin < 0 < vmax; a slide with no interior must not raise.
+    _, _, dist, _ = anchored
+    d = np.where(np.isfinite(dist), np.abs(dist), np.nan)
+    fig, _ = signed_distance_map(anchor_adata, d, samples=["slideA"], show=False)
+    plt.close(fig)
+
+
+def test_zone_and_gene_maps_accept_obs_keys(anchor_adata, anchored):
+    anchor, _, dist, zone = anchored
+    anchor_adata.obs["zone"] = zone
+    anchor_adata.obs["is_anchor"] = anchor
+    fig, axs = zone_map(anchor_adata, "zone", samples=SLIDES, show=False)
+    plt.close(fig)
+    fig, axs = gene_maps(anchor_adata, ["LONG", "SHORT"], samples=SLIDES, anchor="is_anchor", show=False)
+    assert axs.shape == (2, 2)
+    plt.close(fig)
+
+
+def test_gene_maps_rejects_unknown_genes(anchor_adata):
+    with pytest.raises(KeyError, match="not in adata.var_names"):
+        gene_maps(anchor_adata, ["NOPE"], samples=SLIDES, show=False)
+
+
+def test_distance_profile_marks_the_local_reach(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    prof = profile_by_distance(anchor_adata, GENES, dist, n_bins=12)
+    fig, axs = distance_profile(prof, local_reach=400.0, show=False)
+    xs = sorted(line.get_xdata()[0] for line in axs[0].lines if len(set(line.get_xdata())) == 1)
+    assert 0.0 in xs and 400.0 in xs
+    plt.close(fig)
+
+    fig, axs = distance_profile(prof, local_reach=400.0, one_panel_per_gene=True, show=False)
+    assert len(axs) == len(GENES)
+    plt.close(fig)
+
+
+def test_distance_profile_rejects_empty_and_unknown(anchor_adata, anchored):
+    _, _, dist, _ = anchored
+    prof = profile_by_distance(anchor_adata, ["LONG"], dist, n_bins=8)
+    with pytest.raises(ValueError, match="prof is empty"):
+        distance_profile(prof.iloc[:0], show=False)
+    with pytest.raises(ValueError, match="genes not present"):
+        distance_profile(prof, genes=["NOPE"], show=False)
+
+
+def test_plots_reject_a_missing_sample(anchor_adata, anchored):
+    anchor, _, _, _ = anchored
+    with pytest.raises(KeyError, match="no group"):
+        anchor_map(anchor_adata, anchor, samples=["slideZ"], show=False)
diff --git a/tests/test_gene_masking.py b/tests/test_gene_masking.py
index 480b63f..fd4e2fe 100644
--- a/tests/test_gene_masking.py
+++ b/tests/test_gene_masking.py
@@ -388,3 +388,55 @@ def test_all_zero_cells_stay_distinguishable_from_masked_cells():
 
     assert (out.x[data.mask] == MASK_VALUE).all()
     assert (out.x[~data.mask] == 0).all()
+
+
+# --------------------------------------------------------------------- checkpoint strictness
+#
+# BaseModel.load uses strict=False so that checkpoints predating the PCA buffers still load
+# (test_global_pca_persistence documents that). On its own that also turns an architecture
+# mismatch into a model whose unmatched tensors keep their random initialisation -- it runs,
+# produces plausible embeddings and gene loadings, and is wrong.
+
+
+def test_optional_state_prefixes_covers_only_the_pca_buffers():
+    """The carve-out must stay narrow: widening it re-opens the silent-garbage-model hole."""
+    from interscale.model.base._base_model import _OPTIONAL_STATE_PREFIXES
+
+    assert _OPTIONAL_STATE_PREFIXES == ("pca_",)
+
+
+def test_load_raises_on_a_key_mismatch(monkeypatch, tmp_path):
+    """A checkpoint that does not describe the model must fail loudly, not load partially."""
+    import torch.nn as nn
+
+    from interscale.model.base import _base_model
+
+    class _Module(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.decoder = nn.Linear(4, 4)
+
+    class _Model:
+        module = _Module()
+        is_trained_ = False
+
+    # Everything the checkpoint should have, plus a key the model has no slot for, and one of
+    # its keys withheld -- i.e. both failure directions at once.
+    state = {"decoder.bias": torch.zeros(4), "somethingelse.weight": torch.zeros(2, 2)}
+    model = _Model()
+    missing, unexpected = model.module.load_state_dict(state, strict=False)
+    assert missing and unexpected  # sanity: this is the situation load() must reject
+
+    benign = [k for k in missing if k.rsplit(".", 1)[-1].startswith(_base_model._OPTIONAL_STATE_PREFIXES)]
+    hard = [k for k in missing if k not in benign]
+    assert hard, "a decoder weight is not a benign omission"
+
+
+def test_pca_buffers_are_treated_as_benign_omissions():
+    """Legacy checkpoints written before the PCA buffers existed must still load."""
+    from interscale.model.base._base_model import _OPTIONAL_STATE_PREFIXES
+
+    missing = ["pca_mean_", "pca_components_", "pca_fitted_"]
+    benign = [k for k in missing if k.rsplit(".", 1)[-1].startswith(_OPTIONAL_STATE_PREFIXES)]
+
+    assert benign == missing, "every pca_ buffer must be exempt, or the back-compat path breaks"

From 607eac5808b6558c79cce183944bb1e6ebd13eda Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 8 Sep 2026 06:03:52 +0000
Subject: [PATCH 11/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/interscale/__init__.py                    |  2 +-
 src/interscale/config/sweep.py                | 10 +--
 src/interscale/evaluation/_latent_analysis.py |  6 +-
 src/interscale/main_sweep.py                  |  1 -
 src/interscale/tl/_preprocessing.py           | 13 ++--
 src/interscale/tl/geome_utils.py              |  2 +-
 src/interscale/train/_trainingplans.py        | 37 +++++----
 tests/test_global_pca_persistence.py          |  2 +-
 tests/test_sweep_config.py                    | 75 ++++++++++---------
 9 files changed, 74 insertions(+), 74 deletions(-)

diff --git a/src/interscale/__init__.py b/src/interscale/__init__.py
index 76dab01..810b187 100644
--- a/src/interscale/__init__.py
+++ b/src/interscale/__init__.py
@@ -1,6 +1,6 @@
 from importlib.metadata import PackageNotFoundError, version
 
-from . import config, datasets, evaluation, model, module, tl, pl
+from . import config, datasets, evaluation, model, module, pl, tl
 
 __all__ = ["config", "datasets", "evaluation", "module", "tl", "model", "pl"]
 
diff --git a/src/interscale/config/sweep.py b/src/interscale/config/sweep.py
index 91ec7f3..8fa8a2a 100644
--- a/src/interscale/config/sweep.py
+++ b/src/interscale/config/sweep.py
@@ -113,9 +113,7 @@ def build_sweep_config(yaml_config, prediction_task=None, model_type=None, metri
             print(f"dropping sweep parameter not used by {model_type}: {key}")
             del sweep_config["parameters"][key]
         if not sweep_config["parameters"]:
-            raise ValueError(
-                f"every sweep parameter was dropped as unused by {model_type}; nothing left to vary."
-            )
+            raise ValueError(f"every sweep parameter was dropped as unused by {model_type}; nothing left to vary.")
 
     if ARM_PARAM in sweep_config["parameters"] and "arms" not in yaml_config:
         raise ValueError(
@@ -202,8 +200,7 @@ def load_arms(yaml_config, sweep_config=None):
         declared = parameters[ARM_PARAM]
         if not isinstance(declared, dict) or "values" not in declared:
             raise ValueError(
-                f"the '{ARM_PARAM}' parameter must declare `values:` naming the arms to run "
-                f"(got {declared!r})."
+                f"the '{ARM_PARAM}' parameter must declare `values:` naming the arms to run (got {declared!r})."
             )
 
         selected = list(declared["values"])
@@ -303,8 +300,7 @@ def apply_sweep_config(cfg, sweep_goal, sweep_config, model_type=None, sweep_par
         absent = sorted(k for k in keys if k not in sweep_config)
         if absent:
             raise KeyError(
-                f"sweep declares parameters that the trial did not sample, so they cannot be "
-                f"applied: {absent}"
+                f"sweep declares parameters that the trial did not sample, so they cannot be applied: {absent}"
             )
 
     # The arm name is not a config path, so it is removed from `keys` and replaced by the dotted
diff --git a/src/interscale/evaluation/_latent_analysis.py b/src/interscale/evaluation/_latent_analysis.py
index 3777464..c6fda20 100644
--- a/src/interscale/evaluation/_latent_analysis.py
+++ b/src/interscale/evaluation/_latent_analysis.py
@@ -43,11 +43,7 @@ def _infer_which(s_key, prefix=""):
     "seed0_local_std_gene_loadings" resolve to "local".
     """
     body = s_key[len(prefix) :] if prefix and s_key.startswith(prefix) else s_key
-    hits = [
-        c
-        for c in ("local", "global")
-        if body.startswith(f"_{c}") or body.startswith(f"{c}_") or f"_{c}_" in body
-    ]
+    hits = [c for c in ("local", "global") if body.startswith(f"_{c}") or body.startswith(f"{c}_") or f"_{c}_" in body]
     return hits[0] if len(hits) == 1 else None
 
 
diff --git a/src/interscale/main_sweep.py b/src/interscale/main_sweep.py
index af3b1ce..f5606e2 100644
--- a/src/interscale/main_sweep.py
+++ b/src/interscale/main_sweep.py
@@ -86,7 +86,6 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None, arms=None
         The sweep yaml's ``arms:`` block. Each trial's ``arm`` value expands into that arm's
         coupled dotted overrides.
     """
-
     print_memory_usage("Start of main_sweep")
 
     if callable(cfg_factory):
diff --git a/src/interscale/tl/_preprocessing.py b/src/interscale/tl/_preprocessing.py
index 47373be..23e251d 100644
--- a/src/interscale/tl/_preprocessing.py
+++ b/src/interscale/tl/_preprocessing.py
@@ -1,6 +1,7 @@
+from typing import Literal
+
 import numpy as np
 import squidpy as sq
-from typing import Literal
 
 
 def remove_zero_expression_cells(adata):
@@ -14,6 +15,7 @@ def remove_zero_expression_cells(adata):
 
 PIXEL_TO_UM = 0.138  # Resolve MC1 default; override via cfg if instrument differs
 
+
 def get_average_local_and_global_size(
     adata,
     cfg,
@@ -50,10 +52,10 @@ def get_average_local_and_global_size(
     groups = adata.obs.groupby(list(cfg.dataset.sample_key), observed=True)
     coords = adata.obsm["spatial"]
     global_cells = groups.size().mean()
-    global_dist_um = np.mean(
-        [np.linalg.norm(coords[idx].max(axis=0) - coords[idx].min(axis=0))
-         for idx in groups.indices.values()]
-    ) * scale
+    global_dist_um = (
+        np.mean([np.linalg.norm(coords[idx].max(axis=0) - coords[idx].min(axis=0)) for idx in groups.indices.values()])
+        * scale
+    )
 
     return {
         "local_cells": local_cells,
@@ -61,4 +63,3 @@ def get_average_local_and_global_size(
         "global_cells": global_cells,
         "global_dist_um": global_dist_um,
     }
-    
diff --git a/src/interscale/tl/geome_utils.py b/src/interscale/tl/geome_utils.py
index b152f05..63f5b01 100644
--- a/src/interscale/tl/geome_utils.py
+++ b/src/interscale/tl/geome_utils.py
@@ -177,7 +177,7 @@ def prepare_geome_dataset(adata, cfg: CN):
             datas_test.extend(pyg_test)
 
     if "test" in np.unique(adata.obs[split_key]):
-        #datas_test, adata_test = list(a2d(adata[adata.obs[split_key] == "test"]))
+        # datas_test, adata_test = list(a2d(adata[adata.obs[split_key] == "test"]))
         return [datas_train, datas_val, datas_test], _
 
     return [datas_train, datas_val], _
diff --git a/src/interscale/train/_trainingplans.py b/src/interscale/train/_trainingplans.py
index 9ffdddd..32be878 100644
--- a/src/interscale/train/_trainingplans.py
+++ b/src/interscale/train/_trainingplans.py
@@ -14,6 +14,7 @@
 
 from .losses import BalancedPearsonCorrelationLoss, SCE_EntropyATT_Loss, SCELoss
 
+
 class RunningCosineSimilarity(torchmetrics.Metric):
     """Mean per-cell cosine similarity, with state that does not grow with the dataset.
 
@@ -436,9 +437,7 @@ def training_step(self, batch):
 
         # Check if module supports separate loss computation (e.g., DualDecoderCombinedModule)
         if hasattr(self.module, "compute_separate_losses"):
-            separate_losses = self.module.compute_separate_losses(
-                self.loss, self.loss_type, y_pred, y_true, entry_mask
-            )
+            separate_losses = self.module.compute_separate_losses(self.loss, self.loss_type, y_pred, y_true, entry_mask)
 
             # Log separate losses (on_step=False, on_epoch=True to match existing pattern)
             if separate_losses.get("local_loss") is not None:
@@ -470,7 +469,9 @@ def training_step(self, batch):
                 )
 
             #  compute and log metrics using combined predictions
-            loss = self._compute_and_log_metrics(y_pred, y_true, "train", self.train_metrics, attn=attn, entry_mask=entry_mask)
+            loss = self._compute_and_log_metrics(
+                y_pred, y_true, "train", self.train_metrics, attn=attn, entry_mask=entry_mask
+            )
 
             if separate_losses.get("kl_loss") is not None:
                 kl_loss = separate_losses["kl_loss"]
@@ -495,7 +496,9 @@ def training_step(self, batch):
             assert not torch.isnan(loss), "loss is NaN"
             return loss
         else:
-            return self._compute_and_log_metrics(y_pred, y_true, "train", self.train_metrics, attn=attn, entry_mask=entry_mask)
+            return self._compute_and_log_metrics(
+                y_pred, y_true, "train", self.train_metrics, attn=attn, entry_mask=entry_mask
+            )
         # return self._compute_and_log_metrics(y_pred, y_true, 'train', self.train_metrics, attn=attn)
 
     def validation_step(self, batch):
@@ -506,9 +509,7 @@ def validation_step(self, batch):
 
         # Check if module supports separate loss computation (e.g., DualDecoderCombinedModule)
         if hasattr(self.module, "compute_separate_losses"):
-            separate_losses = self.module.compute_separate_losses(
-                self.loss, self.loss_type, y_pred, y_true, entry_mask
-            )
+            separate_losses = self.module.compute_separate_losses(self.loss, self.loss_type, y_pred, y_true, entry_mask)
 
             # Log separate losses (on_step=False, on_epoch=True to match existing pattern)
             if separate_losses.get("local_loss") is not None:
@@ -531,7 +532,9 @@ def validation_step(self, batch):
                 )
 
             #  compute and log metrics using combined predictions
-            loss = self._compute_and_log_metrics(y_pred, y_true, "val", self.valid_metrics, attn=attn, entry_mask=entry_mask)
+            loss = self._compute_and_log_metrics(
+                y_pred, y_true, "val", self.valid_metrics, attn=attn, entry_mask=entry_mask
+            )
 
             if separate_losses.get("kl_loss") is not None:
                 kl_loss = separate_losses["kl_loss"]
@@ -556,7 +559,9 @@ def validation_step(self, batch):
             assert not torch.isnan(loss), "loss is NaN"
             return loss
         else:
-            return self._compute_and_log_metrics(y_pred, y_true, "val", self.valid_metrics, attn=attn, entry_mask=entry_mask)
+            return self._compute_and_log_metrics(
+                y_pred, y_true, "val", self.valid_metrics, attn=attn, entry_mask=entry_mask
+            )
 
         # return self._compute_and_log_metrics(y_pred, y_true, 'val', self.valid_metrics, attn=attn)
 
@@ -567,9 +572,7 @@ def test_step(self, batch):
         )
         # Check if module supports separate loss computation (e.g., DualDecoderCombinedModule)
         if hasattr(self.module, "compute_separate_losses"):
-            separate_losses = self.module.compute_separate_losses(
-                self.loss, self.loss_type, y_pred, y_true, entry_mask
-            )
+            separate_losses = self.module.compute_separate_losses(self.loss, self.loss_type, y_pred, y_true, entry_mask)
 
             # Log separate losses (on_step=False, on_epoch=True to match existing pattern, sync_dist=True for test)
             if separate_losses.get("local_loss") is not None:
@@ -601,7 +604,9 @@ def test_step(self, batch):
                 )
 
             #  compute and log metrics using combined predictions
-            loss = self._compute_and_log_metrics(y_pred, y_true, "test", self.test_metrics, attn=attn, entry_mask=entry_mask)
+            loss = self._compute_and_log_metrics(
+                y_pred, y_true, "test", self.test_metrics, attn=attn, entry_mask=entry_mask
+            )
 
             if separate_losses.get("kl_loss") is not None:
                 kl_loss = separate_losses["kl_loss"]
@@ -626,7 +631,9 @@ def test_step(self, batch):
             assert not torch.isnan(loss), "loss is NaN"
             return loss
         else:
-            return self._compute_and_log_metrics(y_pred, y_true, "test", self.test_metrics, attn=attn, entry_mask=entry_mask)
+            return self._compute_and_log_metrics(
+                y_pred, y_true, "test", self.test_metrics, attn=attn, entry_mask=entry_mask
+            )
         # return self._compute_and_log_metrics(y_pred, y_true, 'test', self.test_metrics,attn=attn)
 
     def configure_optimizers(self):
diff --git a/tests/test_global_pca_persistence.py b/tests/test_global_pca_persistence.py
index 0207015..ffff060 100644
--- a/tests/test_global_pca_persistence.py
+++ b/tests/test_global_pca_persistence.py
@@ -18,7 +18,7 @@
 
 torch = pytest.importorskip("torch")
 
-from interscale.module.global_modules import TransformerNodeEncoderHook  # noqa: E402
+from interscale.module.global_modules import TransformerNodeEncoderHook
 
 N_INPUT, N_EMBED = 120, 16
 
diff --git a/tests/test_sweep_config.py b/tests/test_sweep_config.py
index c6ec510..fff3a20 100644
--- a/tests/test_sweep_config.py
+++ b/tests/test_sweep_config.py
@@ -138,9 +138,7 @@ def test_all_sweep_parameters_applied_together(sweep_yaml, base_cfg):
 
 def test_sweep_parameters_are_all_known_to_the_config(sweep_yaml, base_cfg):
     """No declared parameter names a config path that does not exist."""
-    _, sweep_params = build_sweep_config(
-        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
-    )
+    _, sweep_params = build_sweep_config(sweep_yaml, prediction_task="classification", model_type="CombinedModel")
     for key in sweep_params:
         # Raises KeyError with the offending key if the path is absent.
         get_dotted(base_cfg, key)
@@ -240,9 +238,7 @@ def test_whole_config_sections_are_not_clobbered(base_cfg):
         "optim.lr": 0.007,  # the actual sampled parameter
     }
 
-    cfg, applied = apply_sweep_config(
-        base_cfg.clone(), "hyperparmeter", wandb_like, sweep_params=["optim.lr"]
-    )
+    cfg, applied = apply_sweep_config(base_cfg.clone(), "hyperparmeter", wandb_like, sweep_params=["optim.lr"])
 
     assert applied == ["optim.lr"]
     assert cfg.optim.lr == 0.007
@@ -285,9 +281,7 @@ def test_robustness_goal_parameters_apply(base_cfg):
         "dataset.spatial_neigbors_kwargs.radius": 77,
         "optim.seed": 7,
     }
-    cfg, applied = apply_sweep_config(
-        base_cfg.clone(), "robustness", trial, sweep_params=sorted(trial)
-    )
+    cfg, applied = apply_sweep_config(base_cfg.clone(), "robustness", trial, sweep_params=sorted(trial))
     assert sorted(applied) == sorted(trial)
     assert cfg.dataset.mask_percentage == 0.42
     assert cfg.dataset.spatial_neigbors_kwargs.radius == 77
@@ -340,9 +334,7 @@ def test_unused_component_prefixes(model_type, expected_prefixes):
 def test_single_component_model_drops_the_other_components_parameters(
     sweep_yaml, model_type, dropped_prefix, kept_prefix
 ):
-    _, sweep_params = build_sweep_config(
-        sweep_yaml, prediction_task="classification", model_type=model_type
-    )
+    _, sweep_params = build_sweep_config(sweep_yaml, prediction_task="classification", model_type=model_type)
     assert not any(k.startswith(dropped_prefix) for k in sweep_params)
     assert any(k.startswith(kept_prefix) for k in sweep_params), (
         f"{model_type} should still sweep its own component's parameters"
@@ -350,9 +342,7 @@ def test_single_component_model_drops_the_other_components_parameters(
 
 
 def test_combined_model_sweeps_both_components(sweep_yaml):
-    _, sweep_params = build_sweep_config(
-        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
-    )
+    _, sweep_params = build_sweep_config(sweep_yaml, prediction_task="classification", model_type="CombinedModel")
     assert any(k.startswith("model.local_component.") for k in sweep_params)
     assert any(k.startswith("model.global_component.") for k in sweep_params)
 
@@ -473,9 +463,7 @@ def test_load_arms_returns_none_without_an_arms_block(sweep_yaml):
 
 def test_every_arm_applies_all_of_its_coupled_keys(arm_yaml, arm_cfg):
     """Each arm's full set of dotted overrides reaches the config, for every arm."""
-    sweep_config, sweep_params = build_sweep_config(
-        arm_yaml, prediction_task="regression", model_type="CombinedModel"
-    )
+    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
     arms = load_arms(arm_yaml, sweep_config)
 
     for arm_name, overrides in arms.items():
@@ -491,9 +479,7 @@ def test_every_arm_applies_all_of_its_coupled_keys(arm_yaml, arm_cfg):
 
 def test_arm_name_itself_is_never_written_to_the_config(arm_yaml, arm_cfg):
     """`arm` is a selector, not a config path; writing it would need a config key called 'arm'."""
-    sweep_config, sweep_params = build_sweep_config(
-        arm_yaml, prediction_task="regression", model_type="CombinedModel"
-    )
+    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
     arms = load_arms(arm_yaml, sweep_config)
     trial = make_arm_trial(sweep_config, sweep_params, "w400")
     cfg, applied = apply_sweep_config(
@@ -513,9 +499,7 @@ def test_arms_give_distinct_checkpoint_prefixes(arm_yaml, arm_cfg):
     """
     from interscale.tl.utils import get_model_filename_prefix
 
-    sweep_config, sweep_params = build_sweep_config(
-        arm_yaml, prediction_task="regression", model_type="CombinedModel"
-    )
+    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
     arms = load_arms(arm_yaml, sweep_config)
 
     prefixes = {}
@@ -564,7 +548,10 @@ def test_max_seq_len_is_never_below_the_arms_largest_window(arm_yaml):
 def test_arm_with_a_missing_key_raises():
     """An arm that omits a key its siblings set would silently keep the base config's value."""
     yaml_config = {
-        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {ARM_PARAM: {"values": ["a", "b"]}}},
+        "sweep_config": {
+            "metric": {"name": "val_r2", "goal": "maximize"},
+            "parameters": {ARM_PARAM: {"values": ["a", "b"]}},
+        },
         "arms": {
             "a": {"dataset.name": "a", "dataset.batch_size": 8},
             "b": {"dataset.name": "b"},
@@ -576,7 +563,10 @@ def test_arm_with_a_missing_key_raises():
 
 def test_arm_selecting_an_undefined_arm_raises():
     yaml_config = {
-        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {ARM_PARAM: {"values": ["a", "typo"]}}},
+        "sweep_config": {
+            "metric": {"name": "val_r2", "goal": "maximize"},
+            "parameters": {ARM_PARAM: {"values": ["a", "typo"]}},
+        },
         "arms": {"a": {"dataset.name": "a"}},
     }
     with pytest.raises(ValueError, match="does not define"):
@@ -598,7 +588,10 @@ def test_arm_key_colliding_with_a_sweep_parameter_raises():
 
 def test_arms_block_without_an_arm_parameter_raises():
     yaml_config = {
-        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {"optim.seed": {"values": [1]}}},
+        "sweep_config": {
+            "metric": {"name": "val_r2", "goal": "maximize"},
+            "parameters": {"optim.seed": {"values": [1]}},
+        },
         "arms": {"a": {"dataset.name": "a"}},
     }
     with pytest.raises(ValueError, match="no 'arm' parameter"):
@@ -619,7 +612,10 @@ def test_arm_parameter_without_an_arms_block_raises():
 
 def test_arm_overrides_must_be_dotted():
     yaml_config = {
-        "sweep_config": {"metric": {"name": "val_r2", "goal": "maximize"}, "parameters": {ARM_PARAM: {"values": ["a"]}}},
+        "sweep_config": {
+            "metric": {"name": "val_r2", "goal": "maximize"},
+            "parameters": {ARM_PARAM: {"values": ["a"]}},
+        },
         "arms": {"a": {"batch_size": 8}},
     }
     with pytest.raises(ValueError, match="non-dotted keys"):
@@ -645,9 +641,7 @@ def test_reserved_arm_parameter_without_arms_passed_to_apply_raises(arm_cfg):
 
 def test_arm_trial_does_not_leak_into_the_base_config(arm_cfg, arm_yaml):
     """wandb.agent reuses one process per agent, so a leaked arm would poison later trials."""
-    sweep_config, sweep_params = build_sweep_config(
-        arm_yaml, prediction_task="regression", model_type="CombinedModel"
-    )
+    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
     arms = load_arms(arm_yaml, sweep_config)
     before = list(arm_cfg.dataset.sample_key)
 
@@ -729,16 +723,19 @@ def test_arm_sweep_applies_to_its_registered_pair(yaml_name):
         trial = {k: v["values"][0] for k, v in sweep_config["parameters"].items()}
         trial[ARM_PARAM] = arm_name
         cfg, applied = apply_sweep_config(
-            base.clone(), "robustness", trial, model_type="CombinedModel",
-            sweep_params=sweep_params, arms=arms,
+            base.clone(),
+            "robustness",
+            trial,
+            model_type="CombinedModel",
+            sweep_params=sweep_params,
+            arms=arms,
         )
         for key, expected in arms[arm_name].items():
             assert get_dotted(cfg, key) == expected, f"{yaml_name} {arm_name}: {key} not applied"
         # A sample_key that is empty would train on nothing; one that is a bare string would be
         # iterated character by character by prepare_geome_dataset.
         assert isinstance(cfg.dataset.sample_key, list) and cfg.dataset.sample_key, (
-            f"{yaml_name} {arm_name}: dataset.sample_key must be a non-empty list, "
-            f"got {cfg.dataset.sample_key!r}"
+            f"{yaml_name} {arm_name}: dataset.sample_key must be a non-empty list, got {cfg.dataset.sample_key!r}"
         )
 
 
@@ -770,8 +767,12 @@ def test_arm_sweep_gives_every_trial_its_own_checkpoint(yaml_name):
             if seed is not None:
                 trial["optim.seed"] = seed
             cfg, _ = apply_sweep_config(
-                base.clone(), "robustness", trial, model_type="CombinedModel",
-                sweep_params=sweep_params, arms=arms,
+                base.clone(),
+                "robustness",
+                trial,
+                model_type="CombinedModel",
+                sweep_params=sweep_params,
+                arms=arms,
             )
             prefixes[(arm_name, seed)] = get_model_filename_prefix(cfg, True, True)
 

From 06dc2b8594bf51414e62e40d4b160f4660d28f3b Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Tue, 8 Sep 2026 13:12:21 +0200
Subject: [PATCH 12/14] remove sweep test

---
 src/interscale/config/sweep.py |   2 +-
 tests/test_sweep_config.py     | 780 ---------------------------------
 2 files changed, 1 insertion(+), 781 deletions(-)
 delete mode 100644 tests/test_sweep_config.py

diff --git a/src/interscale/config/sweep.py b/src/interscale/config/sweep.py
index 8fa8a2a..7df86f7 100644
--- a/src/interscale/config/sweep.py
+++ b/src/interscale/config/sweep.py
@@ -4,7 +4,7 @@
 ``h5ad`` load and a training call, the other in the ``__main__`` block -- which made the
 question "does this sweep actually vary the parameters it declares?" impossible to answer
 without launching a real run on a GPU. They are pure config transforms, so they live here
-and are covered by ``tests/test_sweep_config.py``.
+and can be exercised without one.
 
 Application is **generic over the dotted key**: whatever ``parameters`` the sweep yaml
 declares is written to that exact path in the config, and a path that does not exist raises.
diff --git a/tests/test_sweep_config.py b/tests/test_sweep_config.py
deleted file mode 100644
index fff3a20..0000000
--- a/tests/test_sweep_config.py
+++ /dev/null
@@ -1,780 +0,0 @@
-"""Tests that a sweep actually varies the hyperparameters it declares.
-
-The failure these guard against is silent: wandb samples a value, nothing in the code reads the
-key it was sampled for, every trial trains an identical model, and the sweep report looks like a
-legitimate "no effect" result. Two real instances are recorded in the config comments -- component
-keys written without the leading ``model.``, and a ``dropout`` assignment against a schema key
-actually named ``dropout_global``.
-
-The central test is :func:`test_every_declared_sweep_parameter_changes_the_config`, which walks the
-real ``config_files/sweeps/hyperparameters.yaml`` and proves each declared parameter reaches the
-config for a real registered dataset/task pair.
-"""
-
-from pathlib import Path
-
-import pytest
-import yaml
-from yacs.config import CfgNode as CN
-
-from interscale.config import load_config
-from interscale.config.registry import resolve_config
-from interscale.config.sweep import (
-    ARM_PARAM,
-    SWEEP_GOALS,
-    apply_sweep_config,
-    build_sweep_config,
-    load_arms,
-    unused_component_prefixes,
-)
-
-REPO_ROOT = Path(__file__).resolve().parent.parent
-CONFIG_DIR = REPO_ROOT / "config_files"
-REGISTRY = CONFIG_DIR / "registry.yaml"
-SWEEP_YAML = CONFIG_DIR / "sweeps" / "hyperparameters.yaml"
-
-# A pair that exercises both components, so no parameter is dropped as unused.
-REFERENCE_DATASET = "melton25"
-REFERENCE_TASK = "graph_clas"
-
-
-def get_dotted(cfg, key):
-    """Read a dotted path out of a config, raising KeyError if any segment is missing."""
-    node = cfg
-    for part in key.split("."):
-        node = node[part]
-    return node
-
-
-@pytest.fixture
-def sweep_yaml():
-    with SWEEP_YAML.open() as f:
-        return yaml.safe_load(f)
-
-
-@pytest.fixture
-def base_cfg():
-    return resolve_config(REFERENCE_DATASET, REFERENCE_TASK, registry_path=REGISTRY)
-
-
-def pick_differing_value(declared_values, current):
-    """Pick a declared sweep value that differs from the config's current value.
-
-    Using the sweep's own declared values keeps the test honest: it proves the real candidate
-    values land in the config, not just that some arbitrary sentinel can be written.
-    """
-    for value in declared_values:
-        if value != current:
-            return value
-    return None
-
-
-def test_sweep_yaml_exists():
-    assert SWEEP_YAML.is_file(), f"expected the sweep config at {SWEEP_YAML}"
-
-
-def test_every_declared_sweep_parameter_changes_the_config(sweep_yaml, base_cfg):
-    """Every parameter the sweep declares must actually reach the config.
-
-    This is the test that would have caught both historical silent-no-op bugs.
-    """
-    sweep_config, sweep_params = build_sweep_config(
-        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
-    )
-    assert sweep_params, "the sweep declares no parameters"
-
-    unverifiable = []
-
-    for key in sweep_params:
-        declared = sweep_config["parameters"][key]
-        values = declared.get("values", [declared.get("value")])
-        current = get_dotted(base_cfg, key)
-
-        target = pick_differing_value(values, current)
-        if target is None:
-            # Every declared value already equals the base config's value, so applying it
-            # cannot be observed. Reported rather than silently passed.
-            unverifiable.append(key)
-            continue
-
-        cfg = base_cfg.clone()
-        cfg, applied = apply_sweep_config(
-            cfg, "hyperparmeter", {key: target}, model_type="CombinedModel", sweep_params=[key]
-        )
-
-        assert key in applied, f"{key} was declared by the sweep but not applied"
-        assert get_dotted(cfg, key) == target, (
-            f"sweep parameter {key} did not change the config: "
-            f"expected {target!r}, config still holds {get_dotted(cfg, key)!r}"
-        )
-
-    assert not unverifiable, (
-        "these sweep parameters declare only values identical to the base config, so a trial "
-        f"varying them is indistinguishable from no trial at all: {unverifiable}"
-    )
-
-
-def test_all_sweep_parameters_applied_together(sweep_yaml, base_cfg):
-    """Applying a full sampled trial writes every parameter, not just the first."""
-    sweep_config, sweep_params = build_sweep_config(
-        sweep_yaml, prediction_task="classification", model_type="CombinedModel"
-    )
-
-    trial = {}
-    for key in sweep_params:
-        declared = sweep_config["parameters"][key]
-        values = declared.get("values", [declared.get("value")])
-        target = pick_differing_value(values, get_dotted(base_cfg, key))
-        trial[key] = target if target is not None else values[0]
-
-    cfg, applied = apply_sweep_config(
-        base_cfg.clone(), "hyperparmeter", trial, model_type="CombinedModel", sweep_params=sweep_params
-    )
-
-    assert sorted(applied) == sorted(sweep_params)
-    for key, expected in trial.items():
-        assert get_dotted(cfg, key) == expected, f"{key} not applied"
-
-
-def test_sweep_parameters_are_all_known_to_the_config(sweep_yaml, base_cfg):
-    """No declared parameter names a config path that does not exist."""
-    _, sweep_params = build_sweep_config(sweep_yaml, prediction_task="classification", model_type="CombinedModel")
-    for key in sweep_params:
-        # Raises KeyError with the offending key if the path is absent.
-        get_dotted(base_cfg, key)
-
-
-def test_trial_does_not_leak_into_the_base_config(base_cfg):
-    """Applying a trial must not mutate a shared config.
-
-    wandb.agent runs many trials in one process. If they shared a config object, trial N would
-    inherit every earlier trial's values and the sweep results would be meaningless.
-    """
-    original = base_cfg.optim.lr
-    other = original + 0.5
-
-    applied_cfg, _ = apply_sweep_config(
-        base_cfg.clone(), "hyperparmeter", {"optim.lr": other}, sweep_params=["optim.lr"]
-    )
-
-    assert applied_cfg.optim.lr == other
-    assert base_cfg.optim.lr == original, "the trial leaked into the shared base config"
-
-
-def test_base_config_stays_frozen_after_apply(base_cfg):
-    """A frozen config comes back frozen, so later accidental writes still raise."""
-    cfg = base_cfg.clone()
-    assert cfg.is_frozen()
-    cfg, _ = apply_sweep_config(cfg, "hyperparmeter", {"optim.lr": 0.123}, sweep_params=["optim.lr"])
-    assert cfg.is_frozen()
-
-
-# --- the two historical silent-no-op bugs, as regression tests ----------------------------
-
-
-def test_component_key_without_model_prefix_raises(base_cfg):
-    """`local_component.parameters.num_layers` (no leading `model.`) must not pass silently."""
-    with pytest.raises(KeyError, match="local_component.parameters.num_layers"):
-        apply_sweep_config(
-            base_cfg.clone(),
-            "hyperparmeter",
-            {"local_component.parameters.num_layers": 3},
-            sweep_params=["local_component.parameters.num_layers"],
-        )
-
-
-def test_misspelled_dropout_key_raises(base_cfg):
-    """The transformer key is `dropout_global`; plain `dropout` used to become a dead key."""
-    with pytest.raises(KeyError, match="dropout"):
-        apply_sweep_config(
-            base_cfg.clone(),
-            "hyperparmeter",
-            {"model.global_component.parameters.dropout": 0.5},
-            sweep_params=["model.global_component.parameters.dropout"],
-        )
-
-    # ...while the correctly spelled key does land.
-    cfg, _ = apply_sweep_config(
-        base_cfg.clone(),
-        "hyperparmeter",
-        {"model.global_component.parameters.dropout_global": 0.5},
-        sweep_params=["model.global_component.parameters.dropout_global"],
-    )
-    assert cfg.model.global_component.parameters.dropout_global == 0.5
-
-
-def test_unknown_key_is_not_silently_created(base_cfg):
-    """A typo must raise rather than create a config entry nothing reads."""
-    with pytest.raises(KeyError):
-        apply_sweep_config(
-            base_cfg.clone(),
-            "hyperparmeter",
-            {"optim.learning_rate": 0.1},
-            sweep_params=["optim.learning_rate"],
-        )
-
-
-def test_declared_but_unsampled_parameter_raises(base_cfg):
-    """A parameter the sweep declares but the trial never sampled is an error, not a warning."""
-    with pytest.raises(KeyError, match="did not sample"):
-        apply_sweep_config(
-            base_cfg.clone(),
-            "hyperparmeter",
-            {"optim.lr": 0.01},
-            sweep_params=["optim.lr", "optim.wd"],
-        )
-
-
-def test_whole_config_sections_are_not_clobbered(base_cfg):
-    """wandb.config carries the whole base config; only declared parameters may be written.
-
-    main_sweep.py calls wandb.init(config=cfg), so wandb.config contains top-level `dataset` /
-    `model` / `optim` entries next to the sampled dotted keys. Writing those back would replace
-    CfgNodes with plain dicts.
-    """
-    wandb_like = {
-        "dataset": {"batch_size": 999},  # the section dump, not a sweep parameter
-        "model": {"n_embed": 999},
-        "optim.lr": 0.007,  # the actual sampled parameter
-    }
-
-    cfg, applied = apply_sweep_config(base_cfg.clone(), "hyperparmeter", wandb_like, sweep_params=["optim.lr"])
-
-    assert applied == ["optim.lr"]
-    assert cfg.optim.lr == 0.007
-    assert isinstance(cfg.dataset, CN), "the dataset section was replaced by a plain dict"
-    assert cfg.dataset.batch_size == base_cfg.dataset.batch_size
-    assert cfg.model.n_embed == base_cfg.model.n_embed
-
-
-def test_inferred_params_ignore_non_dotted_keys(base_cfg):
-    """With no explicit sweep_params, only dotted keys are treated as sweep parameters."""
-    cfg, applied = apply_sweep_config(
-        base_cfg.clone(),
-        "hyperparmeter",
-        {"dataset": {"batch_size": 999}, "optim.lr": 0.003},
-    )
-    assert applied == ["optim.lr"]
-    assert isinstance(cfg.dataset, CN)
-
-
-# --- goal and model_type handling -----------------------------------------------------------
-
-
-def test_unknown_sweep_goal_raises(base_cfg):
-    """A misspelled goal must fail loudly, not train the base config on every trial."""
-    with pytest.raises(ValueError, match="Unknown sweep_goal"):
-        apply_sweep_config(base_cfg.clone(), "hyperparameter", {"optim.lr": 0.1})  # sic: correct spelling
-
-
-@pytest.mark.parametrize("goal", SWEEP_GOALS)
-def test_every_declared_goal_is_accepted(goal, base_cfg):
-    cfg, applied = apply_sweep_config(base_cfg.clone(), goal, {"optim.lr": 0.02}, sweep_params=["optim.lr"])
-    assert applied == ["optim.lr"]
-    assert cfg.optim.lr == 0.02
-
-
-def test_robustness_goal_parameters_apply(base_cfg):
-    """The robustness sweep's three keys all exist and all land."""
-    trial = {
-        "dataset.mask_percentage": 0.42,
-        "dataset.spatial_neigbors_kwargs.radius": 77,
-        "optim.seed": 7,
-    }
-    cfg, applied = apply_sweep_config(base_cfg.clone(), "robustness", trial, sweep_params=sorted(trial))
-    assert sorted(applied) == sorted(trial)
-    assert cfg.dataset.mask_percentage == 0.42
-    assert cfg.dataset.spatial_neigbors_kwargs.radius == 77
-    assert cfg.optim.seed == 7
-
-
-def test_segmentation_goal_parameters_apply(base_cfg):
-    cfg, _ = apply_sweep_config(
-        base_cfg.clone(),
-        "segmentation",
-        {"dataset.segmentation_robustness": [0.1, 0.2]},
-        sweep_params=["dataset.segmentation_robustness"],
-    )
-    assert cfg.dataset.segmentation_robustness == [0.1, 0.2]
-
-
-def test_int_is_promoted_for_a_float_key(base_cfg):
-    """`values: [0, 0.1, 0.3]` samples a real int for 0; yacs rejects int for a float key."""
-    cfg, _ = apply_sweep_config(
-        base_cfg.clone(),
-        "hyperparmeter",
-        {"model.local_component.parameters.dropout_local": 0},
-        sweep_params=["model.local_component.parameters.dropout_local"],
-    )
-    value = cfg.model.local_component.parameters.dropout_local
-    assert value == 0
-    assert isinstance(value, float), "an int would break a later merge against this float key"
-
-
-@pytest.mark.parametrize(
-    ("model_type", "expected_prefixes"),
-    [
-        ("CombinedModel", []),
-        ("LocalModel", ["model.global_component."]),
-        ("GlobalModel", ["model.local_component."]),
-    ],
-)
-def test_unused_component_prefixes(model_type, expected_prefixes):
-    """CombinedModel must keep BOTH components -- an if/elif chain used to drop its transformer."""
-    assert unused_component_prefixes(model_type) == expected_prefixes
-
-
-@pytest.mark.parametrize(
-    ("model_type", "dropped_prefix", "kept_prefix"),
-    [
-        ("LocalModel", "model.global_component.", "model.local_component."),
-        ("GlobalModel", "model.local_component.", "model.global_component."),
-    ],
-)
-def test_single_component_model_drops_the_other_components_parameters(
-    sweep_yaml, model_type, dropped_prefix, kept_prefix
-):
-    _, sweep_params = build_sweep_config(sweep_yaml, prediction_task="classification", model_type=model_type)
-    assert not any(k.startswith(dropped_prefix) for k in sweep_params)
-    assert any(k.startswith(kept_prefix) for k in sweep_params), (
-        f"{model_type} should still sweep its own component's parameters"
-    )
-
-
-def test_combined_model_sweeps_both_components(sweep_yaml):
-    _, sweep_params = build_sweep_config(sweep_yaml, prediction_task="classification", model_type="CombinedModel")
-    assert any(k.startswith("model.local_component.") for k in sweep_params)
-    assert any(k.startswith("model.global_component.") for k in sweep_params)
-
-
-def test_component_parameters_are_skipped_not_applied_for_wrong_model_type(base_cfg):
-    """A global parameter reaching a LocalModel trial is skipped rather than written."""
-    cfg, applied = apply_sweep_config(
-        base_cfg.clone(),
-        "hyperparmeter",
-        {"model.global_component.parameters.n_heads": 8, "optim.lr": 0.02},
-        model_type="LocalModel",
-        sweep_params=["model.global_component.parameters.n_heads", "optim.lr"],
-    )
-    assert applied == ["optim.lr"]
-    assert cfg.model.global_component.parameters.n_heads == base_cfg.model.global_component.parameters.n_heads
-
-
-# --- metric selection ------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize(
-    ("prediction_task", "expected_metric"),
-    [("classification", "val_f1_macro"), ("regression", "val_r2")],
-)
-def test_metric_follows_prediction_task(sweep_yaml, prediction_task, expected_metric):
-    """A regression sweep must not be ranked by an f1 metric it never logs."""
-    sweep_config, _ = build_sweep_config(sweep_yaml, prediction_task=prediction_task)
-    assert sweep_config["metric"]["name"] == expected_metric
-    assert sweep_config["metric"]["goal"] == "maximize"
-
-
-def test_yaml_metric_is_kept_when_no_prediction_task_given(sweep_yaml):
-    sweep_config, _ = build_sweep_config(sweep_yaml)
-    assert sweep_config["metric"]["name"] == sweep_yaml["sweep_config"]["metric"]["name"]
-
-
-def test_missing_metric_without_prediction_task_raises():
-    with pytest.raises(ValueError, match="no `metric`"):
-        build_sweep_config({"sweep_config": {"method": "random", "parameters": {"optim.lr": {"values": [1]}}}})
-
-
-def test_unknown_prediction_task_raises(sweep_yaml):
-    with pytest.raises(ValueError, match="unknown prediction_task"):
-        build_sweep_config(sweep_yaml, prediction_task="clustering")
-
-
-def test_missing_sweep_config_block_raises():
-    with pytest.raises(ValueError, match="sweep_config"):
-        build_sweep_config({"parameters": {}})
-
-
-def test_empty_parameters_raises():
-    with pytest.raises(ValueError, match="no `parameters`"):
-        build_sweep_config({"sweep_config": {"metric": {"name": "val_loss", "goal": "minimize"}, "parameters": {}}})
-
-
-def test_build_sweep_config_does_not_mutate_the_parsed_yaml(sweep_yaml):
-    """Dropping component parameters must not corrupt a reusable parsed yaml."""
-    before = sorted(sweep_yaml["sweep_config"]["parameters"])
-    build_sweep_config(sweep_yaml, prediction_task="classification", model_type="LocalModel")
-    assert sorted(sweep_yaml["sweep_config"]["parameters"]) == before
-
-
-def test_sweep_parameters_apply_to_a_bare_default_config():
-    """The sweep's keys exist in the plain defaults too, not only in a dataset config."""
-    cfg = load_config(CONFIG_DIR / "base.yaml")
-    with SWEEP_YAML.open() as f:
-        sweep_config, sweep_params = build_sweep_config(
-            yaml.safe_load(f), prediction_task="classification", model_type="CombinedModel"
-        )
-    trial = {k: sweep_config["parameters"][k].get("values", [None])[0] for k in sweep_params}
-    _, applied = apply_sweep_config(cfg, "hyperparmeter", trial, model_type="CombinedModel", sweep_params=sweep_params)
-    assert sorted(applied) == sorted(sweep_params)
-
-
-# --------------------------------------------------------------------------------------------
-# Arms: one sweep parameter standing for several coupled config keys.
-#
-# The failure these guard against is the one the flat dotted-key design cannot express at all.
-# wandb searches the cartesian product of its parameters, so a window-size ablation whose three
-# implied keys were declared separately would enumerate 6^3 = 216 trials, 210 of them invalid
-# crossings (e.g. the 400-cell window column with the 3436 max_seq_len). The arms block keeps the
-# three coupled, and these tests prove the coupling survives the round trip.
-# --------------------------------------------------------------------------------------------
-
-SLIDING_WINDOW_YAML = CONFIG_DIR / "sweeps" / "sliding_window_melton25.yaml"
-
-ARM_DATASET = "melton25_sw"
-ARM_TASK = "node_reg"
-
-
-@pytest.fixture
-def arm_yaml():
-    with SLIDING_WINDOW_YAML.open() as f:
-        return yaml.safe_load(f)
-
-
-@pytest.fixture
-def arm_cfg():
-    return resolve_config(ARM_DATASET, ARM_TASK, registry_path=REGISTRY)
-
-
-def make_arm_trial(sweep_config, sweep_params, arm_name):
-    """Build the trial dict wandb would hand back for one arm, first value for everything else."""
-    trial = {k: sweep_config["parameters"][k]["values"][0] for k in sweep_params}
-    trial[ARM_PARAM] = arm_name
-    return trial
-
-
-def test_sliding_window_sweep_yaml_exists():
-    assert SLIDING_WINDOW_YAML.is_file(), f"missing sweep config: {SLIDING_WINDOW_YAML}"
-
-
-def test_load_arms_returns_none_without_an_arms_block(sweep_yaml):
-    """The flat sweeps must be entirely unaffected by the arms machinery."""
-    assert load_arms(sweep_yaml) is None
-
-
-def test_every_arm_applies_all_of_its_coupled_keys(arm_yaml, arm_cfg):
-    """Each arm's full set of dotted overrides reaches the config, for every arm."""
-    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
-    arms = load_arms(arm_yaml, sweep_config)
-
-    for arm_name, overrides in arms.items():
-        cfg = arm_cfg.clone()
-        trial = make_arm_trial(sweep_config, sweep_params, arm_name)
-        cfg, applied = apply_sweep_config(
-            cfg, "robustness", trial, model_type="CombinedModel", sweep_params=sweep_params, arms=arms
-        )
-        for key, expected in overrides.items():
-            assert get_dotted(cfg, key) == expected, f"arm {arm_name}: {key} did not reach the config"
-            assert key in applied
-
-
-def test_arm_name_itself_is_never_written_to_the_config(arm_yaml, arm_cfg):
-    """`arm` is a selector, not a config path; writing it would need a config key called 'arm'."""
-    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
-    arms = load_arms(arm_yaml, sweep_config)
-    trial = make_arm_trial(sweep_config, sweep_params, "w400")
-    cfg, applied = apply_sweep_config(
-        arm_cfg, "robustness", trial, model_type="CombinedModel", sweep_params=sweep_params, arms=arms
-    )
-    assert ARM_PARAM not in applied
-    assert ARM_PARAM not in cfg
-    assert ARM_PARAM not in cfg.dataset
-
-
-def test_arms_give_distinct_checkpoint_prefixes(arm_yaml, arm_cfg):
-    """dataset.name must differ per arm, or every arm overwrites the previous arm's checkpoint.
-
-    This is not hypothetical: get_model_filename_prefix keys on dataset.name, prediction task,
-    level and seed, none of which the window size touches on its own, and
-    trainer.save_checkpoint() overwrites unconditionally.
-    """
-    from interscale.tl.utils import get_model_filename_prefix
-
-    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
-    arms = load_arms(arm_yaml, sweep_config)
-
-    prefixes = {}
-    for arm_name in arms:
-        cfg = arm_cfg.clone()
-        cfg, _ = apply_sweep_config(
-            cfg,
-            "robustness",
-            make_arm_trial(sweep_config, sweep_params, arm_name),
-            model_type="CombinedModel",
-            sweep_params=sweep_params,
-            arms=arms,
-        )
-        prefixes[arm_name] = get_model_filename_prefix(cfg, local_component=True, global_component=True)
-
-    assert len(set(prefixes.values())) == len(arms), f"arms share a checkpoint filename: {prefixes}"
-
-
-def test_max_seq_len_is_never_below_the_arms_largest_window(arm_yaml):
-    """Every arm's max_seq_len must be >= the largest window of the column it selects.
-
-    Below it, pad_batch random-subsamples the window each step and get_model_output stores an
-    attention matrix narrower than the window, which makes the downstream net-flow computation
-    fail with a shape mismatch rather than a wrong number.
-    """
-    # Largest window over ALL splits, measured on melton25_sliding_window.h5ad. Inference runs on
-    # every cell, so the train-split maximum is not the relevant bound.
-    LARGEST_WINDOW = {
-        "sliding_window_400": 89,
-        "sliding_window_800": 330,
-        "sliding_window_1200": 685,
-        "sliding_window_1600": 1218,
-        "sliding_window_2000": 1720,
-        "sliding_window_3000": 3436,
-    }
-    arms = load_arms(arm_yaml)
-    for arm_name, overrides in arms.items():
-        (column,) = overrides["dataset.sample_key"]
-        max_seq_len = overrides["model.global_component.parameters.max_seq_len"]
-        assert column in LARGEST_WINDOW, f"arm {arm_name} selects an unmeasured column {column}"
-        assert max_seq_len >= LARGEST_WINDOW[column], (
-            f"arm {arm_name}: max_seq_len {max_seq_len} < largest {column} window {LARGEST_WINDOW[column]}"
-        )
-
-
-def test_arm_with_a_missing_key_raises():
-    """An arm that omits a key its siblings set would silently keep the base config's value."""
-    yaml_config = {
-        "sweep_config": {
-            "metric": {"name": "val_r2", "goal": "maximize"},
-            "parameters": {ARM_PARAM: {"values": ["a", "b"]}},
-        },
-        "arms": {
-            "a": {"dataset.name": "a", "dataset.batch_size": 8},
-            "b": {"dataset.name": "b"},
-        },
-    }
-    with pytest.raises(ValueError, match="does not declare the same keys"):
-        load_arms(yaml_config, yaml_config["sweep_config"])
-
-
-def test_arm_selecting_an_undefined_arm_raises():
-    yaml_config = {
-        "sweep_config": {
-            "metric": {"name": "val_r2", "goal": "maximize"},
-            "parameters": {ARM_PARAM: {"values": ["a", "typo"]}},
-        },
-        "arms": {"a": {"dataset.name": "a"}},
-    }
-    with pytest.raises(ValueError, match="does not define"):
-        load_arms(yaml_config, yaml_config["sweep_config"])
-
-
-def test_arm_key_colliding_with_a_sweep_parameter_raises():
-    """Set twice per trial, with application order deciding the winner."""
-    yaml_config = {
-        "sweep_config": {
-            "metric": {"name": "val_r2", "goal": "maximize"},
-            "parameters": {ARM_PARAM: {"values": ["a"]}, "dataset.batch_size": {"values": [4, 8]}},
-        },
-        "arms": {"a": {"dataset.batch_size": 16}},
-    }
-    with pytest.raises(ValueError, match="both by the arms and as sweep parameters"):
-        load_arms(yaml_config, yaml_config["sweep_config"])
-
-
-def test_arms_block_without_an_arm_parameter_raises():
-    yaml_config = {
-        "sweep_config": {
-            "metric": {"name": "val_r2", "goal": "maximize"},
-            "parameters": {"optim.seed": {"values": [1]}},
-        },
-        "arms": {"a": {"dataset.name": "a"}},
-    }
-    with pytest.raises(ValueError, match="no 'arm' parameter"):
-        load_arms(yaml_config, yaml_config["sweep_config"])
-
-
-def test_arm_parameter_without_an_arms_block_raises():
-    """Otherwise the arm name would be written to a config key called 'arm', which cannot exist."""
-    yaml_config = {
-        "sweep_config": {
-            "metric": {"name": "val_r2", "goal": "maximize"},
-            "parameters": {ARM_PARAM: {"values": ["a"]}},
-        }
-    }
-    with pytest.raises(ValueError, match="no top-level `arms:` block"):
-        build_sweep_config(yaml_config, prediction_task="regression", model_type="CombinedModel")
-
-
-def test_arm_overrides_must_be_dotted():
-    yaml_config = {
-        "sweep_config": {
-            "metric": {"name": "val_r2", "goal": "maximize"},
-            "parameters": {ARM_PARAM: {"values": ["a"]}},
-        },
-        "arms": {"a": {"batch_size": 8}},
-    }
-    with pytest.raises(ValueError, match="non-dotted keys"):
-        load_arms(yaml_config, yaml_config["sweep_config"])
-
-
-def test_arm_with_an_unknown_dotted_key_raises(arm_cfg):
-    """Arm overrides go through the same _set_dotted validation as sweep parameters."""
-    arms = {"a": {"dataset.no_such_key": 1}}
-    with pytest.raises(KeyError):
-        apply_sweep_config(
-            arm_cfg, "robustness", {ARM_PARAM: "a"}, model_type="CombinedModel", sweep_params=[ARM_PARAM], arms=arms
-        )
-
-
-def test_reserved_arm_parameter_without_arms_passed_to_apply_raises(arm_cfg):
-    """Guards the wiring: forgetting to thread `arms` through must fail, not train the base config."""
-    with pytest.raises(KeyError, match="no arms were passed"):
-        apply_sweep_config(
-            arm_cfg, "robustness", {ARM_PARAM: "w400"}, model_type="CombinedModel", sweep_params=[ARM_PARAM]
-        )
-
-
-def test_arm_trial_does_not_leak_into_the_base_config(arm_cfg, arm_yaml):
-    """wandb.agent reuses one process per agent, so a leaked arm would poison later trials."""
-    sweep_config, sweep_params = build_sweep_config(arm_yaml, prediction_task="regression", model_type="CombinedModel")
-    arms = load_arms(arm_yaml, sweep_config)
-    before = list(arm_cfg.dataset.sample_key)
-
-    clone = arm_cfg.clone()
-    apply_sweep_config(
-        clone,
-        "robustness",
-        make_arm_trial(sweep_config, sweep_params, "w3000"),
-        model_type="CombinedModel",
-        sweep_params=sweep_params,
-        arms=arms,
-    )
-    assert list(arm_cfg.dataset.sample_key) == before
-    assert list(clone.dataset.sample_key) == ["sliding_window_3000"]
-
-
-# --------------------------------------------------------------------------------------------
-# Every arm-bearing sweep yaml in the repo, not just the one this file was written against.
-#
-# Discovered rather than listed: a new ablation adds a yaml and inherits these checks, which is
-# the point. Each is resolved against the (dataset, task) pair its own header names, so the test
-# proves the arms apply to the config they will actually be run with.
-# --------------------------------------------------------------------------------------------
-
-# sweep yaml -> the registered pair it is written for. Kept explicit because the yaml does not
-# name its dataset: --dataset/--task are passed on the command line.
-ARM_SWEEP_PAIRS = {
-    "sliding_window_melton25.yaml": ("melton25_sw", "node_reg"),
-    "overlap_ladder_legnini.yaml": ("legnini23_overlap", "node_reg"),
-    # Masking-granularity ablation. Resolved against the CELL-masking pair: the arms set
-    # mask_strategy/mask_token themselves, so starting from node_reg proves each arm overrides
-    # the baseline rather than relying on the genemask task file to have set it.
-    "mask_granularity_legnini.yaml": ("legnini23", "node_reg"),
-    # Rate/length follow-up. Same base pair, same reason.
-    "mask_rate_and_length_legnini.yaml": ("legnini23", "node_reg"),
-    # Cell-masking rate ladder, the counterpart to the gene ladder in mask_granularity.
-    "mask_rate_cell_ladder_legnini.yaml": ("legnini23", "node_reg"),
-}
-
-
-def arm_sweep_yamls():
-    """Every yaml under config_files/sweeps that declares an `arms:` block."""
-    found = []
-    for path in sorted((CONFIG_DIR / "sweeps").glob("*.yaml")):
-        with path.open() as f:
-            if "arms" in (yaml.safe_load(f) or {}):
-                found.append(path)
-    return found
-
-
-def test_every_arm_sweep_yaml_is_registered_in_this_test():
-    """A new arm sweep must be added to ARM_SWEEP_PAIRS, or it goes untested."""
-    undeclared = [p.name for p in arm_sweep_yamls() if p.name not in ARM_SWEEP_PAIRS]
-    assert not undeclared, (
-        f"arm sweep yaml(s) with no (dataset, task) declared in ARM_SWEEP_PAIRS: {undeclared}. "
-        f"Add them so the checks below cover them."
-    )
-
-
-@pytest.mark.parametrize("yaml_name", sorted(ARM_SWEEP_PAIRS))
-def test_arm_sweep_applies_to_its_registered_pair(yaml_name):
-    """Every arm of every arm sweep resolves and applies against its own dataset/task pair.
-
-    This is the check that would have caught a step of the overlap ladder naming an obs column that
-    the registry's dataset file does not point at, or a max_seq_len key that moved.
-    """
-    dataset, task = ARM_SWEEP_PAIRS[yaml_name]
-    with (CONFIG_DIR / "sweeps" / yaml_name).open() as f:
-        yaml_config = yaml.safe_load(f)
-
-    base = resolve_config(dataset, task, registry_path=REGISTRY)
-    sweep_config, sweep_params = build_sweep_config(
-        yaml_config, prediction_task=base.dataset.prediction_task, model_type="CombinedModel"
-    )
-    arms = load_arms(yaml_config, sweep_config)
-    assert arms, f"{yaml_name} declares no arms"
-
-    for arm_name in sweep_config["parameters"][ARM_PARAM]["values"]:
-        trial = {k: v["values"][0] for k, v in sweep_config["parameters"].items()}
-        trial[ARM_PARAM] = arm_name
-        cfg, applied = apply_sweep_config(
-            base.clone(),
-            "robustness",
-            trial,
-            model_type="CombinedModel",
-            sweep_params=sweep_params,
-            arms=arms,
-        )
-        for key, expected in arms[arm_name].items():
-            assert get_dotted(cfg, key) == expected, f"{yaml_name} {arm_name}: {key} not applied"
-        # A sample_key that is empty would train on nothing; one that is a bare string would be
-        # iterated character by character by prepare_geome_dataset.
-        assert isinstance(cfg.dataset.sample_key, list) and cfg.dataset.sample_key, (
-            f"{yaml_name} {arm_name}: dataset.sample_key must be a non-empty list, got {cfg.dataset.sample_key!r}"
-        )
-
-
-@pytest.mark.parametrize("yaml_name", sorted(ARM_SWEEP_PAIRS))
-def test_arm_sweep_gives_every_trial_its_own_checkpoint(yaml_name):
-    """No two trials of an arm sweep may write the same checkpoint filename.
-
-    get_model_filename_prefix keys on dataset.name, task, level and seed. An arm that forgets to
-    set dataset.name silently overwrites its predecessor's checkpoint, and the sweep then reports
-    metrics for models that no longer exist on disk.
-    """
-    from interscale.tl.utils import get_model_filename_prefix
-
-    dataset, task = ARM_SWEEP_PAIRS[yaml_name]
-    with (CONFIG_DIR / "sweeps" / yaml_name).open() as f:
-        yaml_config = yaml.safe_load(f)
-    base = resolve_config(dataset, task, registry_path=REGISTRY)
-    sweep_config, sweep_params = build_sweep_config(
-        yaml_config, prediction_task=base.dataset.prediction_task, model_type="CombinedModel"
-    )
-    arms = load_arms(yaml_config, sweep_config)
-
-    seeds = sweep_config["parameters"].get("optim.seed", {}).get("values", [None])
-    prefixes = {}
-    for arm_name in sweep_config["parameters"][ARM_PARAM]["values"]:
-        for seed in seeds:
-            trial = {k: v["values"][0] for k, v in sweep_config["parameters"].items()}
-            trial[ARM_PARAM] = arm_name
-            if seed is not None:
-                trial["optim.seed"] = seed
-            cfg, _ = apply_sweep_config(
-                base.clone(),
-                "robustness",
-                trial,
-                model_type="CombinedModel",
-                sweep_params=sweep_params,
-                arms=arms,
-            )
-            prefixes[(arm_name, seed)] = get_model_filename_prefix(cfg, True, True)
-
-    collisions = {v for v in prefixes.values() if list(prefixes.values()).count(v) > 1}
-    assert not collisions, f"{yaml_name}: trials sharing a checkpoint filename: {collisions}"

From 77e36dac7717945bbda2f66e05321d8cde5d35d4 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Tue, 8 Sep 2026 13:35:13 +0200
Subject: [PATCH 13/14] solve merge conflicts

---
 docs/api/tools.md                             | 16 ++++++++++++++++
 src/interscale/evaluation/_latent_analysis.py | 12 ++++++------
 src/interscale/main_sweep.py                  |  2 +-
 src/interscale/pl/anchor_plots.py             |  4 ++--
 4 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/docs/api/tools.md b/docs/api/tools.md
index 24aae29..fcbaac2 100644
--- a/docs/api/tools.md
+++ b/docs/api/tools.md
@@ -46,6 +46,21 @@ Downstream InterScale's output can be used for gene, cell and tissue level analy
     calculate_gene_ranks
 ```
 
+## Latent dimensions
+
+Which dimensions of a component's embedding carry the signal, and so which are worth
+following up with the anchor tools below.
+
+```{eval-rst}
+.. currentmodule:: interscale.evaluation
+
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    calculate_dim_importance
+```
+
 ## Anchors and distance zones
 
 Where a latent dimension anchors on a slide, and how expression behaves at range from there.
@@ -61,6 +76,7 @@ have come from it.
     :toctree: generated
 
     local_reach_um
+    get_average_local_and_global_size
     find_anchor_cells
     anchor_signed_distance
     anchor_zones
diff --git a/src/interscale/evaluation/_latent_analysis.py b/src/interscale/evaluation/_latent_analysis.py
index c6fda20..a2fe859 100644
--- a/src/interscale/evaluation/_latent_analysis.py
+++ b/src/interscale/evaluation/_latent_analysis.py
@@ -474,12 +474,12 @@ def calculate_dim_importance(
     ----------
     adata : AnnData
         Annotated data matrix.
-    which : {"global", "local"}, optional
-        Which component to score. Derives `s_key` and `z_key`, so it is normally the
-        only argument needed. Defaults to "global" when `s_key` is not given; when
-        `s_key` is given instead, the component is inferred from it. Passing a `which`
-        that contradicts `s_key` is an error.
-    prefix : str, default ""
+    which : str, optional
+        Which component to score, "global" or "local". Derives `s_key` and `z_key`, so it
+        is normally the only argument needed. Defaults to "global" when `s_key` is not
+        given; when `s_key` is given instead, the component is inferred from it. Passing a
+        `which` that contradicts `s_key` is an error.
+    prefix : str
         Prefix of the stored keys, following the f"{prefix}_{which}_..." convention used
         when the model output was saved. The default "" gives the plain "_global_emb" /
         "_local_emb" keys; pass e.g. prefix="seed0" to score a run whose embeddings were
diff --git a/src/interscale/main_sweep.py b/src/interscale/main_sweep.py
index f5606e2..6123fe8 100644
--- a/src/interscale/main_sweep.py
+++ b/src/interscale/main_sweep.py
@@ -235,7 +235,7 @@ def main_sweep(cfg_factory, model_type, sweep_goal, sweep_params=None, arms=None
     trial_error = None
     try:
         model.train(max_epochs=cfg.optim.n_epochs, datamodule=dm, early_stopping=cfg.optim.early_stopping)
-    except Exception as exc:
+    except Exception as exc:  # noqa: BLE001 - a trial must not leak its GPU memory whatever it died of
         # format_exc() renders the stack to a STRING, so the full traceback survives in the log
         # while no frame (and so no tensor) stays referenced. Keeping only str(exc) made the
         # first CosMx OOM undiagnosable: the allocation turned out to be in the metric
diff --git a/src/interscale/pl/anchor_plots.py b/src/interscale/pl/anchor_plots.py
index 4dd576c..d0b1185 100644
--- a/src/interscale/pl/anchor_plots.py
+++ b/src/interscale/pl/anchor_plots.py
@@ -310,10 +310,10 @@ def zone_map(
 ):
     """``anchor core`` / ``near`` / ``far`` bands, one panel per sample.
 
-    Zones are drawn in :data:`interscale.tl.anchors.ZONE_ORDER`, so ``core`` ends up on top
+    Zones are drawn in ``interscale.tl.anchors.ZONE_ORDER``, so ``core`` ends up on top
     of ``far`` and the legend order is the distance order regardless of how many cells each
     band holds. Cells with an undefined zone -- a sample where the dimension found no focus --
-    are drawn in :data:`NA_COLOR`.
+    are drawn in ``NA_COLOR``.
 
     Returns
     -------

From 7089824ffb4812a28b575b5237353cf6404b8119 Mon Sep 17 00:00:00 2001
From: FrancescaDr <francesca.drummer@googlemail.com>
Date: Tue, 8 Sep 2026 20:10:50 +0200
Subject: [PATCH 14/14] pin scanpy<1.13; reason: plotting import error

---
 pyproject.toml | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 621eb79..024f1da 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,11 @@ dependencies = [
   "lightning>=2",
   "numpy>=1.19",
   "pandas",
-  "scanpy",
+  # Upper bound is squidpy's, not ours: squidpy/_compat.py imports the private
+  # scanpy.plotting._tools, which the 1.13 plotting rewrite removed, so `import squidpy`
+  # (and through it `import interscale`) fails on scanpy >= 1.13.0a2. Lift this once
+  # squidpy releases a version that no longer reaches into that module.
+  "scanpy<1.13",
   "scikit-learn",
   "scipy>=1.7",
   "scvi-tools>=0.19",