Multi-slide pooled fit_stain_reference

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",

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,

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)

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(

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: