diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 8b85530..b24278d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -8,7 +8,7 @@ # Workflow: release.yml Environment: pypi # (A *pending* publisher lets the very first release create the project on PyPI.) # 2. Cut a release whose tag matches the version in pyproject.toml, e.g.: -# gh release create v0.1.1 --generate-notes +# gh release create v0.2.0 --generate-notes # This workflow then builds the sdist + wheel and uploads them automatically. name: Release diff --git a/CHANGELOG.md b/CHANGELOG.md index 76e240b..34c91b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,11 +3,11 @@ Notable changes to raw2features, newest first. This project follows [Semantic Versioning](https://semver.org). -## [0.1.1] - 2026-07-16 +## [0.2.0] - Unreleased -This is a focused correctness, reproducibility, and interoperability release. It keeps -the v0.1 store layout and geometry identity while making resume decisions stricter and -multi-grid workflows fully additive. +This release extends named-channel multiplex support with a `channelwise` strategy, adds +four patch encoders, and includes a set of correctness and release-safety fixes. The v0.1 +store layout and existing brightfield workflows remain unchanged. ### Correctness and reproducibility @@ -33,11 +33,21 @@ multi-grid workflows fully additive. ### Models and interoperability +- Add a pluggable multiplex-strategy layer and a built-in `channelwise` strategy for + applying registered RGB patch encoders to named-channel multiplex tissue images. It + records positional channel identity (including a strict metadata-file fallback), marker + selection and order, normalization, RGB conversion, base-model identity, mean/concat + aggregation, and model-agnostic slide pooling without changing brightfield outputs. +- Bind native multiplex outputs to the complete effective positional panel. Nuclear + masking accepts one unambiguous nuclear channel or a recognized same-stain group, such + as DNA1/DNA2 or repeated DAPI channels. It averages each multi-channel group in float32 + and binds every physical index into grid identity; `--no-seg` remains supported. - Add pinned, forward-validated H0-mini, KEEP, OpenMidnight, and OpenPath patch encoders, with their licenses, access requirements, preprocessing, and weight hashes recorded. - Correct cached OME-Zarr reads for noncanonical spatial axis order. The reader accepts x/y in either order (with an optional channel axis) and warns whenever another - non-singleton axis is reduced to index zero. + non-singleton axis is reduced to index zero. Unnamed `omero.channels` entries now + retain their physical positions instead of shifting later marker identities. - Refresh the development lock to current `ngff-zarr`/`ome-zarr` releases while retaining the existing declared compatibility floor. - Choose the float-pixel `[0,1]` versus `[0,255]` convention once per opened slide instead @@ -56,11 +66,15 @@ multi-grid workflows fully additive. ### Compatibility and migration -- Existing v0.1 stores remain readable; the embeddings schema and geometry-only - `grid_hash` are unchanged. +- Existing v0.1 stores remain readable; the store layout/schema and model-independent + grid-identity contract are unchanged. Multiplex nuclear masking now binds the resolved + physical nuclear-channel indices. Averaging a recognized same-stain group is an + intentional value change from v0.1 and therefore creates a new grid. Compatible + historical native-multiplex grids remain appendable only when the original OME panel + proves that the v0.1 selector used the same single physical channel. - The first request for a model output created before fingerprints were introduced will recompute that model in place. The old header is deliberately not treated as proof that - pre-v0.1.1 revision pins were enforced. + pre-v0.2.0 revision pins were enforced. - Plain local output IDs remain basename-based for v0.1 compatibility. Source comparison prevents accidental reuse, and manifest preflight rejects duplicate derived IDs. - Credentials already persisted by v0.1.0 are not rewritten by this release and should be @@ -102,4 +116,4 @@ What 0.1.0 provides: harness, and six runnable tutorial notebooks. [0.1.0]: https://github.com/CraigMyles/raw2features/releases/tag/v0.1.0 -[0.1.1]: https://github.com/CraigMyles/raw2features/compare/v0.1.0...v0.1.1 +[0.2.0]: https://github.com/CraigMyles/raw2features/compare/v0.1.0...v0.2.0 diff --git a/CITATION.cff b/CITATION.cff index 7e5ec8e..8571b07 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -20,4 +20,4 @@ keywords: - foundation models - feature extraction license: MIT -version: "0.1.1" +version: "0.2.0" diff --git a/README.md b/README.md index 08fb023..b53053f 100644 --- a/README.md +++ b/README.md @@ -43,15 +43,15 @@ store too. level such that embeddings are comparable across slides and datasets. - **Modular implementation.** Reader, segmenter, patcher, patch-embedder, - slide-embedder, and sink implementations are plugin seams exposed through Python - entry points. Entry points add implementations/families; the CLI's patch-model names - come from the bundled provenance registry, while Python callers can inject external - embedder instances through `embed_slide(..., embedders=[...])`. + multiplex-strategy, slide-embedder, and sink implementations are plugin seams exposed + through Python entry points. Entry points add implementations/families; the CLI's + patch-model names come from the bundled provenance registry, while Python callers can + inject external embedder instances through `embed_slide(..., embedders=[...])`. - **FAIR & provenance-first.** Stable models' weights are pinned to an **immutable HuggingFace revision** (or a sha256-pinned URL), with preprocessing sourced from each model's card. Every output records that provenance plus a 1:1 coords↔features mapping, so an embedding is reproducible and traceable to the exact - weights that made it. SEAL is explicitly experimental in v0.1.1: its adapter is + weights that made it. SEAL is explicitly experimental in v0.2.0: its adapter is pinned and verified, but its upstream factory still fetches the frozen base from a mutable revision; see [`MODELS.md`](docs/MODELS.md). @@ -114,6 +114,8 @@ on CPU with no model-access-token. Its figures are pre-rendered on GitHub. **Full guide: [`docs/usage.md`](docs/usage.md)** - every command, what actually happens under the hood (exact MPP, decode-once fan-out, output schema), the rerun-safe / skip-if-complete model, thumbnails, and example SLURM cohort runs. +Named-channel inputs, native multiplex encoders, and the `channelwise` RGB-encoder +strategy are covered in [`docs/MODALITIES.md`](docs/MODALITIES.md). ```bash raw2features info slide.ome.zarr diff --git a/docs/MODALITIES.md b/docs/MODALITIES.md index 0c6ecd8..3e13f1c 100644 --- a/docs/MODALITIES.md +++ b/docs/MODALITIES.md @@ -1,87 +1,94 @@ -# Modalities - brightfield (H&E) + multiplex (spatial proteomics) +# Brightfield and named-channel multiplex tissue images -raw2features is **H&E-first**, but the same pipeline also handles a second modality: -**multiplex** imaging (CODEX / Phenocycler / spatial proteomics), via marker-aware -encoders like KRONOS. H&E runs are completely unchanged. +raw2features supports RGB brightfield images and named-channel multiplex tissue images. +Multiplex sources use native `(H, W, C)` reads and preserve positional channel identity. +Registered RGB encoders can be applied through a multiplex strategy, while native +multiplex encoders consume the channel stack directly. Ordinary brightfield execution is +unchanged. -## How a model declares its modality -Each registry entry carries `modality: brightfield` (default) or `modality: multiplex`. -The runner reads the modality off the requested embedders and routes the per-slide path: +## Execution paths -| stage | brightfield (H&E) | multiplex | -|---|---|---| -| reader | `read_region` → RGB `[H,W,3]` uint8 | `read_region_channels` → native `[H,W,C]` + `channel_names` | -| segmenter | `otsu` (saturation), `canny`, … | `nuclear` - Otsu on the DAPI/Hoechst channel | -| embedder | RGB transform → `[B,3,H,W]` | per-marker norm → `[B,M,H,W]` + marker ids | +Each registry model declares `modality: brightfield` (the default) or +`modality: multiplex`. A strategy wraps a brightfield model without changing its registry +entry. -## Multiplex data → OME-Zarr -Multiplex slides are usually OME-TIFF / TIFF stacks. Convert to a channel-named OME-Zarr: +| stage | brightfield | `channelwise` RGB strategy | native multiplex encoder | +|---|---|---|---| +| reader | `read_region` → RGB `[H,W,3]` uint8 | `read_region_channels` → native `[H,W,C]` | `read_region_channels` → native `[H,W,C]` | +| channel identity | none | selected positional names | positional names resolved by the model | +| segmentation | `otsu`, `canny`, … | one nuclear channel, a recognized same-stain group, or `--no-seg` | one nuclear channel, a recognized same-stain group, or `--no-seg` | +| embedding | one RGB input per patch | one RGB input per selected marker, then mean/concat | model-specific marker stack | + +Panel binding happens before receipt or store completion checks. Native multiplex +fingerprints bind the complete effective positional panel; `channelwise` fingerprints +bind the selected physical channel identities and order (all channels when default +selection is used). The full effective panel remains in source/panel provenance. When +nuclear masking is enabled, the resolved physical nuclear-channel index or same-stain +index group is also part of grid identity. + +## Channel metadata + +The preferred source of channel identity is `omero.channels` in the OME-Zarr metadata. +The reader preserves unnamed entries as empty positional slots rather than shifting later +labels onto the wrong pixels. + +For a source whose labels are absent or incomplete, pass `--channel-names-file` with a +complete ordered panel. UTF-8 `.txt`, `.csv`, and `.tsv` files are accepted, with one +unique name per physical C-axis position. Existing non-empty OME labels are treated as +assertions and must agree with the supplied name at the same index. The override is +in-memory only and never rewrites the source. Identity uses only the resolved names. +Provenance also records whether effective names came from OME metadata or the supplied +file and preserves any differing original OME labels. See +[usage.md](usage.md#rgb-encoders-on-named-channel-multiplex-tissue-images) for examples +and marker-selection rules. + +## Converting a multiplex TIFF + +Multiplex images are often supplied as OME-TIFF or TIFF channel stacks. The included +conversion helper accepts a `(C,Y,X)` TIFF and writes a multiscale OME-Zarr with one +`omero.channels[].label` per channel: ```bash -pip install "raw2features[kronos]" -python scripts/codex_to_omezarr.py SLIDE.tif SLIDE.ome.zarr --markers markers.json --mpp 0.5 +pip install "raw2features[zarr]" tifffile imagecodecs +python scripts/codex_to_omezarr.py SLIDE.tif SLIDE.ome.zarr \ + --markers markers.json --mpp 0.5 ``` -`markers.json` is a list (or `{"raw_markers": [...]}`) of one marker name per channel. The -names are written to `omero.channels[].label`; the reader surfaces them as `channel_names`. - -## KRONOS - the first multiplex model -`kronos` (MahmoodLab **KRONOSv1**) is a panel-agnostic spatial-proteomics encoder. Needs -the `[kronos]` extra (the public `kronos` package). ViT-S/16 → **384-d** patch embedding. - -- **Marker matching - your panel → KRONOS's vocabulary.** KRONOS doesn't ship a - name-resolver, only a vocabulary: the **public** `marker_metadata.csv` (175 markers, - each with a marker id + per-marker mean/std). Mapping *your* channels to it is the - pipeline's job (`Embedder.set_panel`), and KRONOS is panel-agnostic - it embeds whatever - subset of its vocabulary your slide carries. Each `channel_name` is matched: exact name → - common **synonyms** (`pancytokeratin`→`CYTOKERATIN`, `granzymeb`→`GZMB`) → a **CD-number** - in a compound name (`cla_cd162`→`CD162`); `hoechst*`/`dapi` → DAPI; blank/empty cycles - drop silently. -- **The mapping is recorded in the store.** The header's - `panel.kronos` block records the full per-channel resolution under `mapping` - each kept - channel as `{channel, channel_index, kronos_marker, marker_id}` (e.g. - `pancytokeratin → CYTOKERATIN (id 322)`) - plus `kept` / `dropped` / `unmatched` and a - `vocabulary` pointer pinned to the exact `marker_metadata.csv` revision. The SpatialData - export (`INTEROP.md`) surfaces this same map as a tidy, queryable `uns["raw2features_panel"]` - table. -- **Unmatched named channels are surfaced.** A *named* marker that - doesn't resolve (genuinely outside KRONOS's vocabulary, or a synonym we don't yet know) - is dropped, **warned** about (with the `matched/total` coverage), and listed in - `panel.kronos.unmatched`. (Zero matches is a hard error; partial panels are fine.) If an - unmatched marker *is* a KRONOS marker under another name, add it to `_resolve_marker`'s - synonym map. -- Run it like any model: - - ```bash - raw2features embed SLIDE.ome.zarr OUT -m kronos --mpp 0.5 - ``` - - The runner detects the multiplex modality and routes the `nuclear` segmenter + native - N-channel reads automatically. - -**Performance.** A multiplex patch is an M-marker stack (~M× the tokens of a single-channel -ViT). The `kronos` family uses **torch SDPA (flash attention) by default** - no xFormers, -fully portable - which on an A100 (M=41) runs at **~41 patches/s in ~1.3 GB at batch 8**, -vs ~20 p/s / 26 GB with KRONOS's upstream naive-attention fallback (it materialises the full -8k×8k attention matrix). Memory is low and predictable, so the default `--batch-size` -fits comfortably; lower it only on small GPUs or very large panels. (Opt out via the -registry `sdpa: false`.) Validated full-slide on GPU against real CODEX multiplex data - a -56-channel panel (44 markers matched, the rest empty/blank cycles) → 28 tissue patches → -`(28, 384)`, all finite. - -> KRONOS weights are **CC-BY-NC-ND** (non-commercial), like UNI / Virchow2 / CONCH - so its -> embeddings are non-commercial too. See `MODEL_LICENSES.md`. - -### Marker intensity scaling -The transform brings marker intensities to `[0,1]` **by the source dtype** (uint16 → `/65535`, -uint8 → `/255`, float assumed already normalised), then standardises each marker with its -`marker_mean`/`marker_std`. For the uint16 CODEX/Phenocycler case this matches KRONOS's own -feature-extraction pipeline exactly - its loader divides by `max_value` (default `65535`) then -applies the same per-marker `(x − mean) / std`, so the published marker statistics are defined -on the same `[0,1]` scale. Scaling by the dtype range (rather than a hard-coded `65535`) keeps -uint8 / pre-normalised sources correct instead of silently squashing them. - -### Adding more multiplex models -Future marker-conditioned models (e.g. KRONOS successors) are added the same way as any -encoder: a registry entry with `modality: multiplex` + a small family class (or reuse the -`kronos` family). The reader / segmenter / runner routing is shared. +`markers.json` is a list, or an object containing `raw_markers`, `channels`, or `markers`. +The converter requires one name per channel. Prefer a domain converter that preserves all +available OME metadata when one is available; this helper is intentionally small. + +## RGB encoders through `channelwise` + +`channelwise` normalizes selected channels independently, repeats each single-channel +patch across RGB, runs the registered RGB encoder and pooling, then combines the marker +vectors with mean or ordered concatenation. It records the effective panel, normalization +level and values, RGB conversion, base-model output fingerprint, pooling, aggregation, +and output dimension. Model-agnostic slide poolers can consume the result. + +This makes it possible to compare models such as UNI or Virchow2 on named-channel data +without representing them as native multiplex models. The strategy's assumptions and +full CLI contract are documented in [usage.md](usage.md#rgb-encoders-on-named-channel-multiplex-tissue-images). + +## Native multiplex encoders + +Native multiplex models use the same positional-panel plumbing but consume the marker +stack directly. The current built-in example, `kronos`, maps effective source names to +its pinned marker vocabulary, records the kept/dropped physical mapping, and binds the +complete effective panel into its output fingerprint before resume. Installation and +model-specific license/access details are in [MODELS.md](MODELS.md) and +[MODEL_LICENSES.md](MODEL_LICENSES.md). + +```bash +raw2features embed SLIDE.ome.zarr OUT -m kronos --mpp 0.5 +``` + +Use `--channel-names-file` when the source lacks a complete panel, or `--no-seg` when the +entire image should be tiled without a nuclear mask. + +## Extending multiplex support + +The `raw2features.multiplex_strategies` entry point separates panel/config preparation +from slide-specific binding. Future marker-to-RGB mappings or learned channel adapters +can use this contract if there is demand. A new native multiplex model instead declares +`modality: multiplex` and implements the standard embedder panel-binding seam. diff --git a/docs/MODELS.md b/docs/MODELS.md index 20c6205..94e5424 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -59,7 +59,7 @@ read and decide on. raw2features makes no commercial-use determination. | `seal_univ2` ⚠ | seal | 1536 | 224 | 0.5 | conv. | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/SEAL) · [GH](https://github.com/mahmoodlab/SEAL) · [paper](https://arxiv.org/abs/2602.14177) | | `kronos` | kronos | 384 | 224 | -¹² | n/a | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/KRONOS) · [GH](https://github.com/mahmoodlab/KRONOS) · [paper](https://arxiv.org/abs/2506.03373) | -⚠ **SEAL is experimental in v0.1.1 and is outside the exact-weight pinning +⚠ **SEAL is experimental in v0.2.0 and is outside the exact-weight pinning guarantee.** raw2features pins and SHA-256 verifies the SEAL LoRA adapter, freezes the adapter constructor, and records its composite identity. The pinned upstream factory still downloads the frozen CONCH/UNI2-h base from mutable HEAD, however, so diff --git a/docs/SPEC.md b/docs/SPEC.md index 195852c..2ade9ce 100644 --- a/docs/SPEC.md +++ b/docs/SPEC.md @@ -16,9 +16,10 @@ The keywords MUST, SHOULD and MAY are used as in RFC 2119. ## Stability -raw2features is alpha (`0.x`); this format ships with it on **one version line** - package -`0.x`, format `0.x`, schema `0.x` all move together (so today's store says -`schema_version: "0.1"`). To let readers depend on the format now during the 0.x line, the stable parts are scoped explicitly: +raw2features is alpha (`0.x`), but the package release and store schema are versioned +independently. Package v0.2.0 continues to write `schema_version: "0.1"` because this +release does not change the established read contract or layout. To let readers depend on +the format now, the stable parts are scoped explicitly: - **Stable (the read contract).** Patch sets live under `grids//`, each a complete, self-describing grid carrying its header under that group's `raw2features` attr; **within @@ -26,7 +27,7 @@ raw2features is alpha (`0.x`); this format ships with it on **one version line** `coords` are level-0 `(x, y)` pixels carrying `units = "level0_px"`; and a reader **MUST tolerate keys/arrays it does not recognise**. Code written against these will keep working across `0.x`. -- **Still revisable during `0.x`.** The exact on-disk array layout, the set of header +- **Still revisable during schema `0.x`.** The exact on-disk array layout, the set of header keys, and the coordinate *convention* (e.g. if/when the store grows a richer, physical coordinate frame - see [Coordinates](#coordinates--relocatability)) may change between `0.x` releases. Each such change bumps `schema_version` and ships an updated JSON Schema @@ -98,7 +99,7 @@ Under each `grids//`: |---|---|---|---| | `grid_index` | (N, 2) | int32 | `(row, col)` of each patch in the regular grid. | | `mask` | (rows, cols) | uint8 | Fraction of each grid cell that is tissue (the mean of the segmentation mask over the cell), stored as `0-255` - divide by 255 for the `[0,1]` value. Index a kept patch's cell via `grid_index[i]`. Absent when tiling without segmentation. | -| `slide/` | (1, dim) | float32 | A slide-level vector (e.g. from a slide encoder or pooling), derived from this grid's features. | +| `slide/` | (1, dim) | float32 | A slide-level vector derived from this grid's features. Ordinary outputs use the encoder name; a model-agnostic pool over a strategy-derived multiplex array uses `__` so several recipes can coexist. | ## Optional QC layer (per grid) @@ -152,7 +153,14 @@ A JSON object; the normative definition is the packaged JSON Schema (see anisotropy), and `level0_translation_um` (the source's level-0 translation/origin in µm, or `null` when it carries none). These are recorded **values**, not an NGFF transform object; they let a consumer re-express `coords` in the source's physical - frame (see [Coordinates](#coordinates--relocatability)). + frame (see [Coordinates](#coordinates--relocatability)). For a named-channel + multiplex source, optional `channel_names` records the complete effective positional + C-axis mapping and `channel_names_source` says whether it came from + `ome_zarr_omero_channels` or a user-supplied channel-names file. When that file fills + incomplete metadata, `omero_channel_names` preserves the original positional labels, + including blank entries. The structured source and channel-name fields do not include + the file path. As with other non-secret command arguments, `provenance.cli` retains the + invocation and can therefore include that path. - `patching` - **this grid's** geometry: `target_mpp`, `achieved_mpp`, `patch_px`, `level0_patch` (patch side in level-0 px), `level0_step`, `read_level`, `step_out_px`, `n_patches`, `grid_shape` (`[rows, cols]`), `coords_convention` (`"level0_xy"`). @@ -175,19 +183,38 @@ A JSON object; the normative definition is the packaged JSON Schema (see loader-contract version. Composite models include every component (SEAL records its pinned adapter and its experimental, upstream-unpinned base separately). The same record is written on `features/` only after every feature row succeeds; that - array attribute is the completion commit marker. -- `grid_hash` - a hash of the patch geometry only. Two grids with the same `grid_hash` - share an identical grid, so feature arrays from different runs are row-comparable. - Model identity deliberately remains separate: a changed output fingerprint replaces - only that model array without renaming or rebuilding the grid. + array attribute is the completion commit marker. A strategy-derived multiplex output + additionally records its strategy name/version and complete resolved marker, + normalization, input-conversion, pooling, and aggregation recipe. Its fingerprint + embeds the complete base-encoder output fingerprint, binding the derived output + transitively to the exact base weights and preprocessing. The readable model header + mirrors this under `base_model` and `multiplex`. A native multiplex model instead adds + a `multiplex_panel` binding containing its contract version, physical C axis/count, and + complete effective positional channel list. +- `grid_hash` - a hash of patch-grid construction only: geometry plus any tissue-mask + inputs that can change the retained coordinates. For multiplex nuclear segmentation, + that includes the resolved physical nuclear-channel index or same-stain index group and + combination rule, while unrelated marker names remain model-output identity. Two grids + with the same `grid_hash` share an identical grid, so feature arrays from different + runs are row-comparable. Model identity deliberately remains separate: a changed output + fingerprint replaces only that model array without renaming or rebuilding the grid. - `provenance` - `raw2features_version`, `created_utc`, `cli`, `git_sha`, `host`, `arch`, `platform`, `python`, and `gpu` when applicable. Secret option values and credentials embedded in URIs are redacted from `cli`. Optional grid keys: `segmentation` (segmenter name + parameters), `thumbnail` (thumbnail -metadata when one was written), `slide_embeddings` (provenance for any `slide/`). - -Each `slide/` array carries an analogous `output_fingerprint`, mirrored in +metadata when one was written), `slide_embeddings` (provenance for any +`slide/`), and `panel`. `panel` is keyed by effective model name and is +model/strategy-specific. `channelwise` records the physical count, complete positional +source list, selected order, exclusions, and resolved normalization. Native encoders may +instead record their model-specific kept/dropped mapping and vocabulary audit; their +output fingerprint still binds the complete effective source panel. Readers retain blank +metadata entries at their physical positions. A strategy that selects every channel +requires a unique effective name at every position; an explicit marker selection may +exclude unnamed positions, and a validated positional channel-names file can fill absent +or incomplete labels. + +Each `slide/` array carries an analogous `output_fingerprint`, mirrored in `slide_embeddings`. It includes the selected patch array's fingerprint, so replacing patch features invalidates any derived slide vector, and records the effective fp16/fp32 slide-forward precision for the resolved device. Fingerprint fields are optional in @@ -235,11 +262,12 @@ source is remote or read-only. ## Versioning -`schema_version` is `MAJOR.MINOR` and tracks the 0.x release line (one version line - -package, format and schema move together). Adding optional arrays or header keys is a -MINOR bump; removing or repurposing anything is a MAJOR bump. Each version ships its -JSON Schema as `embeddings_store-.schema.json`. Readers MUST check -`schema_version`, MUST ignore unknown keys/arrays, and SHOULD warn on an unknown MAJOR. +`schema_version` is `MAJOR.MINOR` and is independent of the package version. It changes +when the normative store contract changes, not for every package release or newly emitted +optional extension key. Backward-compatible schema changes increment MINOR; removing or +repurposing part of the read contract increments MAJOR. Each schema version ships its JSON +Schema as `embeddings_store-.schema.json`. Readers MUST check `schema_version`, +MUST ignore unknown keys/arrays, and SHOULD warn on an unknown MAJOR. ## Conformance diff --git a/docs/usage.md b/docs/usage.md index f7a89f8..84d0775 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -41,11 +41,12 @@ OME-Zarr ──reader──▶ segmenter ──▶ patcher ──▶ embedde | **segmenter** | `otsu` (or `--no-seg`) | low-res tissue mask so only tissue patches are kept | | **patcher** | `grid` | tiles a regular grid at an **exact** target MPP | | **embedder** | `resnet50` (repeatable) | patch → feature vector; one or many, run together | +| **multiplex_strategy** | none | optional named-channel adaptation before an RGB patch encoder | | **slide_embedder** | none | patch-feature matrix → optional slide-level vector | | **sink** | `zarr` | writes `coords`/`grid_index`/`mask`/`features/` + provenance | List what's installed for any component: `raw2features list embedders` (or `readers`, -`segmenters`, `patchers`, `sinks`, `slide_embedders`). +`segmenters`, `patchers`, `sinks`, `slide_embedders`, `multiplex_strategies`). --- @@ -76,7 +77,7 @@ on another (e.g. an x86 cluster). Verify with `raw2features version`. | `raw2features export-h5 ` | export features to HDF5 (one `.h5` per model); needs the `[h5]` extra | | `raw2features validate-store ` | check a store conforms to the embeddings-store spec | | `raw2features models` | list the model registry (name, dim, gated, source) | -| `raw2features list ` | list installed plugins for a component type (readers, segmenters, patchers, embedders, sinks, slide_embedders) | +| `raw2features list ` | list installed plugins for a component type (readers, segmenters, patchers, embedders, sinks, slide_embedders, multiplex_strategies) | | `raw2features version` | print the version | Every option below is also visible via `raw2features --help`. @@ -150,7 +151,7 @@ Any token passed on the command line is redacted from stored command provenance. | `--qc` | none | per-patch QC producer(s) writing `grids//qc//` (e.g. `grandqc`); needs the producer's extra (`raw2features[grandqc]`) | | `--qc-stain-norm` | off | normalize the QC input first (`macenko`\|`reinhard`\|`vahadane`; `vahadane` needs `raw2features[stain]`) - for a stain outside the QC model's domain | | `--device` | `auto` | `auto` (best of cuda→mps→cpu) \| `cuda` \| `mps` \| `cpu` | -| `--devices` | = `--device` | opt-in in-process multi-GPU, e.g. `cuda:0,cuda:1`; shards this slide's patches across them and gathers in coord order (output identical to one device). See _In-process multi-GPU_ in §6. | +| `--devices` | = `--device` | opt-in in-process multi-GPU, e.g. `cuda:0,cuda:1`; shards this slide's patches across them and gathers in coord order (output identical to one device). Multi-device execution is not supported with `--multiplex-strategy`; see _In-process multi-GPU_ in §6. | | `--batch-size` | `256` | patches per forward pass (lower it on small GPUs) | | `--amp` | `auto` | precision: `auto` (each model's card precision) \| `fp32` \| `bf16` \| `fp16` | | `--snap-to-level` | off | read a pyramid level natively (no resample); MPP = that level | @@ -225,6 +226,100 @@ byte-like convention, otherwise the slide is treated as normalized. This bounded avoids scanning an entire WSI at open. `read_region_channels` remains native-dtype and unscaled for multiplex models. +### RGB encoders on named-channel multiplex tissue images + +`--multiplex-strategy channelwise` lets a registered RGB patch encoder process an +OME-Zarr image with a `c` axis and positionally matching channel names: + +```bash +raw2features embed multiplex_slide.ome.zarr multiplex_out \ + -m uni --multiplex-strategy channelwise \ + --marker CD3 --marker CK --marker DAPI \ + --multiplex-aggregation mean -s mean +``` + +Each selected channel is processed independently. raw2features resolves a normalization +for that channel, repeats the normalized single-channel patch across RGB, and applies the +encoder's registered preprocessing and pooling (for example CLS for UNI or +CLS+mean-patch for Virchow2). The marker vectors are combined with `mean` (default) or +ordered `concat`; `concat` requires an explicit marker order. Without an explicit +selection, every physical channel must have a unique name and all channels are used in +source order. If metadata contains unnamed positions, supply a complete +`--channel-names-file` or explicitly select the named markers to exclude those positions. +Selected names must resolve unambiguously. No channel name has implicit inclusion or +exclusion semantics for embedding. Nuclear marker names are resolved separately when +segmentation is enabled. + +`--multiplex-normalization percentile` is the default. It records each selected channel's +1st and 99th percentiles from the first image-pyramid level whose longest side is at most +2048 pixels. Set alternative bounds with `--multiplex-percentile-low` and +`--multiplex-percentile-high`. This is a bounded-memory, reproducible raw2features rule +rather than a platform-specific calibration. The level limit can be changed with +`--multiplex-normalization-max-side-px`; a larger value may choose a finer level and use +more memory. If no pyramid level fits the cap, raw2features asks for a coarser level or a +higher cap instead of silently scanning level 0. `--multiplex-normalization dtype` +instead maps the source dtype range to `[0,1]`; floating-point sources must already be +finite in `[0,1]`. Percentiles describe the selected pyramid level, so spatial +downsampling can smooth extrema relative to level 0; the exact level and measured values +are fingerprinted. The selected markers and physical indices, order, +normalization values and source level, RGB conversion, base-model fingerprint and +pooling, aggregation, and output dimension are stored in the grid's model/panel metadata +and output fingerprint. + +The RGB vision-encoder baselines evaluated in the +[public KRONOS paper](https://arxiv.org/html/2506.03373) provide structural inspiration. +This strategy adapts ordinary RGB encoders rather than implementing the marker-aware +KRONOS architecture. Every marker is normalized independently, repeated to RGB, embedded +independently, then aggregated with a fully recorded deterministic recipe. + +Channel identity normally comes from `omero.channels`. If those labels are absent or +incomplete, provide a complete positional panel without rewriting the source: + +```bash +raw2features embed multiplex_slide.ome.zarr multiplex_out \ + -m uni --multiplex-strategy channelwise \ + --channel-names-file channel_names.tsv \ + --marker CD3 --marker CK --marker DAPI +``` + +The UTF-8 `.txt`, `.csv`, or `.tsv` file contains exactly one unique name for every +physical C-axis position, in order. A conventional one-column header is optional. Any +existing non-empty OME labels must agree; conflicts and channel-count mismatches fail +before model loading or output mutation. The file defines physical channel identity, +while repeated `--marker` options select and order a subset. Only the resolved ordered +panel and ordered marker request affect request identity; the file path and serialization +do not. Store provenance also records whether effective labels came from OME metadata or +a channel-names file and, when they differ, preserves the original OME labels. As with +other non-secret arguments, the recorded CLI invocation can include the file path. One +file applies to every slide in an `embed-many` invocation, so heterogeneous panels should +use separate runs. + +Multiplex runs use nuclear-channel tissue masking when segmentation is enabled. They +require one unambiguous nuclear channel or a recognized same-stain group, such as +DNA1/DNA2 or repeated DAPI channels. A multi-channel group is averaged in float32 before +thresholding. This differs from v0.1's first-match behaviour and is bound into a new grid +identity rather than reusing the old mask. +`--no-seg` is fully supported and tiles the complete image, which is useful for +segmentation-free analyses or panels without a suitable nuclear channel; it can also +retain more background and require more compute. + +Patch geometry remains configurable. If an extracted patch does not match the encoder's +registered input size, raw2features resizes it before model preprocessing. Smaller, +cell-scale fields are therefore mechanically supported, but resizing does not add image +detail and the changed physical field of view should be validated for the intended task. + +Strategy-derived outputs use a cohort-stable effective key under `features/`. +Model-agnostic slide poolers (`mean`, `max`, `meanmax`) accept them, while learned slide +encoders tied to another patch model are rejected. The first implementation processes one +device per slide. + +The strategy is a plugin (`raw2features.multiplex_strategies`), with separate +panel/config preparation and slide-specific binding phases. Future versions may add +marker-to-R/G/B mappings or learned channel adapters if there is demand; either can use +the same panel-binding and fingerprint contract. Third-party strategies can accept a +JSON object through `--multiplex-params`; the built-in `channelwise` strategy uses its +explicit options and does not accept additional strategy parameters. + ### Segmentation The default `otsu` segmenter runs Otsu-on-saturation + morphology on a low-res @@ -250,6 +345,10 @@ devices in one process; it never changes the embeddings (it is excluded from the config hash), so a store built with it is identical to a single-device one and resumes interchangeably. +Multi-device `--devices` execution does not currently support `--multiplex-strategy`. +Strategy-derived multiplex runs use one device per slide; prefer `--device` for that +selection. + - **`embed` (one slide) - patch-parallel, latency mode.** `embed … --devices cuda:0,cuda:1` shards the slide's patches into contiguous blocks, embeds each block on its own device (one model copy + reader per device), and **gathers the @@ -335,6 +434,12 @@ reader, models (order-independent), segmenter / --no-seg, --mpp, --patch-size, --stain-norm, --snap-to-level, --mpp-tolerance, --allow-upsample, --amp ``` +For multiplex requests, the hash also covers the strategy and its parameters, ordered +marker selection, normalization recipe and pyramid-level limit, aggregation, the resolved +positional channel names and the resolved physical nuclear-channel indices when masking +is enabled. The `--channel-names-file` path and serialization are not hashed; their +parsed effective names are. + Runtime-only knobs do **not** change the hash: `--device`, `--devices`, `--batch-size`, `--read-workers`, `--read-block`, `--compile`, `--emit-geojson`, `--emit-thumbnail`, `--thumbnail-mpp`, `--max-px`, @@ -353,6 +458,10 @@ actual store's root/grid source provenance all agree. **Pass `--out-dir` and the content-affecting flags you pass to `embed`**, or the check is not target-aware (or the config hash will not match). +For library callers, `verify` also recognizes the direct single-grid receipt written by +the public `run_slide` primitive. Multi-grid requests use `embed_slide` and its whole-run +receipt. + ```bash raw2features verify slide.zarr --receipts-dir receipts/ --out-dir embeddings/ -m uni -m resnet50 --mpp 1.0 --quiet ``` @@ -491,7 +600,7 @@ Notes: ├── grid_index/ # (N,2) int32 (row,col) on the tiling grid ├── mask/ # (n_rows,n_cols) uint8 fraction of each cell that is tissue; /255 for [0,1] (omitted with --no-seg) ├── features// # (N,dim) float16 one array per extractor - └── slide// # (1,dim) float32 optional slide-level vectors + └── slide// # (1,dim) float32 optional slide-level vectors .patches.geojson # optional (--emit-geojson): per-patch polygons for QuPath .thumbnail.png # optional (--emit-thumbnail / thumbnail cmd) @@ -516,13 +625,16 @@ pixel extent), any embedding is invertible to the exact slide region **Each grid's `.zattrs["raw2features"]`** records the full, self-describing header: `schema_version`, `provenance` (version, CLI, git SHA, host, GPU), `thumbnail` (if emitted), `source` (uri, ngff_version, reader, mpp_level0, dimensions/downsamples, source axes, -per-axis physical scale, and source translation/origin), +per-axis physical scale, source translation/origin, and for multiplex sources the +effective positional channel names, their source, and any differing original OME labels), `patching` (target/achieved MPP, patch_px, read_level, read_px, resample, level0_patch, -level0_step, n_patches, grid_shape, `coords_convention: level0_xy`), `segmentation`, and +level0_step, n_patches, grid_shape, `coords_convention: level0_xy`), `segmentation` +(including the resolved nuclear channel binding when used), top-level `panel`, and per-model `models.` (source, embedding_dim, input_size, pooling, mean, std, -`transform_source_url`, license, gated, weights_sha256). The **root** `.zattrs` carries the -shared `source`/`provenance` plus a `grids` index (`key → {mpp, patch_px, n_patches, -models, grid_hash}`) for discovery without opening each grid. +`transform_source_url`, license, gated, weights pins, output fingerprint, and any +multiplex contract). The **root** `.zattrs` carries the shared `source`/`provenance` plus a +`grids` index (`key → {mpp, patch_px, n_patches, models, grid_hash}`) for discovery without +opening each grid. Read it back: diff --git a/notebooks/01_quickstart.ipynb b/notebooks/01_quickstart.ipynb index cc4a766..f086d6c 100644 --- a/notebooks/01_quickstart.ipynb +++ b/notebooks/01_quickstart.ipynb @@ -43,7 +43,7 @@ "IN_COLAB = \"google.colab\" in sys.modules\n", "HAVE_PKG = importlib.util.find_spec(\"raw2features\") is not None\n", "EXTRAS = \"zarr,image,torch,models\"\n", - "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.1.1) for reproducibility\n", + "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.2.0) for reproducibility\n", "if IN_COLAB or not HAVE_PKG:\n", " subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", SPEC])\n", "\n", diff --git a/notebooks/02_visual_walkthrough.ipynb b/notebooks/02_visual_walkthrough.ipynb index bff7f6e..ddda120 100644 --- a/notebooks/02_visual_walkthrough.ipynb +++ b/notebooks/02_visual_walkthrough.ipynb @@ -52,7 +52,7 @@ "HAVE_PKG = importlib.util.find_spec(\"raw2features\") is not None\n", "\n", "EXTRAS = \"zarr,image,torch,models\"\n", - "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.1.1) for reproducibility\n", + "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.2.0) for reproducibility\n", "if IN_COLAB or not HAVE_PKG:\n", " subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", SPEC])\n", "\n", diff --git a/notebooks/03_the_embeddings_store.ipynb b/notebooks/03_the_embeddings_store.ipynb index d58d89d..2c09256 100644 --- a/notebooks/03_the_embeddings_store.ipynb +++ b/notebooks/03_the_embeddings_store.ipynb @@ -44,7 +44,7 @@ "IN_COLAB = \"google.colab\" in sys.modules\n", "HAVE_PKG = importlib.util.find_spec(\"raw2features\") is not None\n", "EXTRAS = \"zarr,image,torch,models\"\n", - "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.1.1) for reproducibility\n", + "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.2.0) for reproducibility\n", "if IN_COLAB or not HAVE_PKG:\n", " subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", SPEC])\n", "\n", diff --git a/notebooks/04_spatialdata.ipynb b/notebooks/04_spatialdata.ipynb index 00684d2..c811876 100644 --- a/notebooks/04_spatialdata.ipynb +++ b/notebooks/04_spatialdata.ipynb @@ -53,9 +53,9 @@ "HAVE_PKG = importlib.util.find_spec(\"raw2features\") is not None\n", "\n", "# SpatialData export needs the [spatialdata] extra (spatialdata + geopandas) on top of\n", - "# the resnet50 stack. After a PyPI release prefer a hard pin: raw2features[...]==0.1.1\n", + "# the resnet50 stack. After a PyPI release prefer a hard pin: raw2features[...]==0.2.0\n", "EXTRAS = \"zarr,image,torch,models,spatialdata\"\n", - "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.1.1) for reproducibility\n", + "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.2.0) for reproducibility\n", "\n", "if IN_COLAB or not HAVE_PKG:\n", " subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", SPEC])\n", diff --git a/notebooks/05_quality_control.ipynb b/notebooks/05_quality_control.ipynb index 38b91e2..44f0c27 100644 --- a/notebooks/05_quality_control.ipynb +++ b/notebooks/05_quality_control.ipynb @@ -51,7 +51,7 @@ "IN_COLAB = \"google.colab\" in sys.modules\n", "HAVE_PKG = importlib.util.find_spec(\"raw2features\") is not None\n", "EXTRAS = \"zarr,image,torch,models,grandqc\"\n", - "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.1.1) for reproducibility\n", + "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.2.0) for reproducibility\n", "if IN_COLAB or not HAVE_PKG:\n", " subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", SPEC])\n", "\n", diff --git a/notebooks/06_cohort_quality_control.ipynb b/notebooks/06_cohort_quality_control.ipynb index 059fd2e..04d823c 100644 --- a/notebooks/06_cohort_quality_control.ipynb +++ b/notebooks/06_cohort_quality_control.ipynb @@ -46,7 +46,7 @@ "IN_COLAB = \"google.colab\" in sys.modules\n", "HAVE_PKG = importlib.util.find_spec(\"raw2features\") is not None\n", "EXTRAS = \"zarr,image,torch,models,grandqc\"\n", - "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.1.1) for reproducibility\n", + "SPEC = f\"raw2features[{EXTRAS}]\" # PyPI; pin a version (==0.2.0) for reproducibility\n", "if IN_COLAB or not HAVE_PKG:\n", " subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", SPEC])\n", "\n", diff --git a/pyproject.toml b/pyproject.toml index cb23d46..b2c2f14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "raw2features" -version = "0.1.1" +version = "0.2.0" description = "Read OME-Zarr whole-slide images and emit patch- and slide-level foundation-model embeddings - storage backend and model independently swappable." readme = "README.md" authors = [{ name = "Craig Myles", email = "me@craig.im" }] @@ -83,6 +83,9 @@ seal = "raw2features.embedders.seal_embedder:SealEmbedder" [project.entry-points."raw2features.sinks"] zarr = "raw2features.sinks.zarr_sink:ZarrSink" +[project.entry-points."raw2features.multiplex_strategies"] +channelwise = "raw2features.multiplex.channelwise:ChannelwiseStrategy" + [project.entry-points."raw2features.slide_embedders"] mean = "raw2features.slide_embedders.pool:MeanPoolSlideEmbedder" max = "raw2features.slide_embedders.pool:MaxPoolSlideEmbedder" diff --git a/scripts/codex_to_omezarr.py b/scripts/codex_to_omezarr.py index c00168d..812b786 100644 --- a/scripts/codex_to_omezarr.py +++ b/scripts/codex_to_omezarr.py @@ -1,9 +1,10 @@ #!/usr/bin/env python -"""Convert a CODEX / multiplex multi-channel TIFF to a multi-channel OME-Zarr (NGFF). +"""Convert a multiplex multi-channel TIFF to a multi-channel OME-Zarr (NGFF). Marker names (one per channel) are written as ``omero.channels[].label`` so the -raw2features channel-aware reader can surface them for the multiplex (KRONOS) path. A -spatial pyramid is built so low-res nuclear-channel segmentation is cheap. +raw2features channel-aware reader can surface them to multiplex strategies and native +multiplex models. A spatial pyramid is built so low-resolution normalization and +nuclear-channel masking are efficient. Usage:: @@ -13,8 +14,8 @@ ``--markers`` is a JSON file: either a bare list of channel names, or an object with a ``raw_markers`` / ``channels`` / ``markers`` key. The channel count must match the TIFF. -Needs the ``[kronos]`` extra (tifffile + imagecodecs): -``pip install "raw2features[kronos]"``. +Needs the OME-Zarr reader stack plus TIFF codecs: +``pip install "raw2features[zarr]" tifffile imagecodecs``. """ from __future__ import annotations diff --git a/src/raw2features/__init__.py b/src/raw2features/__init__.py index 29e2df4..7cbdf2b 100644 --- a/src/raw2features/__init__.py +++ b/src/raw2features/__init__.py @@ -1,8 +1,9 @@ """raw2features - read OME-Zarr WSIs and emit patch-level FM embeddings. The storage backend (reader) and model implementation are independently swappable -via six plugin seams: readers, segmenters, patchers, embedders, sinks, and -slide_embedders (see ``raw2features.core.plugins.SEAMS``). Third parties add an +via seven plugin seams: readers, segmenters, patchers, embedders, sinks, +slide_embedders, and multiplex_strategies (see ``raw2features.core.plugins.SEAMS``). +Third parties add an implementation through the matching ``raw2features.`` entry-point group. For embedders, that entry point is a loader family used by the packaged model registry; it does not add a new CLI ``--model`` name. Programmatic callers can pass @@ -25,7 +26,7 @@ from typing import TYPE_CHECKING -__version__ = "0.1.1" +__version__ = "0.2.0" __all__ = [ "__version__", diff --git a/src/raw2features/cli/_validation.py b/src/raw2features/cli/_validation.py index 72921b3..f060e45 100644 --- a/src/raw2features/cli/_validation.py +++ b/src/raw2features/cli/_validation.py @@ -2,7 +2,12 @@ from __future__ import annotations +import csv +import json import math +import os +import unicodedata +from typing import Any import typer @@ -40,6 +45,95 @@ def validate_positive_int(value: int | None, param_hint: str) -> None: raise typer.BadParameter("must be greater than zero", param_hint=param_hint) +def parse_channel_names_file(path: str | None) -> list[str]: + """Read a complete positional channel panel from a small text table. + + ``.txt`` files contain one name per line. ``.csv`` and ``.tsv`` files contain + exactly one column and may start with a conventional channel-name header. Blank + lines and ``#`` comments in text files are ignored; the resulting ordered list is + validated against the physical C-axis length after the reader opens. + """ + + if path is None: + return [] + suffix = os.path.splitext(path)[1].casefold() + if suffix not in {".txt", ".csv", ".tsv"}: + raise typer.BadParameter( + "must be a .txt, .csv, or .tsv file", + param_hint="--channel-names-file", + ) + try: + with open(path, encoding="utf-8-sig", newline="") as handle: + if suffix == ".txt": + rows = [ + [line] + for line in handle.read().splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + else: + rows = [ + row + for row in csv.reader( + handle, + delimiter="," if suffix == ".csv" else "\t", + strict=True, + ) + if row + ] + except (OSError, UnicodeError, csv.Error) as exc: + raise typer.BadParameter( + f"could not read file ({exc})", param_hint="--channel-names-file" + ) from exc + + header_names = { + "channel", + "channel_name", + "channel name", + "marker", + "marker_name", + "marker name", + "name", + "label", + } + if rows and len(rows[0]) == 1 and rows[0][0].strip().casefold() in header_names: + rows = rows[1:] + if not rows: + raise typer.BadParameter( + "must contain at least one channel name", + param_hint="--channel-names-file", + ) + + names: list[str] = [] + for index, row in enumerate(rows, start=1): + if len(row) != 1: + raise typer.BadParameter( + f"row {index} must contain exactly one column", + param_hint="--channel-names-file", + ) + name = row[0].strip() + if not name: + raise typer.BadParameter( + f"row {index} has an empty channel name", + param_hint="--channel-names-file", + ) + if any(ord(character) < 32 or ord(character) == 127 for character in name): + raise typer.BadParameter( + f"row {index} contains a control character", + param_hint="--channel-names-file", + ) + names.append(name) + identities = [unicodedata.normalize("NFKC", name).casefold() for name in names] + duplicates = sorted( + {identity for identity in identities if identities.count(identity) > 1} + ) + if duplicates: + raise typer.BadParameter( + "channel names must be unique after Unicode/case normalization", + param_hint="--channel-names-file", + ) + return names + + def validate_geometry( *, mpp: float | None = None, @@ -53,3 +147,33 @@ def validate_geometry( validate_positive_int(patch_size, "--patch-size") validate_positive_int(step, "--step") validate_positive_float(source_mpp, "--source-mpp") + + +def validate_multiplex_percentiles(low: float, high: float) -> None: + """Require a finite, ordered percentile interval within ``[0, 100]``.""" + + if not (math.isfinite(low) and math.isfinite(high) and 0 <= low < high <= 100): + raise typer.BadParameter( + "must be finite and satisfy 0 <= low < high <= 100", + param_hint="--multiplex-percentile-low/--multiplex-percentile-high", + ) + + +def parse_json_object(value: str | None, param_hint: str) -> dict[str, Any]: + """Parse a finite JSON object for a namespaced plugin configuration.""" + + if value is None: + return {} + try: + parsed = json.loads(value, parse_constant=lambda token: (_ for _ in ()).throw( + ValueError(f"non-finite JSON constant {token}") + )) + except (json.JSONDecodeError, ValueError) as exc: + raise typer.BadParameter( + f"must be a valid finite JSON object ({exc})", param_hint=param_hint + ) from exc + if not isinstance(parsed, dict) or any(not isinstance(key, str) for key in parsed): + raise typer.BadParameter( + "must be a JSON object with string keys", param_hint=param_hint + ) + return parsed diff --git a/src/raw2features/cli/embed.py b/src/raw2features/cli/embed.py index ae5a444..3c7e256 100644 --- a/src/raw2features/cli/embed.py +++ b/src/raw2features/cli/embed.py @@ -11,7 +11,15 @@ from raw2features.pipeline.runner import RunConfig, embed_slide from raw2features.viz import DEFAULT_THUMBNAIL_MPP -from ._validation import validate_amp, validate_batch_size, validate_geometry +from ._validation import ( + parse_channel_names_file, + parse_json_object, + validate_amp, + validate_batch_size, + validate_geometry, + validate_multiplex_percentiles, + validate_positive_int, +) def embed( @@ -60,6 +68,58 @@ def embed( "Changes the " "features -- use separate output dirs for with/without-norm experiments.", ), + multiplex_strategy: str | None = typer.Option( + None, + "--multiplex-strategy", + help="Apply a named multiplex strategy (initially: channelwise) around each " + "ordinary patch encoder.", + ), + multiplex_markers: list[str] = typer.Option( + [], + "--marker", + help="Multiplex marker to include; repeat in the required order. Default: all " + "named channels in source order (concat requires an explicit list).", + ), + channel_names_file: str | None = typer.Option( + None, + "--channel-names-file", + help="UTF-8 .txt/.csv/.tsv with exactly one ordered name per physical C-axis " + "position. Supplies missing labels and verifies existing labels; combine " + "with --marker to select or order a subset.", + ), + multiplex_normalization: str = typer.Option( + "percentile", + "--multiplex-normalization", + help="Per-marker intensity normalization (default: percentile over a " + "deterministically selected whole-image pyramid level).", + ), + multiplex_percentile_low: float = typer.Option( + 1.0, + "--multiplex-percentile-low", + help="Lower whole-image-level percentile for per-marker normalization.", + ), + multiplex_percentile_high: float = typer.Option( + 99.0, + "--multiplex-percentile-high", + help="Upper whole-image-level percentile for per-marker normalization.", + ), + multiplex_normalization_max_side_px: int = typer.Option( + 2048, + "--multiplex-normalization-max-side-px", + help="Longest side allowed for the deterministic whole-image normalization " + "level (default: 2048; larger may select a finer level and use more RAM).", + ), + multiplex_aggregation: str = typer.Option( + "mean", + "--multiplex-aggregation", + help="Marker embedding aggregation: mean | concat.", + ), + multiplex_params: str | None = typer.Option( + None, + "--multiplex-params", + help="JSON object of content parameters for a third-party multiplex strategy. " + "The built-in channelwise strategy uses its explicit options instead.", + ), device: str = typer.Option( "auto", "--device", help="auto | cuda | mps | cpu (auto = best available)" ), @@ -165,9 +225,14 @@ def embed( """ validate_amp(amp) validate_batch_size(batch_size) - validate_geometry( - mpp=mpp, patch_size=patch_size, step=step, source_mpp=source_mpp + validate_geometry(mpp=mpp, patch_size=patch_size, step=step, source_mpp=source_mpp) + validate_positive_int( + multiplex_normalization_max_side_px, + "--multiplex-normalization-max-side-px", ) + validate_multiplex_percentiles(multiplex_percentile_low, multiplex_percentile_high) + strategy_params = parse_json_object(multiplex_params, "--multiplex-params") + channel_names_override = parse_channel_names_file(channel_names_file) if hf_token: os.environ["HF_TOKEN"] = hf_token os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token @@ -215,6 +280,15 @@ def embed( tissue_threshold=tissue_threshold, features_dtype=features_dtype, stain_norm=stain_norm, + multiplex_strategy=multiplex_strategy, + multiplex_markers=list(multiplex_markers), + multiplex_normalization=multiplex_normalization, + multiplex_percentile_low=multiplex_percentile_low, + multiplex_percentile_high=multiplex_percentile_high, + multiplex_normalization_max_side_px=multiplex_normalization_max_side_px, + multiplex_aggregation=multiplex_aggregation, + multiplex_strategy_params=strategy_params, + channel_names_override=channel_names_override, snap_to_level=snap_to_level, mpp_tolerance=mpp_tolerance, allow_upsample=allow_upsample, diff --git a/src/raw2features/cli/embed_many.py b/src/raw2features/cli/embed_many.py index 43c5937..1fddf4e 100644 --- a/src/raw2features/cli/embed_many.py +++ b/src/raw2features/cli/embed_many.py @@ -31,12 +31,22 @@ ) from raw2features.pipeline.runner import ( RunConfig, + _validate_multiplex_slide_encoders, embed_slide, load_embedders, + resolve_multiplex_source_config, slide_id_from_path, ) -from ._validation import validate_amp, validate_batch_size, validate_geometry +from ._validation import ( + parse_channel_names_file, + parse_json_object, + validate_amp, + validate_batch_size, + validate_geometry, + validate_multiplex_percentiles, + validate_positive_int, +) class _WorkerStartupError(RuntimeError): @@ -92,6 +102,58 @@ def embed_many( "--stain-norm", help="Stain-normalize each patch before embedding (macenko|reinhard|vahadane).", ), + multiplex_strategy: str | None = typer.Option( + None, + "--multiplex-strategy", + help="Apply a named multiplex strategy (initially: channelwise) around each " + "ordinary patch encoder.", + ), + multiplex_markers: list[str] = typer.Option( + [], + "--marker", + help="Multiplex marker to include; repeat in the required order. Default: all " + "named channels in source order (concat requires an explicit list).", + ), + channel_names_file: str | None = typer.Option( + None, + "--channel-names-file", + help="UTF-8 .txt/.csv/.tsv with exactly one ordered name per physical C-axis " + "position. Applies to every slide in this command and verifies any existing " + "labels; combine with --marker to select/order a subset.", + ), + multiplex_normalization: str = typer.Option( + "percentile", + "--multiplex-normalization", + help="Per-marker intensity normalization (default: percentile over a " + "deterministically selected whole-image pyramid level).", + ), + multiplex_percentile_low: float = typer.Option( + 1.0, + "--multiplex-percentile-low", + help="Lower whole-image-level percentile for per-marker normalization.", + ), + multiplex_percentile_high: float = typer.Option( + 99.0, + "--multiplex-percentile-high", + help="Upper whole-image-level percentile for per-marker normalization.", + ), + multiplex_normalization_max_side_px: int = typer.Option( + 2048, + "--multiplex-normalization-max-side-px", + help="Longest side allowed for the deterministic whole-image normalization " + "level (default: 2048; larger may select a finer level and use more RAM).", + ), + multiplex_aggregation: str = typer.Option( + "mean", + "--multiplex-aggregation", + help="Marker embedding aggregation: mean | concat.", + ), + multiplex_params: str | None = typer.Option( + None, + "--multiplex-params", + help="JSON object of content parameters for a third-party multiplex strategy. " + "The built-in channelwise strategy uses its explicit options instead.", + ), slide_encoder: list[str] = typer.Option( [], "--slide-encoder", @@ -179,9 +241,14 @@ def embed_many( """Warm worker: build the models once and embed this shard's slides.""" validate_amp(amp) validate_batch_size(batch_size) - validate_geometry( - mpp=mpp, patch_size=patch_size, step=step, source_mpp=source_mpp + validate_geometry(mpp=mpp, patch_size=patch_size, step=step, source_mpp=source_mpp) + validate_positive_int( + multiplex_normalization_max_side_px, + "--multiplex-normalization-max-side-px", ) + validate_multiplex_percentiles(multiplex_percentile_low, multiplex_percentile_high) + strategy_params = parse_json_object(multiplex_params, "--multiplex-params") + channel_names_override = parse_channel_names_file(channel_names_file) from raw2features.slide_embedders.model_registry import ( validate_slide_encoder_names, ) @@ -257,7 +324,7 @@ def embed_many( else: models = list(feature_extractor) - from raw2features.embedders.model_registry import resolve_geometry + from raw2features.embedders.model_registry import build_embedder, resolve_geometry groups = resolve_geometry(models, mpp, patch_size, geometry_config) if not groups: @@ -276,6 +343,15 @@ def embed_many( tissue_threshold=tissue_threshold, features_dtype=features_dtype, stain_norm=stain_norm, + multiplex_strategy=multiplex_strategy, + multiplex_markers=list(multiplex_markers), + multiplex_normalization=multiplex_normalization, + multiplex_percentile_low=multiplex_percentile_low, + multiplex_percentile_high=multiplex_percentile_high, + multiplex_normalization_max_side_px=multiplex_normalization_max_side_px, + multiplex_aggregation=multiplex_aggregation, + multiplex_strategy_params=strategy_params, + channel_names_override=channel_names_override, slide_encoders=list(slide_encoder), qc=list(qc), qc_stain_norm=qc_stain_norm, @@ -294,6 +370,42 @@ def embed_many( cli = sanitize_argv(sys.argv) device_list = cfg.device_list() + # Validate every positional panel before a warm worker loads model weights. This is + # especially important for one --channel-names-file applied to a whole cohort: a + # count or metadata conflict must not surface only after an expensive gated load. + try: + strategy = None + base_probes = [] + if cfg.multiplex_strategy is not None: + from raw2features.multiplex import build_strategy + + if len(device_list) != 1: + raise ValueError( + "multiplex strategies are currently single-device; pass one " + "--device" + ) + _validate_multiplex_slide_encoders(cfg.slide_encoders) + strategy = build_strategy(cfg.multiplex_strategy) + base_probes = [build_embedder(name) for name in cfg.models] + for row in shard: + source_cfg = resolve_multiplex_source_config( + row["path"], _with_source_mpp(cfg, row) + ) + if strategy is not None: + for base in base_probes: + strategy.prepare( + base_embedder=base, + channel_names=source_cfg.resolved_channel_names, + channel_count=len(source_cfg.resolved_channel_names), + config=source_cfg, + ) + except (KeyError, OSError, TypeError, ValueError) as exc: + typer.echo( + redact_uri_credentials(f"Error: multiplex panel preflight failed: {exc}"), + err=True, + ) + raise typer.Exit(1) from exc + t0 = time.time() try: if len(device_list) > 1: diff --git a/src/raw2features/cli/main.py b/src/raw2features/cli/main.py index 24d569f..857a1ad 100644 --- a/src/raw2features/cli/main.py +++ b/src/raw2features/cli/main.py @@ -56,7 +56,7 @@ def list_plugins( component: str = typer.Argument( ..., help="Component type to list: readers, segmenters, patchers, embedders, " - "sinks, slide_embedders.", + "sinks, slide_embedders, multiplex_strategies.", ), ) -> None: """List the available plugins of a given component type.""" diff --git a/src/raw2features/cli/slide_embed.py b/src/raw2features/cli/slide_embed.py index 256ecfc..dab837c 100644 --- a/src/raw2features/cli/slide_embed.py +++ b/src/raw2features/cli/slide_embed.py @@ -71,6 +71,7 @@ def slide_embed( encode_slide_embedding, resolve_slide_grid, slide_embedding_is_complete, + slide_output_key, write_slide_embedding, ) from raw2features.slide_embedders.model_registry import ( @@ -114,12 +115,18 @@ def slide_embed( typer.echo(f"Error: {exc}", err=True) raise typer.Exit(1) from exc + # Strategy-derived multiplex pools are patch-qualified so several marker + # recipes can coexist on the same grid without replacing one another. + output_name = slide_output_key( + group, slide_model_name, selected_patch_model + ) # Skip only an output produced from the requested patch model. if not force and slide_embedding_is_complete( group, slide_model_name, patch_model=selected_patch_model, device=device, + output_name=output_name, ): typer.echo( f"{slide_model_name} [{selected_grid}]: already complete (skipping)" @@ -151,6 +158,7 @@ def slide_embed( slide_model_name, encoding.vector, encoding.provenance, + output_name=output_name, ) typer.echo( f"{slide_model_name} [{selected_grid}]: done " diff --git a/src/raw2features/cli/verify.py b/src/raw2features/cli/verify.py index 83cd409..39538f9 100644 --- a/src/raw2features/cli/verify.py +++ b/src/raw2features/cli/verify.py @@ -10,6 +10,7 @@ import typer +from raw2features.core.uris import redact_uri_credentials from raw2features.pipeline.receipt import ( canonical_source_uri, is_complete, @@ -17,10 +18,19 @@ from raw2features.pipeline.runner import ( RunConfig, expected_model_contracts, + resolve_multiplex_output_contracts, + resolve_multiplex_source_config, slide_id_from_path, ) -from ._validation import validate_amp, validate_geometry +from ._validation import ( + parse_channel_names_file, + parse_json_object, + validate_amp, + validate_geometry, + validate_multiplex_percentiles, + validate_positive_int, +) def verify( @@ -54,6 +64,52 @@ def verify( stain_norm: str | None = typer.Option( None, "--stain-norm", help="Must match the embed run (content-hash field)." ), + multiplex_strategy: str | None = typer.Option( + None, + "--multiplex-strategy", + help="Multiplex strategy used by the embed run (initially: channelwise).", + ), + multiplex_markers: list[str] = typer.Option( + [], + "--marker", + help="Multiplex marker; repeat in the exact order used by the embed run.", + ), + channel_names_file: str | None = typer.Option( + None, + "--channel-names-file", + help="UTF-8 .txt/.csv/.tsv with exactly one ordered name per physical C-axis " + "position; must resolve to the same effective panel used by embed.", + ), + multiplex_normalization: str = typer.Option( + "percentile", + "--multiplex-normalization", + help="Per-marker whole-image-level normalization used by the embed run.", + ), + multiplex_percentile_low: float = typer.Option( + 1.0, + "--multiplex-percentile-low", + help="Lower normalization percentile used by the embed run.", + ), + multiplex_percentile_high: float = typer.Option( + 99.0, + "--multiplex-percentile-high", + help="Upper normalization percentile used by the embed run.", + ), + multiplex_normalization_max_side_px: int = typer.Option( + 2048, + "--multiplex-normalization-max-side-px", + help="Whole-image normalization level limit used by the embed run.", + ), + multiplex_aggregation: str = typer.Option( + "mean", + "--multiplex-aggregation", + help="Marker aggregation used by the embed run: mean | concat.", + ), + multiplex_params: str | None = typer.Option( + None, + "--multiplex-params", + help="JSON content parameters used by a third-party multiplex strategy.", + ), amp: str = typer.Option("auto", "--amp"), device: str = typer.Option( "auto", @@ -70,9 +126,14 @@ def verify( ) -> None: """Exit 0 if the slide is already complete & output-validated, else exit 1.""" validate_amp(amp) - validate_geometry( - mpp=mpp, patch_size=patch_size, step=step, source_mpp=source_mpp + validate_geometry(mpp=mpp, patch_size=patch_size, step=step, source_mpp=source_mpp) + validate_positive_int( + multiplex_normalization_max_side_px, + "--multiplex-normalization-max-side-px", ) + validate_multiplex_percentiles(multiplex_percentile_low, multiplex_percentile_high) + strategy_params = parse_json_object(multiplex_params, "--multiplex-params") + channel_names_override = parse_channel_names_file(channel_names_file) from raw2features.core.device import resolve_device from raw2features.embedders.model_registry import resolve_geometry from raw2features.pipeline.runner import resolve_run @@ -116,12 +177,32 @@ def verify( tissue_threshold=tissue_threshold, features_dtype=features_dtype, stain_norm=stain_norm, + multiplex_strategy=multiplex_strategy, + multiplex_markers=list(multiplex_markers), + multiplex_normalization=multiplex_normalization, + multiplex_percentile_low=multiplex_percentile_low, + multiplex_percentile_high=multiplex_percentile_high, + multiplex_normalization_max_side_px=multiplex_normalization_max_side_px, + multiplex_aggregation=multiplex_aggregation, + multiplex_strategy_params=strategy_params, + channel_names_override=channel_names_override, snap_to_level=snap_to_level, mpp_tolerance=mpp_tolerance, allow_upsample=allow_upsample, amp=amp, device=resolve_device(device), ) + try: + cfg = resolve_multiplex_source_config(slide, cfg) + except (OSError, ValueError) as exc: + if not quiet: + typer.echo( + redact_uri_credentials( + f"multiplex panel is invalid: {exc}; incomplete" + ), + err=True, + ) + raise typer.Exit(code=1) from exc # Hash exactly as embed_slide does, so verify matches the embed that wrote it. _, group_cfgs, run_hash = resolve_run(cfg, mpp, patch_size, geometry_config) expected_source = canonical_source_uri(slide) @@ -145,24 +226,66 @@ def verify( if expected_path is not None else None ) - contracts = expected_model_contracts(cfg) - expected_grid_models = { - group_cfg.grid_hash(): list(group_cfg.models) for group_cfg in group_cfgs - } + if cfg.multiplex_strategy is None: + contracts = expected_model_contracts(cfg) + expected_grid_models = { + group_cfg.grid_hash(): list(group_cfg.models) for group_cfg in group_cfgs + } + else: + contracts = {} + expected_grid_models = {} + for group_cfg in group_cfgs: + effective_cfg, group_contracts = resolve_multiplex_output_contracts( + slide, + group_cfg, + device=cfg.device, + ) + for name, contract in group_contracts.items(): + if name in contracts and contracts[name] != contract: + raise ValueError( + f"multiplex effective model key {name!r} resolved to " + "different contracts within one request" + ) + contracts[name] = contract + expected_grid_models[group_cfg.grid_hash()] = list(effective_cfg.models) compatible_grid_hashes = { group_cfg.grid_hash(): group_cfg.compatible_legacy_grid_hashes() for group_cfg in group_cfgs } + compatible_grid_segmenters = { + group_cfg.grid_hash(): group_cfg.compatible_legacy_grid_segmenters() + for group_cfg in group_cfgs + } + allow_hashless_legacy_grids = { + group_cfg.grid_hash(): group_cfg.allows_hashless_legacy_grid() + for group_cfg in group_cfgs + } + completion_kwargs = { + "expected_source_uri": expected_source, + "expected_output_uri": expected_output, + "expected_model_contracts": contracts, + "expected_grid_models": expected_grid_models, + "compatible_grid_hashes": compatible_grid_hashes, + "compatible_grid_segmenters": compatible_grid_segmenters, + "allow_hashless_legacy_grids": allow_hashless_legacy_grids, + } complete = is_complete( receipts_dir, slide_id, run_hash, - expected_source_uri=expected_source, - expected_output_uri=expected_output, - expected_model_contracts=contracts, - expected_grid_models=expected_grid_models, - compatible_grid_hashes=compatible_grid_hashes, + **completion_kwargs, ) + if not complete and len(group_cfgs) == 1: + # ``run_slide`` is the public single-grid primitive and historically writes + # its direct RunConfig hash, while ``embed``/``embed_slide`` write the whole- + # request hash above. Accept either only for a genuine one-grid request; both + # paths still validate source, output, live arrays, and output fingerprints. + complete = is_complete( + receipts_dir, + slide_id, + group_cfgs[0].content_hash(), + **completion_kwargs, + ) if not quiet: typer.echo(f"{slide_id}: {'complete' if complete else 'incomplete'}") raise typer.Exit(code=0 if complete else 1) diff --git a/src/raw2features/core/plugins.py b/src/raw2features/core/plugins.py index 44afbfe..7d1602f 100644 --- a/src/raw2features/core/plugins.py +++ b/src/raw2features/core/plugins.py @@ -1,6 +1,7 @@ """Plugin discovery for the swappable seams. -Each seam (readers, segmenters, patchers, embedders, sinks, slide_embedders) +Each seam (readers, segmenters, patchers, embedders, sinks, slide_embedders, +multiplex_strategies) is resolved by name from two sources, in order: 1. An in-process registry populated by the ``@register`` decorator (used by @@ -21,7 +22,15 @@ # group name -> {plugin name -> object/class} _REGISTRY: dict[str, dict[str, Any]] = {} -SEAMS = ("readers", "segmenters", "patchers", "embedders", "sinks", "slide_embedders") +SEAMS = ( + "readers", + "segmenters", + "patchers", + "embedders", + "sinks", + "slide_embedders", + "multiplex_strategies", +) def register(seam: str, name: str): diff --git a/src/raw2features/core/uris.py b/src/raw2features/core/uris.py index f3a71ab..dcc89dd 100644 --- a/src/raw2features/core/uris.py +++ b/src/raw2features/core/uris.py @@ -14,6 +14,8 @@ import os import posixpath import re +from collections.abc import Mapping +from typing import Any from urllib.parse import unquote, unquote_plus, urlsplit, urlunsplit _STORE_SUFFIXES = (".embeddings.zarr", ".ome.zarr", ".zarr") @@ -200,6 +202,53 @@ def redact_uri_credentials(text: str) -> str: return _QUALIFIED_URI_IN_TEXT.sub(_redact_uri_match, redacted) +def redact_metadata_credentials(value: Any) -> Any: + """Recursively make untrusted plugin metadata safe to persist. + + URI userinfo and signed/auth query values are removed from every string. Values + under unambiguously credential-bearing keys are replaced even when they are plain + tokens rather than URLs. Tuples become JSON-equivalent lists. + """ + + if isinstance(value, Mapping): + safe = {} + for key, item in value.items(): + name = str(key) + normalized = re.sub(r"[-_]", "", name).casefold() + safe[name] = ( + "" + if normalized in _NORMALIZED_GENERIC_AUTH_QUERY_KEYS + else redact_metadata_credentials(item) + ) + return safe + if isinstance(value, (list, tuple)): + return [redact_metadata_credentials(item) for item in value] + if isinstance(value, str): + return redact_uri_credentials(value) + return value + + +def redact_metadata_uri_credentials(value: Any) -> Any: + """Recursively sanitize URI strings without interpreting semantic key names. + + Model constructors and strategy contracts may legitimately use names such as + ``token`` or ``signature``. Those values affect output identity and must not be + mistaken for credentials merely because of their key. This narrower helper strips + credentials only when they occur inside a URI value. + """ + + if isinstance(value, Mapping): + return { + str(key): redact_metadata_uri_credentials(item) + for key, item in value.items() + } + if isinstance(value, (list, tuple)): + return [redact_metadata_uri_credentials(item) for item in value] + if isinstance(value, str): + return redact_uri_credentials(value) + return value + + def slide_id_from_source(source: str) -> str: """Derive a readable output ID, fingerprinting collision-prone sources. diff --git a/src/raw2features/embedders/base.py b/src/raw2features/embedders/base.py index b3e94c5..fd24203 100644 --- a/src/raw2features/embedders/base.py +++ b/src/raw2features/embedders/base.py @@ -144,6 +144,10 @@ class ModelSpec: # exists, else the arXiv DataCite DOI (10.48550/arXiv.*); None only for open-weights # releases with no publication. Flows into the output store's model provenance. doi: str | None = None + # Resolved strategy contract for a derived multiplex output. Ordinary registry + # models leave this unset, preserving their v0.1 fingerprints byte-for-byte. Kept + # last so adding it does not shift any pre-existing positional constructor fields. + multiplex: dict[str, Any] | None = None @property def extract_px(self) -> int: @@ -200,6 +204,21 @@ def transform_signature(self) -> tuple: self.spec.interpolation, ) + @property + def transform_input_dtype(self) -> str: + """Host-patch contract accepted by :meth:`transform_batch`. + + Every built-in RGB embedder currently consumes HWC ``uint8`` patches. A + plugin whose transform genuinely consumes normalized floating-point input + may override this with ``"float32_0_1"``. Multiplex adapters consult this + explicit seam instead of assuming that quantizing to uint8 is universally + valid. It is intentionally not part of ordinary brightfield fingerprints, + preserving their existing output identity; a multiplex strategy records the + resolved value in its own contract. + """ + + return "uint8" + @abstractmethod def load( self, diff --git a/src/raw2features/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index 3f9ead9..e6974f7 100644 --- a/src/raw2features/embedders/fingerprint.py +++ b/src/raw2features/embedders/fingerprint.py @@ -11,6 +11,7 @@ import hashlib import json import posixpath +import re from collections.abc import Mapping from copy import deepcopy from typing import Any @@ -29,6 +30,39 @@ OUTPUT_FINGERPRINT_VERSION = 1 OUTPUT_FINGERPRINT_ALGORITHM = "sha256" +# Plain values under these unambiguous names are runtime credentials, not output +# semantics. Ambiguous names such as ``token``, ``signature``, ``auth``, ``key``, and +# ``secret`` are intentionally excluded because model constructors legitimately use +# them (for example ``token='cls'``). URI query sanitization remains broader and +# provider-aware in ``source_uri``. +_HIGH_CONFIDENCE_CREDENTIAL_KEYS = frozenset( + { + "accesstoken", + "apikey", + "authorization", + "awsaccesskeyid", + "awssecretaccesskey", + "bearertoken", + "clientsecret", + "credential", + "credentials", + "googleaccessid", + "hftoken", + "idtoken", + "oauthtoken", + "password", + "passwd", + "privatekey", + "refreshtoken", + "sastoken", + "securitytoken", + "secretaccesskey", + "sharedaccesssignature", + "subscriptionkey", + "xapikey", + } +) + # Construction inputs in SEAL's pinned fork that affect the image encoder. The # loader consumes this object directly as well as fingerprinting it, preventing a # CWD-local conf/config.yaml or conf/user/.yaml from silently changing output. @@ -152,7 +186,16 @@ def _credential_free(value: Any) -> Any: """Recursively remove URI credentials before hashing or persistence.""" if isinstance(value, Mapping): - return {str(key): _credential_free(item) for key, item in value.items()} + safe = {} + for key, item in value.items(): + name = str(key) + normalized = re.sub(r"[-_]", "", name).casefold() + safe[name] = ( + "" + if normalized in _HIGH_CONFIDENCE_CREDENTIAL_KEYS + else _credential_free(item) + ) + return safe if isinstance(value, (list, tuple)): return [_credential_free(item) for item in value] if isinstance(value, str): @@ -391,6 +434,60 @@ def _patch_constructor(spec: ModelSpec) -> dict[str, Any]: def patch_output_fingerprint(spec: ModelSpec, resolved_amp: str) -> dict[str, Any]: """Fingerprint the complete persisted-output contract for one patch model.""" + if spec.multiplex is not None: + contract = deepcopy(spec.multiplex) + base_fingerprint = contract.pop("base_output_fingerprint", None) + if not valid_output_fingerprint(base_fingerprint): + raise ValueError( + f"Multiplex output {spec.name!r} has no valid base-model output " + "fingerprint. Bind the strategy after resolving the base encoder." + ) + strategy = contract.get("strategy") + contract_version = contract.get("contract_version") + if not isinstance(strategy, str) or not strategy.strip(): + raise ValueError( + f"Multiplex output {spec.name!r} has no valid strategy name." + ) + if ( + isinstance(contract_version, bool) + or not isinstance(contract_version, int) + or contract_version <= 0 + ): + raise ValueError( + f"Multiplex output {spec.name!r} has no positive integer " + "strategy contract version." + ) + base_amp = ( + base_fingerprint.get("payload", {}) + .get("output", {}) + .get("resolved_amp") + ) + if base_amp != resolved_amp: + raise ValueError( + f"Multiplex output {spec.name!r} resolved AMP {resolved_amp!r} " + f"does not match its base-model fingerprint ({base_amp!r})." + ) + return make_output_fingerprint( + { + "kind": "patch_features", + "model": spec.name, + "loader": { + "family": "multiplex_strategy", + "contract_version": contract_version, + "strategy": strategy, + }, + "base_output_fingerprint": base_fingerprint, + "multiplex": contract, + "output": { + "pooling": contract.get("aggregation", {}).get("name"), + "embedding_dim": int(spec.embedding_dim), + "reg_tokens": 0, + "modality": "multiplex", + "resolved_amp": resolved_amp, + }, + } + ) + payload: dict[str, Any] = { "kind": "patch_features", "model": spec.name, diff --git a/src/raw2features/embedders/kronos_embedder.py b/src/raw2features/embedders/kronos_embedder.py index 050c3d0..bc2fcd1 100644 --- a/src/raw2features/embedders/kronos_embedder.py +++ b/src/raw2features/embedders/kronos_embedder.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: # pragma: no cover import torch + def _to_unit_interval(arr: np.ndarray, dtype) -> np.ndarray: """Scale a multiplex stack to ``[0, 1]`` by its source dtype. @@ -226,15 +227,17 @@ def set_panel(self, channel_names: list[str] | None) -> dict: kept.append(name) # Explicit, retrievable record of how this source channel was identified # to KRONOS: original name + position -> KRONOS marker name + id. - mapping.append({ - "channel": name, - "channel_index": i, - "kronos_marker": canonical, - "marker_id": int(mid), - }) + mapping.append( + { + "channel": name, + "channel_index": i, + "kronos_marker": canonical, + "marker_id": int(mid), + } + ) else: dropped.append(name) - if key is not None: # a named marker (not a blank/empty cycle) missed + if key: # a named marker (not a blank/empty cycle) missed unmatched.append(name) if not idx: raise ValueError("no channels matched the KRONOS marker vocabulary") diff --git a/src/raw2features/embedders/registry.yaml b/src/raw2features/embedders/registry.yaml index 18166a1..8b93ebe 100644 --- a/src/raw2features/embedders/registry.yaml +++ b/src/raw2features/embedders/registry.yaml @@ -1084,7 +1084,7 @@ plip: # frozen base. Needs the [seal] extra + the backbone's own package + both gates. seal_embedder.py. seal_conch: family: seal - experimental: true # adapter pinned+verified; frozen base is upstream-managed HEAD in v0.1.1 + experimental: true # adapter pinned+verified; frozen base is upstream-managed HEAD in v0.2.0 source: "conch" # SEAL backbone key; checkpoint seal_conch_vision.pth in MahmoodLab/SEAL recommended_mpp: 0.5 # CONCH convention (unsourced); SEAL keeps the base's scale embedding_dim: 512 # CONCH base dim, unchanged by the LoRA @@ -1102,7 +1102,7 @@ seal_conch: gated: true inference_amp: fp32 notes: >- - EXPERIMENTAL in v0.1.1: the SEAL adapter is revision-pinned and SHA-256 verified, but the + EXPERIMENTAL in v0.2.0: the SEAL adapter is revision-pinned and SHA-256 verified, but the upstream SEAL factory still downloads the frozen CONCH base from mutable HEAD, so this model is outside the exact-weight pinning guarantee. SEAL-CONCH (MahmoodLab) -- CONCH ViT-B/16 + a spatial-transcriptomics LoRA, 512-d. The HF @@ -1114,7 +1114,7 @@ seal_conch: seal_univ2: family: seal - experimental: true # adapter pinned+verified; frozen base is upstream-managed HEAD in v0.1.1 + experimental: true # adapter pinned+verified; frozen base is upstream-managed HEAD in v0.2.0 source: "univ2" # checkpoint seal_univ2_vision.pth recommended_mpp: 0.5 # UNI2-h convention (unsourced) embedding_dim: 1536 # UNI2-h base dim @@ -1132,7 +1132,7 @@ seal_univ2: gated: true inference_amp: fp32 notes: >- - EXPERIMENTAL in v0.1.1: the SEAL adapter is revision-pinned and SHA-256 verified, but the + EXPERIMENTAL in v0.2.0: the SEAL adapter is revision-pinned and SHA-256 verified, but the upstream SEAL factory still downloads the frozen UNI2-h base from mutable HEAD, so this model is outside the exact-weight pinning guarantee. SEAL-UNIv2 (MahmoodLab) -- UNI2-h ViT-H/14 + a spatial-transcriptomics LoRA, 1536-d. Same diff --git a/src/raw2features/embedders/seal_embedder.py b/src/raw2features/embedders/seal_embedder.py index dad3bc2..b227ce0 100644 --- a/src/raw2features/embedders/seal_embedder.py +++ b/src/raw2features/embedders/seal_embedder.py @@ -6,7 +6,7 @@ image-only loader (``ModelMixin.get_img_model``) - the gene/omics model is not loaded. Output dim = the base encoder's (512 for CONCH, 1536 for UNI2-h). -SEAL is **experimental in raw2features v0.1.1**. The LoRA adapter is revision-pinned, +SEAL is **experimental in raw2features v0.2.0**. The LoRA adapter is revision-pinned, SHA-256 verified, and loaded with a frozen constructor contract, but SEAL's upstream factory still fetches the frozen CONCH/UNI2-h base from mutable Hugging Face HEAD. The persisted composite fingerprint records that limitation explicitly; SEAL is outside the @@ -92,7 +92,7 @@ def load( # Resolve and verify the adapter before either deserialising it or paying to # construct its large frozen base. The base itself remains upstream-managed - # and unpinned, which is why SEAL is marked experimental in v0.1.1. + # and unpinned, which is why SEAL is marked experimental in v0.2.0. ckpt_path = download_pinned_hf_file( f"hf-hub:{_SEAL_REPO}", self.spec.weights_filename or f"seal_{backbone}_vision.pth", diff --git a/src/raw2features/export/spatialdata.py b/src/raw2features/export/spatialdata.py index 3f7ee24..553afbe 100644 --- a/src/raw2features/export/spatialdata.py +++ b/src/raw2features/export/spatialdata.py @@ -32,6 +32,7 @@ from __future__ import annotations +import json import os import numpy as np @@ -88,7 +89,7 @@ def _marker_field(entry: dict): KRONOS records it under ``kronos_marker``; a future multiplex model may use another ``*_marker`` key. Falls back through the known names so the table stays generic. """ - for k in ("marker", "kronos_marker"): + for k in ("marker", "kronos_marker", "canonical_name", "source_name"): if k in entry: return entry[k] for k in entry: @@ -114,8 +115,10 @@ def _panel_dataframe(panel: dict): rows.append( { "model": model, - "channel": e.get("channel"), - "channel_index": e.get("channel_index"), + "channel": e.get("channel", e.get("source_name")), + "channel_index": e.get( + "channel_index", e.get("source_index") + ), "marker": _marker_field(e), "marker_id": e.get("marker_id"), } @@ -143,25 +146,66 @@ def _panel_summary(panel: dict) -> dict: } +def _columnarize_records(value): + """Recursively convert record lists to columns that AnnData can round-trip.""" + + if isinstance(value, dict): + return {key: _columnarize_records(item) for key, item in value.items()} + if isinstance(value, list) and value and all( + isinstance(item, dict) for item in value + ): + keys = list(dict.fromkeys(key for item in value for key in item)) + return { + key: [_columnarize_records(item.get(key)) for item in value] + for key in keys + } + if isinstance(value, list): + return [_columnarize_records(item) for item in value] + return value + + def _uns_safe_header(header: dict) -> dict: - """Header copy whose multiplex panel ``mapping`` is columnarised for anndata. + """Copy multiplex records into an AnnData-safe recursive columnar form. - anndata's zarr writer turns a list-of-dicts into an array of stringified dict reprs; - converting each model's ``mapping`` to a dict-of-lists keeps it as plain arrays that - round-trip and stay reconstructable. Shallow-copies (model cards stay untouched). + AnnData's zarr writer turns a list of dictionaries into stringified reprs. This + preserves top-level panel mappings, nested normalization bounds, and the mirrored + model strategy/fingerprint records as reconstructable arrays without changing + ordinary brightfield model headers. """ panel = header.get("panel") if not panel: return header safe_panel = {} for model, p in panel.items(): - p = dict(p or {}) - mapping = p.get("mapping") - if isinstance(mapping, list) and mapping and isinstance(mapping[0], dict): - keys = list(mapping[0].keys()) - p["mapping"] = {k: [e.get(k) for e in mapping] for k in keys} - safe_panel[model] = p - return {**header, "panel": safe_panel} + safe_panel[model] = _columnarize_records(dict(p or {})) + result = {**header, "panel": safe_panel} + models = header.get("models") + if isinstance(models, dict): + safe_models = dict(models) + for model, metadata in models.items(): + if isinstance(metadata, dict) and metadata.get("multiplex") is not None: + # The fingerprint digest authenticates the exact nested payload. + # Columnarising its record lists would mutate that payload while + # retaining the old digest, making the exported contract invalid. + # Canonical JSON is AnnData-safe, lossless, and directly verifiable; + # only the separate human-readable multiplex metadata is reshaped. + fingerprint = metadata.get("output_fingerprint") + readable = { + key: value + for key, value in metadata.items() + if key != "output_fingerprint" + } + safe_metadata = _columnarize_records(readable) + if fingerprint is not None: + safe_metadata["output_fingerprint"] = json.dumps( + fingerprint, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + safe_models[model] = safe_metadata + result["models"] = safe_models + return result def _coordinate_systems(sx, sy, translation): diff --git a/src/raw2features/multiplex/__init__.py b/src/raw2features/multiplex/__init__.py new file mode 100644 index 0000000..1397d65 --- /dev/null +++ b/src/raw2features/multiplex/__init__.py @@ -0,0 +1,17 @@ +"""Named-channel strategies for applying RGB encoders to multiplex images.""" + +from .base import ( + BoundMultiplexStrategy, + MultiplexStrategy, + PreparedMultiplexStrategy, +) +from .channelwise import ChannelwiseStrategy +from .registry import build_strategy + +__all__ = [ + "BoundMultiplexStrategy", + "ChannelwiseStrategy", + "MultiplexStrategy", + "PreparedMultiplexStrategy", + "build_strategy", +] diff --git a/src/raw2features/multiplex/base.py b/src/raw2features/multiplex/base.py new file mode 100644 index 0000000..b2f1c5c --- /dev/null +++ b/src/raw2features/multiplex/base.py @@ -0,0 +1,143 @@ +"""Interfaces for adapting ordinary RGB patch encoders to multiplex images. + +A multiplex strategy has two phases. ``prepare`` validates the source marker panel +and resolves every configuration choice that is shared across a cohort. ``bind`` +then attaches slide-specific inputs (for example resolved intensity percentiles). +Keeping those phases separate lets output array names remain stable across slides +while the complete per-slide production contract is still fingerprinted. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from typing import Any + +from raw2features.embedders.base import Embedder + + +class BoundMultiplexStrategy(Embedder, ABC): + """An embedder-compatible, fully resolved multiplex production contract.""" + + @property + @abstractmethod + def name(self) -> str: + """Effective feature-array key for the strategy-derived output.""" + + @property + @abstractmethod + def embedding_dim(self) -> int: + """Width of the aggregated patch representation.""" + + @property + @abstractmethod + def strategy_metadata(self) -> dict[str, Any]: + """JSON-safe, fully resolved strategy provenance.""" + + @property + @abstractmethod + def panel_metadata(self) -> dict[str, Any]: + """JSON-safe positional marker-panel provenance for the store header.""" + + @property + @abstractmethod + def transform_signature(self) -> tuple[Any, ...]: + """Identity of every strategy and slide-specific transform input.""" + + @abstractmethod + def set_panel(self, channel_names: Sequence[str | None] | None) -> dict[str, Any]: + """Verify that source marker metadata still matches the bound panel.""" + + @abstractmethod + def transform_batch(self, patches_hwc: list[Any], device: str) -> Any: + """Adapt multiplex HWC patches to a marker-aware base-model batch.""" + + @abstractmethod + def multiplex_fingerprint_payload( + self, base_output_fingerprint: Mapping[str, Any] + ) -> dict[str, Any]: + """Return the strategy block for the model-output fingerprint. + + The complete, self-validating base fingerprint is required (a digest alone + is insufficient). The runner supplies it after resolving the base encoder's + effective AMP and any load-time preprocessing contract. + """ + + +class PreparedMultiplexStrategy(ABC): + """A panel/config-bound strategy awaiting slide-specific resolved inputs.""" + + @property + @abstractmethod + def effective_model_key(self) -> str: + """Cohort-stable output key, independent of per-slide measured values.""" + + @property + @abstractmethod + def config_metadata(self) -> dict[str, Any]: + """JSON-safe strategy configuration used to derive the output key.""" + + @property + @abstractmethod + def context_max_side_px(self) -> int: + """Maximum side of the whole-slide context level requested from the reader.""" + + @property + @abstractmethod + def base_inputs_per_patch(self) -> int: + """Number of ordinary base-encoder inputs produced for one source patch.""" + + @abstractmethod + def resolve_slide_context( + self, + image_hwc: Any, + *, + level: int, + source_dtype: Any, + ) -> Any: + """Resolve slide-specific state before completion checks or array creation.""" + + @abstractmethod + def bind(self, resolved: Any) -> BoundMultiplexStrategy: + """Attach slide-specific inputs and return an embedder-compatible object.""" + + +class MultiplexStrategy(ABC): + """Pluggable strategy for using an RGB encoder on named multiplex channels. + + Future implementations can use the same seam for learned N-to-3 adapters or + marker-triplet pseudo-RGB mappings. A strategy owns its configuration and + fingerprint schema; the runner only asks it to prepare and bind. + """ + + name: str + contract_version: int + + @abstractmethod + def prepare( + self, + *, + base_embedder: Any, + channel_names: Sequence[str | None], + channel_count: int, + config: Mapping[str, Any] | Any | None = None, + ) -> PreparedMultiplexStrategy: + """Validate a source panel and resolve the cohort-level strategy config.""" + + def bind( + self, + *, + base_embedder: Any, + channel_names: Sequence[str | None], + channel_count: int, + resolved: Any, + config: Mapping[str, Any] | Any | None = None, + ) -> BoundMultiplexStrategy: + """Convenience wrapper around ``prepare(...).bind(resolved)``.""" + + return self.prepare( + base_embedder=base_embedder, + channel_names=channel_names, + channel_count=channel_count, + config=config, + ).bind(resolved) diff --git a/src/raw2features/multiplex/channelwise.py b/src/raw2features/multiplex/channelwise.py new file mode 100644 index 0000000..87e55f0 --- /dev/null +++ b/src/raw2features/multiplex/channelwise.py @@ -0,0 +1,759 @@ +"""Channel-wise adaptation of ordinary RGB encoders to named multiplex images. + +The RGB vision-encoder baselines evaluated in the public KRONOS paper provide +structural inspiration. This strategy adapts ordinary RGB encoders rather than +implementing the marker-aware KRONOS architecture. Every marker is normalized +independently, repeated to RGB, embedded independently, then aggregated with a fully +recorded deterministic recipe. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import unicodedata +from collections.abc import Mapping +from dataclasses import replace +from typing import Any + +import numpy as np + +from raw2features.core.plugins import register +from raw2features.embedders.base import ModelSpec + +from .base import BoundMultiplexStrategy, MultiplexStrategy, PreparedMultiplexStrategy + + +def _value(config: Any, name: str, default: Any) -> Any: + if config is None: + return default + if isinstance(config, dict): + return config.get(name, default) + return getattr(config, name, default) + + +def _identity(name: str) -> str: + return unicodedata.normalize("NFKC", str(name)).strip().casefold() + + +def _json_digest(value: dict[str, Any]) -> str: + encoded = json.dumps( + value, sort_keys=True, separators=(",", ":"), ensure_ascii=True + ) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def _require_safe_model_segment(name: Any) -> str: + value = str(name) + if ( + not value + or value in {".", ".."} + or any(character in value for character in ("/", "\\", "\x00")) + ): + raise ValueError(f"base model name {value!r} is not a safe Zarr path segment") + return value + + +class PreparedChannelwise(PreparedMultiplexStrategy): + @property + def context_max_side_px(self) -> int: + return self._context_max_side_px + + @property + def base_inputs_per_patch(self) -> int: + return len(self.marker_records) + + def __init__(self, base_embedder: Any, channel_names, channel_count: int, config): + if getattr(base_embedder, "modality", "brightfield") != "brightfield": + raise ValueError( + f"channelwise requires an ordinary RGB encoder; " + f"{base_embedder.name!r} is {base_embedder.modality!r}" + ) + if len(channel_names) != int(channel_count): + raise ValueError( + "multiplex marker metadata does not match the physical C axis: " + f"omero.channels has {len(channel_names)} entries but C={channel_count}" + ) + + base_name = _require_safe_model_segment(base_embedder.name) + panel = ["" if value is None else str(value) for value in channel_names] + self.channel_count = int(channel_count) + self.source_panel = panel + self.source_panel_records = [ + { + "source_index": index, + "source_name": name, + "canonical_name": unicodedata.normalize("NFKC", name).strip(), + } + for index, name in enumerate(panel) + ] + requested = list(_value(config, "multiplex_markers", []) or []) + requested = [str(value).strip() for value in requested] + if any(not value for value in requested): + raise ValueError("selected multiplex marker names must be non-empty") + + positions: dict[str, list[int]] = {} + for index, source_name in enumerate(panel): + identity = _identity(source_name) + if identity: + positions.setdefault(identity, []).append(index) + + if requested: + records = [] + seen: set[str] = set() + for requested_name in requested: + key = _identity(requested_name) + if key in seen: + raise ValueError( + f"multiplex marker {requested_name!r} was requested more " + "than once" + ) + seen.add(key) + matches = positions.get(key, []) + if not matches: + raise ValueError( + f"requested multiplex marker {requested_name!r} is not present" + ) + if len(matches) != 1: + raise ValueError( + f"requested multiplex marker {requested_name!r} is ambiguous; " + f"it matches source channels {matches}" + ) + index = matches[0] + records.append( + { + "source_index": index, + "source_name": panel[index], + "canonical_name": unicodedata.normalize( + "NFKC", panel[index] + ).strip(), + } + ) + else: + unnamed = [index for index, name in enumerate(panel) if not _identity(name)] + if unnamed: + raise ValueError( + "default marker selection requires every channel to be named; " + f"unnamed source channel indices: {unnamed}. Pass --marker " + "explicitly to exclude them." + ) + duplicate = { + key: indices for key, indices in positions.items() if len(indices) > 1 + } + if duplicate: + raise ValueError( + "default marker selection is ambiguous because normalized marker " + f"names are duplicated: {duplicate}" + ) + records = [ + { + "source_index": index, + "source_name": name, + "canonical_name": unicodedata.normalize("NFKC", name).strip(), + } + for index, name in enumerate(panel) + ] + + normalization = str( + _value(config, "multiplex_normalization", "percentile") + ).lower() + if normalization not in {"percentile", "dtype"}: + raise ValueError( + "unknown multiplex normalization " + f"{normalization!r}; available: percentile, dtype" + ) + aggregation = str(_value(config, "multiplex_aggregation", "mean")).lower() + if aggregation not in {"mean", "concat"}: + raise ValueError( + f"unknown multiplex aggregation {aggregation!r}; " + "available: mean, concat" + ) + if aggregation == "concat" and not requested: + raise ValueError( + "channelwise concat requires an explicit ordered --marker list" + ) + low = float(_value(config, "multiplex_percentile_low", 1.0)) + high = float(_value(config, "multiplex_percentile_high", 99.0)) + if not (math.isfinite(low) and math.isfinite(high) and 0 <= low < high <= 100): + raise ValueError( + "multiplex percentiles must satisfy 0 <= low < high <= 100" + ) + max_side = _value(config, "multiplex_normalization_max_side_px", 2048) + if isinstance(max_side, bool) or not isinstance(max_side, int) or max_side <= 0: + raise ValueError( + "multiplex_normalization_max_side_px must be a positive integer" + ) + strategy_params = dict(_value(config, "multiplex_strategy_params", {}) or {}) + if strategy_params: + raise ValueError( + "channelwise does not accept --multiplex-params; use its explicit " + "marker, normalization, and aggregation options" + ) + + self.base_embedder = base_embedder + self.marker_records = records + self.selected_indices = [int(r["source_index"]) for r in records] + self.normalization = normalization + self.percentile_low = low + self.percentile_high = high + self.aggregation = aggregation + self._context_max_side_px = max_side + batch_limit = _value(config, "batch_size", 256) + if ( + isinstance(batch_limit, bool) + or not isinstance(batch_limit, int) + or batch_limit <= 0 + ): + raise ValueError("channelwise requires a positive integer batch_size") + self.base_image_batch_limit = batch_limit + input_contract = str(getattr(base_embedder, "transform_input_dtype", "uint8")) + if input_contract not in {"uint8", "float32_0_1"}: + raise ValueError( + f"channelwise does not support base transform input contract " + f"{input_contract!r}; expected 'uint8' or 'float32_0_1'" + ) + self.transform_input_dtype = input_contract + # The effective key represents cohort-level configuration, not how a marker + # happened to be capitalised in one source's metadata. Exact source spelling + # and indices remain in each slide's full fingerprint/provenance below. + marker_request = { + "mode": "explicit" if requested else "all_named_source_order", + "markers": [_identity(record["canonical_name"]) for record in records], + } + quantization = ( + { + "dtype": "uint8", + "scale": 255, + "clipping": "unit_interval", + "rounding": "numpy.rint_ties_to_even", + } + if input_contract == "uint8" + else { + "dtype": "float32", + "range": "unit_interval", + "clipping": "unit_interval", + "rounding": "none", + } + ) + self._config_metadata = { + "strategy": "channelwise", + "contract_version": ChannelwiseStrategy.contract_version, + "strategy_params": strategy_params, + "base_model": base_name, + "base_pooling": base_embedder.spec.pooling, + "marker_request": marker_request, + "normalization": { + "name": normalization, + "scope": "whole_slide_downsample_level", + "nonfinite_policy": "error", + "statistics_compute_dtype": "float64", + "applied_bounds_dtype": "float32", + "patch_compute_dtype": "float32", + "level_selection": { + "rule": "first_pyramid_level_with_max_side_lte", + "max_side_px": self.context_max_side_px, + }, + **( + {"lower_percentile": low, "upper_percentile": high} + if normalization == "percentile" + else {} + ), + }, + "channel_expansion": "repeat_single_channel_to_rgb", + "execution_policy": { + "configured_batch_unit": "base_encoder_rgb_inputs", + "patch_batch_rule": ( + "max(1,floor(configured_batch_size/selected_marker_count))" + ), + "base_transform_and_forward_chunking": ( + "at_most_configured_batch_size_rgb_images" + ), + }, + "encoder_input_contract": input_contract, + "quantization": quantization, + "encoder_output": { + "pooling": base_embedder.spec.pooling, + "dimension": int(base_embedder.embedding_dim), + }, + "aggregation": {"name": aggregation}, + } + suffix = _json_digest(self._config_metadata)[:16] + self._effective_model_key = f"{base_name}__channelwise_{aggregation}_{suffix}" + + @property + def effective_model_key(self) -> str: + return self._effective_model_key + + @property + def config_metadata(self) -> dict[str, Any]: + return json.loads(json.dumps(self._config_metadata)) + + def resolve_slide_context( + self, image_hwc: np.ndarray, *, level: int, source_dtype: Any + ) -> dict[str, Any]: + image = np.asarray(image_hwc) + if image.ndim != 3 or image.shape[2] != self.channel_count: + raise ValueError( + "normalization image must be HWC with exactly the source C-axis " + f"extent ({self.channel_count}); got shape {image.shape}" + ) + dtype = np.dtype(source_dtype) + if image.dtype != dtype: + raise ValueError( + f"normalization image dtype {image.dtype} does not match declared " + f"source dtype {dtype}" + ) + resolved = [] + for record in self.marker_records: + index = int(record["source_index"]) + values = image[..., index].astype(np.float64, copy=False) + if not np.isfinite(values).all(): + raise ValueError( + f"marker {record['source_name']!r} contains non-finite pixels; " + "channelwise normalization has nonfinite_policy='error'" + ) + if self.normalization == "percentile": + low, high = np.percentile( + values, + [self.percentile_low, self.percentile_high], + method="linear", + ) + elif np.issubdtype(dtype, np.bool_): + low, high = (0.0, 1.0) + elif np.issubdtype(dtype, np.integer): + info = np.iinfo(dtype) + low, high = (float(info.min), float(info.max)) + else: + observed_low = float(values.min()) + observed_high = float(values.max()) + if observed_low < 0.0 or observed_high > 1.0: + raise ValueError( + "dtype normalization for floating-point multiplex images " + "requires finite values in [0, 1]; observed " + f"[{observed_low:g}, {observed_high:g}] for marker " + f"{record['source_name']!r}" + ) + low, high = (0.0, 1.0) + if not (math.isfinite(float(low)) and math.isfinite(float(high))): + raise ValueError( + f"marker {record['source_name']!r} normalization is not finite" + ) + if float(high) <= float(low): + raise ValueError( + f"marker {record['source_name']!r} has degenerate normalization " + f"range ({float(low):g}, {float(high):g})" + ) + # Patch normalization runs in float32. Persist the exact cast bounds that + # are applied rather than higher-precision statistics that never reach the + # encoder input. + low = float(np.float32(low)) + high = float(np.float32(high)) + if high <= low: + raise ValueError( + f"marker {record['source_name']!r} normalization range collapses " + "when represented as float32" + ) + resolved.append( + { + **record, + "low": float(low), + "high": float(high), + } + ) + return { + "name": self.normalization, + "scope": "whole_slide_downsample_level", + "level_selection": { + "rule": "first_pyramid_level_with_max_side_lte", + "max_side_px": self.context_max_side_px, + }, + "source_level": int(level), + "source_level_shape_hwc": [ + int(image.shape[0]), + int(image.shape[1]), + int(image.shape[2]), + ], + "source_dtype": dtype.str, + "nonfinite_policy": "error", + "clipping": "unit_interval", + "statistics_compute_dtype": "float64", + "applied_bounds_dtype": "float32", + "patch_compute_dtype": "float32", + "percentile_method": ( + "numpy_linear" if self.normalization == "percentile" else None + ), + "lower_percentile": ( + self.percentile_low if self.normalization == "percentile" else None + ), + "upper_percentile": ( + self.percentile_high if self.normalization == "percentile" else None + ), + "resolved": resolved, + } + + def bind(self, resolved: Any) -> BoundChannelwise: + return BoundChannelwise(self, resolved) + + +class BoundChannelwise(BoundMultiplexStrategy): + def __init__(self, prepared: PreparedChannelwise, resolved: dict[str, Any]): + self.prepared = prepared + self.base_embedder = prepared.base_embedder + self._normalization = json.loads(json.dumps(resolved)) + if self._normalization.get("name") != prepared.normalization: + raise ValueError( + "resolved normalization method does not match the prepared strategy" + ) + if self._normalization.get("nonfinite_policy") != "error": + raise ValueError( + "resolved normalization must retain nonfinite_policy='error'" + ) + expected_contract = { + "scope": "whole_slide_downsample_level", + "level_selection": { + "rule": "first_pyramid_level_with_max_side_lte", + "max_side_px": prepared.context_max_side_px, + }, + "clipping": "unit_interval", + "statistics_compute_dtype": "float64", + "applied_bounds_dtype": "float32", + "patch_compute_dtype": "float32", + } + if any( + self._normalization.get(field) != value + for field, value in expected_contract.items() + ): + raise ValueError( + "resolved normalization execution contract does not match " + "channelwise v1" + ) + level = self._normalization.get("source_level") + if isinstance(level, bool) or not isinstance(level, int) or level < 0: + raise ValueError("resolved normalization has no valid source level") + shape = self._normalization.get("source_level_shape_hwc") + if ( + not isinstance(shape, list) + or len(shape) != 3 + or any( + isinstance(value, bool) or not isinstance(value, int) for value in shape + ) + or any(value <= 0 for value in shape) + or shape[2] != prepared.channel_count + ): + raise ValueError( + "resolved normalization has no valid source-level HWC shape" + ) + if prepared.normalization == "percentile": + if ( + self._normalization.get("percentile_method") != "numpy_linear" + or self._normalization.get("lower_percentile") + != prepared.percentile_low + or self._normalization.get("upper_percentile") + != prepared.percentile_high + ): + raise ValueError( + "resolved percentile recipe does not match the prepared strategy" + ) + elif any( + self._normalization.get(field) is not None + for field in ( + "percentile_method", + "lower_percentile", + "upper_percentile", + ) + ): + raise ValueError("dtype normalization must not carry a percentile recipe") + try: + self._source_dtype = np.dtype(self._normalization["source_dtype"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError( + "resolved normalization has no valid source dtype" + ) from exc + rows = self._normalization.get("resolved", []) + if len(rows) != len(prepared.marker_records): + raise ValueError("resolved normalization does not match selected markers") + lows: list[float] = [] + highs: list[float] = [] + for expected, row in zip(prepared.marker_records, rows, strict=True): + if not isinstance(row, Mapping) or any( + row.get(field) != expected[field] + for field in ("source_index", "source_name", "canonical_name") + ): + raise ValueError( + "resolved normalization marker order/identity does not match " + f"the prepared panel at {expected!r}" + ) + try: + low = float(row["low"]) + high = float(row["high"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError( + f"resolved normalization for {expected['source_name']!r} " + "has no numeric low/high range" + ) from exc + if not (math.isfinite(low) and math.isfinite(high) and high > low): + raise ValueError( + f"resolved normalization for {expected['source_name']!r} must " + "have finite low < high" + ) + lows.append(low) + highs.append(high) + self._low = np.asarray(lows, dtype=np.float32) + self._high = np.asarray(highs, dtype=np.float32) + base_dim = int(self.base_embedder.embedding_dim) + self._embedding_dim = ( + base_dim * len(rows) if prepared.aggregation == "concat" else base_dim + ) + spec: ModelSpec = replace( + self.base_embedder.spec, + name=prepared.effective_model_key, + embedding_dim=self._embedding_dim, + modality="multiplex", + ) + super().__init__(spec) + self._device = getattr(self.base_embedder, "_device", "cpu") + self._dtype = getattr(self.base_embedder, "_dtype", None) + # The runner sets ownership after binding. Caller-injected warm embedders are + # deliberately never owned by this wrapper and must survive across slides. + self._raw2features_owns_base = False + self._base_load_attempted_by_wrapper = False + + @property + def name(self) -> str: + return self.spec.name + + @property + def modality(self) -> str: + return "multiplex" + + @property + def embedding_dim(self) -> int: + return self._embedding_dim + + @property + def transform_signature(self) -> tuple[Any, ...]: + # Never share a transformed tensor across a different panel/statistics recipe. + return ("multiplex", self.name, _json_digest(self.strategy_metadata)) + + @property + def strategy_metadata(self) -> dict[str, Any]: + input_conversion = ( + "quantize_uint8" + if self.prepared.transform_input_dtype == "uint8" + else "cast_float32_unit_interval" + ) + return { + **self.prepared.config_metadata, + "markers": json.loads(json.dumps(self.prepared.marker_records)), + "normalization": json.loads(json.dumps(self._normalization)), + "output_dimension": self.embedding_dim, + "operation_order": [ + "select_marker", + "normalize_to_unit_interval", + input_conversion, + "repeat_to_rgb", + "base_encoder_preprocessing", + "base_encoder_output", + "aggregate_markers", + ], + } + + @property + def panel_metadata(self) -> dict[str, Any]: + selected = {int(row["source_index"]) for row in self.prepared.marker_records} + return { + "strategy": "channelwise", + "physical_channel_count": self.prepared.channel_count, + "n_markers": len(selected), + "kept": [row["source_name"] for row in self.prepared.marker_records], + "excluded": json.loads( + json.dumps( + [ + row + for row in self.prepared.source_panel_records + if int(row["source_index"]) not in selected + ] + ) + ), + "source_channels": json.loads( + json.dumps(self.prepared.source_panel_records) + ), + "mapping": json.loads(json.dumps(self.prepared.marker_records)), + "normalization": json.loads(json.dumps(self._normalization)), + } + + def multiplex_fingerprint_payload( + self, base_output_fingerprint: Mapping[str, Any] + ) -> dict[str, Any]: + if not isinstance(base_output_fingerprint, Mapping): + raise ValueError( + "channelwise fingerprint requires the full base fingerprint" + ) + return { + **self.strategy_metadata, + "base_output_fingerprint": dict(base_output_fingerprint), + } + + def set_panel(self, channel_names) -> dict[str, Any]: + values = [] if channel_names is None else channel_names + panel = ["" if value is None else str(value) for value in values] + if panel != self.prepared.source_panel: + raise ValueError( + "source marker panel changed after channelwise strategy binding; " + f"expected {self.prepared.source_panel!r}, got {panel!r}" + ) + return self.panel_metadata + + def load(self, device="cuda", dtype=None, compile=False): + if self._raw2features_owns_base: + # Mark before loading so a partially failed load is still cleaned up by + # the runner's finally block. + self._base_load_attempted_by_wrapper = True + self.base_embedder.load(device=device, dtype=dtype, compile=compile) + self._device = device + self._dtype = dtype + return self + + def unload(self) -> None: + if self._raw2features_owns_base and self._base_load_attempted_by_wrapper: + self.base_embedder.unload() + self._base_load_attempted_by_wrapper = False + + def transform_batch(self, patches_hwc: list[np.ndarray], device: str): + if not patches_hwc: + raise ValueError("channelwise transform requires at least one patch") + indices = self.prepared.selected_indices + rgb: list[np.ndarray] = [] + for patch_index, patch in enumerate(patches_hwc): + patch = np.asarray(patch) + if patch.ndim != 3 or int(patch.shape[2]) != self.prepared.channel_count: + raise ValueError( + f"multiplex patch {patch_index} must be HWC with " + f"C={self.prepared.channel_count}; got shape {patch.shape}" + ) + if patch.dtype != self._source_dtype: + raise ValueError( + f"multiplex patch {patch_index} dtype {patch.dtype} does not " + f"match bound source dtype {self._source_dtype}" + ) + selected = patch[..., indices].astype(np.float32) + if not np.isfinite(selected).all(): + raise ValueError( + f"multiplex patch {patch_index} contains non-finite selected " + "marker pixels; channelwise normalization has " + "nonfinite_policy='error'" + ) + if ( + self.prepared.normalization == "dtype" + and np.issubdtype(self._source_dtype, np.floating) + and (float(selected.min()) < 0.0 or float(selected.max()) > 1.0) + ): + raise ValueError( + f"multiplex patch {patch_index} violates float dtype " + "normalization: selected marker pixels must be in [0, 1]" + ) + unit = (selected - self._low) / (self._high - self._low) + unit = np.clip(unit, 0.0, 1.0) + if self.prepared.transform_input_dtype == "uint8": + encoder_input = np.rint(unit * 255.0).astype(np.uint8) + else: + encoder_input = unit.astype(np.float32, copy=False) + for marker in range(encoder_input.shape[2]): + rgb.append(np.repeat(encoder_input[..., marker, None], 3, axis=2)) + chunks = [ + self.base_embedder.transform_batch( + rgb[start : start + self.prepared.base_image_batch_limit], device + ) + for start in range(0, len(rgb), self.prepared.base_image_batch_limit) + ] + transformed = chunks[0] + if len(chunks) > 1: + try: + import torch + + if torch.is_tensor(transformed): + transformed = torch.cat(chunks, dim=0) + else: + transformed = np.concatenate(chunks, axis=0) + except ImportError: + transformed = np.concatenate(chunks, axis=0) + if int(transformed.shape[0]) != len(rgb): + raise ValueError( + "base transform output row count does not match expanded marker " + f"images ({int(transformed.shape[0])} != {len(rgb)})" + ) + return transformed.reshape( + len(patches_hwc), len(indices), *transformed.shape[1:] + ) + + def embed_batch(self, batch): + import torch + + if getattr(batch, "ndim", 0) < 3: + raise ValueError( + "channelwise transformed batch must have batch, marker, and base " + "transform dimensions" + ) + batch_size, markers = int(batch.shape[0]), int(batch.shape[1]) + expected_markers = len(self.prepared.selected_indices) + if batch_size <= 0 or markers != expected_markers: + raise ValueError( + "channelwise transformed batch marker count does not match the " + f"bound panel ({markers} != {expected_markers})" + ) + flat = batch.reshape(batch_size * markers, *batch.shape[2:]) + embedded_parts = [ + self.base_embedder.embed_batch( + flat[start : start + self.prepared.base_image_batch_limit] + ).float() + for start in range( + 0, + int(flat.shape[0]), + self.prepared.base_image_batch_limit, + ) + ] + embedded = torch.cat(embedded_parts, dim=0) + expected_rows = batch_size * markers + expected_dim = int(self.base_embedder.embedding_dim) + if embedded.ndim != 2 or tuple(embedded.shape) != ( + expected_rows, + expected_dim, + ): + raise ValueError( + "base encoder output does not match the bound channelwise contract: " + f"expected ({expected_rows}, {expected_dim}), got " + f"{tuple(embedded.shape)}" + ) + marker_embeddings = embedded.reshape(batch_size, markers, -1) + aggregation = self.prepared.aggregation + if aggregation == "mean": + output = marker_embeddings.mean(dim=1) + elif aggregation == "concat": + output = marker_embeddings.reshape(batch_size, -1) + else: + raise RuntimeError( + f"unbound channelwise aggregation {aggregation!r}; " + "expected mean or concat" + ) + return output.float().cpu() + + +@register("multiplex_strategies", "channelwise") +class ChannelwiseStrategy(MultiplexStrategy): + name = "channelwise" + contract_version = 1 + + def prepare( + self, + *, + base_embedder, + channel_names, + channel_count: int, + config=None, + ) -> PreparedChannelwise: + if channel_names is None: + raise ValueError( + "channelwise requires positional marker names in omero.channels" + ) + return PreparedChannelwise(base_embedder, channel_names, channel_count, config) diff --git a/src/raw2features/multiplex/registry.py b/src/raw2features/multiplex/registry.py new file mode 100644 index 0000000..5522fab --- /dev/null +++ b/src/raw2features/multiplex/registry.py @@ -0,0 +1,20 @@ +"""Build multiplex strategies through the public plugin seam.""" + +from __future__ import annotations + +from raw2features.core import plugins + +from .base import MultiplexStrategy + + +def build_strategy(name: str) -> MultiplexStrategy: + """Instantiate a registered multiplex strategy by name.""" + + cls = plugins.get("multiplex_strategies", name) + strategy = cls() + if not isinstance(strategy, MultiplexStrategy): + raise TypeError( + f"multiplex strategy {name!r} returned {type(strategy).__name__}, " + "not a MultiplexStrategy" + ) + return strategy diff --git a/src/raw2features/patcher/grid.py b/src/raw2features/patcher/grid.py index 069de49..1d3be1f 100644 --- a/src/raw2features/patcher/grid.py +++ b/src/raw2features/patcher/grid.py @@ -136,12 +136,17 @@ def resample_patch(patch_hwc: np.ndarray, out_px: int) -> np.ndarray: import cv2 interp = cv2.INTER_AREA if out_px < max(h, w) else cv2.INTER_LINEAR - # cv2.resize handles at most 4 channels; multiplex stacks (C>4) resize per channel - # (the equal-size fast path above means native-MPP reads never reach here). - if patch_hwc.ndim == 3 and patch_hwc.shape[2] > 4: + # cv2.resize handles at most 4 channels and drops the trailing singleton axis + # for HWC C=1. Multiplex stacks outside the safe 2..4 range therefore resize per + # channel so both channel identity and the HWC contract survive exact-MPP reads. + if patch_hwc.ndim == 3 and patch_hwc.shape[2] not in {2, 3, 4}: return np.stack( - [cv2.resize(patch_hwc[:, :, c], (out_px, out_px), interpolation=interp) - for c in range(patch_hwc.shape[2])], + [ + cv2.resize( + patch_hwc[:, :, c], (out_px, out_px), interpolation=interp + ) + for c in range(patch_hwc.shape[2]) + ], axis=-1, ) return cv2.resize(patch_hwc, (out_px, out_px), interpolation=interp) diff --git a/src/raw2features/pipeline/receipt.py b/src/raw2features/pipeline/receipt.py index d23cedb..e454e6f 100644 --- a/src/raw2features/pipeline/receipt.py +++ b/src/raw2features/pipeline/receipt.py @@ -16,9 +16,9 @@ from raw2features.core.uris import source_uri # The embeddings-store format version, written into every store header and fed into the -# config hash. Tracks the 0.x release line while the format is alpha (one version line: -# package 0.x == format 0.x == schema 0.x). Bumps with each release that changes the -# format; the matching JSON Schema ships under schema/embeddings_store-.json +# config hash. It is independent of the package version and changes only when the +# normative store contract changes; package v0.2.0 therefore continues to write schema +# 0.1. The matching JSON Schema ships under schema/embeddings_store-.json. SCHEMA_VERSION = "0.1" @@ -297,14 +297,19 @@ def validate_store_models( expected_model_contracts: dict[str, dict] | None = None, expected_grid_models: dict[str, list[str]] | None = None, compatible_grid_hashes: dict[str, tuple[str, ...]] | None = None, + compatible_grid_segmenters: dict[str, dict[str, str]] | None = None, + allow_hashless_legacy_grids: dict[str, bool] | None = None, ) -> bool: """True iff the requested model outputs are complete in their requested grids. ``expected_grid_models`` binds each requested model set to its authoritative full ``grid_hash``. A valid same-named model in another grid must not satisfy that - requirement. ``compatible_grid_hashes`` retains the explicitly-supported legacy - AMP hashes. A hashless store is accepted only for the unambiguous case of one - requested grid and one stored grid. + requirement. ``compatible_grid_hashes`` retains explicitly-supported legacy + identities. ``compatible_grid_segmenters`` can require live segmentation evidence + for an otherwise ambiguous alias. A hashless store is accepted only for the + unambiguous one-grid case, when its per-grid policy permits it, and, when evidence + is required, a matching live header. Missing policy entries retain the historical + permissive fallback for API compatibility. When no per-grid mapping is supplied, the historical cross-grid union behaviour is retained for API compatibility. Production resume/verify callers supply the map. @@ -349,6 +354,9 @@ def validate_store_models( ) used: set[str] = set() for expected_hash, requested in expected_grid_models.items(): + segmenter_requirements = (compatible_grid_segmenters or {}).get( + expected_hash, {} + ) candidates = tuple( dict.fromkeys( ( @@ -363,15 +371,41 @@ def validate_store_models( # current grid; the current one must win instead of making the # receipt permanently ambiguous. for candidate in candidates: + required_segmenter = segmenter_requirements.get(candidate) matches = [ key for key in keys - if key not in used and stored_hashes[key] == candidate + if key not in used + and stored_hashes[key] == candidate + and ( + required_segmenter is None + or ( + dict( + root[GRIDS][key].attrs.get("raw2features", {}) + ).get("segmentation") + or {} + ).get("segmenter") + == required_segmenter + ) ] if matches: break - if not matches and sole_hashless: - matches = [keys[0]] + hashless_allowed = (allow_hashless_legacy_grids or {}).get( + expected_hash, True + ) + if not matches and sole_hashless and hashless_allowed: + required = set(segmenter_requirements.values()) + if not required or ( + len(required) == 1 + and ( + dict( + root[GRIDS][keys[0]].attrs.get("raw2features", {}) + ).get("segmentation") + or {} + ).get("segmenter") + == next(iter(required)) + ): + matches = [keys[0]] if len(matches) != 1: return False key = matches[0] @@ -422,6 +456,8 @@ def is_complete( expected_model_contracts: dict[str, dict] | None = None, expected_grid_models: dict[str, list[str]] | None = None, compatible_grid_hashes: dict[str, tuple[str, ...]] | None = None, + compatible_grid_segmenters: dict[str, dict[str, str]] | None = None, + allow_hashless_legacy_grids: dict[str, bool] | None = None, ) -> bool: """True iff a request-bound, output-validated ``complete`` receipt exists. @@ -468,4 +504,6 @@ def is_complete( expected_model_contracts=expected_model_contracts, expected_grid_models=expected_grid_models, compatible_grid_hashes=compatible_grid_hashes, + compatible_grid_segmenters=compatible_grid_segmenters, + allow_hashless_legacy_grids=allow_hashless_legacy_grids, ) diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index a063c00..600d635 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -6,6 +6,7 @@ from __future__ import annotations +import json import math import os import queue @@ -14,23 +15,30 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from dataclasses import dataclass, field, replace -from typing import ClassVar +from typing import Any, ClassVar from raw2features.benchmark.profiler import null_profiler from raw2features.core import plugins, provenance from raw2features.core.geometry import Point, Region, Size from raw2features.core.store import grid_key from raw2features.core.uris import ( + redact_metadata_credentials, + redact_metadata_uri_credentials, redact_uri_credentials, slide_id_from_source, source_uri, ) from raw2features.embedders.fingerprint import ( expected_patch_outputs, + make_output_fingerprint, patch_output_fingerprint, resolved_patch_amp, ) -from raw2features.embedders.model_registry import build_embedder, resolve_geometry +from raw2features.embedders.model_registry import ( + build_embedder, + get_spec, + resolve_geometry, +) from raw2features.patcher.grid import GridPatcher, resample_patch from raw2features.sinks.zarr_sink import ( ZarrSink, @@ -74,6 +82,30 @@ def _positive_int(value, *, field: str) -> int: return integer +def _registered_models_need_nuclear(models: list[str]) -> bool: + """Whether any known native patch model consumes a multiplex marker stack.""" + + for name in models: + try: + if get_spec(name).modality == "multiplex": + return True + except KeyError: + # Programmatic external models are resolved from their injected specs in + # resolve_run/run_slide, where that information is available. + continue + return False + + +def _models_need_nuclear(models: list[str], specs: dict | None = None) -> bool: + provided = specs or {} + for name in models: + spec = provided.get(name) + if spec is not None and getattr(spec, "modality", "brightfield") == "multiplex": + return True + missing = [name for name in models if name not in provided] + return _registered_models_need_nuclear(missing) + + @dataclass class RunConfig: """Everything needed to embed one slide. Content-affecting fields feed the @@ -144,6 +176,30 @@ class RunConfig: # GrandQC artifact model scale in µm/px ("1.0"/"1.5"/"2.0"): coarser is less # focus-sensitive; raise it for a softer-scanned cohort that over-reads as blur. qc_artifact_mpp: str = "1.5" + # Optional strategy that adapts ordinary RGB patch encoders to a named-channel + # multiplex source. Appended after every v0.1 field to preserve RunConfig's public + # positional constructor. These affect model output/receipts, never grid geometry. + multiplex_strategy: str | None = None + multiplex_markers: list[str] = field(default_factory=list) + multiplex_normalization: str = "percentile" + multiplex_percentile_low: float = 1.0 + multiplex_percentile_high: float = 99.0 + multiplex_aggregation: str = "mean" + # JSON-safe, namespaced content parameters for third-party strategy plugins. + # The built-in channelwise strategy uses its explicit typed fields instead. + multiplex_strategy_params: dict[str, Any] = field(default_factory=dict) + multiplex_normalization_max_side_px: int = 2048 + # A user-supplied complete positional panel, normally parsed from + # ``--channel-names-file``. It is an input to source binding rather than the + # persisted identity itself; ``resolved_channel_names`` below is the authoritative + # panel after validating it against the physical C axis and any existing labels. + channel_names_override: list[str] = field(default_factory=list) + resolved_channel_names: list[str] = field(default_factory=list) + resolved_nuclear_channel_indices: list[int] = field(default_factory=list) + # Original positional OME labels before a channel-names-file fills blanks. This is + # runtime migration evidence only: effective names above determine current output, + # while the original list lets us prove whether a v0.1 selector used the same index. + resolved_original_channel_names: list[str] = field(default_factory=list) def __post_init__(self) -> None: """Reject invalid geometry before hashing, model loading, or store writes.""" @@ -154,6 +210,105 @@ def __post_init__(self) -> None: self.step_px = _positive_int(self.step_px, field="step_px") if self.source_mpp is not None: self.source_mpp = _positive_float(self.source_mpp, field="source_mpp") + if self.multiplex_strategy is not None: + self.multiplex_strategy = self.multiplex_strategy.strip().lower() or None + if not self.no_seg and ( + self.multiplex_strategy is not None + or _registered_models_need_nuclear(self.models) + ): + # Execution has always routed multiplex inputs to the nuclear segmenter. + # Canonicalise the configured value too so grid identity cannot alias an + # Otsu brightfield grid whose coordinates differ. + self.segmenter = "nuclear" + self.multiplex_markers = [str(name).strip() for name in self.multiplex_markers] + self.channel_names_override = [ + str(name).strip() for name in self.channel_names_override + ] + if any(not name for name in self.channel_names_override): + raise ValueError("channel_names_override entries must be non-empty") + self.resolved_channel_names = [ + "" if name is None else str(name) for name in self.resolved_channel_names + ] + self.resolved_original_channel_names = [ + "" if name is None else str(name) + for name in self.resolved_original_channel_names + ] + self.resolved_nuclear_channel_indices = list( + self.resolved_nuclear_channel_indices + ) + if any( + isinstance(index, bool) or not isinstance(index, int) or index < 0 + for index in self.resolved_nuclear_channel_indices + ): + raise ValueError( + "resolved_nuclear_channel_indices must contain non-negative integers" + ) + if len(set(self.resolved_nuclear_channel_indices)) != len( + self.resolved_nuclear_channel_indices + ): + raise ValueError("resolved_nuclear_channel_indices must be unique") + if not isinstance(self.multiplex_strategy_params, dict) or any( + not isinstance(key, str) for key in self.multiplex_strategy_params + ): + raise ValueError( + "multiplex_strategy_params must be a mapping with string keys" + ) + try: + encoded_params = json.dumps( + self.multiplex_strategy_params, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + self.multiplex_strategy_params = json.loads(encoded_params) + except (TypeError, ValueError) as exc: + raise ValueError( + "multiplex_strategy_params must contain only finite JSON values" + ) from exc + if self.multiplex_strategy is None: + orphaned = bool( + self.multiplex_markers or self.multiplex_strategy_params + ) or any( + ( + str(self.multiplex_normalization).lower() != "percentile", + self.multiplex_percentile_low != 1.0, + self.multiplex_percentile_high != 99.0, + str(self.multiplex_aggregation).lower() != "mean", + self.multiplex_normalization_max_side_px != 2048, + ) + ) + if orphaned: + raise ValueError( + "multiplex marker/normalization/aggregation options require " + "multiplex_strategy (for example 'channelwise')" + ) + # Ordinary brightfield requests deliberately do not validate or hash the + # dormant strategy defaults. Their historical identity and execution path + # remain byte-for-byte unchanged. + return + if any(not name for name in self.multiplex_markers): + raise ValueError("multiplex marker names must be non-empty") + self.multiplex_normalization = str(self.multiplex_normalization).lower() + if self.multiplex_normalization not in {"percentile", "dtype"}: + raise ValueError( + "multiplex_normalization must be one of: percentile, dtype" + ) + self.multiplex_aggregation = str(self.multiplex_aggregation).lower() + if self.multiplex_aggregation not in {"mean", "concat"}: + raise ValueError("multiplex_aggregation must be one of: mean, concat") + low = float(self.multiplex_percentile_low) + high = float(self.multiplex_percentile_high) + if not (math.isfinite(low) and math.isfinite(high) and 0 <= low < high <= 100): + raise ValueError( + "multiplex percentiles must be finite and satisfy " + "0 <= low < high <= 100" + ) + self.multiplex_percentile_low = low + self.multiplex_percentile_high = high + self.multiplex_normalization_max_side_px = _positive_int( + self.multiplex_normalization_max_side_px, + field="multiplex_normalization_max_side_px", + ) # Single source of truth for which fields affect the output (and thus the # config hash) vs. runtime-only knobs. ``test_config_integrity`` asserts these @@ -170,6 +325,16 @@ def __post_init__(self) -> None: "tissue_threshold", "features_dtype", "stain_norm", + "multiplex_strategy", + "multiplex_markers", + "multiplex_normalization", + "multiplex_percentile_low", + "multiplex_percentile_high", + "multiplex_aggregation", + "multiplex_strategy_params", + "multiplex_normalization_max_side_px", + "resolved_channel_names", + "resolved_nuclear_channel_indices", "snap_to_level", "mpp_tolerance", "allow_upsample", @@ -191,6 +356,8 @@ def __post_init__(self) -> None: "qc", "qc_stain_norm", "qc_artifact_mpp", + "channel_names_override", + "resolved_original_channel_names", ) def _hash_payload( @@ -199,12 +366,46 @@ def _hash_payload( include_models: bool, include_amp: bool = True, models: list[str] | None = None, + segmenter_override: str | None = None, ) -> str: # Build the hashed payload from _CONTENT_FIELDS, applying the few value # transforms inline. config_hash sorts keys, so insertion order is # irrelevant -- the bytes are identical to listing the dict by hand. payload: dict = {} for name in self._CONTENT_FIELDS: + if name == "resolved_channel_names": + # The complete panel changes model outputs/receipts, but only the + # resolved nuclear position below is a geometry input. + if not self.resolved_channel_names or not include_models: + continue + if name == "resolved_nuclear_channel_indices": + if not self.resolved_nuclear_channel_indices or self.no_seg: + continue + payload["nuclear_channel_binding"] = { + "contract_version": 2, + "source_indices": self.resolved_nuclear_channel_indices, + "combination": ( + "identity" + if len(self.resolved_nuclear_channel_indices) == 1 + else "float32_mean" + ), + } + continue + if name.startswith("multiplex_"): + # Multiplex strategy settings belong to the model output/receipt, + # not the geometry. Omit defaults entirely for ordinary RGB runs so + # all v0.1 content hashes remain byte-for-byte stable. + if not include_models or self.multiplex_strategy is None: + continue + if ( + name + in { + "multiplex_percentile_low", + "multiplex_percentile_high", + } + and self.multiplex_normalization != "percentile" + ): + continue if name == "amp" and not include_amp: continue if name == "models": @@ -214,7 +415,15 @@ def _hash_payload( continue value = getattr(self, name) if name == "segmenter": - value = "none" if self.no_seg else value + value = ( + "none" + if self.no_seg + else segmenter_override + if segmenter_override is not None + else "nuclear" + if self.multiplex_strategy is not None + else value + ) elif name == "step_px": value = self.step_px or self.patch_px elif name == "source_mpp" and self.source_mpp is None: @@ -245,16 +454,167 @@ def grid_hash(self) -> str: def legacy_grid_hash(self) -> str: """Pre-fingerprint grid identity, retained solely to append to v0.1 stores.""" - return self._hash_payload(include_models=False, include_amp=True) + return replace( + self, + resolved_channel_names=[], + resolved_nuclear_channel_indices=[], + )._hash_payload(include_models=False, include_amp=True) def compatible_legacy_grid_hashes(self) -> tuple[str, ...]: - """Old grid identities, preferring this request's configured AMP first.""" + """Safe old grid identities, preferring configured AMP first. + + Native-multiplex v0.1 grids are compatible only when the old selector provably + used the same singleton physical channel. Both its nuclear-valued hash family + and its historical Otsu-valued aliases additionally require a stored header + proving nuclear segmentation. Brightfield and strategy requests retain their + existing compatibility families. + """ amps = (self.amp, "auto", "fp32", "bf16", "fp16") - return tuple( - dict.fromkeys(replace(self, amp=value).legacy_grid_hash() for value in amps) + configs = [ + replace( + self, + amp=value, + resolved_channel_names=[], + resolved_nuclear_channel_indices=[], + ) + for value in amps + ] + hashes = [config.legacy_grid_hash() for config in configs] + if self._is_native_multiplex_grid_request() and not ( + self._v01_nuclear_selector_matches_current() + ): + # A v0.1 native-multiplex grid may have thresholded a different physical + # channel after blank labels were dropped or ambiguous names were matched + # by substring. Its header still says ``nuclear``, so segmenter evidence + # cannot make this family safe without the positional-index proof. + hashes = [] + otsu_aliases = self._native_multiplex_otsu_grid_aliases() + if otsu_aliases: + # Prefer the newest no-AMP historical identity, then retain older + # AMP-bearing variants for pre-fingerprint stores. + hashes[0:0] = otsu_aliases + return tuple(dict.fromkeys(hashes)) + + def compatible_legacy_grid_segmenters(self) -> dict[str, str]: + """Segmenter evidence required before accepting ambiguous legacy hashes. + + Every accepted native-multiplex legacy family requires a live grid header + proving that nuclear segmentation produced its coordinates. The positional + selector check in :meth:`compatible_legacy_grid_hashes` remains independently + load-bearing because old nuclear headers did not record the physical index. + """ + + if not self._is_native_multiplex_grid_request() or not ( + self._v01_nuclear_selector_matches_current() + ): + return {} + nuclear_hashes = [ + replace( + self, + amp=value, + resolved_channel_names=[], + resolved_nuclear_channel_indices=[], + ).legacy_grid_hash() + for value in (self.amp, "auto", "fp32", "bf16", "fp16") + ] + return { + grid_hash: "nuclear" + for grid_hash in ( + *nuclear_hashes, + *self._native_multiplex_otsu_grid_aliases(), + ) + } + + def allows_hashless_legacy_grid(self) -> bool: + """Whether a sole grid without identity may be reused for this request. + + The historical brightfield fallback remains available. A derived multiplex + strategy cannot predate this feature, and a native multiplex request whose + old selector cannot be proven equivalent must not reuse coordinates whose + physical nuclear-channel binding is unknowable. + """ + + if self.multiplex_strategy is not None: + return False + return not self._is_native_multiplex_grid_request() or ( + self._v01_nuclear_selector_matches_current() ) + def _is_native_multiplex_grid_request(self) -> bool: + """Whether legacy reuse concerns a segmented native multiplex model.""" + + return ( + not self.no_seg + and self.multiplex_strategy is None + and ( + _registered_models_need_nuclear(self.models) + # An injected native multiplex embedder is not present in the built-in + # registry, but source resolution still binds its positional panel. + or (self.segmenter == "nuclear" and bool(self.resolved_channel_names)) + ) + ) + + def _native_multiplex_otsu_grid_aliases(self) -> tuple[str, ...]: + if ( + not self._is_native_multiplex_grid_request() + or not self._v01_nuclear_selector_matches_current() + ): + return () + amps = (self.amp, "auto", "fp32", "bf16", "fp16") + configs = [ + replace( + self, + amp=value, + resolved_channel_names=[], + resolved_nuclear_channel_indices=[], + ) + for value in amps + ] + legacy_base = replace( + self, + resolved_channel_names=[], + resolved_nuclear_channel_indices=[], + ) + hashes = [ + legacy_base._hash_payload( + include_models=False, + include_amp=False, + segmenter_override="otsu", + ) + ] + hashes.extend( + config._hash_payload( + include_models=False, + include_amp=True, + segmenter_override="otsu", + ) + for config in configs + ) + return tuple(dict.fromkeys(hashes)) + + def _v01_nuclear_selector_matches_current(self) -> bool: + """Whether v0.1 provably thresholded the same one physical channel. + + v0.1 dropped blank OME channel entries, then selected the first remaining + label containing one of four substrings. Reusing its grid is safe only when + emulating those exact rules on the original metadata yields the same singleton + physical index selected now. A DNA1/DNA2 mean or any uncertain panel gets a new + grid while the historical grid remains readable. + """ + + if ( + len(self.resolved_nuclear_channel_indices) != 1 + or not self.resolved_original_channel_names + ): + return False + legacy_names = [name for name in self.resolved_original_channel_names if name] + for legacy_index, name in enumerate(legacy_names): + lowered = name.lower() + if any(alias in lowered for alias in ("dapi", "hoechst", "hochst", "dna")): + return legacy_index == self.resolved_nuclear_channel_indices[0] + return False + def device_list(self) -> list[str]: """The devices to run on, in order. ``--devices`` overrides the single ``--device``; the default (``devices is None``) is ``[self.device]``, so the @@ -292,6 +652,388 @@ def _amp_dtype(amp: str): return getattr(torch, _AMP.get(amp, "float32")) +def _multiplex_normalization_level(reader, max_side: int = 2048) -> int: + """Highest-resolution whole-slide level bounded to a practical scan size.""" + + for level, dimensions in enumerate(reader.level_dimensions): + if max(dimensions.width, dimensions.height) <= max_side: + return level + smallest = reader.level_dimensions[-1] + raise ValueError( + "multiplex slide-level normalization requires a pyramid level whose " + f"maximum side is <= {max_side} px; the smallest available level is " + f"{smallest.width}x{smallest.height}. Add a downsampled level, or set " + "--multiplex-normalization-max-side-px to at least " + f"{max(smallest.width, smallest.height)} (which may use more RAM)." + ) + + +def resolve_multiplex_source_config( + slide_path: str, + cfg: RunConfig, + *, + model_specs: dict | None = None, +) -> RunConfig: + """Bind the physical C axis and its positional names before completion checks. + + The returned config contains the resolved panel, which feeds model-output identity + and (when nuclear masking is enabled) grid identity. A supplied override may fill + absent or blank OME metadata, but the reader rejects conflicts with existing labels. + """ + + native_multiplex = _models_need_nuclear(cfg.models, model_specs) + if ( + cfg.channel_names_override + and cfg.multiplex_strategy is None + and not native_multiplex + ): + raise ValueError( + "--channel-names-file is only valid with a native multiplex model or " + "--multiplex-strategy" + ) + if cfg.multiplex_strategy is None and not native_multiplex: + return cfg + + reader_cls = plugins.get("readers", cfg.reader) + with reader_cls(slide_path) as reader: + original_names = getattr(reader, "original_channel_names", None) + original_names = [ + "" if value is None else str(value) for value in (original_names or []) + ] + if cfg.channel_names_override: + reader.apply_channel_names(cfg.channel_names_override) + if not getattr(reader, "has_channel_axis", False): + raise ValueError( + "multiplex embedding requires a source with an explicit channel axis" + ) + names = reader.channel_names + if names is None: + raise ValueError( + "multiplex embedding requires positional channel names in " + "omero.channels or --channel-names-file" + ) + names = ["" if value is None else str(value) for value in names] + if len(names) != int(reader.channel_count): + raise ValueError( + "multiplex channel metadata does not match the physical C axis: " + f"found {len(names)} names but C={reader.channel_count}" + ) + if cfg.resolved_channel_names and names != cfg.resolved_channel_names: + raise ValueError( + "source channel names changed after multiplex request resolution; " + f"expected {cfg.resolved_channel_names!r}, got {names!r}" + ) + if ( + cfg.resolved_original_channel_names + and original_names != cfg.resolved_original_channel_names + ): + raise ValueError( + "source OME channel metadata changed after multiplex request resolution" + ) + nuclear_indices: list[int] = [] + if not cfg.no_seg: + from raw2features.segmenters.nuclear import NuclearSegmenter + + nuclear_indices = [ + index for index, _name in NuclearSegmenter().nuclear_channels(names) + ] + return replace( + cfg, + resolved_channel_names=names, + resolved_nuclear_channel_indices=nuclear_indices, + resolved_original_channel_names=original_names, + ) + + +def _validate_multiplex_slide_encoders(names: list[str]) -> None: + """Allow only model-agnostic poolers over strategy-derived patch features.""" + + from raw2features.slide_embedders.model_registry import get_slide_spec + + incompatible = [ + name for name in names if get_slide_spec(name).patch_encoder != "any" + ] + if incompatible: + raise ValueError( + "multiplex strategy outputs may be passed only to model-agnostic slide " + f"poolers (mean, max, meanmax); learned slide encoder(s) {incompatible} " + "are tied to their native registered patch encoders" + ) + + +def _resolve_multiplex_request( + slide_path: str, + cfg: RunConfig, + *, + embedders: list | None, + device: str, + load_models: bool, +) -> tuple[list, RunConfig, dict[str, dict]]: + """Resolve source-dependent strategy wrappers and their exact contracts. + + Metadata, marker selection, and whole-slide normalization are resolved before + optional registry models are loaded. ``verify`` calls the same helper with + ``load_models=False``; embedding calls it with ``True``. This keeps the effective + feature keys and per-slide fingerprints identical between both entry points. + """ + + from raw2features.multiplex import BoundMultiplexStrategy, build_strategy + + # Resolve the plugin before constructing even an unloaded base model, so an + # unknown strategy cannot trigger an expensive gated-model operation first. + strategy = build_strategy(cfg.multiplex_strategy) + if embedders is None: + base_embedders = [build_embedder(name) for name in cfg.models] + registry_owned = True + else: + by_name = {embedder.name: embedder for embedder in embedders} + if len(by_name) != len(embedders): + raise ValueError("loaded base embedders contain duplicate model names") + missing = [name for name in cfg.models if name not in by_name] + if missing: + raise ValueError( + "loaded base embedders do not match the requested multiplex models; " + f"missing {missing}" + ) + base_embedders = [by_name[name] for name in cfg.models] + registry_owned = False + + reader_cls = plugins.get("readers", cfg.reader) + with reader_cls(slide_path) as reader: + if cfg.channel_names_override: + reader.apply_channel_names(cfg.channel_names_override) + if not getattr(reader, "has_channel_axis", False): + raise ValueError( + "multiplex strategy requires a source with an explicit channel axis" + ) + names = reader.channel_names + if names is None: + raise ValueError( + "multiplex strategy requires marker identity in omero.channels" + ) + names = ["" if value is None else str(value) for value in names] + if cfg.resolved_channel_names and names != cfg.resolved_channel_names: + raise ValueError( + "source channel names changed after multiplex request resolution" + ) + prepared = [ + strategy.prepare( + base_embedder=embedder, + channel_names=names, + channel_count=reader.channel_count, + config=cfg, + ) + for embedder in base_embedders + ] + for item in prepared: + if ( + cfg.multiplex_strategy_params + and item.config_metadata.get("strategy_params") + != cfg.multiplex_strategy_params + ): + raise ValueError( + f"multiplex strategy {cfg.multiplex_strategy!r} did not bind " + "multiplex_strategy_params into its static config metadata" + ) + context_sides = {int(item.context_max_side_px) for item in prepared} + if len(context_sides) != 1: + raise ValueError( + "multiplex strategy preparations disagree on slide-context level" + ) + max_side = context_sides.pop() + expansion_factors = [item.base_inputs_per_patch for item in prepared] + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in expansion_factors + ): + raise ValueError( + "multiplex strategy base_inputs_per_patch must be a positive integer" + ) + # One native whole-image read supplies deterministic slide-level context for + # every base encoder. The chosen level and actual HWC shape are persisted. + level = _multiplex_normalization_level(reader, max_side=max_side) + dimensions = reader.level_dimensions[level] + image = reader.read_region_channels( + Region(level, Point(0, 0), Size(dimensions.width, dimensions.height)) + ) + resolved = [ + item.resolve_slide_context( + image, + level=level, + source_dtype=image.dtype, + ) + for item in prepared + ] + + wrapped = [] + contracts: dict[str, dict] = {} + for base, item, normalization in zip( + base_embedders, prepared, resolved, strict=True + ): + bound = item.bind(normalization) + if not isinstance(bound, BoundMultiplexStrategy): + raise TypeError( + f"multiplex strategy {cfg.multiplex_strategy!r} returned " + f"{type(bound).__name__}, not a BoundMultiplexStrategy" + ) + effective_name = bound.name + if ( + not isinstance(effective_name, str) + or not effective_name + or effective_name in {".", ".."} + or any(character in effective_name for character in ("/", "\\", "\x00")) + ): + raise ValueError( + f"multiplex strategy produced unsafe effective model key " + f"{effective_name!r}; keys must be single Zarr path segments" + ) + # Registry-owned bases remain unloaded through store completion checks. + # _select_embedders invokes the normal lifecycle only if writing is needed. + bound._raw2features_deferred_load = bool(load_models and registry_owned) + bound._raw2features_owns_base = bool(registry_owned) + resolved_amp = _resolve_amp(cfg, base.spec, device) + base_fingerprint = patch_output_fingerprint(base.spec, resolved_amp) + multiplex_contract = bound.multiplex_fingerprint_payload(base_fingerprint) + bound.spec = replace(bound.spec, multiplex=multiplex_contract) + contract = { + "embedding_dim": int(bound.embedding_dim), + "output_fingerprint": patch_output_fingerprint(bound.spec, resolved_amp), + } + if effective_name in contracts: + raise ValueError( + "multiplex strategy produced duplicate effective key " + f"{effective_name!r}" + ) + contracts[effective_name] = contract + wrapped.append(bound) + + # The recursive primitive must stay on the pre-existing brightfield/native- + # multiplex path. Reset every strategy knob as well as the selector so RunConfig + # cannot interpret them as orphaned options; their full values are already bound + # transitively into each effective model key and fingerprint above. + effective_cfg = replace( + cfg, + models=[embedder.name for embedder in wrapped], + segmenter="nuclear" if not cfg.no_seg else cfg.segmenter, + batch_size=max( + 1, + cfg.batch_size // max(expansion_factors), + ), + multiplex_strategy=None, + multiplex_markers=[], + multiplex_normalization="percentile", + multiplex_percentile_low=1.0, + multiplex_percentile_high=99.0, + multiplex_aggregation="mean", + multiplex_strategy_params={}, + multiplex_normalization_max_side_px=2048, + ) + return wrapped, effective_cfg, contracts + + +def resolve_multiplex_output_contracts( + slide_path: str, + cfg: RunConfig, + *, + embedders: list | None = None, + device: str | None = None, +) -> tuple[RunConfig, dict[str, dict]]: + """Read-only source resolution used by ``verify`` and external schedulers.""" + + if cfg.multiplex_strategy is None: + raise ValueError("no multiplex strategy is configured") + cfg = resolve_multiplex_source_config(slide_path, cfg) + _validate_multiplex_slide_encoders(cfg.slide_encoders) + _wrapped, effective_cfg, contracts = _resolve_multiplex_request( + slide_path, + cfg, + embedders=embedders, + device=device or cfg.device, + load_models=False, + ) + return effective_cfg, contracts + + +def _run_slide_with_multiplex_strategy( + slide_path: str, + out_dir: str, + cfg: RunConfig, + *, + receipts_dir: str | None, + cli: str | None, + embedders: list | None, + embedder_factory, + force: bool, + profiler, + allow_hashless_legacy: bool, + devices: list[str], +) -> dict: + """Bind one source-dependent strategy, then use the unchanged single-grid path. + + This deliberately sits in front of (rather than inside) the existing brightfield + orchestration. Ordinary RGB requests never enter it, preserving their fingerprints, + completion ordering, and resume semantics byte-for-byte. + """ + + if len(devices) != 1: + raise ValueError( + "multiplex strategies are currently single-device; pass one --device. " + "The marker panel and slide-level normalization must be bound identically " + "before patch workers can be sharded safely." + ) + if embedder_factory is not None and embedders is None: + # A custom model remains supported by passing its loaded embedder explicitly. + raise ValueError( + "multiplex strategies do not support an unprobed embedder_factory; pass " + "the loaded base embedder through embedders= on the single-device path" + ) + + _validate_multiplex_slide_encoders(cfg.slide_encoders) + expected_source = canonical_source_uri(slide_path) + if expected_source is None: + raise ValueError("Source URI is malformed and cannot be compared safely.") + try: + slide_id = slide_id_from_path(slide_path) + except Exception: # noqa: BLE001 - suppress credential-bearing parser errors + raise ValueError( + "Source URI is malformed and no safe output ID can be derived." + ) from None + out_path = os.path.join(out_dir, f"{slide_id}.embeddings.zarr") + if os.path.exists(out_path) and not force: + _assert_store_source(out_path, expected_source) + + wrapped, effective_cfg, contracts = _resolve_multiplex_request( + slide_path, + cfg, + embedders=embedders, + device=devices[0], + load_models=True, + ) + try: + return run_slide( + slide_path, + out_dir, + effective_cfg, + receipts_dir=receipts_dir, + cli=cli, + embedders=wrapped, + embedder_factory=None, + force=force, + profiler=profiler, + # A derived-strategy grid cannot predate this feature. Never treat an + # unidentifiable hashless brightfield store as compatible with it. + allow_hashless_legacy=False, + _expected_model_contracts=contracts, + _receipt_config_hash=cfg.content_hash(), + ) + finally: + # Registry-created bases are local to this call. Injected bases belong to + # warm-worker/library callers and their wrappers intentionally ignore unload. + for embedder in wrapped: + if getattr(embedder, "_raw2features_owns_base", False): + embedder.unload() + + def run_slide( slide_path: str, out_dir: str, @@ -305,6 +1047,8 @@ def run_slide( profiler=None, allow_hashless_legacy: bool = True, _expected_model_contracts: dict[str, dict] | None = None, + _model_specs: dict | None = None, + _receipt_config_hash: str | None = None, ) -> dict: """Embed one slide into exactly one configured grid. Returns a summary dict. @@ -340,6 +1084,16 @@ def run_slide( validate_slide_encoder_names(cfg.slide_encoders) + injected_specs = dict(_model_specs or {}) + injected_specs.update( + {embedder.name: embedder.spec for embedder in (embedders or [])} + ) + cfg = resolve_multiplex_source_config( + slide_path, + cfg, + model_specs=injected_specs, + ) + import torch # noqa: F401 - imported so a torch-less env fails clearly, here from raw2features.core.device import resolve_device @@ -347,6 +1101,22 @@ def run_slide( cfg.device = resolve_device(cfg.device) # "auto" -> cuda/mps/cpu (idempotent) devices = cfg.device_list() multi_device = len(devices) > 1 + if cfg.multiplex_strategy is not None: + return _run_slide_with_multiplex_strategy( + slide_path, + out_dir, + cfg, + receipts_dir=receipts_dir, + cli=cli, + embedders=embedders, + embedder_factory=embedder_factory, + force=force, + profiler=profiler, + allow_hashless_legacy=allow_hashless_legacy, + devices=devices, + ) + if not cfg.no_seg and _models_need_nuclear(cfg.models, injected_specs): + cfg = replace(cfg, segmenter="nuclear") expected_source = canonical_source_uri(slide_path) if expected_source is None: raise ValueError("Source URI is malformed and cannot be compared safely.") @@ -356,8 +1126,11 @@ def run_slide( raise ValueError( "Source URI is malformed and no safe output ID can be derived." ) from None - content_hash = cfg.content_hash() + content_hash = _receipt_config_hash or cfg.content_hash() grid_hash = cfg.grid_hash() + allow_hashless_for_request = ( + allow_hashless_legacy and cfg.allows_hashless_legacy_grid() + ) out_path = os.path.join(out_dir, f"{slide_id}.embeddings.zarr") expected_output = f"file://{os.path.abspath(out_path)}" store_exists = os.path.exists(out_path) and not force @@ -397,9 +1170,18 @@ def run_slide( expected_model_contracts=model_contracts, expected_grid_models={grid_hash: list(cfg.models)}, compatible_grid_hashes={grid_hash: cfg.compatible_legacy_grid_hashes()}, + compatible_grid_segmenters={ + grid_hash: cfg.compatible_legacy_grid_segmenters() + }, + allow_hashless_legacy_grids={grid_hash: allow_hashless_for_request}, ) ): - return {"slide_id": slide_id, "status": "skipped", "reason": "already complete"} + return { + "slide_id": slide_id, + "status": "skipped", + "reason": "already complete", + "requested_models": list(cfg.models), + } present_valid: list[str] = [] grid_key_existing: str | None = None @@ -410,7 +1192,8 @@ def run_slide( cfg.models, expected_model_contracts=model_contracts, compatible_grid_hashes=cfg.compatible_legacy_grid_hashes(), - allow_hashless_legacy=allow_hashless_legacy, + compatible_grid_segmenters=cfg.compatible_legacy_grid_segmenters(), + allow_hashless_legacy=allow_hashless_for_request, require_mask=not cfg.no_seg, ) # append: a grid of THIS geometry already exists -> add the missing models to it. @@ -477,6 +1260,7 @@ def run_slide( "status": "skipped", "reason": "all requested models and auxiliary outputs already present", "models_present": present_valid, + "requested_models": list(cfg.models), "grid": grid_key_existing, "geojson": geojson_path, "thumbnail": thumb_meta, @@ -513,8 +1297,10 @@ def embedder_factory(device: str) -> list: ) ) run_contracts = {name: model_contracts[name] for name in models_to_do} - _assert_loaded_model_contracts(run_embedders, run_contracts) - # Modality of this run: multiplex (marker stacks, e.g. KRONOS) routes the + _assert_loaded_model_contracts( + run_embedders, run_contracts, cfg.resolved_channel_names + ) + # Modality of this run: a marker stack routes the # nuclear segmenter + N-channel reads; brightfield is the RGB default. multiplex = any( getattr(e, "modality", "brightfield") == "multiplex" @@ -525,12 +1311,21 @@ def embedder_factory(device: str) -> list: # native channels, so multiplex on >1 device would embed wrong data. # Reject it cleanly rather than produce silently-corrupt features. raise ValueError( - "multiplex models (e.g. kronos) are not supported on the " + "multiplex models are not supported on the " "multi-device path; run with a single --device (the panel is " "bound per-slide and channels are read natively only there)." ) reader_cls = plugins.get("readers", cfg.reader) with reader_cls(slide_path) as reader: + if cfg.channel_names_override: + reader.apply_channel_names(cfg.channel_names_override) + if cfg.resolved_channel_names and ( + list(reader.channel_names or []) != cfg.resolved_channel_names + ): + raise ValueError( + "source channel names changed after multiplex request " + "resolution" + ) _warn_channel_collapse(reader, multiplex) # An uncalibrated source (no axis unit -> reader.mpp is None) needs an # explicit physical pixel size; apply --source-mpp when given, else the @@ -547,7 +1342,9 @@ def embedder_factory(device: str) -> list: panel_meta: dict[str, dict] = {} if multiplex: for e in run_embedders: - panel_meta[e.name] = e.set_panel(reader.channel_names) + panel_meta[e.name] = redact_metadata_credentials( + e.set_panel(reader.channel_names) + ) if append: ( n, @@ -563,6 +1360,10 @@ def embedder_factory(device: str) -> list: key=grid_key_existing, new_model_dims={e.name: e.embedding_dim for e in run_embedders}, new_model_meta=_models_header(run_embedders, run_contracts), + new_panel_meta=panel_meta, + new_source_meta=( + _source_channel_metadata(reader) if multiplex else None + ), replace_models=models_to_do, ) else: @@ -763,6 +1564,16 @@ def embedder_factory(device: str) -> list: if qc_to_do or thumbnail_to_do: reader_cls = plugins.get("readers", cfg.reader) with reader_cls(slide_path) as reader: + if cfg.channel_names_override: + reader.apply_channel_names(cfg.channel_names_override) + if cfg.resolved_channel_names and ( + list(reader.channel_names or []) + != cfg.resolved_channel_names + ): + raise ValueError( + "source channel names changed after multiplex request " + "resolution" + ) if ( reader.mpp is None and cfg.source_mpp is not None @@ -823,7 +1634,11 @@ def embedder_factory(device: str) -> list: if encoders_here: with prof.stage("slide_encode"): slide_results = _run_slide_encoders( - sink, encoders_here, cfg.device, available + sink, + encoders_here, + cfg.device, + available, + preferred_patch_models=list(cfg.models), ) final_dims = sink.feature_dims() @@ -832,6 +1647,12 @@ def embedder_factory(device: str) -> list: elapsed = time.time() - started if receipts_dir: + receipt_models = ( + sorted(cfg.models) + if _receipt_config_hash is not None + else sorted(final_dims) + ) + receipt_dims = {model: final_dims[model] for model in receipt_models} write_receipt( receipts_dir, Receipt( @@ -840,10 +1661,14 @@ def embedder_factory(device: str) -> list: source_uri=source_uri(slide_path), output_uri=sink.uri, reader=cfg.reader, - models=sorted(final_dims), - config_hash=cfg.content_hash(sorted(final_dims)), + models=receipt_models, + config_hash=( + _receipt_config_hash + if _receipt_config_hash is not None + else cfg.content_hash(receipt_models) + ), n_patches=n, - model_dims=final_dims, + model_dims=receipt_dims, started_utc=provenance.now_utc_iso(), finished_utc=provenance.now_utc_iso(), elapsed_s=round(elapsed, 2), @@ -858,6 +1683,7 @@ def embedder_factory(device: str) -> list: "models": sorted(final_dims), "models_added": models_to_do, "models_skipped": present_valid, + "requested_models": list(cfg.models), "grid": actual_grid_key, "slide_embeddings": slide_results, "output_uri": sink.uri, @@ -943,10 +1769,19 @@ def resolve_run( ) if not groups: raise ValueError("at least one model/extraction is required") - group_cfgs = [ - replace(cfg, models=list(g.models), target_mpp=g.mpp, patch_px=g.patch_px) - for g in groups - ] + group_cfgs = [] + for group in groups: + group_cfg = replace( + cfg, + models=list(group.models), + target_mpp=group.mpp, + patch_px=group.patch_px, + ) + if not group_cfg.no_seg and _models_need_nuclear( + list(group.models), model_specs + ): + group_cfg = replace(group_cfg, segmenter="nuclear") + group_cfgs.append(group_cfg) return groups, group_cfgs, _run_content_hash(group_cfgs) @@ -1009,14 +1844,20 @@ def embed_slide( "Source URI is malformed and no safe output ID can be derived." ) from None model_specs = {embedder.name: embedder.spec for embedder in (embedders or [])} + multiplex_request = cfg.multiplex_strategy is not None probed_contracts: dict[str, dict] | None = None - if not model_specs and embedder_factory is not None: + if not multiplex_request and not model_specs and embedder_factory is not None: # A factory can be the only source for a programmatic external model. Probe it # once before geometry resolution, then pass the resulting contracts into each # run_slide call so the same loaded copy is not probed a second time. model_specs, probed_contracts = _probe_factory_contracts( cfg, embedder_factory, devices ) + cfg = resolve_multiplex_source_config( + slide_path, + cfg, + model_specs=model_specs, + ) groups, group_cfgs, run_hash = resolve_run( cfg, requested_mpp, @@ -1027,12 +1868,20 @@ def embed_slide( out_path = os.path.join(out_dir, f"{slide_id}.embeddings.zarr") grids = {grid_key(g.mpp, g.patch_px): list(g.models) for g in groups} expected_output = f"file://{os.path.abspath(out_path)}" - model_contracts = probed_contracts or _expected_contracts_with_factory_probe( - cfg, - embedders, - embedder_factory, - devices=devices, - ) + model_contracts = None + if not multiplex_request: + model_contracts = probed_contracts or _expected_contracts_with_factory_probe( + cfg, + embedders, + embedder_factory, + devices=devices, + ) + model_contracts = _bind_native_multiplex_panel_contracts( + model_contracts, + cfg.models, + model_specs, + cfg.resolved_channel_names, + ) expected_grid_models = { group_cfg.grid_hash(): list(group_cfg.models) for group_cfg in group_cfgs } @@ -1040,12 +1889,21 @@ def embed_slide( group_cfg.grid_hash(): group_cfg.compatible_legacy_grid_hashes() for group_cfg in group_cfgs } + compatible_grid_segmenters = { + group_cfg.grid_hash(): group_cfg.compatible_legacy_grid_segmenters() + for group_cfg in group_cfgs + } + allow_hashless_legacy_grids = { + group_cfg.grid_hash(): group_cfg.allows_hashless_legacy_grid() + for group_cfg in group_cfgs + } if os.path.exists(out_path) and not force: _assert_store_source(out_path, expected_source) if ( - not force + not multiplex_request + and not force and receipts_dir and not _runtime_aux_requested(cfg) and is_complete( @@ -1057,9 +1915,18 @@ def embed_slide( expected_model_contracts=model_contracts, expected_grid_models=expected_grid_models, compatible_grid_hashes=compatible_grid_hashes, + compatible_grid_segmenters=compatible_grid_segmenters, + allow_hashless_legacy_grids=allow_hashless_legacy_grids, ) ): - grids = _stored_grid_summary(out_path, group_cfgs) or grids + grids = ( + _stored_grid_summary( + out_path, + group_cfgs, + allow_hashless_legacy_grids=allow_hashless_legacy_grids, + ) + or grids + ) return { "slide_id": slide_id, "status": "skipped", @@ -1085,13 +1952,22 @@ def embed_slide( # requested geometry, but cannot identify which of several grids it # represents. A multi-grid request therefore leaves it untouched. allow_hashless_legacy=(len(group_cfgs) == 1), - _expected_model_contracts={ - model: model_contracts[model] for model in gc.models + _expected_model_contracts=( + None + if model_contracts is None + else {model: model_contracts[model] for model in gc.models} + ), + _model_specs={ + model: model_specs[model] + for model in gc.models + if model in model_specs }, ) ) grids = { - (result.get("grid") or grid_key(group.mpp, group.patch_px)): list(group.models) + (result.get("grid") or grid_key(group.mpp, group.patch_px)): list( + result.get("requested_models", group.models) + ) for group, result in zip(groups, results, strict=True) } status = "skipped" if all(r["status"] == "skipped" for r in results) else "complete" @@ -1184,12 +2060,54 @@ def expected_model_contracts( """ specs = {e.name: e.spec for e in (embedders or [])} - return expected_patch_outputs( + contracts = expected_patch_outputs( cfg.models, cfg.amp, cfg.device if device is None else device, specs=specs, ) + return _bind_native_multiplex_panel_contracts( + contracts, + cfg.models, + specs, + cfg.resolved_channel_names, + ) + + +def _bind_native_multiplex_panel_contracts( + contracts: dict[str, dict], + models: list[str], + specs: dict, + channel_names: list[str], +) -> dict[str, dict]: + """Bind a positional source panel into every native multiplex fingerprint. + + Strategy-derived multiplex specs already carry their selected-panel contract in + ``spec.multiplex``. Native models otherwise have a static weights/loader + fingerprint, even though their output changes when channel identity changes; this + transitive wrapper closes that completion gap before a receipt or array is accepted. + """ + + if not channel_names: + return dict(contracts) + bound = {name: dict(contract) for name, contract in contracts.items()} + for name in models: + spec = specs.get(name) or get_spec(name) + if ( + getattr(spec, "modality", "brightfield") != "multiplex" + or getattr(spec, "multiplex", None) is not None + ): + continue + fingerprint = bound[name]["output_fingerprint"] + payload = json.loads(json.dumps(fingerprint["payload"])) + payload["multiplex_panel"] = { + "binding_contract_version": 1, + "channel_axis": "c", + "physical_channel_count": len(channel_names), + "effective_channel_names": list(channel_names), + } + bound[name]["output_fingerprint"] = make_output_fingerprint(payload) + return bound def _expected_contracts_for_devices( @@ -1264,7 +2182,7 @@ def _probe_factory_contracts( probe = factory(worker_devices[0]) try: contracts = _expected_contracts_for_devices(cfg, probe, worker_devices) - _assert_loaded_model_contracts(probe, contracts) + _assert_loaded_model_contracts(probe, contracts, cfg.resolved_channel_names) specs = {embedder.name: embedder.spec for embedder in probe} return specs, contracts finally: @@ -1285,7 +2203,9 @@ def _amp_label(dtype) -> str: }.get(dtype, "fp32") -def _loaded_model_contracts(embedders: list) -> dict[str, dict]: +def _loaded_model_contracts( + embedders: list, channel_names: list[str] | None = None +) -> dict[str, dict]: """Derive contracts from loaded specs and their effective execution precision.""" contracts: dict[str, dict] = {} @@ -1300,13 +2220,22 @@ def _loaded_model_contracts(embedders: list) -> dict[str, dict]: "embedding_dim": int(embedder.embedding_dim), "output_fingerprint": patch_output_fingerprint(embedder.spec, effective), } - return contracts + return _bind_native_multiplex_panel_contracts( + contracts, + [embedder.name for embedder in embedders], + {embedder.name: embedder.spec for embedder in embedders}, + list(channel_names or []), + ) -def _assert_loaded_model_contracts(embedders: list, expected: dict[str, dict]) -> None: +def _assert_loaded_model_contracts( + embedders: list, + expected: dict[str, dict], + channel_names: list[str] | None = None, +) -> None: """Fail before store mutation if loaded model copies differ from provenance.""" - actual = _loaded_model_contracts(embedders) + actual = _loaded_model_contracts(embedders, channel_names) if set(actual) != set(expected): missing = sorted(set(expected) - set(actual)) extra = sorted(set(actual) - set(expected)) @@ -1382,7 +2311,7 @@ def _warn_channel_collapse(reader, multiplex: bool) -> None: """Warn when a brightfield run silently drops channels of a multi-channel source. The RGB read path keeps the first three channels, so feeding a brightfield model a - source whose ``omero`` lists more than three channels (a multiplex/CODEX stack) + source whose ``omero`` lists more than three channels (a multiplex stack) discards the rest with no diagnostic - exactly the kind of garbage-in the embeddings would hide. The reverse (a multiplex model on a source with no marker panel) already fails loudly when the panel is bound, so it needs no guard here. @@ -1396,7 +2325,8 @@ def _warn_channel_collapse(reader, multiplex: bool) -> None: warnings.warn( f"source has {len(names)} channels ({', '.join(names[:4])}…) but a " "brightfield (RGB) model keeps only the first three; the rest are dropped. " - "If this is a multiplex slide, run a multiplex model (e.g. kronos).", + "If this is a multiplex slide, use --multiplex-strategy or a native " + "multiplex model.", stacklevel=2, ) @@ -1415,7 +2345,21 @@ def _select_embedders( can mix precisions across models. """ if embedders is not None: - return [e for e in embedders if e.name in names] + selected = [e for e in embedders if e.name in names] + for embedder in selected: + if not getattr(embedder, "_raw2features_deferred_load", False): + continue + base = getattr(embedder, "base_embedder", embedder) + _warn_scale_mismatch(cfg, base.spec) + target_device = device or cfg.device + dtype = _amp_dtype(_resolve_amp(cfg, base.spec, target_device)) + embedder.load( + device=target_device, + dtype=dtype, + compile=cfg.compile, + ) + embedder._raw2features_deferred_load = False + return selected return _build_embedders_on(cfg, names, device or cfg.device) @@ -1430,12 +2374,27 @@ def _segment(reader, cfg: RunConfig, segmenter_name: str | None = None): name = segmenter_name or cfg.segmenter seg = plugins.get("segmenters", name)() tissue = seg.segment(reader) - return tissue, { + metadata = { "segmenter": name, "tissue_threshold": cfg.tissue_threshold, "seg_level": tissue.level, "seg_downsample": tissue.downsample, } + if name == "nuclear" and hasattr(seg, "nuclear_channels"): + channels = seg.nuclear_channels(reader.channel_names) + indices = [int(index) for index, _channel_name in channels] + names = [str(channel_name) for _index, channel_name in channels] + metadata["nuclear_channel_indices"] = indices + metadata["nuclear_channel_names"] = names + metadata["nuclear_channel_combination"] = ( + "identity" if len(indices) == 1 else "float32_mean" + ) + metadata["channel_binding_contract_version"] = 2 + if len(indices) == 1: + # Retain the convenient singular fields for existing readers. + metadata["nuclear_channel_index"] = indices[0] + metadata["nuclear_channel_name"] = names[0] + return tissue, metadata def _decode_one(reader, x, y, read_level, read_px, patch_px, multichannel=False): @@ -1889,7 +2848,9 @@ def _worker(idx: int, device: str, lo: int, hi: int) -> None: try: embedders = embedder_factory(device) if expected_contracts is not None: - _assert_loaded_model_contracts(embedders, expected_contracts) + _assert_loaded_model_contracts( + embedders, expected_contracts, cfg.resolved_channel_names + ) collector = _FeatureCollector(hi - lo, model_dims, cfg.features_dtype) # Own reader per worker: per-reader chunk cache, thread-safe zarr reads. with reader_cls(slide_path) as reader: @@ -1945,6 +2906,7 @@ def _inspect_store( *, expected_model_contracts: dict[str, dict] | None = None, compatible_grid_hashes: tuple[str, ...] = (), + compatible_grid_segmenters: dict[str, str] | None = None, allow_hashless_legacy: bool = True, require_mask: bool = False, ): @@ -2008,18 +2970,45 @@ def _result_for(k: str): valid.append(model) return (k, n, valid) + def _recorded_segmenter(k: str) -> str | None: + try: + header = dict(root[GRIDS][k].attrs.get("raw2features", {})) + segmentation = header.get("segmentation") or {} + name = segmentation.get("segmenter") + return str(name) if name is not None else None + except Exception: # noqa: BLE001 - missing evidence cannot satisfy a guard + return None + # Prefer the new identity, then the exact requested legacy AMP, then other legal - # legacy AMP values. A hashless hand-built/legacy grid is the last-resort match. + # legacy AMP values. Ambiguous aliases additionally require live segmenter + # evidence. A hashless hand-built/legacy grid is the last-resort match. + requirements = compatible_grid_segmenters or {} candidates = tuple(dict.fromkeys((expected_grid_hash, *compatible_grid_hashes))) for candidate in candidates: + required_segmenter = requirements.get(candidate) for k in keys: stored = dict(root[GRIDS][k].attrs.get("raw2features", {})).get("grid_hash") - if stored == candidate and (result := _result_for(k)) is not None: + if ( + stored == candidate + and ( + required_segmenter is None + or _recorded_segmenter(k) == required_segmenter + ) + and (result := _result_for(k)) is not None + ): return result if allow_hashless_legacy and len(keys) == 1: k = keys[0] stored = dict(root[GRIDS][k].attrs.get("raw2features", {})).get("grid_hash") - if stored is None and (result := _result_for(k)) is not None: + required = set(requirements.values()) + evidence_matches = not required or ( + len(required) == 1 and _recorded_segmenter(k) == next(iter(required)) + ) + if ( + stored is None + and evidence_matches + and (result := _result_for(k)) is not None + ): return result return (None, 0, []) @@ -2320,7 +3309,10 @@ def _write_grid_thumbnail( def _stored_grid_summary( - out_path: str, group_cfgs: list[RunConfig] + out_path: str, + group_cfgs: list[RunConfig], + *, + allow_hashless_legacy_grids: dict[str, bool] | None = None, ) -> dict[str, list[str]]: """Resolve requested grid hashes to their actual (possibly suffixed) labels.""" @@ -2339,10 +3331,25 @@ def _stored_grid_summary( used: set[str] = set() for cfg in group_cfgs: candidates = (cfg.grid_hash(), *cfg.compatible_legacy_grid_hashes()) + requirements = cfg.compatible_legacy_grid_segmenters() match = None for candidate in dict.fromkeys(candidates): + required_segmenter = requirements.get(candidate) hits = [ - key for key in keys if key not in used and stored[key] == candidate + key + for key in keys + if key not in used + and stored[key] == candidate + and ( + required_segmenter is None + or ( + dict(root[GRIDS][key].attrs.get("raw2features", {})).get( + "segmentation" + ) + or {} + ).get("segmenter") + == required_segmenter + ) ] if len(hits) == 1: match = hits[0] @@ -2354,8 +3361,19 @@ def _stored_grid_summary( and len(group_cfgs) == 1 and len(keys) == 1 and stored[keys[0]] is None + and (allow_hashless_legacy_grids or {}).get(cfg.grid_hash(), True) ): - match = keys[0] + required = set(requirements.values()) + recorded = ( + dict(root[GRIDS][keys[0]].attrs.get("raw2features", {})).get( + "segmentation" + ) + or {} + ).get("segmenter") + if not required or ( + len(required) == 1 and recorded == next(iter(required)) + ): + match = keys[0] if match is None: return {} used.add(match) @@ -2433,6 +3451,8 @@ def _run_slide_encoders( slide_encoder_names: list[str], device: str, available_patch_models: list[str], + *, + preferred_patch_models: list[str] | None = None, ) -> dict[str, str]: """Run slide-level encoders on patch features already written to *sink*. @@ -2444,23 +3464,36 @@ def _run_slide_encoders( encode_slide_embedding, resolve_slide_patch_model, slide_embedding_is_complete, + slide_output_key, ) results: dict[str, str] = {} for slide_model_name in slide_encoder_names: + from raw2features.slide_embedders.model_registry import get_slide_spec + + preferred = list(preferred_patch_models or []) + explicit_patch_model = ( + preferred[0] + if get_slide_spec(slide_model_name).patch_encoder == "any" + and len(preferred) == 1 + else None + ) patch_model = resolve_slide_patch_model( sink._group, slide_model_name, + patch_model=explicit_patch_model, available_patch_models=available_patch_models, ) + output_name = slide_output_key(sink._group, slide_model_name, patch_model) if slide_embedding_is_complete( sink._group, slide_model_name, patch_model=patch_model, device=device, + output_name=output_name, ): - results[slide_model_name] = f"slide/{slide_model_name}" + results[slide_model_name] = f"slide/{output_name}" continue encoding = encode_slide_embedding( @@ -2476,8 +3509,9 @@ def _run_slide_encoders( slide_model_name, encoding.vector, encoding.provenance, + output_name=output_name, ) - results[slide_model_name] = f"slide/{slide_model_name}" + results[slide_model_name] = f"slide/{output_name}" return results @@ -2499,6 +3533,7 @@ def _run_slide_encoders_from_store( encode_slide_embedding, resolve_slide_grid, slide_embedding_is_complete, + slide_output_key, write_slide_embedding, ) @@ -2517,15 +3552,15 @@ def _run_slide_encoders_from_store( # compatible pre-existing grid was found; model-load/encode errors below # still propagate rather than being mistaken for a missing grid. continue + output_name = slide_output_key(group, slide_model_name, patch_model) if slide_embedding_is_complete( group, slide_model_name, patch_model=patch_model, device=device, + output_name=output_name, ): - results[slide_model_name] = ( - f"grids/{selected_grid}/slide/{slide_model_name}" - ) + results[slide_model_name] = f"grids/{selected_grid}/slide/{output_name}" continue encoding = encode_slide_embedding( group, @@ -2540,9 +3575,10 @@ def _run_slide_encoders_from_store( slide_model_name, encoding.vector, encoding.provenance, + output_name=output_name, ) wrote = True - results[slide_model_name] = f"grids/{selected_grid}/slide/{slide_model_name}" + results[slide_model_name] = f"grids/{selected_grid}/slide/{output_name}" if wrote: try: @@ -2566,7 +3602,7 @@ def _models_header( } for e in embedders } - return { + headers = { e.name: { "source": e.spec.source, "embedding_dim": e.spec.embedding_dim, @@ -2589,6 +3625,23 @@ def _models_header( } for e in embedders } + for e in embedders: + if e.spec.multiplex is None: + continue + # Mirror only the already credential-sanitized canonical fingerprint payload; + # plugin-supplied in-memory metadata is untrusted persistence input. + fingerprint = contracts[e.name]["output_fingerprint"] + multiplex = dict(fingerprint["payload"]["multiplex"]) + headers[e.name]["base_model"] = multiplex.get("base_model") + headers[e.name]["multiplex"] = multiplex + headers[e.name]["pooling"] = multiplex.get("aggregation", {}).get("name") + # Programmatic embedders may supply private model URLs. Sanitize URI values without + # interpreting semantic constructor keys such as ``token`` as authentication fields; + # those keys must remain intact for complete output identity. + return { + name: redact_metadata_uri_credentials(metadata) + for name, metadata in headers.items() + } def _build_header( @@ -2645,9 +3698,26 @@ def _build_header( "segmentation": seg_meta, "models": _models_header(embedders, model_contracts), } + header["source"].update(_source_channel_metadata(reader)) # Multiplex marker-panel resolution (per model): which of the slide's channels # matched the model's marker vocabulary and which were dropped. Absent for # brightfield runs. Part of the reproducibility record for a multiplex embedding. if panel_meta: header["panel"] = panel_meta return header + + +def _source_channel_metadata(reader) -> dict: + """Effective positional channel provenance for a named-channel source.""" + + channel_names = getattr(reader, "channel_names", None) + if channel_names is None: + return {} + metadata = { + "channel_names": list(channel_names), + "channel_names_source": getattr(reader, "channel_names_source", None), + } + original_names = getattr(reader, "original_channel_names", None) + if original_names != channel_names: + metadata["omero_channel_names"] = original_names + return metadata diff --git a/src/raw2features/readers/base.py b/src/raw2features/readers/base.py index a1384da..77a3d05 100644 --- a/src/raw2features/readers/base.py +++ b/src/raw2features/readers/base.py @@ -83,14 +83,57 @@ def read_region(self, region: Region) -> np.ndarray: # -- multi-channel / multiplex (optional; brightfield readers need not) -- @property def channel_names(self) -> list[str] | None: - """Channel/marker names for multiplex sources (e.g. CODEX), else ``None``. + """Positional names for named-channel multiplex sources, else ``None``. ``None`` means a plain RGB brightfield source (H&E) - the default. A multiplex reader returns one name per channel (the marker panel), which the - multiplex embedders (e.g. KRONOS) map to per-marker normalisation + IDs. + multiplex strategies and native multiplex embedders bind to physical channels. + Readers preserve unnamed positional metadata as ``""`` rather than removing + the slot and shifting every later marker away from its physical channel. """ return None + @property + def has_channel_axis(self) -> bool: + """Whether the source exposes a positional channel axis. + + Brightfield readers without a native multi-channel representation retain the + default ``False``. Multi-channel readers override this alongside + :attr:`channel_count` and :meth:`read_region_channels`. + """ + + return False + + @property + def channel_count(self) -> int: + """Physical channel count, independent of optional channel-name metadata.""" + + return 1 + + @property + def channel_names_source(self) -> str | None: + """Where :attr:`channel_names` came from, when the reader can report it.""" + + return None + + @property + def original_channel_names(self) -> list[str] | None: + """Unmodified positional names read from source metadata, when available.""" + + return self.channel_names + + def apply_channel_names(self, channel_names: list[str]) -> None: + """Apply a complete positional channel-name override. + + Multi-channel readers may implement this for sources whose native metadata is + absent or incomplete. The default fails explicitly so an override can never be + accepted while the reader silently continues with another panel. + """ + + raise NotImplementedError( + f"{type(self).__name__} does not support channel-name overrides" + ) + def read_region_channels(self, region: Region) -> np.ndarray: """Native multi-channel read -> ``(h, w, C)`` in the source dtype. diff --git a/src/raw2features/readers/omezarr.py b/src/raw2features/readers/omezarr.py index ae0c0fe..a674c3e 100644 --- a/src/raw2features/readers/omezarr.py +++ b/src/raw2features/readers/omezarr.py @@ -156,6 +156,8 @@ def __init__(self, path: str) -> None: self._ngff_version: str | None = None self._ome_version: str | None = None # NGFF 0.5+ keeps version in `ome` self._channel_names: list[str] | None = None + self._original_channel_names: list[str] | None = None + self._channel_names_source: str | None = None # Self-description of the source coordinate frame (captured as plain VALUES # so a coordinate-systems emitter is a re-encode, not a # re-extraction). axes = the NGFF axis order; axis_units = per-axis unit @@ -265,6 +267,12 @@ def open(self) -> OmeZarrReader: ] self._float_to_uint8_scale = self._infer_float_to_uint8_scale(self._arrays[0]) self._channel_names = self._read_omero_channels(ms_path) + self._original_channel_names = ( + list(self._channel_names) if self._channel_names is not None else None + ) + self._channel_names_source = ( + "ome_zarr_omero_channels" if self._channel_names is not None else None + ) self._warn_plane_collapse(self._arrays[0].shape) # type: ignore[attr-defined] self._chunk_cache.clear() self._open = True @@ -366,6 +374,95 @@ def apply_source_mpp(self, mpp_level0: float) -> None: def channel_names(self) -> list[str] | None: return list(self._channel_names) if self._channel_names else None + @property + def channel_names_source(self) -> str | None: + return self._channel_names_source + + @property + def original_channel_names(self) -> list[str] | None: + return ( + list(self._original_channel_names) + if self._original_channel_names is not None + else None + ) + + def apply_channel_names(self, channel_names: list[str]) -> None: + """Fill missing positional names without silently relabelling a channel. + + The supplied panel must cover the complete physical C axis. Existing non-empty + ``omero.channels`` labels remain authoritative and must agree after Unicode, + whitespace, and case normalization; absent/blank entries may be supplied. + """ + + import unicodedata + + if not self._open: + raise RuntimeError( + "reader is not open; use `with OmeZarrReader(path) as r:`" + ) + if not self.has_channel_axis: + raise ValueError( + "--channel-names-file requires an OME-Zarr source with a C axis" + ) + names = [str(value).strip() for value in channel_names] + if len(names) != self.channel_count: + raise ValueError( + "channel-name override must contain exactly one name per physical " + f"channel: file has {len(names)} names but C={self.channel_count}" + ) + if any(not name for name in names): + raise ValueError("channel-name override contains an empty channel name") + + declared = list(self._channel_names or []) + if len(declared) > self.channel_count: + raise ValueError( + "OME-Zarr omero.channels has more entries than the physical C axis " + f"({len(declared)} > {self.channel_count})" + ) + + def identity(value: str) -> str: + return unicodedata.normalize("NFKC", value).strip().casefold() + + identities = [identity(name) for name in names] + if len(set(identities)) != len(identities): + raise ValueError( + "channel-name override must use a unique name for every physical " + "channel (disambiguate repeated assays, for example DAPI_cycle1)" + ) + for index, existing in enumerate(declared): + if identity(existing) and identity(existing) != identity(names[index]): + raise ValueError( + "channel-name override conflicts with OME-Zarr metadata at " + f"C index {index}: {existing!r} != {names[index]!r}" + ) + self._channel_names = names + self._channel_names_source = "channel_names_file" + + @property + def has_channel_axis(self) -> bool: + """Whether the source declares a positional NGFF channel axis.""" + + return "c" in self._dims + + @property + def channel_count(self) -> int: + """Number of source channels, independent of ``omero`` label metadata. + + Multiplex strategy binding compares this physical C-axis extent with the + positional ``channel_names`` metadata before selecting markers. Keeping the + values separate makes an incomplete ``omero.channels`` block detectable + instead of silently treating its metadata length as the image channel count. + A source without a C axis is a single-channel plane. + """ + + if "c" not in self._dims: + return 1 + if not self._arrays: + raise RuntimeError( + "reader is not open; use `with OmeZarrReader(path) as r:`" + ) + return int(self._arrays[0].shape[self._dims.index("c")]) # type: ignore[attr-defined] + @property def ngff_version(self) -> str | None: return self._ngff_version @@ -447,7 +544,15 @@ def read_region_channels(self, region: Region) -> np.ndarray: return np.ascontiguousarray(block) def _read_omero_channels(self, ms_path: str) -> list[str] | None: - """Marker/channel names from NGFF ``omero.channels`` labels, or None (RGB).""" + """Positional names from NGFF ``omero.channels``, or ``None`` when absent. + + Every metadata entry occupies one slot, including an unnamed entry (stored as + ``""``). Dropping unnamed entries would shift every later label away from its + physical C-axis index and could silently assign pixels to the wrong marker. + Strategy binding validates the metadata length, selected names, and duplicate + identities against :attr:`channel_count`; this reader only preserves the + source's positional metadata faithfully. + """ import zarr try: @@ -455,8 +560,15 @@ def _read_omero_channels(self, ms_path: str) -> list[str] | None: except Exception: # noqa: BLE001 - no omero block -> plain RGB source return None channels = (attrs.get("omero") or {}).get("channels") or [] - names = [c.get("label") or c.get("name") for c in channels] - names = [n for n in names if n] + names: list[str] = [] + for channel in channels: + if not isinstance(channel, dict): + names.append("") + continue + value = channel.get("label") + if value is None or value == "": + value = channel.get("name") + names.append("" if value is None else str(value)) return names or None # -- internals ----------------------------------------------------------- diff --git a/src/raw2features/segmenters/nuclear.py b/src/raw2features/segmenters/nuclear.py index b822463..19dd15d 100644 --- a/src/raw2features/segmenters/nuclear.py +++ b/src/raw2features/segmenters/nuclear.py @@ -1,14 +1,16 @@ -"""Nuclear-channel tissue segmenter for MULTIPLEX slides (CODEX/Phenocycler). +"""Nuclear-channel tissue segmenter for named-channel multiplex slides. Multiplex slides have no H&E saturation/colour to threshold; instead, tissue is where the nuclear stain (DAPI / Hoechst) is present. This segmenter finds the nuclear marker -channel by name (via the reader's ``channel_names``), reads it at a low-res level, and -runs Otsu + morphology - the multiplex analogue of the default Otsu-on-saturation -segmenter. Pure OpenCV + numpy. +channel(s) by name (via the reader's ``channel_names``), reads them at a low-res level, +and runs Otsu + morphology - the multiplex analogue of the default +Otsu-on-saturation segmenter. Pure OpenCV + numpy. """ from __future__ import annotations +import re +import unicodedata from typing import TYPE_CHECKING import numpy as np @@ -23,11 +25,15 @@ from raw2features.readers.base import WSISource _NUCLEAR_ALIASES = ("dapi", "hoechst", "hochst", "dna") +# Numbered DNA channels occur both as plain labels (``DNA1``) and as delimited +# tokens in IMC exports (``191Ir_DNA1`` / ``Ir193_DNA2``). Keep the boundaries +# strict so biomarkers and prose containing ``DNA`` never match. +_DNA_NUMBERED_RE = re.compile(r"(? None: self.seg_mpp = seg_mpp self.morph_kernel = morph_kernel - self.nuclear_aliases = tuple(a.lower() for a in nuclear_aliases) + self.nuclear_aliases = tuple( + unicodedata.normalize("NFKC", a).strip().casefold() for a in nuclear_aliases + ) def _pick_level(self, reader: WSISource) -> int: return nearest_level(reader.mpp, reader.level_downsamples(), self.seg_mpp) - def _nuclear_index(self, channel_names: list[str] | None) -> int: + def _nuclear_matches( + self, channel_names: list[str] | None + ) -> list[tuple[int, str]]: if not channel_names: raise ValueError( "NuclearSegmenter needs a multiplex reader with channel_names " - "(no DAPI/Hoechst channel to threshold)" + "(no recognized DAPI/Hoechst/DNA channel to threshold)" ) + matches: list[tuple[int, str]] = [] for i, name in enumerate(channel_names): - low = name.lower() - if any(alias in low for alias in self.nuclear_aliases): - return i + normalized = unicodedata.normalize("NFKC", str(name)).strip().casefold() + match_kind = None + for alias in self.nuclear_aliases: + if alias == "dna": + if normalized == "dna": + match_kind = "dna" + else: + numbered_dna = _DNA_NUMBERED_RE.search(normalized) + if numbered_dna is not None: + match_kind = f"dna{numbered_dna.group(1)}" + elif ( + re.search(rf"(? list[int]: + matches = self._nuclear_matches(channel_names) + if len(matches) == 1: + return [matches[0][0]] + # Repeated acquisitions of one stain (for example DAPI1/DAPI2 in CODEX) + # represent one nuclear signal. Numbered DNA1/DNA2 channels are likewise + # one family and may contain a canonical pair or repeated acquisitions. + # Never silently mix distinct stains such as DAPI and Hoechst/DNA. + family_by_kind = { + "hochst": "hoechst", + "dna1": "dna-numbered", + "dna2": "dna-numbered", + } + families = {family_by_kind.get(kind, kind) for _, kind in matches} + if len(families) == 1: + return [index for index, _ in matches] + indices = [index for index, _ in matches] raise ValueError( - f"no nuclear channel (aliases {self.nuclear_aliases}) found in the panel" + "nuclear segmentation requires one nuclear stain family; matched mixed " + f"families at physical C indices {indices}. Use --no-seg or provide " + "unambiguous channel names." ) + def _nuclear_index(self, channel_names: list[str] | None) -> int: + """Return the first bound index (legacy private helper).""" + + return self._nuclear_indices(channel_names)[0] + + def nuclear_channels( + self, channel_names: list[str] | None + ) -> list[tuple[int, str]]: + """Return physical nuclear-channel bindings in source order.""" + + indices = self._nuclear_indices(channel_names) + assert channel_names is not None + return [(index, channel_names[index]) for index in indices] + + def nuclear_channel(self, channel_names: list[str] | None) -> tuple[int, str]: + """Return the first physical binding (kept for source compatibility).""" + + return self.nuclear_channels(channel_names)[0] + + @staticmethod + def _combine_nuclear_channels(block: np.ndarray, indices: list[int]) -> np.ndarray: + """Return one signal; repeated channels are averaged in float32.""" + + if len(indices) == 1: + # Preserve the v0.1 single-channel arithmetic exactly; OpenCV + # normalizes in the source dtype before the established uint8 cast. + return np.asarray(block[..., indices[0]]) + selected = np.asarray(block[..., indices], dtype=np.float32) + return selected.mean(axis=2, dtype=np.float32) + def segment(self, reader: WSISource) -> TissueMask: import cv2 - idx = self._nuclear_index(reader.channel_names) + indices = self._nuclear_indices(reader.channel_names) level = self._pick_level(reader) dim: Size = reader.level_dimensions[level] block = reader.read_region_channels( Region(level=level, location=Point(0, 0), size=Size(dim.width, dim.height)) ) - nuclear = block[:, :, idx] + nuclear = self._combine_nuclear_channels(block, indices) # Scale the (possibly uint16) nuclear channel to 8-bit for Otsu. chan8 = cv2.normalize(nuclear, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) _, binary = cv2.threshold(chan8, 0, 255, cv2.THRESH_OTSU + cv2.THRESH_BINARY) diff --git a/src/raw2features/sinks/zarr_sink.py b/src/raw2features/sinks/zarr_sink.py index bd6d9f5..289ef84 100644 --- a/src/raw2features/sinks/zarr_sink.py +++ b/src/raw2features/sinks/zarr_sink.py @@ -14,6 +14,7 @@ import os import secrets import stat +import unicodedata import numpy as np @@ -236,6 +237,29 @@ def create( else: # use_consolidated=False: read live metadata (a prior run consolidated it). root = zarr.open_group(path, mode="r+", use_consolidated=False) + # Validate shared source metadata before deleting or rebuilding a target + # grid. + rh = dict(root.attrs.get("raw2features", {})) + rh.setdefault("schema_version", header.get("schema_version")) + if "source" not in rh and "source" in header: + rh["source"] = header["source"] + elif "source" in header: + source_update = { + name: header["source"][name] + for name in ( + "channel_names", + "channel_names_source", + "omero_channel_names", + ) + if name in header["source"] + } + if source_update: + rh["source"] = self._merge_source_metadata( + rh.get("source", {}), source_update + ) + for k in ("provenance", "segmentation", "thumbnail"): + if k not in rh and k in header: + rh[k] = header[k] rebuild = _unappendable_grid_to_rebuild( root, grid, @@ -254,11 +278,6 @@ def create( actual_grid = rebuild else: actual_grid = _select_grid_label(root, grid, header.get("grid_hash")) - rh = dict(root.attrs.get("raw2features", {})) - rh.setdefault("schema_version", header.get("schema_version")) - for k in ("source", "provenance", "segmentation", "thumbnail"): - if k not in rh and k in header: - rh[k] = header[k] grids = dict(rh.get("grids", {})) grids[actual_grid] = _grids_index_entry(header, model_dims, n_patches) rh["grids"] = grids @@ -321,6 +340,8 @@ def open_append( key: str | None = None, new_model_dims: dict[str, int], new_model_meta: dict | None = None, + new_panel_meta: dict | None = None, + new_source_meta: dict | None = None, replace_models: list[str] | tuple[str, ...] | set[str] = (), ) -> int: """Open an existing store's grid ``r+`` to add feature arrays, no clobber. @@ -331,7 +352,10 @@ def open_append( recreated; callers use that only after the old output fails its current model contract. Coords/grid/mask and unrelated feature arrays are left untouched. ``key=None`` targets the sole grid. Returns the grid's patch count. Pass empty - ``new_model_dims`` to open for slide-only work. + ``new_model_dims`` to open for slide-only work. ``new_panel_meta`` is merged by + model key alongside ``new_model_meta``. ``new_source_meta`` may add or fill the + effective positional channel metadata in both grid and root headers, but cannot + relabel an already named physical channel. """ import zarr @@ -345,6 +369,16 @@ def open_append( root = zarr.open_group(path, mode="r+", use_consolidated=False) g = open_grid(root, key) actual_key = key if key is not None else grid_keys(root)[0] + gh = dict(g.attrs.get("raw2features", {})) + rh = None + if new_source_meta: + gh["source"] = self._merge_source_metadata( + gh.get("source", {}), new_source_meta + ) + rh = dict(root.attrs.get("raw2features", {})) + rh["source"] = self._merge_source_metadata( + rh.get("source", {}), new_source_meta + ) feats = g["features"] existing = list(feats.keys()) self._dtype = feats[existing[0]].dtype if existing else "float16" @@ -371,10 +405,18 @@ def open_append( a.attrs["role"] = "features" a.attrs["model"] = model added.append(model) - if new_model_meta: - gh = dict(g.attrs.get("raw2features", {})) - gh.setdefault("models", {}).update(new_model_meta) + if new_model_meta or new_panel_meta or new_source_meta: + if new_model_meta: + models = dict(gh.get("models", {})) + models.update(new_model_meta) + gh["models"] = models + if new_panel_meta: + panel = dict(gh.get("panel", {})) + panel.update(new_panel_meta) + gh["panel"] = panel g.attrs["raw2features"] = gh + if rh is not None: + root.attrs["raw2features"] = rh if added: self._update_root_models(root, actual_key, list(feats.keys())) self._root = root @@ -382,6 +424,63 @@ def open_append( self._active_key = actual_key return n + @staticmethod + def _merge_source_metadata(existing: dict, update: dict) -> dict: + """Fill positional channel metadata without changing an established label.""" + + merged = dict(existing or {}) + new_names = update.get("channel_names") + if new_names is None: + return merged + new_names = ["" if value is None else str(value) for value in new_names] + old_names = merged.get("channel_names") + if old_names is None: + merged.update(update) + merged["channel_names"] = new_names + return merged + old_names = ["" if value is None else str(value) for value in old_names] + if len(old_names) > len(new_names): + raise ValueError( + "stored source channel metadata has more positional entries than " + f"the effective physical panel ({len(old_names)} > {len(new_names)})" + ) + old_names.extend([""] * (len(new_names) - len(old_names))) + + def identity(value: str) -> str: + return unicodedata.normalize("NFKC", value).strip().casefold() + + completed = [] + filled = False + for index, (old_name, new_name) in enumerate( + zip(old_names, new_names, strict=True) + ): + old_identity = identity(old_name) + new_identity = identity(new_name) + if old_identity and old_identity != new_identity: + raise ValueError( + "stored source channel metadata conflicts at physical C index " + f"{index}: {old_name!r} != {new_name!r}" + ) + if old_identity: + completed.append(old_name) + else: + completed.append(new_name) + filled = filled or bool(new_identity) + merged["channel_names"] = completed + if filled: + merged["channel_names_source"] = update.get("channel_names_source") + else: + merged.setdefault( + "channel_names_source", update.get("channel_names_source") + ) + if ( + "omero_channel_names" in update + and "omero_channel_names" not in merged + and update["omero_channel_names"] != completed + ): + merged["omero_channel_names"] = update["omero_channel_names"] + return merged + @staticmethod def _drop_slide_dependents(group, patch_models: set[str]) -> None: """Remove slide vectors derived from feature arrays being replaced.""" @@ -475,13 +574,21 @@ def write_slide_embedding( slide_model: str, vector: np.ndarray, provenance: dict, + *, + output_name: str | None = None, ) -> None: """Write a single slide-level vector to the grid's ``slide//``.""" if self._group is None: raise RuntimeError("sink not created; call create() first") from raw2features.slide_embedders.encoding import write_slide_embedding - write_slide_embedding(self._group, slide_model, vector, provenance) + write_slide_embedding( + self._group, + slide_model, + vector, + provenance, + output_name=output_name, + ) def write_qc( self, diff --git a/src/raw2features/slide_embedders/encoding.py b/src/raw2features/slide_embedders/encoding.py index b4a55d3..def9a85 100644 --- a/src/raw2features/slide_embedders/encoding.py +++ b/src/raw2features/slide_embedders/encoding.py @@ -22,6 +22,34 @@ class SlideEncoding: provenance: dict +def slide_output_key(group, slide_model: str, patch_model: str) -> str: + """Storage key for a slide output, preserving multiplex coexistence. + + Historical brightfield/native outputs remain ``slide/`` exactly. + Strategy-derived multiplex patch arrays can coexist on one geometry grid, so a + model-agnostic pool is qualified by that cohort-stable effective patch key and + cannot overwrite a pool derived from another marker/aggregation recipe. + """ + + try: + fingerprint = dict(group["features"][patch_model].attrs).get( + "output_fingerprint" + ) + payload = fingerprint.get("payload", {}) + strategy_derived = ( + payload.get("loader", {}).get("family") == "multiplex_strategy" + ) + except (AttributeError, KeyError, TypeError): + strategy_derived = False + if not strategy_derived: + return slide_model + if "/" in patch_model: + raise ValueError( + f"multiplex patch model key {patch_model!r} is not a safe zarr path segment" + ) + return f"{slide_model}__{patch_model}" + + def _validated_patch_fingerprint(group, patch_model: str) -> dict: """Return a committed, complete patch-output fingerprint. @@ -93,11 +121,13 @@ def slide_embedding_is_complete( *, patch_model: str | None = None, device: str = "cpu", + output_name: str | None = None, ) -> bool: """Return whether ``slide/`` is valid for the requested patch model.""" - if "slide" not in group or slide_model not in group["slide"]: + stored_name = output_name or slide_model + if "slide" not in group or stored_name not in group["slide"]: return False - array = group["slide"][slide_model] + array = group["slide"][stored_name] attrs = dict(array.attrs) stored_patch_model = attrs.get("patch_encoder") if attrs.get("role") != "slide_embedding" or not stored_patch_model: @@ -146,7 +176,7 @@ def slide_embedding_is_complete( return False header = dict(group.attrs.get("raw2features", {})) - mirrored = header.get("slide_embeddings", {}).get(slide_model, {}) + mirrored = header.get("slide_embeddings", {}).get(stored_name, {}) if not isinstance(mirrored, dict): return False if mirrored.get("patch_encoder") != stored_patch_model: @@ -348,26 +378,32 @@ def write_slide_embedding( slide_model: str, vector: np.ndarray, provenance: dict, + *, + output_name: str | None = None, ) -> None: """Replace ``slide/`` with one vector and refresh grid metadata.""" + stored_name = output_name or slide_model + stored_provenance = dict(provenance) + if stored_name != slide_model: + stored_provenance["slide_encoder"] = slide_model slide_group = group.require_group("slide") - if slide_model in slide_group: - del slide_group[slide_model] + if stored_name in slide_group: + del slide_group[stored_name] vector_2d = np.asarray(vector, dtype=np.float32).reshape(1, -1) array = slide_group.create_array( - slide_model, + stored_name, shape=vector_2d.shape, chunks=vector_2d.shape, dtype="float32", ) array[:] = vector_2d array.attrs["role"] = "slide_embedding" - for key, value in provenance.items(): + for key, value in stored_provenance.items(): array.attrs[key] = value header = dict(group.attrs.get("raw2features", {})) slide_metadata = dict(header.get("slide_embeddings", {})) - slide_metadata[slide_model] = dict(provenance) + slide_metadata[stored_name] = stored_provenance header["slide_embeddings"] = slide_metadata group.attrs["raw2features"] = header diff --git a/tests/test_additive_embed.py b/tests/test_additive_embed.py index 66b8286..387e04b 100644 --- a/tests/test_additive_embed.py +++ b/tests/test_additive_embed.py @@ -25,6 +25,7 @@ RunConfig, _assert_store_source, _inspect_store, + _stored_grid_summary, embed_slide, run_slide, ) @@ -47,6 +48,7 @@ def _make_store( grid_hash="GH", fill=1.0, source="file:///source/s.zarr", + segmenter=None, ): """Build a minimal *.embeddings.zarr by hand (no embedding needed).""" os.makedirs(out_dir, exist_ok=True) @@ -66,6 +68,8 @@ def _make_store( } if grid_hash is not None: header["grid_hash"] = grid_hash + if segmenter is not None: + header["segmentation"] = {"segmenter": segmenter} g.attrs["raw2features"] = header root.attrs["raw2features"] = { "source": {"uri": source}, @@ -129,6 +133,162 @@ def test_inspect_store_legacy_without_grid_hash_is_compatible(tmp_path): assert valid == ["a"] +def test_inspect_store_reuses_native_multiplex_pre_v011_grid_identity(tmp_path): + """KRONOS used nuclear execution but persisted the default Otsu hash in v0.1.""" + + native = RunConfig( + models=["kronos"], + resolved_channel_names=["DAPI", "CD3"], + resolved_nuclear_channel_indices=[0], + resolved_original_channel_names=["DAPI", "CD3"], + ) + legacy_otsu = RunConfig(models=["uni"]).grid_hash() + path = _make_store( + str(tmp_path / "out"), + models=(("kronos", 4),), + grid_hash=legacy_otsu, + segmenter="nuclear", + ) + + key, n_patches, valid = _inspect_store( + path, + native.grid_hash(), + ["kronos"], + compatible_grid_hashes=native.compatible_legacy_grid_hashes(), + compatible_grid_segmenters=native.compatible_legacy_grid_segmenters(), + ) + + assert key == "mpp1_px64" + assert n_patches == 12 + assert valid == ["kronos"] + + +def test_inspect_store_rejects_unsafe_native_nuclear_legacy_hash(tmp_path): + """A nuclear header cannot prove which physical channel v0.1 thresholded.""" + + native = RunConfig( + models=["kronos"], + resolved_channel_names=["", "DAPI", "CD3"], + resolved_nuclear_channel_indices=[1], + resolved_original_channel_names=["", "DAPI", "CD3"], + ) + legacy_nuclear = native.legacy_grid_hash() + path = _make_store( + str(tmp_path / "out"), + models=(("kronos", 4),), + grid_hash=legacy_nuclear, + segmenter="nuclear", + ) + + key, n_patches, valid = _inspect_store( + path, + native.grid_hash(), + ["kronos"], + compatible_grid_hashes=native.compatible_legacy_grid_hashes(), + compatible_grid_segmenters=native.compatible_legacy_grid_segmenters(), + ) + + assert key is None + assert n_patches == 0 + assert valid == [] + + +def test_inspect_and_summary_reject_unsafe_native_hashless_grid(tmp_path): + """No hash means the legacy physical nuclear-channel binding is unknowable.""" + + native = RunConfig( + models=["kronos"], + resolved_channel_names=["", "DAPI", "CD3"], + resolved_nuclear_channel_indices=[1], + resolved_original_channel_names=["", "DAPI", "CD3"], + ) + path = _make_store( + str(tmp_path / "out"), + models=(("kronos", 4),), + grid_hash=None, + segmenter="nuclear", + ) + policy = native.allows_hashless_legacy_grid() + + assert not policy + assert _inspect_store( + path, + native.grid_hash(), + ["kronos"], + compatible_grid_hashes=native.compatible_legacy_grid_hashes(), + compatible_grid_segmenters=native.compatible_legacy_grid_segmenters(), + allow_hashless_legacy=policy, + ) == (None, 0, []) + assert ( + _stored_grid_summary( + path, + [native], + allow_hashless_legacy_grids={native.grid_hash(): policy}, + ) + == {} + ) + + +def test_inspect_store_requires_nuclear_evidence_for_native_nuclear_legacy_hash( + tmp_path, +): + native = RunConfig( + models=["kronos"], + resolved_channel_names=["DAPI", "CD3"], + resolved_nuclear_channel_indices=[0], + resolved_original_channel_names=["DAPI", "CD3"], + ) + legacy_nuclear = native.legacy_grid_hash() + path = _make_store( + str(tmp_path / "out"), + models=(("kronos", 4),), + grid_hash=legacy_nuclear, + segmenter="otsu", + ) + + key, n_patches, valid = _inspect_store( + path, + native.grid_hash(), + ["kronos"], + compatible_grid_hashes=native.compatible_legacy_grid_hashes(), + compatible_grid_segmenters=native.compatible_legacy_grid_segmenters(), + ) + + assert key is None + assert n_patches == 0 + assert valid == [] + + +def test_inspect_store_rejects_brightfield_grid_with_same_native_legacy_hash( + tmp_path, +): + native = RunConfig( + models=["kronos"], + resolved_channel_names=["DAPI", "CD3"], + resolved_nuclear_channel_indices=[0], + resolved_original_channel_names=["DAPI", "CD3"], + ) + legacy_otsu = RunConfig(models=["uni"]).grid_hash() + path = _make_store( + str(tmp_path / "out"), + models=(("kronos", 4),), + grid_hash=legacy_otsu, + segmenter="otsu", + ) + + key, n_patches, valid = _inspect_store( + path, + native.grid_hash(), + ["kronos"], + compatible_grid_hashes=native.compatible_legacy_grid_hashes(), + compatible_grid_segmenters=native.compatible_legacy_grid_segmenters(), + ) + + assert key is None + assert n_patches == 0 + assert valid == [] + + @pytest.mark.skipif(not _TORCH, reason="torch not installed") def test_same_label_different_grid_hash_uses_deterministic_suffix( synthetic_ngff, tmp_path diff --git a/tests/test_channelwise_runner.py b/tests/test_channelwise_runner.py new file mode 100644 index 0000000..abb5bda --- /dev/null +++ b/tests/test_channelwise_runner.py @@ -0,0 +1,646 @@ +"""End-to-end multiplex strategy tests over the real runner/store seams.""" + +from __future__ import annotations + +import os +from dataclasses import replace +from importlib import import_module + +import numpy as np +import pytest +import zarr +from typer.testing import CliRunner + +from conftest import MockEmbedder +from raw2features.cli.main import app +from raw2features.core.store import open_grid +from raw2features.pipeline.receipt import canonical_source_uri, is_complete +from raw2features.pipeline.runner import ( + RunConfig, + _multiplex_normalization_level, + embed_slide, + resolve_multiplex_output_contracts, + resolve_multiplex_source_config, + resolve_run, + run_slide, +) + + +def _cfg(*, markers=("CD3", "DAPI"), aggregation="mean", slide=True): + return RunConfig( + models=["mock"], + target_mpp=0.5, + patch_px=64, + step_px=64, + no_seg=True, + features_dtype="float32", + device="cpu", + batch_size=2, + read_workers=1, + multiplex_strategy="channelwise", + multiplex_markers=list(markers), + multiplex_normalization="percentile", + multiplex_aggregation=aggregation, + slide_encoders=["mean"] if slide else [], + ) + + +def _root_and_grid(out_dir, summary): + path = os.path.join(out_dir, "multiplex.embeddings.zarr") + root = zarr.open_group(path, mode="r", use_consolidated=False) + return root, open_grid(root, summary["grid"]) + + +def test_channelwise_runner_writes_self_describing_patch_and_slide_outputs( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + out = str(tmp_path / "out") + summary = run_slide( + synthetic_multiplex_ngff, + out, + _cfg(), + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + + assert summary["status"] == "complete" + assert len(summary["requested_models"]) == 1 + effective = summary["requested_models"][0] + assert effective.startswith("mock__channelwise_mean_") + assert len(effective.rsplit("_", 1)[1]) == 16 + + _root, group = _root_and_grid(out, summary) + features = group["features"][effective] + assert features.shape[1] == 4 + assert np.isfinite(np.asarray(features[:])).all() + + fingerprint = dict(features.attrs)["output_fingerprint"] + payload = fingerprint["payload"] + assert payload["output"]["modality"] == "multiplex" + assert payload["multiplex"]["base_model"] == "mock" + assert payload["multiplex"]["base_pooling"] == "cls" + assert payload["multiplex"]["normalization"]["source_level"] == 0 + assert [ + row["source_name"] for row in payload["multiplex"]["normalization"]["resolved"] + ] == ["CD3", "DAPI"] + + header = dict(group.attrs["raw2features"]) + assert header["segmentation"] == {"segmenter": "none"} + assert "mask" not in group + model_header = header["models"][effective] + assert model_header["pooling"] == "mean" + assert model_header["output_fingerprint"]["payload"]["output"]["pooling"] == "mean" + assert model_header["multiplex"]["base_pooling"] == "cls" + panel = header["panel"][effective] + assert [row["source_name"] for row in panel["source_channels"]] == [ + "DAPI", + "CD3", + "CD8", + "CD20", + "FOXP3", + ] + assert [row["source_name"] for row in panel["excluded"]] == [ + "CD8", + "CD20", + "FOXP3", + ] + slide_key = f"mean__{effective}" + assert tuple(group["slide"][slide_key].shape) == (1, 4) + assert dict(group["slide"][slide_key].attrs)["patch_encoder"] == effective + assert header["slide_embeddings"][slide_key]["slide_encoder"] == "mean" + + +def test_dna1_dna2_pair_is_bound_into_grid_and_segmentation_provenance( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + source = zarr.open_group(synthetic_multiplex_ngff, mode="a", use_consolidated=False) + omero = dict(source.attrs["omero"]) + omero["channels"] = [ + {"label": name} + for name in ("DNA1(Ir191)", "DNA2(Ir193)", "CD3", "CD20", "FOXP3") + ] + source.attrs["omero"] = omero + + cfg = replace(_cfg(markers=("CD3",), slide=False), no_seg=False) + summary = run_slide( + synthetic_multiplex_ngff, + str(tmp_path / "out"), + cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + + _root, group = _root_and_grid(str(tmp_path / "out"), summary) + segmentation = dict(group.attrs["raw2features"])["segmentation"] + assert segmentation["nuclear_channel_indices"] == [0, 1] + assert segmentation["nuclear_channel_names"] == [ + "DNA1(Ir191)", + "DNA2(Ir193)", + ] + assert segmentation["nuclear_channel_combination"] == "float32_mean" + assert segmentation["channel_binding_contract_version"] == 2 + + resolved = resolve_multiplex_source_config( + synthetic_multiplex_ngff, + cfg, + model_specs={"mock": MockEmbedder(name="mock").spec}, + ) + assert resolved.resolved_nuclear_channel_indices == [0, 1] + assert resolved.grid_hash() == dict(group.attrs["raw2features"])["grid_hash"] + + +def test_channelwise_append_keeps_both_panel_contracts_and_slide_pools( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + out = str(tmp_path / "out") + first = run_slide( + synthetic_multiplex_ngff, + out, + _cfg(markers=("CD3", "DAPI")), + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + second = run_slide( + synthetic_multiplex_ngff, + out, + _cfg(markers=("DAPI", "CD3")), + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + first_key = first["requested_models"][0] + second_key = second["requested_models"][0] + assert first_key != second_key + assert first["grid"] == second["grid"] + + _root, group = _root_and_grid(out, second) + assert {first_key, second_key} <= set(group["features"].keys()) + header = dict(group.attrs["raw2features"]) + assert {first_key, second_key} <= set(header["panel"]) + assert {f"mean__{first_key}", f"mean__{second_key}"} <= set(group["slide"].keys()) + + +def test_channel_names_override_enables_incomplete_ome_metadata( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + root = zarr.open_group(synthetic_multiplex_ngff, mode="r+") + root.attrs["omero"] = {"channels": [{"label": "DAPI"}, {}]} + cfg = _cfg(slide=False) + cfg.channel_names_override = ["DAPI", "CD3", "CD8", "CD20", "FOXP3"] + + summary = run_slide( + synthetic_multiplex_ngff, + str(tmp_path / "out"), + cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + _root, group = _root_and_grid(str(tmp_path / "out"), summary) + header = dict(group.attrs["raw2features"]) + assert header["source"]["channel_names"] == cfg.channel_names_override + assert header["source"]["channel_names_source"] == "channel_names_file" + assert header["source"]["omero_channel_names"] == ["DAPI", ""] + + +def test_channel_names_override_is_reapplied_for_aux_thumbnail( + synthetic_multiplex_ngff, tmp_path, monkeypatch +): + pytest.importorskip("torch") + root = zarr.open_group(synthetic_multiplex_ngff, mode="r+") + root.attrs["omero"] = {"channels": [{"label": "DAPI"}, {}]} + cfg = replace(_cfg(slide=False), no_seg=False) + cfg.channel_names_override = ["DAPI", "CD3", "CD8", "CD20", "FOXP3"] + out = str(tmp_path / "out") + + run_slide( + synthetic_multiplex_ngff, + out, + cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + runner = import_module("raw2features.pipeline.runner") + monkeypatch.setattr( + runner, + "_embed_patches", + lambda *args, **kwargs: pytest.fail("auxiliary run re-embedded patches"), + ) + run_slide( + synthetic_multiplex_ngff, + out, + replace(cfg, emit_thumbnail=True), + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + + assert os.path.isfile(os.path.join(out, "multiplex.thumbnail.png")) + assert os.path.isfile(os.path.join(out, "multiplex.thumbnail.overlay.png")) + + +@pytest.mark.parametrize("multiplex_mpp", [0.5, 1.0]) +def test_multiplex_addition_fills_shared_source_channel_provenance( + synthetic_multiplex_ngff, tmp_path, multiplex_mpp +): + pytest.importorskip("torch") + source = zarr.open_group(synthetic_multiplex_ngff, mode="r+") + source.attrs["omero"] = {"channels": [{"label": "DAPI"}, {}]} + out = str(tmp_path / "out") + brightfield_cfg = RunConfig( + models=["bright"], + target_mpp=0.5, + patch_px=64, + step_px=64, + no_seg=True, + features_dtype="float32", + device="cpu", + batch_size=2, + read_workers=1, + ) + run_slide( + synthetic_multiplex_ngff, + out, + brightfield_cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="bright")], + ) + + multiplex_cfg = replace(_cfg(slide=False), target_mpp=multiplex_mpp) + multiplex_cfg.channel_names_override = [ + "DAPI", + "CD3", + "CD8", + "CD20", + "FOXP3", + ] + resolved_cfg = resolve_multiplex_source_config( + synthetic_multiplex_ngff, + multiplex_cfg, + model_specs={"mock": MockEmbedder(name="mock").spec}, + ) + assert resolved_cfg.resolved_original_channel_names == ["DAPI", ""] + summary = run_slide( + synthetic_multiplex_ngff, + out, + multiplex_cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + + root, group = _root_and_grid(out, summary) + for header in ( + dict(root.attrs["raw2features"]), + dict(group.attrs["raw2features"]), + ): + assert header["source"]["channel_names"] == ( + multiplex_cfg.channel_names_override + ) + assert header["source"]["channel_names_source"] == "channel_names_file" + assert header["source"]["omero_channel_names"] == ["DAPI", ""] + + +def test_panel_source_origin_does_not_change_channelwise_output_fingerprint( + synthetic_multiplex_ngff, +): + base = MockEmbedder(dim=4, input_size=64, name="mock") + metadata_cfg = resolve_multiplex_source_config( + synthetic_multiplex_ngff, _cfg(slide=False), model_specs={"mock": base.spec} + ) + override_cfg = _cfg(slide=False) + override_cfg.channel_names_override = ["DAPI", "CD3", "CD8", "CD20", "FOXP3"] + override_cfg = resolve_multiplex_source_config( + synthetic_multiplex_ngff, override_cfg, model_specs={"mock": base.spec} + ) + assert metadata_cfg.resolved_original_channel_names == [ + "DAPI", + "CD3", + "CD8", + "CD20", + "FOXP3", + ] + assert override_cfg.resolved_original_channel_names == [ + "DAPI", + "CD3", + "CD8", + "CD20", + "FOXP3", + ] + _, metadata_contracts = resolve_multiplex_output_contracts( + synthetic_multiplex_ngff, + metadata_cfg, + embedders=[base], + device="cpu", + ) + _, override_contracts = resolve_multiplex_output_contracts( + synthetic_multiplex_ngff, + override_cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + device="cpu", + ) + first = next(iter(metadata_contracts.values()))["output_fingerprint"] + second = next(iter(override_contracts.values()))["output_fingerprint"] + assert first == second + + +@pytest.mark.parametrize("strategy_first", [False, True]) +def test_brightfield_otsu_and_channelwise_nuclear_grids_never_alias( + synthetic_multiplex_ngff, tmp_path, strategy_first +): + pytest.importorskip("torch") + out = str(tmp_path / "out") + brightfield = RunConfig( + models=["mock"], + target_mpp=0.5, + patch_px=64, + step_px=64, + segmenter="otsu", + features_dtype="float32", + device="cpu", + batch_size=2, + read_workers=1, + ) + strategy = replace(_cfg(slide=False), no_seg=False) + requests = [ + (strategy, MockEmbedder(dim=4, input_size=64, name="mock")), + (brightfield, MockEmbedder(dim=4, input_size=64, name="mock")), + ] + if not strategy_first: + requests.reverse() + + summaries = [ + run_slide(synthetic_multiplex_ngff, out, cfg, embedders=[emb]) + for cfg, emb in requests + ] + + assert summaries[0]["grid"] != summaries[1]["grid"] + root = zarr.open_group( + os.path.join(out, "multiplex.embeddings.zarr"), + mode="r", + use_consolidated=False, + ) + headers = [ + dict(root["grids"][key].attrs["raw2features"]) + for key in root["grids"].group_keys() + ] + assert {header["segmentation"]["segmenter"] for header in headers} == { + "otsu", + "nuclear", + } + assert len({header["grid_hash"] for header in headers}) == 2 + + +def test_complete_channelwise_rerun_does_not_load_registry_base_model( + synthetic_multiplex_ngff, tmp_path, monkeypatch +): + pytest.importorskip("torch") + loads = [] + unloads = [] + + class TrackingEmbedder(MockEmbedder): + def load(self, device="cpu", dtype=None, compile=False): + loads.append((device, dtype, compile)) + return super().load(device=device, dtype=dtype, compile=compile) + + def unload(self): + unloads.append(True) + return super().unload() + + monkeypatch.setattr( + "raw2features.pipeline.runner.build_embedder", + lambda name: TrackingEmbedder(dim=4, input_size=64, name=name), + ) + cfg = _cfg(slide=False) + out = str(tmp_path / "out") + first = run_slide(synthetic_multiplex_ngff, out, cfg) + second = run_slide(synthetic_multiplex_ngff, out, cfg) + + assert first["status"] == "complete" + assert second["status"] == "skipped" + assert len(loads) == 1 + assert len(unloads) == 1 + + +def test_channelwise_does_not_unload_a_caller_owned_warm_base( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + unloads = [] + + class TrackingEmbedder(MockEmbedder): + def unload(self): + unloads.append(True) + return super().unload() + + base = TrackingEmbedder(dim=4, input_size=64, name="mock") + run_slide( + synthetic_multiplex_ngff, + str(tmp_path / "out-a"), + _cfg(slide=False), + embedders=[base], + ) + run_slide( + synthetic_multiplex_ngff, + str(tmp_path / "out-b"), + _cfg(slide=False), + embedders=[base], + ) + + assert unloads == [] + + +def test_effective_patch_batch_preserves_base_rgb_batch_budget( + synthetic_multiplex_ngff, +): + cfg = _cfg(markers=("DAPI", "CD3"), slide=False) + cfg.batch_size = 16 + effective, _contracts = resolve_multiplex_output_contracts( + synthetic_multiplex_ngff, + cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + device="cpu", + ) + assert effective.batch_size == 8 + + +def test_normalization_refuses_an_unbounded_single_level_read(): + class HugeSingleLevel: + level_dimensions = [ + type("Dimensions", (), {"width": 50_000, "height": 40_000})() + ] + + with pytest.raises(ValueError) as exc_info: + _multiplex_normalization_level(HugeSingleLevel(), max_side=2048) + message = str(exc_info.value) + assert "Add a downsampled level" in message + assert "--multiplex-normalization-max-side-px" in message + assert "50000" in message + assert "more RAM" in message + + +def test_normalization_level_honours_the_configured_maximum_side(): + class Pyramid: + level_dimensions = [ + type("Dimensions", (), {"width": side, "height": side})() + for side in (4096, 2048, 1024) + ] + + assert _multiplex_normalization_level(Pyramid(), max_side=4096) == 0 + assert _multiplex_normalization_level(Pyramid(), max_side=2048) == 1 + assert _multiplex_normalization_level(Pyramid(), max_side=1024) == 2 + + +def test_high_level_receipt_validates_against_source_resolved_effective_contract( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + out = str(tmp_path / "out") + receipts = str(tmp_path / "receipts") + cfg = _cfg(slide=False) + summary = embed_slide( + synthetic_multiplex_ngff, + out, + cfg, + requested_mpp=0.5, + requested_patch_px=64, + receipts_dir=receipts, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + + resolved_cfg = resolve_multiplex_source_config( + synthetic_multiplex_ngff, + cfg, + model_specs={"mock": MockEmbedder(dim=4, input_size=64).spec}, + ) + _groups, group_cfgs, run_hash = resolve_run( + resolved_cfg, + requested_mpp=0.5, + requested_patch_px=64, + model_specs={"mock": MockEmbedder(dim=4, input_size=64).spec}, + ) + contracts = {} + expected_grid_models = {} + for group_cfg in group_cfgs: + effective_cfg, group_contracts = resolve_multiplex_output_contracts( + synthetic_multiplex_ngff, + group_cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + device="cpu", + ) + contracts.update(group_contracts) + expected_grid_models[group_cfg.grid_hash()] = effective_cfg.models + + assert is_complete( + receipts, + "multiplex", + run_hash, + expected_source_uri=canonical_source_uri(synthetic_multiplex_ngff), + expected_output_uri=summary["output_uri"], + expected_model_contracts=contracts, + expected_grid_models=expected_grid_models, + compatible_grid_hashes={ + group_cfg.grid_hash(): group_cfg.compatible_legacy_grid_hashes() + for group_cfg in group_cfgs + }, + ) + + +def test_verify_accepts_a_direct_single_grid_channelwise_receipt( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + from raw2features.embedders.model_registry import get_spec + + spec = get_spec("resnet50") + base = MockEmbedder( + dim=spec.embedding_dim, + input_size=spec.input_size, + name=spec.name, + ) + base.spec = spec + cfg = RunConfig( + models=["resnet50"], + target_mpp=0.5, + patch_px=64, + step_px=64, + no_seg=True, + features_dtype="float32", + device="cpu", + batch_size=2, + read_workers=1, + multiplex_strategy="channelwise", + multiplex_markers=["CD3", "DAPI"], + ) + out = str(tmp_path / "out") + receipts = str(tmp_path / "receipts") + summary = run_slide( + synthetic_multiplex_ngff, + out, + cfg, + receipts_dir=receipts, + embedders=[base], + ) + second = run_slide( + synthetic_multiplex_ngff, + out, + cfg, + receipts_dir=receipts, + embedders=[base], + ) + result = CliRunner().invoke( + app, + [ + "verify", + synthetic_multiplex_ngff, + "--receipts-dir", + receipts, + "--out-dir", + out, + "-m", + "resnet50", + "--mpp", + "0.5", + "--patch-size", + "64", + "--step", + "64", + "--no-seg", + "--features-dtype", + "float32", + "--device", + "cpu", + "--multiplex-strategy", + "channelwise", + "--marker", + "CD3", + "--marker", + "DAPI", + "--quiet", + ], + ) + + assert summary["status"] == "complete" + assert second["status"] == "skipped" + assert second["reason"] == "already complete" + assert result.exit_code == 0, result.output + + +def test_channelwise_rejects_learned_slide_encoder_before_embedding( + synthetic_multiplex_ngff, tmp_path +): + pytest.importorskip("torch") + cfg = _cfg(slide=False) + cfg.slide_encoders = ["titan"] + with pytest.raises(ValueError, match="model-agnostic slide poolers"): + run_slide( + synthetic_multiplex_ngff, + str(tmp_path / "out"), + cfg, + embedders=[MockEmbedder(dim=4, input_size=64, name="mock")], + ) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"multiplex_markers": ["CD3"]}, + {"multiplex_normalization": "dtype"}, + {"multiplex_aggregation": "attention"}, + {"multiplex_normalization_max_side_px": 1024}, + ], +) +def test_brightfield_rejects_orphaned_multiplex_options(kwargs): + with pytest.raises(ValueError, match="require multiplex_strategy"): + RunConfig(models=["mock"], **kwargs) diff --git a/tests/test_channelwise_strategy.py b/tests/test_channelwise_strategy.py new file mode 100644 index 0000000..0657fb2 --- /dev/null +++ b/tests/test_channelwise_strategy.py @@ -0,0 +1,629 @@ +"""Focused contract tests for the channel-wise multiplex strategy.""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import numpy as np +import pytest + +from raw2features.embedders.base import ModelSpec +from raw2features.embedders.fingerprint import ( + patch_output_fingerprint, + valid_output_fingerprint, +) +from raw2features.multiplex import build_strategy +from raw2features.multiplex.channelwise import ChannelwiseStrategy +from raw2features.pipeline.runner import _models_header + + +class TinyRGBEmbedder: + """A deterministic two-dimensional RGB encoder with no model weights.""" + + def __init__(self, *, modality: str = "brightfield") -> None: + self.spec = ModelSpec( + name="tiny_rgb", + family="mock", + source="mock://tiny-rgb", + embedding_dim=2, + input_size=1, + pooling="cls", + mean=(0.0, 0.0, 0.0), + std=(1.0, 1.0, 1.0), + transform_source_url="https://example.org/tiny-rgb", + license="MIT", + gated=False, + modality=modality, + ) + self._device = "cpu" + self._dtype = None + self.last_rgb: list[np.ndarray] = [] + + @property + def name(self) -> str: + return self.spec.name + + @property + def modality(self) -> str: + return self.spec.modality + + @property + def embedding_dim(self) -> int: + return self.spec.embedding_dim + + def transform_batch(self, patches_hwc: list[np.ndarray], device: str): + assert device == "cpu" + self.last_rgb = [np.asarray(patch).copy() for patch in patches_hwc] + intensity = np.asarray( + [float(patch[..., 0].mean()) / 255.0 for patch in patches_hwc], + dtype=np.float32, + ) + values = np.stack([intensity, 1.0 - intensity**2], axis=1) + try: + import torch + except ImportError: + return values + return torch.from_numpy(values) + + def embed_batch(self, batch): + return batch.float().cpu() + + +class TinyFloatRGBEmbedder(TinyRGBEmbedder): + @property + def transform_input_dtype(self) -> str: + return "float32_0_1" + + def transform_batch(self, patches_hwc: list[np.ndarray], device: str): + self.last_rgb = [np.asarray(patch).copy() for patch in patches_hwc] + return np.asarray([[float(patch[..., 0].mean()), 0.0] for patch in patches_hwc]) + + +def _config( + markers=("A", "B", "C"), + *, + normalization="dtype", + aggregation="mean", + max_side=2048, +): + return { + "multiplex_markers": list(markers) if markers is not None else [], + "multiplex_normalization": normalization, + "multiplex_percentile_low": 1.0, + "multiplex_percentile_high": 99.0, + "multiplex_aggregation": aggregation, + "multiplex_normalization_max_side_px": max_side, + } + + +def _prepare( + *, + panel=("A", "B", "C"), + channel_count=None, + config=None, + base=None, +): + if channel_count is None and panel is not None: + channel_count = len(panel) + return ChannelwiseStrategy().prepare( + base_embedder=base or TinyRGBEmbedder(), + channel_names=list(panel) if panel is not None else None, + channel_count=channel_count, + config=config or _config(), + ) + + +def _resolved(prepared, *, values=(0, 128, 255), dtype=np.uint8, level=0): + values_array = np.asarray(values) + if prepared.normalization == "percentile": + # Percentiles are slide-wide per marker, so each channel needs a genuine + # intensity distribution rather than the single-pixel dtype fixture. + image = np.stack([values_array, values_array + 20], axis=0) + image = image.reshape(2, 1, -1).astype(dtype) + else: + image = values_array.astype(dtype).reshape(1, 1, -1) + return prepared.resolve_slide_context(image, level=level, source_dtype=dtype) + + +def _bound( + *, + panel=("A", "B", "C"), + config=None, + values=(0, 128, 255), + dtype=np.uint8, + base=None, +): + prepared = _prepare(panel=panel, config=config, base=base) + return prepared.bind(_resolved(prepared, values=values, dtype=dtype)) + + +def _full_fingerprint(bound, resolved_amp="fp32"): + base_fingerprint = patch_output_fingerprint(bound.base_embedder.spec, resolved_amp) + contract = bound.multiplex_fingerprint_payload(base_fingerprint) + spec = replace(bound.spec, multiplex=contract) + return patch_output_fingerprint(spec, resolved_amp) + + +def test_channelwise_is_available_through_the_strategy_registry(): + strategy = build_strategy("channelwise") + assert isinstance(strategy, ChannelwiseStrategy) + assert strategy.contract_version == 1 + + +def test_panel_binding_preserves_requested_order_and_physical_indices(): + prepared = _prepare( + panel=("af", "CD3", "CK", "DAPI"), + config=_config((" dapi ", "CD3")), + ) + + assert prepared.selected_indices == [3, 1] + assert prepared.marker_records == [ + {"source_index": 3, "source_name": "DAPI", "canonical_name": "DAPI"}, + {"source_index": 1, "source_name": "CD3", "canonical_name": "CD3"}, + ] + + resolved = prepared.resolve_slide_context( + np.asarray([[[10, 20, 30, 40]]], dtype=np.uint8), + level=2, + source_dtype=np.uint8, + ) + bound = prepared.bind(resolved) + assert [row["source_index"] for row in bound.strategy_metadata["markers"]] == [ + 3, + 1, + ] + assert bound.panel_metadata["kept"] == ["DAPI", "CD3"] + + +def test_panel_provenance_preserves_exact_source_label_spelling(): + prepared = _prepare( + panel=(" CD3 ", "DAPI"), + config=_config(("CD3",)), + ) + assert prepared.marker_records == [ + { + "source_index": 0, + "source_name": " CD3 ", + "canonical_name": "CD3", + } + ] + + +def test_explicit_selection_can_exclude_unnamed_physical_channels(): + prepared = _prepare( + panel=("DAPI", "", "CD3"), config=_config(("CD3", "DAPI")) + ) + assert prepared.selected_indices == [2, 0] + + +@pytest.mark.parametrize( + ("panel", "channel_count", "config", "message"), + [ + (("A", "B"), 3, _config(("A",)), "does not match the physical C axis"), + (("A", "B"), 2, _config(("missing",)), "is not present"), + (("CD3", " cd3 "), 2, _config(("CD3",)), "is ambiguous"), + (("A", "", "C"), 3, _config(None), "every channel to be named"), + ], +) +def test_invalid_panel_bindings_fail_closed(panel, channel_count, config, message): + with pytest.raises(ValueError, match=message): + _prepare(panel=panel, channel_count=channel_count, config=config) + + +def test_missing_panel_metadata_fails_closed(): + with pytest.raises(ValueError, match="requires positional marker names"): + ChannelwiseStrategy().prepare( + base_embedder=TinyRGBEmbedder(), + channel_names=None, + channel_count=3, + config=_config(), + ) + + +def test_duplicate_requested_marker_fails_after_identity_normalization(): + with pytest.raises(ValueError, match="requested more than once"): + _prepare(panel=("CD3", "DAPI"), config=_config(("CD3", " cd3 "))) + + +def test_non_brightfield_model_cannot_be_used_as_the_base_encoder(): + with pytest.raises(ValueError, match="ordinary RGB encoder"): + _prepare(base=TinyRGBEmbedder(modality="multiplex")) + + +def test_unsafe_base_model_name_fails_before_a_zarr_key_is_derived(): + base = TinyRGBEmbedder() + base.spec = replace(base.spec, name="nested/model") + with pytest.raises(ValueError, match="safe Zarr path segment"): + _prepare(base=base) + + +def test_default_all_marker_panels_have_distinct_cohort_keys(): + first = _bound( + panel=("CD3", "DAPI"), + config=_config(None), + values=(10, 20), + ) + second = _bound( + panel=("CK", "DAPI"), + config=_config(None), + values=(10, 20), + ) + assert first.name != second.name + + +def test_channelwise_rejects_third_party_strategy_params(): + with pytest.raises(ValueError, match="does not accept --multiplex-params"): + _prepare( + config={ + **_config(), + "multiplex_strategy_params": {"adapter": "example"}, + } + ) + + +def test_reorder_exclusion_and_static_recipe_change_the_output_identity(): + variants = [ + _bound(config=_config(("A", "B", "C"))), + _bound(config=_config(("C", "B", "A"))), + _bound(config=_config(("A", "C")), values=(0, 128, 255)), + _bound(config=_config(("A", "B", "C"), aggregation="concat")), + ] + keys = [bound.name for bound in variants] + digests = [_full_fingerprint(bound)["digest"] for bound in variants] + + assert len(set(keys)) == len(keys) + assert len(set(digests)) == len(digests) + + +def test_irrelevant_percentile_bounds_do_not_fragment_dtype_keys(): + dtype_default = _bound(config=_config(normalization="dtype")) + dtype_other_bounds = _bound( + config={ + **_config(normalization="dtype"), + "multiplex_percentile_low": 10.0, + "multiplex_percentile_high": 90.0, + } + ) + assert dtype_default.name == dtype_other_bounds.name + + +def test_normalization_and_aggregation_change_full_patch_fingerprint(): + percentile = _bound( + config=_config(normalization="percentile"), values=(0, 100, 200) + ) + dtype = _bound(config=_config(normalization="dtype"), values=(0, 100, 200)) + concat = _bound( + config=_config(("A", "B", "C"), normalization="dtype", aggregation="concat"), + values=(0, 100, 200), + ) + + assert len( + { + _full_fingerprint(percentile)["digest"], + _full_fingerprint(dtype)["digest"], + _full_fingerprint(concat)["digest"], + } + ) == 3 + + +def test_slide_specific_normalization_moves_fingerprint_not_effective_key(): + first = _bound( + config=_config(normalization="percentile"), values=(0, 100, 200) + ) + second = _bound( + config=_config(normalization="percentile"), values=(20, 120, 220) + ) + + assert first.name == second.name + assert _full_fingerprint(first)["digest"] != _full_fingerprint(second)["digest"] + + +def test_concat_requires_explicit_order_and_expands_the_output_dimension(): + with pytest.raises(ValueError, match="explicit ordered --marker list"): + _prepare(config=_config(None, aggregation="concat")) + + bound = _bound(config=_config(("C", "A"), aggregation="concat")) + assert bound.embedding_dim == 4 + assert bound.spec.embedding_dim == 4 + assert bound.strategy_metadata["aggregation"] == {"name": "concat"} + + +def test_normalization_max_side_changes_static_and_full_output_identity(): + first = _bound(config=_config(max_side=2048)) + second = _bound(config=_config(max_side=1024)) + + assert first.name != second.name + assert _full_fingerprint(first)["digest"] != _full_fingerprint(second)["digest"] + assert ( + second.strategy_metadata["normalization"]["level_selection"]["max_side_px"] + == 1024 + ) + + +def test_channelwise_rejects_unimplemented_attention_aggregation(): + with pytest.raises(ValueError, match="available: mean, concat"): + _prepare(config=_config(aggregation="attention")) + + +@pytest.mark.parametrize("aggregation", ["mean", "concat"]) +def test_aggregation_matches_the_recorded_numerical_recipe(aggregation): + torch = pytest.importorskip("torch") + base = TinyRGBEmbedder() + bound = _bound( + base=base, + config=_config(("A", "B", "C"), aggregation=aggregation), + ) + patch = np.asarray([[[26, 102, 230]]], dtype=np.uint8) + transformed = bound.transform_batch([patch], "cpu") + marker_vectors = transformed[0] + actual = bound.embed_batch(transformed) + + if aggregation == "mean": + expected = marker_vectors.mean(dim=0, keepdim=True) + else: + expected = marker_vectors.reshape(1, -1) + + assert torch.allclose(actual, expected, atol=1e-7, rtol=0) + assert all(np.array_equal(rgb[..., 0], rgb[..., 1]) for rgb in base.last_rgb) + assert all(np.array_equal(rgb[..., 1], rgb[..., 2]) for rgb in base.last_rgb) + assert [int(rgb[0, 0, 0]) for rgb in base.last_rgb] == [26, 102, 230] + + +def test_many_markers_chunk_base_transform_and_forward_to_batch_budget(): + torch = pytest.importorskip("torch") + + class TrackingEmbedder(TinyRGBEmbedder): + def __init__(self): + super().__init__() + self.transform_batch_sizes = [] + self.forward_batch_sizes = [] + + def transform_batch(self, patches_hwc, device): + self.transform_batch_sizes.append(len(patches_hwc)) + return super().transform_batch(patches_hwc, device) + + def embed_batch(self, batch): + self.forward_batch_sizes.append(int(batch.shape[0])) + return super().embed_batch(batch) + + base = TrackingEmbedder() + config = {**_config(), "batch_size": 4} + bound = _bound(base=base, config=config) + patches = [np.asarray([[[10, 20, 30]]], dtype=np.uint8) for _ in range(3)] + transformed = bound.transform_batch(patches, "cpu") + output = bound.embed_batch(transformed) + + assert tuple(output.shape) == (3, 2) + assert max(base.transform_batch_sizes) <= 4 + assert max(base.forward_batch_sizes) <= 4 + assert sum(base.transform_batch_sizes) == sum(base.forward_batch_sizes) == 9 + assert torch.isfinite(output).all() + + +def test_percentile_quantization_uses_numpy_ties_to_even_rounding(): + base = TinyRGBEmbedder() + prepared = _prepare( + panel=("A",), + config=_config(("A",), normalization="percentile"), + base=base, + ) + resolved = prepared.resolve_slide_context( + np.asarray([[[0.0]], [[1.0]]], dtype=np.float32), + level=0, + source_dtype=np.float32, + ) + resolved["resolved"][0]["low"] = 0.0 + resolved["resolved"][0]["high"] = 1.0 + bound = prepared.bind(resolved) + bound.transform_batch([np.asarray([[[2.5 / 255.0]]], dtype=np.float32)], "cpu") + assert int(base.last_rgb[0][0, 0, 0]) == 2 + + +def test_float_input_transform_receives_unquantized_unit_interval(): + base = TinyFloatRGBEmbedder() + bound = _bound( + base=base, + config=_config(normalization="dtype"), + dtype=np.uint16, + values=(0, 32768, 65535), + ) + bound.transform_batch( + [np.asarray([[[0, 32768, 65535]]], dtype=np.uint16)], "cpu" + ) + + assert all(rgb.dtype == np.float32 for rgb in base.last_rgb) + assert [float(rgb[0, 0, 0]) for rgb in base.last_rgb] == pytest.approx( + [0.0, 32768 / 65535, 1.0] + ) + assert bound.strategy_metadata["encoder_input_contract"] == "float32_0_1" + assert bound.strategy_metadata["quantization"]["rounding"] == "none" + assert "cast_float32_unit_interval" in bound.strategy_metadata["operation_order"] + assert "quantize_uint8" not in bound.strategy_metadata["operation_order"] + + +def test_patch_pixels_recheck_nonfinite_and_float_dtype_range(): + percentile = _bound( + panel=("A",), + config=_config(("A",), normalization="percentile"), + values=(0.0,), + dtype=np.float32, + ) + with pytest.raises(ValueError, match="non-finite selected marker pixels"): + percentile.transform_batch( + [np.asarray([[[np.nan]]], dtype=np.float32)], "cpu" + ) + + dtype = _bound( + panel=("A",), + config=_config(("A",), normalization="dtype"), + values=(0.5,), + dtype=np.float32, + ) + with pytest.raises(ValueError, match=r"must be in \[0, 1\]"): + dtype.transform_batch([np.asarray([[[1.1]]], dtype=np.float32)], "cpu") + + +def test_resolved_normalization_rejects_row_count_or_marker_mismatch(): + prepared = _prepare() + valid = _resolved(prepared) + + with pytest.raises(ValueError, match="does not match selected markers"): + prepared.bind({**valid, "resolved": valid["resolved"][:-1]}) + + wrong = json.loads(json.dumps(valid)) + wrong["resolved"][0]["source_index"] = 2 + with pytest.raises(ValueError, match="does not match"): + prepared.bind(wrong) + + +@pytest.mark.parametrize( + ("low", "high", "message"), + [ + (float("nan"), 1.0, "finite low < high"), + (0.0, float("inf"), "finite low < high"), + (1.0, 1.0, "finite low < high"), + (2.0, 1.0, "finite low < high"), + ], +) +def test_bound_normalization_rejects_nonfinite_and_degenerate_ranges( + low, high, message +): + prepared = _prepare() + resolved = _resolved(prepared) + resolved["resolved"][0]["low"] = low + resolved["resolved"][0]["high"] = high + with pytest.raises(ValueError, match=message): + prepared.bind(resolved) + + +def test_resolving_normalization_rejects_nonfinite_or_constant_marker(): + prepared = _prepare( + panel=("A",), + config=_config(("A",), normalization="percentile"), + ) + with pytest.raises(ValueError, match="non-finite pixels"): + prepared.resolve_slide_context( + np.full((2, 2, 1), np.nan, dtype=np.float32), + level=0, + source_dtype=np.float32, + ) + with pytest.raises(ValueError, match="degenerate normalization"): + prepared.resolve_slide_context( + np.ones((2, 2, 1), dtype=np.float32), + level=0, + source_dtype=np.float32, + ) + + +def test_dtype_normalization_handles_signed_range_and_rejects_unscaled_float(): + signed = _prepare(config=_config(normalization="dtype")) + resolved = signed.resolve_slide_context( + np.zeros((1, 1, 3), dtype=np.int16), + level=0, + source_dtype=np.int16, + ) + assert resolved["resolved"][0]["low"] == np.iinfo(np.int16).min + assert resolved["resolved"][0]["high"] == np.iinfo(np.int16).max + + floating = _prepare(config=_config(normalization="dtype")) + with pytest.raises(ValueError, match=r"requires finite values in \[0, 1\]"): + floating.resolve_slide_context( + np.full((1, 1, 3), 2.0, dtype=np.float32), + level=0, + source_dtype=np.float32, + ) + + +def test_bound_panel_cannot_be_reused_with_changed_source_metadata(): + bound = _bound(panel=("A", "B", "C")) + assert bound.set_panel(["A", "B", "C"])["n_markers"] == 3 + with pytest.raises(ValueError, match="changed after channelwise strategy binding"): + bound.set_panel(["A", "C", "B"]) + + +@pytest.mark.parametrize( + ("patch", "message"), + [ + (np.zeros((2, 2), dtype=np.uint8), "HWC"), + (np.zeros((2, 2, 2), dtype=np.uint8), "C=3"), + (np.zeros((2, 2, 3), dtype=np.uint16), "dtype"), + ], +) +def test_patch_shape_channel_count_and_dtype_must_match_bound_source(patch, message): + bound = _bound(dtype=np.uint8) + with pytest.raises(ValueError, match=message): + bound.transform_batch([patch], "cpu") + + +def test_deriving_channelwise_output_does_not_mutate_brightfield_fingerprint(): + base = TinyRGBEmbedder() + before = patch_output_fingerprint(base.spec, "fp32") + before_bytes = json.dumps(before, sort_keys=True, separators=(",", ":")).encode() + + derived = _bound(base=base) + derived_fingerprint = _full_fingerprint(derived) + + after = patch_output_fingerprint(base.spec, "fp32") + after_bytes = json.dumps(after, sort_keys=True, separators=(",", ":")).encode() + assert after_bytes == before_bytes + assert base.spec.multiplex is None + assert derived_fingerprint != before + + +def test_readable_model_header_redacts_injected_base_credentials(): + base = TinyRGBEmbedder() + base.spec = replace( + base.spec, + source=( + "https://user:password@private.example/model?revision=v1&token=SENTINEL" + ), + transform_source_url=( + "https://private.example/transform.py?api_key=SECOND_SENTINEL" + ), + ) + bound = _bound(base=base) + fingerprint = _full_fingerprint(bound) + bound.spec = replace( + bound.spec, + multiplex=bound.multiplex_fingerprint_payload( + patch_output_fingerprint(base.spec, "fp32") + ), + ) + + header = _models_header( + [bound], + { + bound.name: { + "embedding_dim": bound.embedding_dim, + "output_fingerprint": fingerprint, + } + }, + )[bound.name] + persisted = json.dumps(header, sort_keys=True) + + assert "password" not in persisted + assert "SENTINEL" not in persisted + assert "SECOND_SENTINEL" not in persisted + assert header["source"] == "https://private.example/model?revision=v1" + assert header["transform_source_url"] == "https://private.example/transform.py" + assert valid_output_fingerprint(header["output_fingerprint"]) + + +def test_model_name_that_looks_like_a_secret_key_is_not_redacted(): + base = TinyRGBEmbedder() + base.spec = replace(base.spec, name="token") + fingerprint = patch_output_fingerprint(base.spec, "fp32") + + header = _models_header( + [base], + { + "token": { + "embedding_dim": base.embedding_dim, + "output_fingerprint": fingerprint, + } + }, + ) + + assert isinstance(header["token"], dict) + assert valid_output_fingerprint(header["token"]["output_fingerprint"]) diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py index 24de8a5..101099b 100644 --- a/tests/test_cli_commands.py +++ b/tests/test_cli_commands.py @@ -7,13 +7,16 @@ from __future__ import annotations import os +from dataclasses import replace from importlib import import_module import pytest +import typer from click.utils import strip_ansi from typer.testing import CliRunner from conftest import MockEmbedder, build_ngff_v04 +from raw2features.cli._validation import parse_channel_names_file from raw2features.cli.main import app try: @@ -24,6 +27,37 @@ _TORCH = False +def test_channel_names_file_parsers_are_format_and_path_independent(tmp_path): + txt = tmp_path / "channels.txt" + csv = tmp_path / "panel.csv" + tsv = tmp_path / "renamed.tsv" + txt.write_text("\ufeff# physical order\n CD3 \n\nCK\nDAPI\n") + csv.write_text('channel_name\nCD3\n"CK"\nDAPI\n') + tsv.write_text("marker\nCD3\nCK\nDAPI\n") + + assert parse_channel_names_file(str(txt)) == ["CD3", "CK", "DAPI"] + assert parse_channel_names_file(str(csv)) == ["CD3", "CK", "DAPI"] + assert parse_channel_names_file(str(tsv)) == ["CD3", "CK", "DAPI"] + + +@pytest.mark.parametrize( + ("name", "contents", "message"), + [ + ("empty.txt", "", "at least one"), + ("duplicate.txt", "CD3\ncd3\n", "unique"), + ("columns.csv", "CD3,CK\n", "one column"), + ("blank.csv", "channel_name\n\n", "at least one"), + ], +) +def test_channel_names_file_parser_rejects_ambiguous_inputs( + tmp_path, name, contents, message +): + path = tmp_path / name + path.write_text(contents) + with pytest.raises(typer.BadParameter, match=message): + parse_channel_names_file(str(path)) + + def test_embed_cli_seeds_runconfig_from_resolved_geometry(tmp_path, monkeypatch): """The high-level CLI must not encode "unset" as the real 1.0/224 geometry.""" embed_module = import_module("raw2features.cli.embed") @@ -82,6 +116,393 @@ def fake_contracts(cfg): assert captured["source_mpp"] == pytest.approx(0.37) +def test_verify_cli_passes_native_multiplex_legacy_segmenter_guard( + tmp_path, monkeypatch +): + verify_module = import_module("raw2features.cli.verify") + captured = {} + + monkeypatch.setattr(verify_module, "expected_model_contracts", lambda cfg: {}) + monkeypatch.setattr( + verify_module, + "resolve_multiplex_source_config", + lambda slide, cfg: replace( + cfg, + resolved_channel_names=["DAPI", "CD3"], + resolved_nuclear_channel_indices=[0], + resolved_original_channel_names=["DAPI", "CD3"], + ), + ) + + def fake_complete(*args, **kwargs): + captured.update(kwargs) + return True + + monkeypatch.setattr(verify_module, "is_complete", fake_complete) + result = CliRunner().invoke( + app, + [ + "verify", + str(tmp_path / "slide.zarr"), + "--receipts-dir", + str(tmp_path / "receipts"), + "-f", + "kronos", + "--device", + "cpu", + "--quiet", + ], + ) + + assert result.exit_code == 0, result.output + requirements = captured["compatible_grid_segmenters"] + assert requirements + assert set(next(iter(requirements.values())).values()) == {"nuclear"} + + +def _multiplex_cli_options(channel_names_file=None) -> list[str]: + options = [ + "--multiplex-strategy", + "channelwise", + "--marker", + "CD3", + "--marker", + "CK", + "--multiplex-normalization", + "percentile", + "--multiplex-percentile-low", + "2", + "--multiplex-percentile-high", + "98", + "--multiplex-aggregation", + "concat", + "--multiplex-normalization-max-side-px", + "1024", + ] + if channel_names_file is not None: + options.extend(["--channel-names-file", str(channel_names_file)]) + return options + + +def _assert_multiplex_config(cfg) -> None: + assert cfg.multiplex_strategy == "channelwise" + assert cfg.multiplex_markers == ["CD3", "CK"] + assert cfg.multiplex_normalization == "percentile" + assert cfg.multiplex_percentile_low == pytest.approx(2.0) + assert cfg.multiplex_percentile_high == pytest.approx(98.0) + assert cfg.multiplex_aggregation == "concat" + assert cfg.multiplex_normalization_max_side_px == 1024 + assert cfg.channel_names_override == ["CD3", "CK", "DAPI"] + + +def test_embed_cli_threads_multiplex_options(tmp_path, monkeypatch): + embed_module = import_module("raw2features.cli.embed") + captured = {} + + def fake_embed_slide(slide, out_dir, cfg, **kwargs): + captured["cfg"] = cfg + return {"status": "complete"} + + monkeypatch.setattr(embed_module, "embed_slide", fake_embed_slide) + names = tmp_path / "channels.txt" + names.write_text("CD3\nCK\nDAPI\n") + result = CliRunner().invoke( + app, + [ + "embed", + str(tmp_path / "unused.zarr"), + str(tmp_path / "out"), + *_multiplex_cli_options(names), + ], + ) + + assert result.exit_code == 0, result.output + _assert_multiplex_config(captured["cfg"]) + + +def test_embed_many_cli_threads_multiplex_options(tmp_path, monkeypatch): + embed_many_module = import_module("raw2features.cli.embed_many") + slides = tmp_path / "slides" + slides.mkdir() + (slides / "one.zarr").mkdir() + names = tmp_path / "channels.csv" + names.write_text("channel_name\nCD3\nCK\nDAPI\n") + captured = {} + + def fake_embed_shard_serial(shard, out_dir, cfg, *args): + captured["cfg"] = cfg + return 1, 0, 0 + + monkeypatch.setattr( + embed_many_module, "_embed_shard_serial", fake_embed_shard_serial + ) + monkeypatch.setattr( + embed_many_module, + "resolve_multiplex_source_config", + lambda slide, cfg: replace( + cfg, resolved_channel_names=list(cfg.channel_names_override) + ), + ) + result = CliRunner().invoke( + app, + [ + "embed-many", + str(slides), + str(tmp_path / "out"), + *_multiplex_cli_options(names), + ], + ) + + assert result.exit_code == 0, result.output + _assert_multiplex_config(captured["cfg"]) + + +def test_embed_many_preflights_every_panel_before_starting_worker( + tmp_path, monkeypatch +): + embed_many_module = import_module("raw2features.cli.embed_many") + slides = tmp_path / "slides" + slides.mkdir() + (slides / "one.zarr").mkdir() + (slides / "two.zarr").mkdir() + checked = [] + worker_started = False + + def fake_resolve(slide, cfg): + checked.append(os.path.basename(slide)) + if len(checked) == 2: + raise ValueError("channel count mismatch") + return cfg + + def fake_embed_shard_serial(*args, **kwargs): + nonlocal worker_started + worker_started = True + return 0, 0, 0 + + monkeypatch.setattr( + embed_many_module, "resolve_multiplex_source_config", fake_resolve + ) + monkeypatch.setattr( + embed_many_module, "_embed_shard_serial", fake_embed_shard_serial + ) + result = CliRunner().invoke( + app, + ["embed-many", str(slides), str(tmp_path / "out")], + ) + + assert result.exit_code == 1, result.output + assert checked == ["one.zarr", "two.zarr"] + assert not worker_started + assert "multiplex panel preflight failed" in strip_ansi(result.output) + + +def test_embed_many_rejects_unknown_strategy_before_starting_worker( + tmp_path, monkeypatch +): + embed_many_module = import_module("raw2features.cli.embed_many") + slides = tmp_path / "slides" + slides.mkdir() + (slides / "one.zarr").mkdir() + worker_started = False + + def fake_embed_shard_serial(*args, **kwargs): + nonlocal worker_started + worker_started = True + return 0, 0, 0 + + monkeypatch.setattr( + embed_many_module, "_embed_shard_serial", fake_embed_shard_serial + ) + result = CliRunner().invoke( + app, + [ + "embed-many", + str(slides), + str(tmp_path / "out"), + "--multiplex-strategy", + "not-a-strategy", + ], + ) + + assert result.exit_code == 1, result.output + assert not worker_started + assert "multiplex panel preflight failed" in strip_ansi(result.output) + + +@pytest.mark.parametrize( + ("extra_options", "message"), + [ + (["--slide-encoder", "titan"], "model-agnostic slide poolers"), + (["--devices", "cpu,cpu"], "single-device"), + ], +) +def test_embed_many_rejects_incompatible_strategy_execution_before_worker( + tmp_path, monkeypatch, extra_options, message +): + embed_many_module = import_module("raw2features.cli.embed_many") + slides = tmp_path / "slides" + slides.mkdir() + (slides / "one.zarr").mkdir() + worker_started = False + + def fake_embed_shard_serial(*args, **kwargs): + nonlocal worker_started + worker_started = True + return 0, 0, 0 + + monkeypatch.setattr( + embed_many_module, "_embed_shard_serial", fake_embed_shard_serial + ) + result = CliRunner().invoke( + app, + [ + "embed-many", + str(slides), + str(tmp_path / "out"), + "--multiplex-strategy", + "channelwise", + *extra_options, + ], + ) + + assert result.exit_code == 1, result.output + assert not worker_started + assert message in strip_ansi(result.output) + + +def test_verify_cli_threads_multiplex_options(tmp_path, monkeypatch): + verify_module = import_module("raw2features.cli.verify") + captured = {} + + def fake_resolve(slide, cfg, *, device): + captured["cfg"] = cfg + return cfg, {} + + monkeypatch.setattr( + verify_module, "resolve_multiplex_output_contracts", fake_resolve + ) + monkeypatch.setattr( + verify_module, "resolve_multiplex_source_config", lambda slide, cfg: cfg + ) + monkeypatch.setattr(verify_module, "is_complete", lambda *args, **kwargs: True) + names = tmp_path / "channels.tsv" + names.write_text("marker\nCD3\nCK\nDAPI\n") + result = CliRunner().invoke( + app, + [ + "verify", + str(tmp_path / "slide.zarr"), + "--receipts-dir", + str(tmp_path / "receipts"), + "--device", + "cpu", + "--quiet", + *_multiplex_cli_options(names), + ], + ) + + assert result.exit_code == 0, result.output + _assert_multiplex_config(captured["cfg"]) + + +def test_embed_cli_parses_namespaced_strategy_params(tmp_path, monkeypatch): + embed_module = import_module("raw2features.cli.embed") + captured = {} + + def fake_embed_slide(slide, out_dir, cfg, **kwargs): + captured["cfg"] = cfg + return {"status": "complete"} + + monkeypatch.setattr(embed_module, "embed_slide", fake_embed_slide) + result = CliRunner().invoke( + app, + [ + "embed", + str(tmp_path / "unused.zarr"), + str(tmp_path / "out"), + "--multiplex-strategy", + "third_party", + "--multiplex-params", + '{"adapter":{"width":3},"enabled":true}', + ], + ) + + assert result.exit_code == 0, result.output + assert captured["cfg"].multiplex_strategy_params == { + "adapter": {"width": 3}, + "enabled": True, + } + + +@pytest.mark.parametrize( + "args", + [ + ["embed", "slide.zarr", "out"], + ["embed-many", "slides", "out"], + ["verify", "slide.zarr", "--receipts-dir", "receipts"], + ], +) +def test_commands_reject_invalid_strategy_params_json(args): + result = CliRunner().invoke( + app, + [ + *args, + "--multiplex-strategy", + "third_party", + "--multiplex-params", + '{"bad": NaN}', + ], + ) + output = strip_ansi(result.output) + + assert result.exit_code == 2 + assert "--multiplex-params" in output + assert "finite JSON object" in output + + +@pytest.mark.parametrize( + "args", + [ + ["embed", "slide.zarr", "out"], + ["embed-many", "slides", "out"], + ["verify", "slide.zarr", "--receipts-dir", "receipts"], + ], +) +@pytest.mark.parametrize( + "bounds", + [ + ["--multiplex-percentile-low", "99", "--multiplex-percentile-high", "1"], + ["--multiplex-percentile-low", "nan"], + ["--multiplex-percentile-high", "101"], + ], +) +def test_commands_reject_invalid_multiplex_percentiles(args, bounds): + result = CliRunner().invoke(app, [*args, *bounds]) + output = strip_ansi(result.output) + + assert result.exit_code == 2 + assert "--multiplex-percentile" in output + assert "0 <= low < high <= 100" in output + + +@pytest.mark.parametrize( + "args", + [ + ["embed", "slide.zarr", "out"], + ["embed-many", "slides", "out"], + ["verify", "slide.zarr", "--receipts-dir", "receipts"], + ], +) +def test_commands_reject_nonpositive_multiplex_normalization_max_side(args): + result = CliRunner().invoke( + app, [*args, "--multiplex-normalization-max-side-px", "0"] + ) + output = strip_ansi(result.output) + + assert result.exit_code == 2 + assert "--multiplex-normalization-max-side-px" in output + + @pytest.mark.parametrize( "args", [ @@ -153,9 +574,7 @@ def test_geometry_options_reject_zero_before_work(command, option): def test_mpp_option_rejects_nonfinite_value_before_work(): - result = CliRunner().invoke( - app, ["embed", "slide.zarr", "out", "--mpp", "nan"] - ) + result = CliRunner().invoke(app, ["embed", "slide.zarr", "out", "--mpp", "nan"]) output = strip_ansi(result.output) assert result.exit_code == 2 diff --git a/tests/test_config_integrity.py b/tests/test_config_integrity.py index 50f333f..7586936 100644 --- a/tests/test_config_integrity.py +++ b/tests/test_config_integrity.py @@ -10,7 +10,7 @@ from __future__ import annotations import inspect -from dataclasses import fields +from dataclasses import fields, replace import pytest @@ -83,6 +83,254 @@ def test_grid_hash_is_model_independent_but_geometry_sensitive(): assert base != RunConfig(models=["resnet50"]).content_hash() +def test_multiplex_settings_affect_receipts_but_never_grid_identity(): + # Disable segmentation so only strategy settings vary; the default channelwise + # path deliberately uses nuclear rather than brightfield Otsu geometry. + brightfield = RunConfig(models=["uni"], no_seg=True) + multiplex = RunConfig( + models=["uni"], + no_seg=True, + multiplex_strategy="channelwise", + multiplex_markers=["CD3", "DAPI"], + multiplex_aggregation="mean", + ) + reordered = replace(multiplex, multiplex_markers=["DAPI", "CD3"]) + + assert multiplex.grid_hash() == brightfield.grid_hash() + assert reordered.grid_hash() == brightfield.grid_hash() + assert multiplex.content_hash() != brightfield.content_hash() + assert reordered.content_hash() != multiplex.content_hash() + + +def test_brightfield_identities_are_unchanged_by_multiplex_support(): + """Release guard: ordinary UNI stores/resume keep their exact v0.1 identity.""" + from raw2features.embedders.fingerprint import patch_output_fingerprint + from raw2features.embedders.model_registry import get_spec + + cfg = RunConfig(models=["uni"]) + assert cfg.grid_hash() == "a7410e019eac5b45" + assert cfg.content_hash() == "5543560bf2dc9a3a" + assert patch_output_fingerprint(get_spec("uni"), "fp32")["digest"] == ( + "4f6b83b5f6958108f0f68c24164c6dad7906a297982248a11472a7b1bec46680" + ) + + +def test_native_multiplex_accepts_only_its_historical_otsu_grid_alias(): + """Safe native legacy families require evidence of nuclear execution.""" + + legacy_otsu = RunConfig(models=["uni"]).grid_hash() + native = RunConfig( + models=["kronos"], + resolved_channel_names=["DAPI", "CD3"], + resolved_nuclear_channel_indices=[0], + resolved_original_channel_names=["DAPI", "CD3"], + ) + derived = RunConfig(models=["uni"], multiplex_strategy="channelwise") + legacy_nuclear = native.legacy_grid_hash() + + assert native.segmenter == "nuclear" + assert legacy_otsu in native.compatible_legacy_grid_hashes() + assert legacy_nuclear in native.compatible_legacy_grid_hashes() + assert native.compatible_legacy_grid_segmenters()[legacy_otsu] == "nuclear" + assert native.compatible_legacy_grid_segmenters()[legacy_nuclear] == "nuclear" + assert native.allows_hashless_legacy_grid() + assert legacy_otsu not in derived.compatible_legacy_grid_hashes() + assert not derived.allows_hashless_legacy_grid() + assert RunConfig(models=["uni"]).allows_hashless_legacy_grid() + + +@pytest.mark.parametrize( + ("original_names", "current_indices"), + [ + (["DNA1", "DNA2", "CD3"], [0, 1]), + (["DNA-PKcs", "DAPI", "CD3"], [1]), + (["", "DAPI", "CD3"], [1]), + ], +) +def test_native_multiplex_rejects_unproven_historical_nuclear_grids( + original_names, current_indices +): + legacy_otsu = RunConfig(models=["uni"]).grid_hash() + native = RunConfig( + models=["kronos"], + resolved_channel_names=original_names, + resolved_nuclear_channel_indices=current_indices, + resolved_original_channel_names=original_names, + ) + legacy_nuclear = native.legacy_grid_hash() + + assert legacy_otsu not in native.compatible_legacy_grid_hashes() + assert legacy_nuclear not in native.compatible_legacy_grid_hashes() + assert native.compatible_legacy_grid_segmenters() == {} + assert not native.allows_hashless_legacy_grid() + + +def test_multiplex_settings_move_content_identity_but_not_grid_identity(): + """A strategy changes model output/receipt identity, never patch geometry.""" + base = RunConfig(models=["uni"], no_seg=True) + multiplex = RunConfig( + models=["uni"], + no_seg=True, + multiplex_strategy="channelwise", + multiplex_markers=["CD3", "CK"], + multiplex_normalization="percentile", + multiplex_percentile_low=1.0, + multiplex_percentile_high=99.0, + multiplex_aggregation="mean", + multiplex_normalization_max_side_px=2048, + ) + reordered = RunConfig( + models=["uni"], + no_seg=True, + multiplex_strategy="channelwise", + multiplex_markers=["CK", "CD3"], + ) + + assert multiplex.grid_hash() == base.grid_hash() + assert multiplex.content_hash() != base.content_hash() + assert reordered.content_hash() != multiplex.content_hash() + + +def test_multiplex_actual_nuclear_segmentation_has_distinct_grid_identity(): + brightfield = RunConfig(models=["uni"]) + strategy = RunConfig(models=["uni"], multiplex_strategy="channelwise") + native = RunConfig(models=["kronos"]) + + assert brightfield.segmenter == "otsu" + assert strategy.segmenter == native.segmenter == "nuclear" + assert strategy.grid_hash() == native.grid_hash() + assert strategy.grid_hash() != brightfield.grid_hash() + + +def test_runconfig_preserves_the_v01_positional_constructor_prefix(): + expected_prefix = [ + "models", + "reader", + "segmenter", + "no_seg", + "target_mpp", + "source_mpp", + "patch_px", + "step_px", + "tissue_threshold", + "features_dtype", + "stain_norm", + "snap_to_level", + "mpp_tolerance", + "allow_upsample", + "amp", + ] + assert [field.name for field in fields(RunConfig)[:15]] == expected_prefix + + cfg = RunConfig( + ["uni"], + "omezarr", + "otsu", + True, + 0.5, + None, + 128, + 64, + 0.2, + "float32", + None, + True, + 0.01, + True, + "bf16", + ) + assert cfg.snap_to_level is True + assert cfg.mpp_tolerance == pytest.approx(0.01) + assert cfg.allow_upsample is True + assert cfg.amp == "bf16" + assert cfg.multiplex_strategy is None + + +def test_irrelevant_multiplex_knobs_do_not_move_receipt_identity(): + mean = RunConfig(models=["uni"], multiplex_strategy="channelwise") + dtype = replace(mean, multiplex_normalization="dtype") + changed_bounds = replace( + dtype, + multiplex_percentile_low=10.0, + multiplex_percentile_high=90.0, + ) + assert dtype.content_hash() == changed_bounds.content_hash() + + +def test_channel_identity_affects_outputs_but_only_nuclear_index_affects_grid(): + first = RunConfig( + models=["uni"], + no_seg=True, + multiplex_strategy="channelwise", + resolved_channel_names=["CD3", "DAPI"], + ) + renamed = replace(first, resolved_channel_names=["CK", "DAPI"]) + assert first.content_hash() != renamed.content_hash() + assert first.grid_hash() == renamed.grid_hash() + + segmented = replace( + first, + no_seg=False, + resolved_nuclear_channel_indices=[1], + ) + nonnuclear_rename = replace(segmented, resolved_channel_names=["CK", "DAPI"]) + moved_nucleus = replace(segmented, resolved_nuclear_channel_indices=[0]) + paired_nucleus = replace(segmented, resolved_nuclear_channel_indices=[0, 1]) + different_metadata_origin = replace( + segmented, resolved_original_channel_names=["", "DAPI"] + ) + assert segmented.grid_hash() == nonnuclear_rename.grid_hash() + assert segmented.grid_hash() != moved_nucleus.grid_hash() + assert segmented.grid_hash() != paired_nucleus.grid_hash() + assert segmented.content_hash() == different_metadata_origin.content_hash() + assert segmented.grid_hash() == different_metadata_origin.grid_hash() + + +def test_normalization_max_side_is_content_not_grid_identity(): + first = RunConfig( + models=["uni"], + no_seg=True, + multiplex_strategy="channelwise", + multiplex_normalization_max_side_px=2048, + ) + second = replace(first, multiplex_normalization_max_side_px=1024) + assert first.content_hash() != second.content_hash() + assert first.grid_hash() == second.grid_hash() + + with pytest.raises(ValueError, match="multiplex_normalization_max_side_px"): + replace(first, multiplex_normalization_max_side_px=0) + + +def test_namespaced_strategy_params_are_content_not_grid_identity(): + first = RunConfig( + models=["uni"], + no_seg=True, + multiplex_strategy="third_party", + multiplex_strategy_params={"adapter": {"width": 3}}, + ) + second = replace( + first, + multiplex_strategy_params={"adapter": {"width": 4}}, + ) + assert first.grid_hash() == second.grid_hash() + assert first.content_hash() != second.content_hash() + + +def test_strategy_params_must_be_finite_json_with_string_keys(): + with pytest.raises(ValueError, match="string keys"): + RunConfig( + models=["uni"], + multiplex_strategy="third_party", + multiplex_strategy_params={1: "bad"}, + ) + with pytest.raises(ValueError, match="finite JSON"): + RunConfig( + models=["uni"], + multiplex_strategy="third_party", + multiplex_strategy_params={"value": float("nan")}, + ) + + def test_amp_auto_is_default_and_resolves_per_model(): from raw2features.embedders.model_registry import get_spec from raw2features.pipeline.runner import _resolve_amp @@ -149,6 +397,13 @@ def test_embed_and_embed_many_share_content_option_defaults(): "no_seg", "tissue_threshold", "features_dtype", + "multiplex_strategy", + "multiplex_markers", + "multiplex_normalization", + "multiplex_percentile_low", + "multiplex_percentile_high", + "multiplex_aggregation", + "multiplex_normalization_max_side_px", "amp", "snap_to_level", "mpp_tolerance", diff --git a/tests/test_export_spatialdata.py b/tests/test_export_spatialdata.py index d188dde..912b879 100644 --- a/tests/test_export_spatialdata.py +++ b/tests/test_export_spatialdata.py @@ -7,9 +7,17 @@ from __future__ import annotations +import json +from dataclasses import replace + import numpy as np import pytest +from raw2features.embedders.fingerprint import ( + patch_output_fingerprint, + valid_output_fingerprint, +) +from raw2features.embedders.model_registry import get_spec from raw2features.export.spatialdata import export_spatialdata from raw2features.sinks.zarr_sink import ZarrSink @@ -222,6 +230,107 @@ def test_export_anisotropic_with_origin_translation(tmp_path): } +def test_channelwise_panel_mapping_is_queryable_without_losing_positions(): + pytest.importorskip("pandas") + from raw2features.export.spatialdata import _panel_dataframe + + frame = _panel_dataframe( + { + "uni__channelwise_mean_0123456789abcdef": { + "mapping": [ + { + "source_index": 3, + "source_name": "DAPI", + "canonical_name": "DAPI", + }, + { + "source_index": 1, + "source_name": "CD3", + "canonical_name": "CD3", + }, + ] + } + } + ) + + assert frame["channel"].tolist() == ["DAPI", "CD3"] + assert frame["channel_index"].tolist() == [3, 1] + assert frame["marker"].tolist() == ["DAPI", "CD3"] + + +def test_channelwise_panel_audit_lists_are_anndata_safe(): + from raw2features.export.spatialdata import _uns_safe_header + + base_fingerprint = patch_output_fingerprint(get_spec("uni"), "fp32") + fingerprint = patch_output_fingerprint( + replace( + get_spec("uni"), + multiplex={ + "strategy": "channelwise", + "contract_version": 1, + "base_output_fingerprint": base_fingerprint, + "aggregation": {"name": "mean"}, + "markers": [ + {"source_index": 0, "source_name": "DAPI"}, + {"source_index": 1, "source_name": "CD3"}, + ], + }, + ), + "fp32", + ) + + header = { + "panel": { + "derived": { + "source_channels": [ + {"source_index": 0, "source_name": "DAPI"}, + {"source_index": 1, "source_name": "CD3"}, + ], + "excluded": [{"source_index": 1, "source_name": "CD3"}], + "normalization": { + "resolved": [ + {"source_index": 0, "low": 1.0, "high": 99.0}, + {"source_index": 1, "low": 2.0, "high": 98.0}, + ] + }, + } + }, + "models": { + "derived": { + "output_fingerprint": fingerprint, + "multiplex": { + "markers": [ + {"source_index": 0, "source_name": "DAPI"}, + {"source_index": 1, "source_name": "CD3"}, + ] + } + } + }, + } + panel = _uns_safe_header(header)["panel"]["derived"] + assert panel["source_channels"] == { + "source_index": [0, 1], + "source_name": ["DAPI", "CD3"], + } + assert panel["excluded"] == { + "source_index": [1], + "source_name": ["CD3"], + } + assert panel["normalization"]["resolved"] == { + "source_index": [0, 1], + "low": [1.0, 2.0], + "high": [99.0, 98.0], + } + safe_model = _uns_safe_header(header)["models"]["derived"] + assert safe_model["multiplex"]["markers"] == { + "source_index": [0, 1], + "source_name": ["DAPI", "CD3"], + } + exported_fingerprint = json.loads(safe_model["output_fingerprint"]) + assert exported_fingerprint == fingerprint + assert valid_output_fingerprint(exported_fingerprint) + + def test_export_surfaces_multiplex_marker_panel(tmp_path): """A multiplex run's marker map must survive as a first-class, queryable table - not the stringified ``"{...}"`` blobs anndata makes of a raw list-of-dicts.""" diff --git a/tests/test_multiplex.py b/tests/test_multiplex.py index d1e0901..9f1266e 100644 --- a/tests/test_multiplex.py +++ b/tests/test_multiplex.py @@ -4,6 +4,8 @@ from __future__ import annotations +import warnings + import numpy as np import pytest @@ -14,6 +16,8 @@ def test_reader_exposes_channels_and_native_read(synthetic_multiplex_ngff): with OmeZarrReader(synthetic_multiplex_ngff) as r: assert r.channel_names == ["DAPI", "CD3", "CD8", "CD20", "FOXP3"] + assert r.has_channel_axis is True + assert r.channel_count == 5 reg = Region(level=0, location=Point(0, 0), size=Size(32, 32)) mc = r.read_region_channels(reg) assert mc.shape == (32, 32, 5) and mc.dtype == np.uint16 @@ -24,6 +28,78 @@ def test_reader_exposes_channels_and_native_read(synthetic_multiplex_ngff): def test_reader_channel_names_none_for_rgb(synthetic_ngff): with OmeZarrReader(synthetic_ngff) as r: assert r.channel_names is None # plain RGB -> no marker panel + assert r.has_channel_axis is True + assert r.channel_count == 3 + + +def test_reader_preserves_unnamed_omero_channel_positions( + synthetic_multiplex_ngff, +): + """An unnamed channel must not shift every later marker to the wrong pixels.""" + import zarr + + root = zarr.open_group(synthetic_multiplex_ngff, mode="r+") + root.attrs["omero"] = { + "channels": [ + {"label": "DAPI"}, + {"active": False}, + {"name": "CD8"}, + {"label": "CD20"}, + {"label": "FOXP3"}, + ] + } + + with OmeZarrReader(synthetic_multiplex_ngff) as reader: + assert reader.channel_names == ["DAPI", "", "CD8", "CD20", "FOXP3"] + assert len(reader.channel_names or []) == reader.channel_count == 5 + patch = reader.read_region_channels( + Region(level=0, location=Point(0, 0), size=Size(1, 1)) + ) + + # The label after the unnamed slot is still bound to physical channel index 2. + assert patch[0, 0, 2] == 3000 + + +def test_reader_surfaces_incomplete_omero_panel_without_padding( + synthetic_multiplex_ngff, +): + """The strategy binder can detect metadata-count/C-axis mismatches explicitly.""" + import zarr + + root = zarr.open_group(synthetic_multiplex_ngff, mode="r+") + root.attrs["omero"] = {"channels": [{"label": "DAPI"}, {"label": "CD3"}]} + + with OmeZarrReader(synthetic_multiplex_ngff) as reader: + assert reader.channel_names == ["DAPI", "CD3"] + assert reader.channel_count == 5 + + +def test_channel_name_override_fills_missing_metadata_positionally( + synthetic_multiplex_ngff, +): + import zarr + + root = zarr.open_group(synthetic_multiplex_ngff, mode="r+") + root.attrs["omero"] = {"channels": [{"label": "DAPI"}, {}, {"label": "CD8"}]} + supplied = ["dapi", "CD3", "CD8", "CD20", "FOXP3"] + + with OmeZarrReader(synthetic_multiplex_ngff) as reader: + reader.apply_channel_names(supplied) + assert reader.original_channel_names == ["DAPI", "", "CD8"] + assert reader.channel_names == supplied + assert reader.channel_names_source == "channel_names_file" + + +def test_channel_name_override_rejects_conflicts_counts_and_duplicates( + synthetic_multiplex_ngff, +): + with OmeZarrReader(synthetic_multiplex_ngff) as reader: + with pytest.raises(ValueError, match="C index 1"): + reader.apply_channel_names(["DAPI", "CK", "CD8", "CD20", "FOXP3"]) + with pytest.raises(ValueError, match="file has 2 names but C=5"): + reader.apply_channel_names(["DAPI", "CD3"]) + with pytest.raises(ValueError, match="unique name"): + reader.apply_channel_names(["DAPI", "CD3", "cd3", "CD20", "FOXP3"]) def test_nuclear_segmenter_thresholds_dapi(synthetic_multiplex_ngff): @@ -43,6 +119,101 @@ def test_nuclear_segmenter_errors_without_a_nuclear_channel(synthetic_multiplex_ seg.segment(r) +def test_nuclear_segmenter_combines_repeated_dapi_channels( + synthetic_multiplex_ngff, +): + from raw2features.segmenters.nuclear import NuclearSegmenter + + with OmeZarrReader(synthetic_multiplex_ngff): + assert NuclearSegmenter()._nuclear_indices( + ["DAPI1", "CD3", "DAPI2", "CD20", "FOXP3"] + ) == [0, 2] + + +def test_nuclear_segmenter_does_not_treat_dna_biomarker_as_nuclear_stain(): + from raw2features.segmenters.nuclear import NuclearSegmenter + + segmenter = NuclearSegmenter() + assert segmenter._nuclear_index(["DAPI", "DNA-PKcs", "CD3"]) == 0 + for non_nuclear_name in ["DNA3", "DNA12", "DNA-PKcs", "cDNA1", "DNA damage"]: + with pytest.raises(ValueError, match="no nuclear channel"): + segmenter._nuclear_index([non_nuclear_name, "CD3"]) + + +@pytest.mark.parametrize( + "name", + [ + "DAPI_cycle1", + "Hoechst 33342", + "DNA", + "DNA1", + "DNA 2", + "DNA1(Ir191)", + "191Ir_DNA1", + "Ir193_DNA2", + ], +) +def test_nuclear_segmenter_keeps_intended_nuclear_name_variants(name): + from raw2features.segmenters.nuclear import NuclearSegmenter + + assert NuclearSegmenter()._nuclear_index(["CD3", name]) == 1 + + +def test_nuclear_segmenter_combines_one_dna1_dna2_pair_without_uint16_overflow(): + from raw2features.segmenters.nuclear import NuclearSegmenter + + segmenter = NuclearSegmenter() + names = ["CD3", "DNA1(Ir191)", "DNA2(Ir193)"] + assert segmenter._nuclear_indices(names) == [1, 2] + block = np.zeros((2, 2, 3), dtype=np.uint16) + block[..., 1] = np.uint16(60_000) + block[..., 2] = np.uint16(50_000) + single = segmenter._combine_nuclear_channels(block, [1]) + combined = segmenter._combine_nuclear_channels(block, [1, 2]) + assert single.dtype == np.uint16 + assert np.all(single == 60_000) + assert combined.dtype == np.float32 + assert np.all(combined == 55_000.0) + + +def test_nuclear_segmenter_combines_metal_prefixed_dna_pair_and_repeated_dapi(): + from raw2features.segmenters.nuclear import NuclearSegmenter + + segmenter = NuclearSegmenter() + assert segmenter._nuclear_indices(["CD3", "191Ir_DNA1", "Ir193_DNA2"]) == [ + 1, + 2, + ] + assert segmenter._nuclear_indices(["DAPI1", "CD3", "DAPI2", "DAPI3"]) == [ + 0, + 2, + 3, + ] + + +@pytest.mark.parametrize( + "names", + [ + ["DAPI", "DNA1", "DNA2"], + ["DNA", "DNA2"], + ["DAPI1", "Hoechst1"], + ], +) +def test_nuclear_segmenter_rejects_noncanonical_multiple_nuclear_matches(names): + from raw2features.segmenters.nuclear import NuclearSegmenter + + with pytest.raises(ValueError, match="one nuclear stain family"): + NuclearSegmenter()._nuclear_indices(names) + + +def test_nuclear_segmenter_combines_repeated_numbered_dna_channels(): + from raw2features.segmenters.nuclear import NuclearSegmenter + + assert NuclearSegmenter()._nuclear_indices( + ["DNA1", "DNA1(Ir191)", "Ir193_DNA2"] + ) == [0, 1, 2] + + def test_kronos_spec_is_multiplex(): from raw2features.embedders.model_registry import get_spec @@ -75,12 +246,12 @@ def test_kronos_marker_resolver_handles_synonyms_and_compound_names(): vocab = {"CD4": 0, "CD162": 0, "CYTOKERATIN": 0, "GZMB": 0, "DAPI": 0} cases = { - "cd4": "CD4", # exact + "cd4": "CD4", # exact "pancytokeratin": "CYTOKERATIN", # synonym (pan- cocktail) - "granzymeb": "GZMB", # synonym (gene symbol) - "cla_cd162": "CD162", # compound name -> CD-number token - "cd62l": None, # genuinely absent -> dropped - "emptya488_1": None, # empty cycle -> normed to None + "granzymeb": "GZMB", # synonym (gene symbol) + "cla_cd162": "CD162", # compound name -> CD-number token + "cd62l": None, # genuinely absent -> dropped + "emptya488_1": None, # empty cycle -> normed to None } for raw, expected in cases.items(): assert _resolve_marker(_norm_marker(raw), vocab) == expected, raw @@ -112,6 +283,22 @@ def test_kronos_set_panel_records_retrievable_mapping(): assert "vocabulary" in s # where the canonical names/ids came from +def test_kronos_set_panel_does_not_warn_about_an_empty_channel_name(): + from raw2features.embedders.kronos_embedder import KronosEmbedder + from raw2features.embedders.model_registry import get_spec + + emb = KronosEmbedder(get_spec("kronos")) + emb._vocab = {"DAPI": (4, 0.1, 0.2, "DAPI")} + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + summary = emb.set_panel(["DAPI", ""]) + + assert not caught + assert summary["unmatched"] == [] + assert summary["dropped"] == [""] + + def test_multiplex_scaling_is_dtype_aware(): # The KRONOS input must reach [0,1] for any source dtype, not just uint16 - a # hard-coded /65535 silently squashes uint8 multiplex. (Pure helper, no weights.) diff --git a/tests/test_output_fingerprints.py b/tests/test_output_fingerprints.py index 3271c54..ba87f80 100644 --- a/tests/test_output_fingerprints.py +++ b/tests/test_output_fingerprints.py @@ -17,7 +17,12 @@ valid_output_fingerprint, ) from raw2features.pipeline.receipt import validate_model, validate_store_models -from raw2features.pipeline.runner import RunConfig, embed_slide, run_slide +from raw2features.pipeline.runner import ( + RunConfig, + embed_slide, + expected_model_contracts, + run_slide, +) from raw2features.sinks.zarr_sink import ZarrSink from raw2features.slide_embedders.encoding import slide_embedding_is_complete @@ -29,6 +34,30 @@ _TORCH = False +def test_native_multiplex_output_fingerprint_binds_the_positional_panel(): + first = RunConfig( + models=["kronos"], + no_seg=True, + resolved_channel_names=["CD3", "DAPI"], + device="cpu", + ) + second = replace(first, resolved_channel_names=["DAPI", "CD3"]) + + first_fingerprint = expected_model_contracts(first)["kronos"][ + "output_fingerprint" + ] + second_fingerprint = expected_model_contracts(second)["kronos"][ + "output_fingerprint" + ] + assert first_fingerprint["digest"] != second_fingerprint["digest"] + assert first_fingerprint["payload"]["multiplex_panel"] == { + "binding_contract_version": 1, + "channel_axis": "c", + "physical_channel_count": 2, + "effective_channel_names": ["CD3", "DAPI"], + } + + def _spec(**overrides) -> ModelSpec: values = { "name": "contract_model", @@ -139,6 +168,56 @@ def test_fingerprint_never_persists_uri_credentials_and_ignores_rotation(): assert not valid_output_fingerprint({**left, "payload": leaked_payload}) +def test_semantic_constructor_keys_are_not_mistaken_for_credentials(): + cls_token = _spec(timm_kwargs={"token": "cls", "signature": "cls_pool"}) + mean_token = replace( + cls_token, timm_kwargs={"token": "mean", "signature": "mean_pool"} + ) + + first = patch_output_fingerprint(cls_token, "fp32") + second = patch_output_fingerprint(mean_token, "fp32") + + assert first["digest"] != second["digest"] + assert first["payload"]["loader"]["constructor"]["timm_kwargs"] == { + "token": "cls", + "signature": "cls_pool", + } + assert second["payload"]["loader"]["constructor"]["timm_kwargs"] == { + "token": "mean", + "signature": "mean_pool", + } + + +def test_unambiguous_plaintext_credentials_are_redacted_from_fingerprints(): + first = _spec( + timm_kwargs={ + "api_key": "SECRET_A", + "credentials": "PRIVATE_A", + "token": "cls", + } + ) + second = replace( + first, + timm_kwargs={ + "api_key": "SECRET_B", + "credentials": "PRIVATE_B", + "token": "cls", + }, + ) + + left = patch_output_fingerprint(first, "fp32") + right = patch_output_fingerprint(second, "fp32") + + assert left == right + assert "SECRET" not in repr(left) + assert "PRIVATE" not in repr(left) + assert left["payload"]["loader"]["constructor"]["timm_kwargs"] == { + "api_key": "", + "credentials": "", + "token": "cls", + } + + def test_auto_amp_resolves_and_grid_hash_stays_geometry_only(): spec = _spec(name="x", inference_amp="bf16") automatic = expected_patch_outputs(["x"], "auto", specs={"x": spec})["x"] diff --git a/tests/test_pixel_and_resample_correctness.py b/tests/test_pixel_and_resample_correctness.py index 5b66504..e121d26 100644 --- a/tests/test_pixel_and_resample_correctness.py +++ b/tests/test_pixel_and_resample_correctness.py @@ -86,6 +86,15 @@ def test_resample_patch_downscale_and_noop(): assert same is patch # exact-size read is a no-op (no needless resample) +@pytest.mark.parametrize("dtype", [np.uint16, np.float32]) +def test_resample_patch_preserves_single_named_channel_axis_and_dtype(dtype): + patch = np.arange(32 * 32, dtype=dtype).reshape(32, 32, 1) + out = resample_patch(patch, 17) + + assert out.shape == (17, 17, 1) + assert out.dtype == dtype + + # -- C: full-array + all-zero-tail validation ------------------------------ def test_validate_output_rejects_unwritten_tail(tmp_path): n, dim = 300, 8 # tail (rows 256-299) is beyond the old 256-row window diff --git a/tests/test_receipt.py b/tests/test_receipt.py index 2d51a96..3100e0f 100644 --- a/tests/test_receipt.py +++ b/tests/test_receipt.py @@ -18,6 +18,7 @@ validate_store_models, write_receipt, ) +from raw2features.pipeline.runner import RunConfig from raw2features.sinks.zarr_sink import ZarrSink @@ -84,19 +85,20 @@ def test_is_complete_true_then_false_on_mismatch(tmp_path): n_patches=4, ), ) - assert is_complete( - rec_dir, - "s", - chash, - expected_source_uri="file:///x", - expected_output_uri=uri, - ) is True + assert ( + is_complete( + rec_dir, + "s", + chash, + expected_source_uri="file:///x", + expected_output_uri=uri, + ) + is True + ) # Legacy callers no longer raise, but fail closed without the current source. assert is_complete(rec_dir, "s", chash) is False # wrong config hash -> not complete - assert is_complete( - rec_dir, "s", "other", expected_source_uri="file:///x" - ) is False + assert is_complete(rec_dir, "s", "other", expected_source_uri="file:///x") is False def test_is_complete_false_when_output_missing_model(tmp_path): @@ -116,9 +118,7 @@ def test_is_complete_false_when_output_missing_model(tmp_path): n_patches=4, ), ) - assert is_complete( - rec_dir, "s", "h", expected_source_uri="file:///x" - ) is False + assert is_complete(rec_dir, "s", "h", expected_source_uri="file:///x") is False def test_expected_contract_mapping_must_cover_every_receipt_model(tmp_path): @@ -272,6 +272,76 @@ def test_per_grid_requirements_allow_unambiguous_legacy_hashless_store(tmp_path) ) +def test_unsafe_native_selector_disallows_hashless_receipt_completion(tmp_path): + uri = _make_output(tmp_path / "out") + native = RunConfig( + models=["kronos"], + resolved_channel_names=["", "DAPI", "CD3"], + resolved_nuclear_channel_indices=[1], + resolved_original_channel_names=["", "DAPI", "CD3"], + ) + expected_hash = native.grid_hash() + policy = {expected_hash: native.allows_hashless_legacy_grid()} + + assert not validate_store_models( + uri, + ["m"], + expected_grid_models={expected_hash: ["m"]}, + allow_hashless_legacy_grids=policy, + ) + + receipts = str(tmp_path / "receipts") + write_receipt( + receipts, + Receipt( + slide_id="s", + status="complete", + source_uri="file:///x", + output_uri=uri, + reader="omezarr", + models=["m"], + config_hash="unsafe-native", + n_patches=4, + ), + ) + assert not is_complete( + receipts, + "s", + "unsafe-native", + expected_source_uri="file:///x", + expected_output_uri=uri, + expected_grid_models={expected_hash: ["m"]}, + allow_hashless_legacy_grids=policy, + ) + + +def test_ambiguous_legacy_grid_alias_requires_live_segmenter_evidence(tmp_path): + import zarr + + uri = _make_output(tmp_path / "out") + root = zarr.open_group( + uri.removeprefix("file://"), mode="r+", use_consolidated=False + ) + group = root["grids"]["mpp1_px224"] + header = dict(group.attrs["raw2features"]) + header["grid_hash"] = "legacy-otsu-hash" + header["segmentation"] = {"segmenter": "otsu"} + group.attrs["raw2features"] = header + kwargs = { + "expected_grid_models": {"current-nuclear-hash": ["m"]}, + "compatible_grid_hashes": {"current-nuclear-hash": ("legacy-otsu-hash",)}, + "compatible_grid_segmenters": { + "current-nuclear-hash": {"legacy-otsu-hash": "nuclear"} + }, + } + + assert not validate_store_models(uri, ["m"], **kwargs) + + header["segmentation"] = {"segmenter": "nuclear"} + group.attrs["raw2features"] = header + assert validate_store_models(uri, ["m"], **kwargs) + + def test_is_complete_binds_receipt_to_source_and_requested_output(tmp_path): rec_dir = str(tmp_path / "rec") uri = _make_output(tmp_path / "out") @@ -289,9 +359,7 @@ def test_is_complete_binds_receipt_to_source_and_requested_output(tmp_path): ), ) - assert not is_complete( - rec_dir, "s", "h", expected_source_uri="file:///other" - ) + assert not is_complete(rec_dir, "s", "h", expected_source_uri="file:///other") assert not is_complete( rec_dir, "s", @@ -305,9 +373,7 @@ def test_is_complete_binds_receipt_to_source_and_requested_output(tmp_path): receipt["slide_id"] = "copied-from-another-slide" with open(os.path.join(rec_dir, "s.json"), "w", encoding="utf-8") as fh: json.dump(receipt, fh) - assert not is_complete( - rec_dir, "s", "h", expected_source_uri="file:///x" - ) + assert not is_complete(rec_dir, "s", "h", expected_source_uri="file:///x") def test_is_complete_checks_live_store_source_binding(tmp_path): @@ -336,9 +402,7 @@ def test_is_complete_checks_live_store_source_binding(tmp_path): header["source"] = {"uri": "file:///different"} root[GRIDS][key].attrs["raw2features"] = header - assert not is_complete( - rec_dir, "s", "h", expected_source_uri="file:///x" - ) + assert not is_complete(rec_dir, "s", "h", expected_source_uri="file:///x") def test_is_complete_canonicalises_rotated_remote_credentials(tmp_path): @@ -381,9 +445,7 @@ def test_corrupt_or_non_object_receipt_is_incomplete(tmp_path): path = rec_dir / "s.json" path.write_text('{"status": "complete"') assert read_receipt(str(rec_dir), "s") is None - assert not is_complete( - str(rec_dir), "s", "h", expected_source_uri="file:///x" - ) + assert not is_complete(str(rec_dir), "s", "h", expected_source_uri="file:///x") path.write_text("[]") assert read_receipt(str(rec_dir), "s") is None diff --git a/tests/test_slide_embedders.py b/tests/test_slide_embedders.py index acdf926..80a1fdb 100644 --- a/tests/test_slide_embedders.py +++ b/tests/test_slide_embedders.py @@ -182,6 +182,23 @@ def test_slide_registry_has_titan(): assert spec.doi == "10.1038/s41591-025-03982-3" +def test_native_multiplex_pool_keeps_legacy_slide_output_key(tmp_path): + from raw2features.embedders.model_registry import get_spec + from raw2features.slide_embedders.encoding import slide_output_key + + path = _write_slide_store( + tmp_path, + {"mpp0.5_px224": (224, {"kronos": np.zeros((2, 384), np.float32)})}, + ) + root = zarr.open_group(path, mode="r+", use_consolidated=False) + group = root["grids"]["mpp0.5_px224"] + group["features"]["kronos"].attrs["output_fingerprint"] = ( + patch_output_fingerprint(get_spec("kronos"), "fp32") + ) + + assert slide_output_key(group, "mean", "kronos") == "mean" + + def test_get_slide_spec_unknown_raises(): with pytest.raises(KeyError, match="Unknown slide encoder"): get_slide_spec("does_not_exist") diff --git a/tests/test_source_uris.py b/tests/test_source_uris.py index b61ec93..fbbd717 100644 --- a/tests/test_source_uris.py +++ b/tests/test_source_uris.py @@ -19,6 +19,8 @@ is_qualified_uri, is_remote_uri, join_uri_path, + redact_metadata_credentials, + redact_metadata_uri_credentials, redact_uri_credentials, source_uri, ) @@ -42,6 +44,38 @@ def test_unauthenticated_remote_source_uris_are_preserved_verbatim(uri): assert source_uri(uri) == uri +def test_untrusted_plugin_metadata_is_recursively_credential_free(): + sentinel = "R2F_PLUGIN_SECRET" + safe = redact_metadata_credentials( + { + "checkpoint": { + "url": ( + f"https://user:{sentinel}@weights.example/model.pt?" + f"part=1&token={sentinel}" + ), + "api_key": sentinel, + }, + "records": [{"source": f"s3://bucket/x?X-Amz-Signature={sentinel}"}], + } + ) + assert sentinel not in repr(safe) + assert safe["checkpoint"]["url"] == ( + "https://weights.example/model.pt?part=1" + ) + assert safe["checkpoint"]["api_key"] == "" + + +def test_semantic_metadata_keys_survive_uri_only_redaction(): + safe = redact_metadata_uri_credentials( + { + "constructor": {"token": "cls", "signature": "mean_pool"}, + "url": "https://weights.example/model.pt?part=1&token=SECRET", + } + ) + assert safe["constructor"] == {"token": "cls", "signature": "mean_pool"} + assert safe["url"] == "https://weights.example/model.pt?part=1" + + @pytest.mark.parametrize( ("uri", "expected"), [ diff --git a/uv.lock b/uv.lock index c5b8e65..0e97eab 100644 --- a/uv.lock +++ b/uv.lock @@ -3371,7 +3371,7 @@ wheels = [ [[package]] name = "raw2features" -version = "0.1.1" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "jsonschema" },