diff --git a/src/squidpy/experimental/im/_stain/__init__.py b/src/squidpy/experimental/im/_stain/__init__.py index e0982c7c1..9a6cf8fbd 100644 --- a/src/squidpy/experimental/im/_stain/__init__.py +++ b/src/squidpy/experimental/im/_stain/__init__.py @@ -18,10 +18,6 @@ apply_decomposition, fit_decomposition, ) -from squidpy.experimental.im._stain._mask import ( - absorbance_foreground_mask, - luminosity_foreground_mask, -) from squidpy.experimental.im._stain._normalize import ( estimate_white_point, fit_stain_reference, @@ -48,7 +44,6 @@ "StainFittingError", "StainMethod", "StainFit", - "absorbance_foreground_mask", "apply_decomposition", "apply_reinhard", "complement_third_column", @@ -57,7 +52,6 @@ "fit_reinhard", "fit_stain_reference", "lab_ruderman_to_rgb", - "luminosity_foreground_mask", "reorder_to_canonical", "rgb_to_lab_ruderman", "rgb_to_sda", diff --git a/src/squidpy/experimental/im/_stain/_decomposition.py b/src/squidpy/experimental/im/_stain/_decomposition.py index 42e4e1357..e4dc3b45b 100644 --- a/src/squidpy/experimental/im/_stain/_decomposition.py +++ b/src/squidpy/experimental/im/_stain/_decomposition.py @@ -10,6 +10,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Any import numpy as np @@ -158,19 +159,35 @@ def _max_concentrations(concentrations: np.ndarray) -> np.ndarray: def fit_decomposition( - image_rgb: xr.DataArray, + image_rgb: xr.DataArray | Sequence[xr.DataArray], method: StainMethod, params: Any, white_point: np.ndarray, *, - tissue_mask: np.ndarray | None = None, - image_key: str | None = None, + tissue_mask: np.ndarray | Sequence[np.ndarray | None] | None = None, + image_key: str | Sequence[str | None] | None = None, reference: dict[str, np.ndarray] = RUIFROK_HE, max_angle_deg: float = 45.0, ) -> StainFit: - """Fit a decomposition :class:`StainFit` (stain matrix + max concentrations).""" - od = _tissue_od(image_rgb, white_point, params["beta"], tissue_mask=tissue_mask, image_key=image_key) - matrix = _stain_matrix(od, method, params, image_key=image_key, reference=reference, max_angle_deg=max_angle_deg) + """Fit a decomposition :class:`StainFit` (stain matrix + max concentrations). + + Accepts one image or several: each image's tissue OD is gathered (naming the + image on empty tissue) and stacked into one ``(SUM_N, 3)`` array the fit runs + on. A single image is a pool of one. + """ + das = [image_rgb] if isinstance(image_rgb, xr.DataArray) else list(image_rgb) + masks = ( + [tissue_mask] * len(das) if tissue_mask is None or isinstance(tissue_mask, np.ndarray) else list(tissue_mask) + ) + keys = [image_key] * len(das) if image_key is None or isinstance(image_key, str) else list(image_key) + od = np.vstack( + [ + _tissue_od(da, white_point, params["beta"], tissue_mask=m, image_key=k) + for da, m, k in zip(das, masks, keys, strict=True) + ] + ) + key = keys[0] if len(keys) == 1 else None # the gate names a single image, not a pool + matrix = _stain_matrix(od, method, params, image_key=key, reference=reference, max_angle_deg=max_angle_deg) return StainFit( method=method, stain_matrix=matrix, diff --git a/src/squidpy/experimental/im/_stain/_mask.py b/src/squidpy/experimental/im/_stain/_mask.py index 833e3da18..258907ce4 100644 --- a/src/squidpy/experimental/im/_stain/_mask.py +++ b/src/squidpy/experimental/im/_stain/_mask.py @@ -16,11 +16,6 @@ RUDERMAN_LMS_TO_LAB, RUDERMAN_RGB_TO_LMS, ) -from squidpy.experimental.im._stain._conversion import ( - _check_channel_dim, - rgb_to_lab_ruderman, - rgb_to_sda, -) def _white_luminosity() -> float: @@ -51,31 +46,6 @@ def foreground_mask_from_lab(lab: xr.DataArray, threshold: float) -> xr.DataArra return luminosity <= threshold -def luminosity_foreground_mask(rgb: xr.DataArray, threshold: float) -> xr.DataArray: - """Boolean tissue mask from normalised Ruderman Lab luminosity. - - Pixels darker than ``threshold`` are tissue; brighter (near-white) - pixels are background and excluded from stain-statistic fitting. - - Parameters - ---------- - rgb - Image with a ``"c"`` dimension of length 3. May be numpy- or - dask-backed; the operation stays lazy. - threshold - Cutoff on normalised luminosity in ``[0, 1]``. Pixels with - luminosity ``<= threshold`` are tissue. Semantics follow - HistomicsTK's ``reinhard`` so literature thresholds transfer. - - Returns - ------- - Boolean ``(y, x)`` DataArray: ``True`` = tissue, ``False`` = - background. Lazy if and only if ``rgb`` was lazy. - """ - _check_channel_dim(rgb) - return foreground_mask_from_lab(rgb_to_lab_ruderman(rgb), threshold) - - def as_spatial_mask(mask: np.ndarray, like: xr.DataArray) -> xr.DataArray: """Wrap a ``(y, x)`` boolean array as a DataArray aligned to ``like``'s y/x. @@ -96,28 +66,3 @@ def foreground_mask_from_sda(sda: xr.DataArray, beta: float = 0.15) -> xr.DataAr RGB->SDA conversion once. ``True`` = tissue (mean absorbance ``> beta``). """ return sda.mean(dim="c") > beta - - -def absorbance_foreground_mask(rgb: xr.DataArray, white_point: np.ndarray, beta: float = 0.15) -> xr.DataArray: - """Boolean tissue mask in optical-density (absorbance) space. - - The convention the Macenko/Vahadane fits expect: a pixel is tissue if its - mean absorbance across channels exceeds ``beta``. Near-white background - has near-zero absorbance and is excluded. - - Parameters - ---------- - rgb - Image with a ``"c"`` dimension of length 3. Numpy- or dask-backed. - white_point - Per-channel white point ``I_0`` (shape ``(3,)``), as used by - :func:`~squidpy.experimental.im._stain._conversion.rgb_to_sda`. - beta - Mean-absorbance cutoff. Pixels with mean SDA ``> beta`` are tissue. - - Returns - ------- - Boolean ``(y, x)`` DataArray: ``True`` = tissue. Lazy if ``rgb`` was lazy. - """ - _check_channel_dim(rgb) - return foreground_mask_from_sda(rgb_to_sda(rgb, white_point), beta) diff --git a/src/squidpy/experimental/im/_stain/_normalize.py b/src/squidpy/experimental/im/_stain/_normalize.py index b6dc869bd..d91428da7 100644 --- a/src/squidpy/experimental/im/_stain/_normalize.py +++ b/src/squidpy/experimental/im/_stain/_normalize.py @@ -33,7 +33,12 @@ validate_macenko_params, validate_vahadane_params, ) -from squidpy.experimental.im._stain._reference import StainFit, StainMethod +from squidpy.experimental.im._stain._reference import ( + _DECOMPOSITION_METHODS, + _VALID_METHODS, + StainFit, + StainMethod, +) from squidpy.experimental.im._stain._reinhard import ( apply_reinhard, fit_reinhard, @@ -64,8 +69,6 @@ "macenko": validate_macenko_params, "vahadane": validate_vahadane_params, } -_VALID_METHODS = tuple(_METHOD_PARAMS) -_DECOMPOSITION_METHODS = ("macenko", "vahadane") _CONCENTRATION_CHANNELS = ["hematoxylin", "eosin", "residual"] # Public union accepted by the method_params argument of the dispatchers. @@ -124,6 +127,14 @@ def _resolve_tissue_bool_mask( return mask +def _resolve_fit_inputs( + sdata: sd.SpatialData, *, image_key: str, scale: str, tissue_mask_key: str | None +) -> tuple[xr.DataArray, np.ndarray]: + da = _resolve_image(sdata, image_key=image_key, scale=scale, prefer="coarsest") + validate_rgb_range(da) + return da, _resolve_tissue_bool_mask(sdata, image_key=image_key, fit_da=da, tissue_mask_key=tissue_mask_key) + + def _resolve_output_tissue_mask( sdata: sd.SpatialData, *, image_key: str, target_da: xr.DataArray, tissue_mask_key: str | None ) -> xr.DataArray: @@ -209,21 +220,19 @@ def estimate_white_point( Shape-``(3,)`` white point; pass it as ``white_point`` to :func:`fit_stain_reference` / :meth:`~squidpy.experimental.im.StainFit.decompose`. """ - da = _resolve_image(sdata, image_key=image_key, scale=scale, prefer="coarsest") - validate_rgb_range(da) - tissue_mask = _resolve_tissue_bool_mask(sdata, image_key=image_key, fit_da=da, tissue_mask_key=tissue_mask_key) + da, tissue_mask = _resolve_fit_inputs(sdata, image_key=image_key, scale=scale, tissue_mask_key=tissue_mask_key) return white_point_from_background(da, ~tissue_mask) def fit_stain_reference( sdata: sd.SpatialData, *, - image_key: str, + image_key: str | list[str], method: StainMethod = "macenko", scale: str | Literal["auto"] = "auto", method_params: MethodParams = None, white_point: np.ndarray | None = None, - tissue_mask_key: str | None = None, + tissue_mask_key: str | list[str] | None = None, max_angle_deg: float = 45.0, canonical_reference: Mapping[str, np.ndarray] | None = None, ) -> StainFit: @@ -234,7 +243,12 @@ def fit_stain_reference( sdata SpatialData object containing the image. image_key - Key of the RGB image in ``sdata.images`` to fit on. + Key of the RGB image in ``sdata.images`` to fit on, or a **list of keys** + to fit one reference from the pooled tissue pixels of several images + (e.g. a representative cohort). Pooled images must share a dtype. The + pool is a pool of *pixels*: each image contributes in proportion to its + tissue pixel count at the resolved ``scale``, not equally, so images + whose pyramids bottom out at different sizes carry different weight. method Fitting method: ``"macenko"`` (default) or ``"vahadane"`` (physical stain-matrix decomposition, usable by both :meth:`~squidpy.experimental.im.StainFit.transform` and @@ -259,7 +273,9 @@ def fit_stain_reference( :func:`!detect_tissue`) restricting the fit to tissue pixels. If ``None``, ``f"{image_key}_tissue"`` is used. A tissue mask is **required**: if neither exists, a :class:`KeyError` asks you to - run :func:`!detect_tissue` first. + run :func:`!detect_tissue` first. When ``image_key`` is a list, pass a + list of mask keys **order-matched** to it (or ``None`` for the + ``{key}_tissue`` convention per image). max_angle_deg Tolerance of the H/E sanity gate for the decomposition methods: the fit raises :class:`!StainFittingError` if either recovered stain vector @@ -276,22 +292,45 @@ def fit_stain_reference( The fitted :class:`~squidpy.experimental.im.StainFit`. Nothing is written to ``sdata``. """ if method not in _VALID_METHODS: - raise ValueError(f"Unknown method {method!r}; expected one of {list(_VALID_METHODS)}.") - da = _resolve_image(sdata, image_key=image_key, scale=scale, prefer="coarsest") - validate_rgb_range(da) + raise ValueError(f"Unknown method {method!r}; expected one of {sorted(_VALID_METHODS)}.") + # Normalise to lists: a single image is a pool of one, so both cases share one path. + keys = [image_key] if isinstance(image_key, str) else list(image_key) + if not keys: + raise ValueError("`image_key` is empty; pass at least one image key.") + if len(set(keys)) != len(keys): + dupes = sorted({k for k in keys if keys.count(k) > 1}) + raise ValueError(f"`image_key` has duplicate keys: {dupes}; each image may be pooled only once.") + if tissue_mask_key is None: + mask_keys: list[str | None] = [None] * len(keys) + else: + mask_keys = [tissue_mask_key] if isinstance(tissue_mask_key, str) else list(tissue_mask_key) + if len(mask_keys) != len(keys): + raise ValueError( + f"`tissue_mask_key` has {len(mask_keys)} entries but `image_key` has {len(keys)}; " + "pass one mask key per image (order-matched) or None for the `{image_key}_tissue` convention." + ) + params = resolve_params(method_params, _METHOD_PARAMS[method], validate=_METHOD_VALIDATORS[method]) - tissue_mask = _resolve_tissue_bool_mask(sdata, image_key=image_key, fit_da=da, tissue_mask_key=tissue_mask_key) + resolved = [ + _resolve_fit_inputs(sdata, image_key=k, scale=scale, tissue_mask_key=mk) + for k, mk in zip(keys, mask_keys, strict=True) + ] + das, masks = [d for d, _ in resolved], [m for _, m in resolved] + dtypes = {str(da.dtype) for da in das} + if len(dtypes) != 1: + raise ValueError(f"pooled images must share a dtype; got {sorted(dtypes)}.") + if method == "reinhard": - return fit_reinhard(da, params, tissue_mask=tissue_mask) - bg = default_white_point(da) if white_point is None else np.asarray(white_point, np.float64) + return fit_reinhard(das, params, tissue_mask=masks, image_key=keys) + bg = default_white_point(das[0]) if white_point is None else np.asarray(white_point, np.float64) reference = RUIFROK_HE if canonical_reference is None else dict(canonical_reference) return fit_decomposition( - da, + das, method, params, bg, - tissue_mask=tissue_mask, - image_key=image_key, + tissue_mask=masks, + image_key=keys, reference=reference, max_angle_deg=max_angle_deg, ) @@ -321,9 +360,8 @@ def _normalize_stains( # Source statistics (Reinhard mu/sigma or the decomposition source matrix) # are reduced on a coarse level with a tissue mask; the lazy transform is # then applied to the full-resolution `da`. - fit_rgb = _resolve_image(sdata, image_key=image_key, scale=scale, prefer="coarsest") - validate_rgb_range(fit_rgb) # reject mis-typed source (e.g. 0-255 float) before the dtype-clipped reconstruction - tissue_mask = _resolve_tissue_bool_mask(sdata, image_key=image_key, fit_da=fit_rgb, tissue_mask_key=tissue_mask_key) + # the range check rejects a mis-typed source (e.g. 0-255 float) before the dtype-clipped reconstruction + fit_rgb, tissue_mask = _resolve_fit_inputs(sdata, image_key=image_key, scale=scale, tissue_mask_key=tissue_mask_key) out_dtype = da.dtype if output_dtype is None else np.dtype(output_dtype) # clip range + final cast if reference.method == "reinhard": normalized = apply_reinhard( diff --git a/src/squidpy/experimental/im/_stain/_reference.py b/src/squidpy/experimental/im/_stain/_reference.py index 8cdfa904f..1d75be59c 100644 --- a/src/squidpy/experimental/im/_stain/_reference.py +++ b/src/squidpy/experimental/im/_stain/_reference.py @@ -7,7 +7,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, get_args import numpy as np @@ -19,8 +19,9 @@ from squidpy.experimental.im._stain._normalize import MethodParams StainMethod = Literal["macenko", "vahadane", "reinhard"] -_DECOMPOSITION_METHODS: frozenset[str] = frozenset({"macenko", "vahadane"}) -_VALID_METHODS: frozenset[str] = _DECOMPOSITION_METHODS | {"reinhard"} +#: The method vocabulary, derived from ``StainMethod`` so the two cannot drift. +_VALID_METHODS: frozenset[str] = frozenset(get_args(StainMethod)) +_DECOMPOSITION_METHODS: frozenset[str] = _VALID_METHODS - {"reinhard"} def _coerce_finite(arr: Any, *, shape: tuple[int, ...], name: str) -> np.ndarray: diff --git a/src/squidpy/experimental/im/_stain/_reinhard.py b/src/squidpy/experimental/im/_stain/_reinhard.py index 39902b9b5..efebc013a 100644 --- a/src/squidpy/experimental/im/_stain/_reinhard.py +++ b/src/squidpy/experimental/im/_stain/_reinhard.py @@ -10,6 +10,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Any import numpy as np @@ -24,6 +25,7 @@ ) from squidpy.experimental.im._stain._mask import as_spatial_mask, foreground_mask_from_lab from squidpy.experimental.im._stain._reference import StainFit +from squidpy.experimental.im._stain._validation import StainFittingError from squidpy.types import ReinhardParams # Numerical safeguard against divide-by-zero on flat (constant-colour) @@ -39,33 +41,6 @@ def validate_reinhard_params(params: dict[str, Any]) -> None: raise ValueError(f"`luminosity_threshold` must be in (0, 1], got {params['luminosity_threshold']}.") -def _masked_channel_stats(lab: xr.DataArray, mask: xr.DataArray | None) -> tuple[np.ndarray, np.ndarray]: - """Per-channel mean and std over the spatial dims, tissue pixels only. - - Lazy: the masked mean and std are bundled into one dataset and computed - in a single pass, never materialising the full image. Returns two - shape-``(3,)`` float64 arrays in channel order. - - Raises ``ValueError`` if the mask leaves no tissue pixels in any channel - (the mean would be NaN), with an actionable message. - """ - masked = lab.where(mask) if mask is not None else lab - stats = xr.Dataset( - { - "mu": masked.mean(dim=("y", "x"), skipna=True), - "sigma": masked.std(dim=("y", "x"), skipna=True), - } - ).compute() - mu = np.asarray(stats["mu"].values, dtype=np.float64) - sigma = np.asarray(stats["sigma"].values, dtype=np.float64) - if not (np.all(np.isfinite(mu)) and np.all(np.isfinite(sigma))): - raise ValueError( - "Foreground mask leaves zero tissue pixels in at least one channel; " - "the luminosity_threshold may be too low or the image may be blank." - ) - return mu, sigma - - def _transfer_kernel( x: np.ndarray, *, @@ -89,18 +64,67 @@ def _reinhard_mask(lab: xr.DataArray, params: ReinhardParams, tissue_mask: np.nd return None -def fit_reinhard(image_rgb: xr.DataArray, params: ReinhardParams, *, tissue_mask: np.ndarray | None = None) -> StainFit: - """Fit Reinhard channel statistics on a reference image. - - Converts to Ruderman Lab, computes per-channel ``mu``/``sigma`` over - tissue pixels, and packs them into a ``StainFit(method="reinhard")``. - ``tissue_mask`` (a ``(y, x)`` boolean aligned to ``image_rgb``) selects the - tissue pixels when given; otherwise the ``mask_background`` / - ``luminosity_threshold`` params drive the mask. - """ +def _tissue_lab_moments( + image_rgb: xr.DataArray, params: ReinhardParams, tissue_mask: np.ndarray | None, *, image_key: str | None +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: _check_channel_dim(image_rgb) lab = rgb_to_lab_ruderman(image_rgb) - mu, sigma = _masked_channel_stats(lab, _reinhard_mask(lab, params, tissue_mask)) + mask = _reinhard_mask(lab, params, tissue_mask) + masked = lab.where(mask) if mask is not None else lab + # Accumulate in float64: this is a reduction, not a map, so the float32 + # working dtype is not enough. sigma falls out of E[x^2] - E[x]^2, whose two + # terms are ~5000x the variance they bracket; in float32 that cancellation + # costs ~1% of sigma and makes the answer depend on chunking. The cast is + # elementwise, so it fuses into the per-chunk graph and stays lazy. + wide = masked.astype(np.float64) + stats = xr.Dataset( + { + "n": wide.count(dim=("y", "x")), + "s": wide.sum(dim=("y", "x"), skipna=True), + "s2": (wide**2).sum(dim=("y", "x"), skipna=True), + } + ).compute() + n = np.asarray(stats["n"].values, dtype=np.float64) + if not np.all(n > 0): + raise StainFittingError( + "foreground mask leaves zero tissue pixels in at least one channel; " + "the luminosity_threshold may be too low or the image may be blank.", + image_key=image_key, + ) + return n, np.asarray(stats["s"].values, dtype=np.float64), np.asarray(stats["s2"].values, dtype=np.float64) + + +def _stats_from_moments(n: np.ndarray, s: np.ndarray, s2: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + mu = s / n + return mu, np.sqrt(np.maximum(s2 / n - mu**2, 0.0)) + + +def fit_reinhard( + image_rgb: xr.DataArray | Sequence[xr.DataArray], + params: ReinhardParams, + *, + tissue_mask: np.ndarray | Sequence[np.ndarray | None] | None = None, + image_key: str | Sequence[str | None] | None = None, +) -> StainFit: + """Fit Reinhard channel statistics on one or more reference images. + + Converts to Ruderman Lab and computes per-channel ``mu``/``sigma`` + (population std) over the tissue pixels. Several images pool into one + reference by summing their per-channel moments, so a slide contributes in + proportion to its **tissue pixel count**, not equally. A single image is a + pool of one. ``tissue_mask`` (``(y, x)`` booleans aligned to each image) + selects the tissue pixels when given; otherwise the ``mask_background`` / + ``luminosity_threshold`` params drive the mask. ``image_key`` only names the + image in error messages. + """ + das = [image_rgb] if isinstance(image_rgb, xr.DataArray) else list(image_rgb) + masks = ( + [tissue_mask] * len(das) if tissue_mask is None or isinstance(tissue_mask, np.ndarray) else list(tissue_mask) + ) + keys = [image_key] * len(das) if image_key is None or isinstance(image_key, str) else list(image_key) + moments = [_tissue_lab_moments(da, params, m, image_key=k) for da, m, k in zip(das, masks, keys, strict=True)] + pooled = (np.sum(np.stack(x), axis=0) for x in zip(*moments, strict=True)) + mu, sigma = _stats_from_moments(*pooled) return StainFit(method="reinhard", mu=mu, sigma=sigma) @@ -124,8 +148,8 @@ def apply_reinhard( Lazy if and only if ``image_rgb`` is lazy. """ _check_channel_dim(image_rgb) - fit_lab = rgb_to_lab_ruderman(fit_rgb if fit_rgb is not None else image_rgb) - mu_src, sigma_src = _masked_channel_stats(fit_lab, _reinhard_mask(fit_lab, params, tissue_mask)) + fit_src = fit_rgb if fit_rgb is not None else image_rgb + mu_src, sigma_src = _stats_from_moments(*_tissue_lab_moments(fit_src, params, tissue_mask, image_key=None)) sigma_src = np.maximum(sigma_src, _SIGMA_FLOOR) lab = rgb_to_lab_ruderman(image_rgb) diff --git a/tests/experimental/test_stain_mask.py b/tests/experimental/test_stain_mask.py index c0e70f23b..467e35a67 100644 --- a/tests/experimental/test_stain_mask.py +++ b/tests/experimental/test_stain_mask.py @@ -5,9 +5,10 @@ import pytest import xarray as xr +from squidpy.experimental.im._stain._conversion import rgb_to_lab_ruderman, rgb_to_sda from squidpy.experimental.im._stain._mask import ( - absorbance_foreground_mask, - luminosity_foreground_mask, + foreground_mask_from_lab, + foreground_mask_from_sda, ) _WHITE = np.array([255.0, 255.0, 255.0]) @@ -22,49 +23,49 @@ class TestLuminosityForegroundMask: @pytest.mark.parametrize("chunked", [False, True]) def test_all_white_is_background(self, chunked: bool) -> None: white = np.full((3, 16, 16), 255.0) - mask = luminosity_foreground_mask(_rgb_dataarray(white, chunked=chunked), 0.8) + mask = foreground_mask_from_lab(rgb_to_lab_ruderman(_rgb_dataarray(white, chunked=chunked)), 0.8) assert mask.dims == ("y", "x") assert not bool(mask.values.any()) @pytest.mark.parametrize("chunked", [False, True]) def test_all_black_is_tissue(self, chunked: bool) -> None: black = np.zeros((3, 16, 16)) - mask = luminosity_foreground_mask(_rgb_dataarray(black, chunked=chunked), 0.8) + mask = foreground_mask_from_lab(rgb_to_lab_ruderman(_rgb_dataarray(black, chunked=chunked)), 0.8) assert bool(mask.values.all()) def test_half_split(self) -> None: values = np.zeros((3, 8, 16)) values[:, :, 8:] = 255.0 # right half white - mask = luminosity_foreground_mask(_rgb_dataarray(values, chunked=False), 0.8) + mask = foreground_mask_from_lab(rgb_to_lab_ruderman(_rgb_dataarray(values, chunked=False)), 0.8) assert bool(mask.values[:, :8].all()) assert not bool(mask.values[:, 8:].any()) def test_lazy_in_lazy_out(self) -> None: values = np.full((3, 16, 16), 100.0) - mask = luminosity_foreground_mask(_rgb_dataarray(values, chunked=True), 0.8) + mask = foreground_mask_from_lab(rgb_to_lab_ruderman(_rgb_dataarray(values, chunked=True)), 0.8) assert isinstance(mask.data, da.Array) def test_non_three_channel_raises(self) -> None: values = np.zeros((2, 8, 8)) with pytest.raises(ValueError, match="3-channel RGB"): - luminosity_foreground_mask(xr.DataArray(values, dims=("c", "y", "x")), 0.8) + foreground_mask_from_lab(rgb_to_lab_ruderman(xr.DataArray(values, dims=("c", "y", "x"))), 0.8) class TestAbsorbanceForegroundMask: def test_white_is_background_dark_is_tissue(self) -> None: values = np.full((3, 8, 16), 255.0) values[:, :, 8:] = 30.0 # dark right half = high absorbance = tissue - mask = absorbance_foreground_mask(_rgb_dataarray(values, chunked=False), _WHITE) + mask = foreground_mask_from_sda(rgb_to_sda(_rgb_dataarray(values, chunked=False), _WHITE)) assert mask.dims == ("y", "x") assert not bool(mask.values[:, :8].any()) assert bool(mask.values[:, 8:].all()) def test_lazy_in_lazy_out(self) -> None: values = np.full((3, 16, 16), 50.0) - mask = absorbance_foreground_mask(_rgb_dataarray(values, chunked=True), _WHITE) + mask = foreground_mask_from_sda(rgb_to_sda(_rgb_dataarray(values, chunked=True), _WHITE)) assert isinstance(mask.data, da.Array) def test_non_three_channel_raises(self) -> None: values = np.zeros((2, 8, 8)) with pytest.raises(ValueError, match="3-channel RGB"): - absorbance_foreground_mask(xr.DataArray(values, dims=("c", "y", "x")), _WHITE) + foreground_mask_from_sda(rgb_to_sda(xr.DataArray(values, dims=("c", "y", "x")), _WHITE)) diff --git a/tests/experimental/test_stain_normalize.py b/tests/experimental/test_stain_normalize.py index ae24a14f6..4343b85fc 100644 --- a/tests/experimental/test_stain_normalize.py +++ b/tests/experimental/test_stain_normalize.py @@ -13,14 +13,19 @@ from spatialdata.transformations import Scale, get_transformation, set_transformation import squidpy as sq -from squidpy._params import defaults_of +from squidpy._params import defaults_of, resolve_params from squidpy.experimental.im import ( StainFit, fit_stain_reference, ) +from squidpy.experimental.im._stain._constants import RUIFROK_HE +from squidpy.experimental.im._stain._decomposition import fit_decomposition +from squidpy.experimental.im._stain._normalize import _resolve_image +from squidpy.experimental.im._stain._reinhard import fit_reinhard from squidpy.experimental.im._stain._validation import StainFittingError +from squidpy.experimental.im._stain._white_point import default_white_point from squidpy.experimental.im._utils import get_element_data -from squidpy.types import ReinhardParams +from squidpy.types import MacenkoParams, ReinhardParams from tests.conftest import PlotTester, PlotTesterMeta _ = sdp # registers the `.pl` spatialdata accessor @@ -285,3 +290,110 @@ def test_plot_reinhard_before_after(self, sdata_hne) -> None: _, axes = plt.subplots(1, 2, figsize=(8, 4)) sdata_hne.pl.render_images("hne_shifted").pl.show(ax=axes[0], title="before") sdata_hne.pl.render_images("hne_normalized").pl.show(ax=axes[1], title="after") + + +# --------------------------------------------------------------------------- +# Multi-slide pooled fit (one reference from several images in one sdata) +# --------------------------------------------------------------------------- + + +class TestPooledFit: + @staticmethod + def _he(seed: int, shape: tuple[int, int] = (48, 48), dtype=np.uint8) -> np.ndarray: + """Synthetic H&E from the Ruifrok H/E vectors so macenko/vahadane can fit.""" + rng = np.random.default_rng(seed) + h, w = shape + wmat = np.stack([RUIFROK_HE["hematoxylin"], RUIFROK_HE["eosin"]], axis=1) # (3, 2) + conc = rng.uniform(0.05, 1.3, (h * w, 2)) + rgb = np.clip(255.0 * np.exp(-(conc @ wmat.T)), 0, 255).reshape(h, w, 3).transpose(2, 0, 1) + arr = rgb.astype(np.uint8) + return (arr.astype(np.uint16) * 257) if dtype == np.uint16 else arr + + @classmethod + def _add(cls, sdata: sd.SpatialData, key: str, *, seed: int, dtype=np.uint8, tissue: np.ndarray | None = None): + arr = cls._he(seed=seed, dtype=dtype) + sdata.images[key] = Image2DModel.parse(arr, dims=("c", "y", "x")) + mask = np.ones(arr.shape[-2:], dtype=np.uint32) if tissue is None else tissue + sdata.labels[f"{key}_tissue"] = Labels2DModel.parse(mask, dims=("y", "x")) + + def _cohort(self, n: int = 3) -> tuple[sd.SpatialData, list[str]]: + sdata = sd.SpatialData() + keys = [f"img{i}" for i in range(n)] + for i, k in enumerate(keys): + self._add(sdata, k, seed=i + 1) + return sdata, keys + + @pytest.mark.parametrize("method", ["reinhard", "macenko", "vahadane"]) + def test_pooled_fit(self, method: str) -> None: + sdata, keys = self._cohort(n=3) + pooled = fit_stain_reference(sdata, image_key=keys, method=method) + assert pooled.method == method + assert (pooled.mu.shape == (3,)) if method == "reinhard" else (pooled.stain_matrix.shape == (3, 3)) + # a pool of one is the single-image fit, bit for bit (same code path) + single = fit_stain_reference(sdata, image_key=keys[0], method=method) + one = fit_stain_reference(sdata, image_key=keys[:1], method=method) + # NMF (vahadane) is not bit-reproducible even when seeded (threaded BLAS); the rest is + for attr in ("mu", "sigma") if method == "reinhard" else ("stain_matrix", "max_concentrations"): + if method == "vahadane": + np.testing.assert_allclose(getattr(one, attr), getattr(single, attr), rtol=1e-4) + else: + np.testing.assert_array_equal(getattr(one, attr), getattr(single, attr)) + + @pytest.mark.parametrize("method", ["reinhard", "macenko"]) + def test_pool_of_duplicates_is_single(self, method: str) -> None: + # pooling an image with itself must not move the reference: catches a + # botched concatenate/weighting that `pool of one == single` cannot. + sdata, keys = self._cohort(n=1) + da = _resolve_image(sdata, image_key=keys[0], scale="auto", prefer="coarsest") + mask = np.ones(da.shape[-2:], dtype=bool) + if method == "reinhard": + single = fit_reinhard(da, resolve_params(None, ReinhardParams), tissue_mask=mask) + double = fit_reinhard([da, da], resolve_params(None, ReinhardParams), tissue_mask=[mask, mask]) + attrs = ("mu", "sigma") + else: + params, wp = resolve_params(None, MacenkoParams), default_white_point(da) + single = fit_decomposition(da, method, params, wp, tissue_mask=mask) + double = fit_decomposition([da, da], method, params, wp, tissue_mask=[mask, mask]) + attrs = ("stain_matrix", "max_concentrations") + # reinhard pools moments, so it is exact; macenko's angular percentile + # interpolates between order statistics, so doubling N shifts it slightly. + rtol = 1e-12 if method == "reinhard" else 1e-3 + for attr in attrs: + np.testing.assert_allclose(getattr(double, attr), getattr(single, attr), rtol=rtol, atol=rtol) + + def test_order_matched_non_convention_masks(self) -> None: + # non-convention mask names selecting different halves; swapping the order + # must change the fit (proves order is honoured, not name-matched). + sdata, keys = self._cohort(n=2) + h, w = 48, 48 + top = np.zeros((h, w), np.uint32) + top[: h // 2] = 1 + bot = np.zeros((h, w), np.uint32) + bot[h // 2 :] = 1 + sdata.labels["m_a"] = Labels2DModel.parse(top, dims=("y", "x")) + sdata.labels["m_b"] = Labels2DModel.parse(bot, dims=("y", "x")) + ref = fit_stain_reference(sdata, image_key=keys, method="reinhard", tissue_mask_key=["m_a", "m_b"]) + swapped = fit_stain_reference(sdata, image_key=keys, method="reinhard", tissue_mask_key=["m_b", "m_a"]) + assert not np.allclose(ref.mu, swapped.mu) + + @pytest.mark.parametrize( + ("kwargs", "extra", "match"), + [ + ({"image_key": []}, None, "empty"), + ({"image_key": ["img0", "img0"]}, None, "duplicate"), + ({"image_key": ["img0", "img1"], "tissue_mask_key": "img0_tissue"}, None, "one mask key per image"), + ({"image_key": ["img0", "img1"], "tissue_mask_key": ["img0_tissue"]}, None, "one mask key per image"), + ({"image_key": "img0", "tissue_mask_key": ["img0_tissue", "img1_tissue"]}, None, "one mask key per image"), + ({"image_key": ["img0", "img16"]}, "uint16", "share a dtype"), + ({"image_key": ["img0", "blank"], "method": "macenko"}, "blank", r"^\[blank\] "), + ({"image_key": ["img0", "blank"], "method": "reinhard"}, "blank", r"^\[blank\] "), + ], + ) + def test_validation(self, kwargs: dict, extra: str | None, match: str) -> None: + sdata, _ = self._cohort(n=2) + if extra == "uint16": + self._add(sdata, "img16", seed=9, dtype=np.uint16) + elif extra == "blank": # empty tissue mask -> the error names the slide + self._add(sdata, "blank", seed=9, tissue=np.zeros((48, 48), np.uint32)) + with pytest.raises((ValueError, StainFittingError), match=match): + fit_stain_reference(sdata, **{"method": "macenko", **kwargs}) diff --git a/tests/experimental/test_stain_reinhard.py b/tests/experimental/test_stain_reinhard.py index 628819eb9..9e8f2fc8b 100644 --- a/tests/experimental/test_stain_reinhard.py +++ b/tests/experimental/test_stain_reinhard.py @@ -6,13 +6,16 @@ import xarray as xr from squidpy._params import resolve_params +from squidpy.experimental.im._stain._conversion import rgb_to_lab_ruderman from squidpy.experimental.im._stain._reference import StainFit from squidpy.experimental.im._stain._reinhard import ( _SIGMA_FLOOR, - _masked_channel_stats, + _stats_from_moments, + _tissue_lab_moments, apply_reinhard, fit_reinhard, ) +from squidpy.experimental.im._stain._validation import StainFittingError from squidpy.types import ReinhardParams @@ -33,34 +36,35 @@ def rgb_b() -> np.ndarray: return rng.uniform(60.0, 190.0, size=(3, 32, 32)) -class TestMaskedChannelStats: +class TestTissueLabMoments: + @staticmethod + def _stats(rgb: xr.DataArray, mask: np.ndarray | None) -> tuple[np.ndarray, np.ndarray]: + params = resolve_params({"mask_background": False}, ReinhardParams) + return _stats_from_moments(*_tissue_lab_moments(rgb, params, mask, image_key=None)) + @pytest.mark.parametrize("chunked", [False, True]) def test_matches_numpy(self, chunked: bool) -> None: rng = np.random.default_rng(7) values = rng.uniform(0.0, 10.0, size=(3, 16, 16)) - mask_np = rng.random((16, 16)) > 0.3 - lab = _da(values, chunked=chunked) - mask = xr.DataArray(mask_np, dims=("y", "x")) + mask = rng.random((16, 16)) > 0.3 - mu, sigma = _masked_channel_stats(lab, mask) + mu, sigma = self._stats(_da(values, chunked=chunked), mask) - sel = values[:, mask_np] + sel = np.asarray(rgb_to_lab_ruderman(_da(values, chunked=False)).values)[:, mask] np.testing.assert_allclose(mu, sel.mean(axis=1), atol=1e-6) np.testing.assert_allclose(sigma, sel.std(axis=1), atol=1e-6) def test_numpy_dask_identical(self) -> None: rng = np.random.default_rng(8) values = rng.uniform(0.0, 10.0, size=(3, 16, 16)) - mu_n, sigma_n = _masked_channel_stats(_da(values, chunked=False), None) - mu_d, sigma_d = _masked_channel_stats(_da(values, chunked=True), None) + mu_n, sigma_n = self._stats(_da(values, chunked=False), None) + mu_d, sigma_d = self._stats(_da(values, chunked=True), None) np.testing.assert_allclose(mu_n, mu_d, atol=1e-6) np.testing.assert_allclose(sigma_n, sigma_d, atol=1e-6) def test_empty_mask_raises(self) -> None: - lab = _da(np.ones((3, 8, 8)), chunked=False) - mask = xr.DataArray(np.zeros((8, 8), dtype=bool), dims=("y", "x")) - with pytest.raises(ValueError, match="zero tissue pixels"): - _masked_channel_stats(lab, mask) + with pytest.raises(StainFittingError, match="zero tissue pixels"): + self._stats(_da(np.ones((3, 8, 8)), chunked=False), np.zeros((8, 8), dtype=bool)) class TestFitReinhard: