Skip to content

Review follow-up: dedupe the stain fit reductions - #21

Closed
selmanozleyen wants to merge 2 commits into
stain/pr4-pooled-fitfrom
stain/pr4-pooled-fit-review
Closed

selmanozleyen wants to merge 2 commits into
stain/pr4-pooled-fitfrom
stain/pr4-pooled-fit-review

Conversation

@selmanozleyen

@selmanozleyen selmanozleyen commented Oct 9, 2026 •

Copy link
Copy Markdown
Owner

Review of scverse#1209, applied. Base is that PR's head, so this diff is only the review changes.

Correctness / efficiency

The Reinhard fit stopped being lazy. _tissue_lab_pixels materialised a 3xHxW float64 frame per image to take a mean and a std. The function it replaced documented the opposite ("never materialising the full image"). It now reduces to per-channel (count, sum, sum of squares) in one bundled pass; pooling is an element-wise sum of moments, so nothing is resident. A/B on a 2000x2000 pool of two, dask-chunked:

before   peak 103.2 MiB   sigma [0.10411141 0.17843567 0.04403144]
after    peak  37.9 MiB   sigma [0.10411182 0.17843678 0.04403172]

The reductions have to run in float64. _working_dtype is float32, and a float32 sum of squares over a whole slide loses the variance to cancellation (~1% on L) and stops being chunk-stable. The cast is per-chunk, so it stays lazy. Chunked and eager now agree to 1e-10. The before sigma above is itself slightly off for this reason.

One error convention. StainFittingError(reason, image_key=...) exists precisely for attributing a failure to a slide in a cohort fit, and the decomposition path already used it. The Reinhard path raised a bare ValueError in its own format, so a caller skipping bad slides needed two excepts and two parsers.

Deduplication

_masked_channel_stats is gone. apply_reinhard standardises by the source image's own tissue statistics -- that is a Reinhard fit of the source. So the module held two implementations of the same reduction, and they had already diverged (fit in float64 after the fix above, apply in float32). Both paths now run _tissue_lab_moments -> _stats_from_moments. Net -1 function; apply_reinhard drops its own Lab conversion, stays lazy, and inherits float64 for free.

_resolve_fit_inputs. "Resolve the coarse image, validate its range, resolve its tissue mask" was written out three times (estimate_white_point, fit_stain_reference, _normalize_stains). The validate_rgb_range step is exactly the kind of guard that gets added to one site and not the others.

The method vocabulary. _VALID_METHODS and _DECOMPOSITION_METHODS were each declared twice, in _reference.py and _normalize.py, alongside the StainMethod literal that is the real source of truth. A fourth method meant four edits and a silent half-registration. Now derived once via get_args(StainMethod).

Docs / tests

  • image_key docstring now states the pool is a pool of pixels: each slide weighted by tissue pixel count at the resolved scale, not equally. Verified empirically; worth knowing when scale="auto" resolves per-image.
  • match="blank" could not fail -- it also matched the message's own trailing "may be blank." Now match=r"^\[blank\] ", which only the key produces. pytest.raises((ValueError, RuntimeError)) -> (ValueError, StainFittingError).
  • New test_pool_of_duplicates_is_single: pooling an image with itself must not move the reference. This is what catches a broken concatenate or weighting; pool-of-one == single cannot, since both sides run the identical branch. It found a real property in passing -- exact for Reinhard, ~2e-5 for Macenko, because np.percentile interpolates between order statistics and doubling N moves the interpolation point.
  • The three _masked_channel_stats tests move onto _tissue_lab_moments, so they now exercise the path both callers use.
  • Duplicate-key check is len(set(keys)) and names the offending key; _unique_order_preserving import dropped.

148 stain tests pass locally, including the normalize_stain image baselines -- the float64 shift on the apply path did not move a pixel.

Left alone

  • The das/masks/keys normalisation is still duplicated across the two fits. It is plumbing with no numerics, so divergence costs nothing, and every home for a shared helper is a module it does not belong in.
  • _reinhard_mask and the mask branch in _tissue_od look alike but resolve in different spaces against different params; _mask.py says the two variants exist on purpose.
    (luminosity_foreground_mask / absorbance_foreground_mask were in this list as dead-but-out-of-scope; they are now deleted -- see below.)

Dead code

luminosity_foreground_mask and absorbance_foreground_mask are gone. Neither had a production caller: every fit path composes the conversion and the *_from_lab / *_from_sda primitive itself, precisely so the mask and the statistics read from the same lazy graph and dask materialises the conversion once. Only tests reached the wrappers. Both conversions already guard the channel dim, so the tests keep their coverage by composing the primitives the way the fits do.

@selmanozleyen
selmanozleyen force-pushed the stain/pr4-pooled-fit-review branch from a7a01ee to 7bdae93 Compare October 9, 2026 20:24
- fit_reinhard pools per-channel (count, sum, sum of squares) instead of
  materialising the tissue pixels, so the Reinhard fit is lazy again and a
  pooled cohort never holds a frame resident (103 -> 38 MiB peak on a
  2000x2000 pool of two).
- accumulate those moments in float64. sigma falls out of E[x^2] - E[x]^2,
  whose terms are ~5000x the variance they bracket; in float32 that
  cancellation costs ~1% of sigma and makes the answer depend on chunking
  (eager 0.10272 vs chunked 0.10410 on the same slide).
- drop _masked_channel_stats. apply_reinhard was a second implementation of
  the same reduction and had already diverged from the fit; both now go
  through _tissue_lab_moments + _stats_from_moments.
- raise StainFittingError (with image_key) on empty tissue, matching the
  decomposition path, instead of a bare ValueError in its own format.
- single _resolve_fit_inputs for the coarse image + tissue mask, shared by
  estimate_white_point, fit_stain_reference and _normalize_stains.
- derive the method vocabulary from StainMethod once in _reference.py; drop
  the second copy in _normalize.py.
- document that the pool is a pool of pixels, weighted by tissue pixel count.
- tests: assert the error names the slide (the old pattern also matched the
  message's own trailing "may be blank"), and add pool-of-duplicates
  invariance, which pool-of-one cannot catch.
@selmanozleyen
selmanozleyen force-pushed the stain/pr4-pooled-fit-review branch from 7bdae93 to 3c52bdf Compare October 9, 2026 20:30
luminosity_foreground_mask and absorbance_foreground_mask had no production
caller: every fit path composes the conversion and the *_from_lab/_from_sda
primitives itself, so the mask reads from the same lazy graph the statistics
do, and dask materialises the conversion once. Only tests reached the
wrappers. Both conversions already guard the channel dim, so the tests keep
their coverage by composing the primitives the way the fits do.
@selmanozleyen
selmanozleyen force-pushed the stain/pr4-pooled-fit-review branch from 3c52bdf to 7b12ad0 Compare October 9, 2026 20:32
@selmanozleyen

Copy link
Copy Markdown
Owner Author

Landed directly on scverse#1209 (commits 0b961ed, 7b12ad0). Closing as redundant.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant