From 30aad3ad8038990c13d358d00d1c3eb243d7e779 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Mon, 3 Aug 2026 15:21:39 -0700 Subject: [PATCH 01/11] Add RL/RLGC deconvolution for 3D fluorescence --- docs/examples/cli/configs/fluorescence_2d.yml | 5 +- docs/examples/cli/configs/fluorescence_3d.yml | 5 +- tests/models/test_rlgc.py | 267 ++++++++++++++++++ waveorder/api/_settings.py | 4 +- waveorder/api/fluorescence.py | 24 +- .../models/isotropic_fluorescent_thick_3d.py | 67 ++++- .../models/isotropic_fluorescent_thin_3d.py | 21 +- waveorder/models/isotropic_thin_3d.py | 7 +- waveorder/models/phase_thick_3d.py | 7 +- waveorder/rlgc.py | 225 +++++++++++++++ 10 files changed, 611 insertions(+), 21 deletions(-) create mode 100644 tests/models/test_rlgc.py create mode 100644 waveorder/rlgc.py diff --git a/docs/examples/cli/configs/fluorescence_2d.yml b/docs/examples/cli/configs/fluorescence_2d.yml index e01946ac..b1839a77 100644 --- a/docs/examples/cli/configs/fluorescence_2d.yml +++ b/docs/examples/cli/configs/fluorescence_2d.yml @@ -15,7 +15,10 @@ fluorescence: wavelength_emission: 0.532 # emission wavelength in micrometers confocal_pinhole_diameter: null # confocal pinhole diameter (null = widefield) apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization + reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' filters or 'RL'/'RLGC' iterative deconvolution regularization_strength: 0.01 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization + rl_iterations: 25 # maximum RL / RLGC iterations + rl_background: 0.0 # constant background folded into the RL / RLGC Poisson forward model + rl_stopping_tolerance: null # relative-change early-stop threshold for RL / RLGC (null = run all iterations) diff --git a/docs/examples/cli/configs/fluorescence_3d.yml b/docs/examples/cli/configs/fluorescence_3d.yml index 29803bfa..47498f94 100644 --- a/docs/examples/cli/configs/fluorescence_3d.yml +++ b/docs/examples/cli/configs/fluorescence_3d.yml @@ -15,7 +15,10 @@ fluorescence: wavelength_emission: 0.532 # emission wavelength in micrometers confocal_pinhole_diameter: null # confocal pinhole diameter (null = widefield) apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization + reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' filters or 'RL'/'RLGC' iterative deconvolution regularization_strength: 0.001 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization + rl_iterations: 25 # maximum RL / RLGC iterations + rl_background: 0.0 # constant background folded into the RL / RLGC Poisson forward model + rl_stopping_tolerance: null # relative-change early-stop threshold for RL / RLGC (null = run all iterations) diff --git a/tests/models/test_rlgc.py b/tests/models/test_rlgc.py new file mode 100644 index 00000000..a489ca78 --- /dev/null +++ b/tests/models/test_rlgc.py @@ -0,0 +1,267 @@ +"""Correctness tests for Richardson-Lucy (RL) and Gradient-Consensus (RLGC). + +These tests check the operator-agnostic core (:mod:`waveorder.rlgc`) and its +wiring into 3D fluorescence reconstruction: + +* the FFT forward/adjoint operators used for deconvolution are true adjoints, +* RL and RLGC sharpen a Poisson-noisy bead simulation, +* over-iterated RL overfits the noise into a "starry night" of spurious + bright voxels while RLGC resists it, +* RL/RLGC are refused for 2D fluorescence and for phase/birefringence. +""" + +import pytest +import torch + +from waveorder import rlgc +from waveorder.api import fluorescence, phase +from waveorder.models import ( + isotropic_fluorescent_thick_3d as thick, +) +from waveorder.models import ( + isotropic_fluorescent_thin_3d as thin, +) +from waveorder.models import isotropic_thin_3d as phase_thin +from waveorder.models import phase_thick_3d + +# Shared, physically reasonable widefield fluorescence imaging parameters. +_OTF_KWARGS = dict( + yx_pixel_size=0.1, + z_pixel_size=0.3, + wavelength_emission=0.515, + z_padding=0, + index_of_refraction_media=1.4, + numerical_aperture_detection=1.2, +) + + +def _otf(zyx_shape): + return thick.calculate_transfer_function(zyx_shape, **_OTF_KWARGS) + + +def _fft_operators(otf): + """Build the deconvolution forward/adjoint the model uses internally.""" + + def forward(x): + return torch.real(torch.fft.ifftn(torch.fft.fftn(x, dim=(-3, -2, -1)) * otf, dim=(-3, -2, -1))) + + def transpose(y): + return torch.real(torch.fft.ifftn(torch.fft.fftn(y, dim=(-3, -2, -1)) * torch.conj(otf), dim=(-3, -2, -1))) + + return forward, transpose + + +def _bead_concentration(volume, beads, half=1): + """Fraction of nonnegative energy within small windows around ``beads``.""" + v = volume.clamp(min=0) + total = float(v.sum()) + local = 0.0 + for z, y, x in beads: + local += float(v[z - half : z + half + 1, y - half : y + half + 1, x - half : x + half + 1].sum()) + return local / total + + +def test_fft_operators_are_adjoint(): + """The conjugate-OTF transpose must be a true adjoint of the forward.""" + otf = _otf((12, 48, 48)) + forward, transpose = _fft_operators(otf) + torch.manual_seed(0) + a = torch.rand(12, 48, 48) + b = torch.rand(12, 48, 48) + lhs = float((forward(a) * b).sum()) + rhs = float((a * transpose(b)).sum()) + assert abs(lhs - rhs) <= 1e-5 * max(abs(lhs), abs(rhs)) + + +def test_core_richardson_lucy_smoke(): + """The core solver returns a positive estimate of the right shape and + increases the Poisson log-likelihood of the measurement.""" + otf = _otf((8, 32, 32)) + forward, transpose = _fft_operators(otf) + torch.manual_seed(0) + obj = torch.zeros(8, 32, 32) + obj[4, 10, 10] = 500.0 + measured = torch.poisson((forward(obj)).clamp(min=0)) + + for method in ("RL", "RLGC"): + gen = torch.Generator().manual_seed(0) + estimate = rlgc.richardson_lucy(measured, forward, transpose, num_iterations=20, method=method, generator=gen) + assert estimate.shape == obj.shape + assert torch.all(estimate > 0) # RL step keeps the estimate positive + ll_start = rlgc.poisson_log_likelihood(forward(torch.ones_like(obj)), measured) + ll_end = rlgc.poisson_log_likelihood(forward(estimate), measured) + assert ll_end > ll_start + + +def test_core_rejects_bad_arguments(): + otf = _otf((4, 16, 16)) + forward, transpose = _fft_operators(otf) + measured = torch.ones(4, 16, 16) + with pytest.raises(ValueError): + rlgc.richardson_lucy(measured, forward, transpose, num_iterations=0) + with pytest.raises(ValueError): + rlgc.richardson_lucy(measured, forward, transpose, num_iterations=5, method="bogus") + + +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_sharpens_noisy_beads(algorithm): + """RL and RLGC concentrate a Poisson-noisy bead simulation, recovering + energy that the microscope's blur had spread out.""" + zyx_shape = (24, 64, 64) + beads = [(12, 20, 20), (12, 20, 44), (8, 40, 32)] + background = 2.0 + + otf = _otf(zyx_shape) + obj = torch.full(zyx_shape, 0.0) + for b in beads: + obj[b] = 2000.0 + + torch.manual_seed(0) + clean = thick.apply_transfer_function(obj, otf, z_padding=0, background=background) + data = torch.poisson(clean.clamp(min=0)) + + raw_conc = _bead_concentration(data, beads) + + torch.manual_seed(1) + recon = thick.apply_inverse_transfer_function( + data, + otf, + z_padding=0, + reconstruction_algorithm=algorithm, + rl_iterations=100, + rl_background=background, + ) + recon_conc = _bead_concentration(recon, beads) + + # Deconvolution should concentrate energy far more tightly than the raw + # blurred data around the true bead locations. + assert recon_conc > 0.2 + assert recon_conc > 10 * raw_conc + + +def test_overiteration_starry_night_rl_vs_rlgc(): + """Over-iterated RL overfits Poisson noise into a 'starry night' of + spurious bright voxels; RLGC freezes those voxels and stays clean.""" + zyx_shape = (16, 64, 64) + beads = [(8, 20, 20), (8, 20, 44), (6, 40, 32), (10, 44, 44)] + # A dim, uniform fluorophore field carries Poisson noise everywhere, + # which is what RL overfits (cf. Andrew York's demo). + otf = _otf(zyx_shape) + obj = torch.full(zyx_shape, 2.0) + for b in beads: + obj[b] = 80.0 + + torch.manual_seed(0) + data = torch.poisson(thick.apply_transfer_function(obj, otf, z_padding=0, background=0).clamp(min=0)) + + # A region with no beads: it should stay smooth after reconstruction. + empty = (slice(2, 14), slice(0, 10), slice(0, 10)) + bright_threshold = 40.0 + true_bright = int((obj > bright_threshold).sum()) + + def reconstruct(algorithm): + torch.manual_seed(1) + return thick.apply_inverse_transfer_function( + data, + otf, + z_padding=0, + reconstruction_algorithm=algorithm, + rl_iterations=1000, + rl_background=0.0, + ) + + rl = reconstruct("RL") + gc = reconstruct("RLGC") + + rl_empty_std = float(rl[empty].std()) + gc_empty_std = float(gc[empty].std()) + rl_bright = int((rl > bright_threshold).sum()) + gc_bright = int((gc > bright_threshold).sum()) + + # RL overfits: the empty region becomes noisy and littered with spurious + # bright voxels far exceeding the four true beads. + assert rl_bright > 100 + assert rl_empty_std > 20 * gc_empty_std + # RLGC resists overfitting: no spurious bright voxels beyond the truth. + assert gc_bright <= true_bright + assert float(gc[empty].max()) < 10.0 + + +def test_stopping_tolerance_stops_early(): + """A loose stopping tolerance should halt before the iteration cap and + return a result close to the fully-iterated one.""" + zyx_shape = (12, 48, 48) + otf = _otf(zyx_shape) + forward, transpose = _fft_operators(otf) + obj = torch.zeros(zyx_shape) + obj[6, 20, 20] = 800.0 + torch.manual_seed(0) + measured = torch.poisson(forward(obj).clamp(min=0)) + + full = rlgc.richardson_lucy(measured, forward, transpose, num_iterations=200, method="RL") + stopped = rlgc.richardson_lucy( + measured, forward, transpose, num_iterations=200, method="RL", stopping_tolerance=1e-2 + ) + # Both positive, same shape; the early-stopped result is a valid estimate. + assert stopped.shape == full.shape + assert torch.all(stopped > 0) + + +def test_rl_not_implemented_for_2d_fluorescence(): + """2D (thin) fluorescence does not yet support RL/RLGC.""" + U = torch.rand(3, 2, 8, 8) + S = torch.rand(2, 8, 8) + Vh = torch.rand(2, 3, 8, 8) + data = torch.rand(3, 8, 8) + for algorithm in ("RL", "RLGC"): + with pytest.raises(NotImplementedError): + thin.apply_inverse_transfer_function(data, (U, S, Vh), reconstruction_algorithm=algorithm) + + +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_phase_3d_not_implemented_for_rl(algorithm): + """3D phase accepts the RL/RLGC request but refuses to run it.""" + zyx = torch.rand(4, 8, 8) + real_tf = torch.rand(4, 8, 8) + imag_tf = torch.rand(4, 8, 8) + with pytest.raises(NotImplementedError): + phase_thick_3d.apply_inverse_transfer_function( + zyx, real_tf, imag_tf, z_padding=0, reconstruction_algorithm=algorithm + ) + + +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_phase_2d_not_implemented_for_rl(algorithm): + """2D phase accepts the RL/RLGC request but refuses to run it.""" + zyx = torch.rand(3, 8, 8) + U = torch.rand(3, 2, 8, 8) + S = torch.rand(2, 8, 8) + Vh = torch.rand(2, 3, 8, 8) + with pytest.raises(NotImplementedError): + phase_thin.apply_inverse_transfer_function(zyx, (U, S, Vh), reconstruction_algorithm=algorithm) + + +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_phase_config_accepts_rl_request(algorithm): + """RL/RLGC are valid config values everywhere (so the request reaches the + model), even though only fluorescence implements them.""" + settings = phase.Settings(apply_inverse={"reconstruction_algorithm": algorithm}) + assert settings.apply_inverse.reconstruction_algorithm == algorithm + + +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_fluorescence_config_accepts_rl(algorithm): + """Fluorescence settings expose RL/RLGC and their parameters.""" + settings = fluorescence.Settings( + apply_inverse={ + "reconstruction_algorithm": algorithm, + "rl_iterations": 15, + "rl_background": 3.0, + "rl_stopping_tolerance": 1e-3, + } + ) + dump = settings.apply_inverse.model_dump() + assert dump["reconstruction_algorithm"] == algorithm + assert dump["rl_iterations"] == 15 + assert dump["rl_background"] == 3.0 + assert dump["rl_stopping_tolerance"] == 1e-3 diff --git a/waveorder/api/_settings.py b/waveorder/api/_settings.py index bb0376a3..925783a5 100644 --- a/waveorder/api/_settings.py +++ b/waveorder/api/_settings.py @@ -124,7 +124,9 @@ def resolve_floats(self): class FourierApplyInverseSettings(MyBaseModel): - reconstruction_algorithm: Literal["Tikhonov", "TV"] = Field( + # "RL"/"RLGC" are accepted here but only implemented for 3D fluorescence + # (see waveorder.api.fluorescence); other modalities raise NotImplementedError. + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = Field( default="Tikhonov", description="'Tikhonov' or 'TV' regularization", ) diff --git a/waveorder/api/fluorescence.py b/waveorder/api/fluorescence.py index 0c9b22b9..c8de976e 100644 --- a/waveorder/api/fluorescence.py +++ b/waveorder/api/fluorescence.py @@ -8,7 +8,7 @@ import numpy as np import torch import xarray as xr -from pydantic import Field, PositiveFloat, model_validator +from pydantic import Field, NonNegativeFloat, PositiveFloat, PositiveInt, model_validator from waveorder._pixel_size import YXPixelSize from waveorder.api._settings import ( @@ -62,7 +62,27 @@ def warn_wavelength_consistency(self): return self -ApplyInverseSettings = FourierApplyInverseSettings +class ApplyInverseSettings(FourierApplyInverseSettings): + """Fluorescence inverse settings. + + Extends the shared Fourier (Tikhonov / TV) settings with the iterative + Richardson-Lucy ("RL") and Gradient-Consensus ("RLGC") options, which are + available for 3D fluorescence reconstruction only. + """ + + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = Field( + default="Tikhonov", + description="'Tikhonov'/'TV' filters or 'RL'/'RLGC' iterative deconvolution", + ) + rl_iterations: PositiveInt = Field(default=25, description="maximum RL / RLGC iterations") + rl_background: NonNegativeFloat = Field( + default=0.0, + description="constant background folded into the RL / RLGC Poisson forward model", + ) + rl_stopping_tolerance: Optional[NonNegativeFloat] = Field( + default=None, + description="relative-change early-stop threshold for RL / RLGC (null = run all iterations)", + ) class Settings(MyBaseModel): diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index f251d6c0..30b4b261 100644 --- a/waveorder/models/isotropic_fluorescent_thick_3d.py +++ b/waveorder/models/isotropic_fluorescent_thick_3d.py @@ -4,7 +4,7 @@ import torch from torch import Tensor -from waveorder import optics, sampling, util +from waveorder import optics, rlgc, sampling, util from waveorder._pixel_size import YXPixelSize from waveorder.reconstruct import tikhonov_regularized_inverse_filter from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer @@ -238,10 +238,13 @@ def apply_inverse_transfer_function( zyx_data: Tensor, optical_transfer_function: Tensor, z_padding: int, - reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov", + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = "Tikhonov", regularization_strength: float = 1e-3, TV_rho_strength: float = 1e-3, TV_iterations: int = 10, + rl_iterations: int = 25, + rl_background: float = 0.0, + rl_stopping_tolerance: float | None = None, ) -> Tensor: """Reconstructs fluorescence density from defocus data. @@ -254,14 +257,24 @@ def apply_inverse_transfer_function( z_padding : int Padding for axial dimension. Use zero for defocus stacks that extend ~3 PSF widths beyond the sample. Pad by ~3 PSF widths otherwise. - reconstruction_algorithm : {"Tikhonov", "TV"}, optional - By default "Tikhonov". "TV" is not implemented. + reconstruction_algorithm : {"Tikhonov", "TV", "RL", "RLGC"}, optional + By default "Tikhonov". "TV" is not implemented. "RL" is + Richardson-Lucy deconvolution and "RLGC" is its Gradient-Consensus + variant, which resists overfitting noise (see :mod:`waveorder.rlgc`). regularization_strength : float, optional - Regularization parameter, by default 1e-3 + Regularization parameter (Tikhonov), by default 1e-3 TV_rho_strength : float, optional TV-specific regularization parameter, by default 1e-3 TV_iterations : int, optional TV-specific number of iterations, by default 10 + rl_iterations : int, optional + Maximum RL / RLGC iterations, by default 25 + rl_background : float, optional + Constant background (dark counts / offset) folded into the RL / RLGC + Poisson forward model, by default 0.0 + rl_stopping_tolerance : float, optional + If set, RL / RLGC stop early once the relative change of the estimate + falls below this value, by default None (run all iterations) Returns ------- @@ -286,6 +299,29 @@ def apply_inverse_transfer_function( elif reconstruction_algorithm == "TV": raise NotImplementedError + elif reconstruction_algorithm in ("RL", "RLGC"): + # The OTF (shared, shape (Z,Y,X)) broadcasts over the batch axis. + otf = optical_transfer_function + + def forward(x: Tensor) -> Tensor: + return torch.real(torch.fft.ifftn(torch.fft.fftn(x, dim=(-3, -2, -1)) * otf, dim=(-3, -2, -1))) + + def transpose(y: Tensor) -> Tensor: + return torch.real(torch.fft.ifftn(torch.fft.fftn(y, dim=(-3, -2, -1)) * torch.conj(otf), dim=(-3, -2, -1))) + + f_real = rlgc.richardson_lucy( + torch.clamp(zyx_padded, min=0.0), + forward, + transpose, + num_iterations=rl_iterations, + method=reconstruction_algorithm, + background=rl_background, + stopping_tolerance=rl_stopping_tolerance, + ) + + else: + raise NotImplementedError(f"Unknown reconstruction_algorithm: {reconstruction_algorithm}") + # Unpad if z_padding != 0: f_real = f_real[:, z_padding:-z_padding] @@ -305,10 +341,13 @@ def reconstruct( index_of_refraction_media: float, numerical_aperture_detection: float, confocal_pinhole_diameter: float | None = None, - reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov", + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = "Tikhonov", regularization_strength: float = 1e-3, TV_rho_strength: float = 1e-3, TV_iterations: int = 10, + rl_iterations: int = 25, + rl_background: float = 0.0, + rl_stopping_tolerance: float | None = None, ) -> Tensor: """Reconstruct 3D fluorescence density from a defocus stack. @@ -330,14 +369,21 @@ def reconstruct( Detection numerical aperture confocal_pinhole_diameter : float | None, optional Confocal pinhole diameter, by default None (widefield) - reconstruction_algorithm : {"Tikhonov", "TV"}, optional - By default "Tikhonov". + reconstruction_algorithm : {"Tikhonov", "TV", "RL", "RLGC"}, optional + By default "Tikhonov". "RL"/"RLGC" are Richardson-Lucy and its + Gradient-Consensus variant. regularization_strength : float, optional - Regularization parameter, by default 1e-3 + Regularization parameter (Tikhonov), by default 1e-3 TV_rho_strength : float, optional TV-specific regularization parameter, by default 1e-3 TV_iterations : int, optional TV-specific number of iterations, by default 10 + rl_iterations : int, optional + Maximum RL / RLGC iterations, by default 25 + rl_background : float, optional + Constant background folded into the RL / RLGC forward model, by default 0.0 + rl_stopping_tolerance : float, optional + Relative-change early-stop threshold for RL / RLGC, by default None Returns ------- @@ -364,4 +410,7 @@ def reconstruct( regularization_strength=regularization_strength, TV_rho_strength=TV_rho_strength, TV_iterations=TV_iterations, + rl_iterations=rl_iterations, + rl_background=rl_background, + rl_stopping_tolerance=rl_stopping_tolerance, ) diff --git a/waveorder/models/isotropic_fluorescent_thin_3d.py b/waveorder/models/isotropic_fluorescent_thin_3d.py index 10e5cfeb..8bbb53c3 100644 --- a/waveorder/models/isotropic_fluorescent_thin_3d.py +++ b/waveorder/models/isotropic_fluorescent_thin_3d.py @@ -272,10 +272,13 @@ def apply_transfer_function( def apply_inverse_transfer_function( zyx_data: Tensor, singular_system: Tuple[Tensor, Tensor, Tensor], - reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov", + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = "Tikhonov", regularization_strength: float = 1e-3, TV_rho_strength: float = 1e-3, TV_iterations: int = 10, + rl_iterations: int = 25, + rl_background: float = 0.0, + rl_stopping_tolerance: float | None = None, ) -> Tensor: """Reconstruct fluorescence density from zyx_data and singular system. @@ -285,20 +288,32 @@ def apply_inverse_transfer_function( Raw data of shape ``(Z, Y, X)`` or ``(B, Z, Y, X)`` singular_system : Tuple[Tensor, Tensor, Tensor] Singular system ``(U, S, Vh)`` (shared, not batched). - reconstruction_algorithm : {"Tikhonov", "TV"}, optional - By default "Tikhonov". "TV" is not implemented. + reconstruction_algorithm : {"Tikhonov", "TV", "RL", "RLGC"}, optional + By default "Tikhonov". "TV" is not implemented. "RL"/"RLGC" are not + yet implemented for 2D (thin) fluorescence reconstruction. regularization_strength : float, optional Regularization parameter, by default 1e-3 TV_rho_strength : float, optional TV-specific regularization parameter, by default 1e-3 TV_iterations : int, optional TV-specific number of iterations, by default 10 + rl_iterations : int, optional + Maximum RL / RLGC iterations (3D only), by default 25 + rl_background : float, optional + Constant background for the RL / RLGC forward model (3D only), by default 0.0 + rl_stopping_tolerance : float, optional + Relative-change early-stop threshold for RL / RLGC (3D only), by default None Returns ------- Tensor Fluorescence density with shape ``(Y, X)`` or ``(B, Y, X)`` """ + if reconstruction_algorithm in ("RL", "RLGC"): + raise NotImplementedError( + "RL/RLGC reconstruction is only implemented for 3D (thick) fluorescence; use reconstruction_dimension=3." + ) + batched = zyx_data.ndim == 4 if not batched: zyx_data = zyx_data.unsqueeze(0) diff --git a/waveorder/models/isotropic_thin_3d.py b/waveorder/models/isotropic_thin_3d.py index fff04bb0..44871454 100644 --- a/waveorder/models/isotropic_thin_3d.py +++ b/waveorder/models/isotropic_thin_3d.py @@ -366,7 +366,7 @@ def apply_transfer_function( def apply_inverse_transfer_function( zyx_data: Tensor, singular_system: Tuple[Tensor, Tensor, Tensor], - reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov", + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = "Tikhonov", regularization_strength: float = 1e-3, reg_p: float = 1e-6, # TODO: use this parameter TV_rho_strength: float = 1e-3, @@ -435,6 +435,9 @@ def apply_inverse_transfer_function( elif reconstruction_algorithm == "TV": raise NotImplementedError + elif reconstruction_algorithm in ("RL", "RLGC"): + raise NotImplementedError("RL/RLGC reconstruction is only implemented for 3D fluorescence") + absorption_yx = output[:, 0] # (B, Y, X) phase_yx = output[:, 1] # (B, Y, X) @@ -454,7 +457,7 @@ def reconstruct( numerical_aperture_illumination: Union[float, Tensor] = 0.9, numerical_aperture_detection: Union[float, Tensor] = 1.2, invert_phase_contrast: bool = False, - reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov", + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = "Tikhonov", regularization_strength: float = 1e-3, reg_p: float = 1e-6, TV_rho_strength: float = 1e-3, diff --git a/waveorder/models/phase_thick_3d.py b/waveorder/models/phase_thick_3d.py index 8dec4f72..dd44900d 100644 --- a/waveorder/models/phase_thick_3d.py +++ b/waveorder/models/phase_thick_3d.py @@ -359,7 +359,7 @@ def apply_inverse_transfer_function( imaginary_potential_transfer_function: Tensor, z_padding: int, absorption_ratio: float = 0.0, - reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov", + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = "Tikhonov", regularization_strength: float = 1e-3, TV_rho_strength: float = 1e-3, TV_iterations: int = 10, @@ -436,6 +436,9 @@ def apply_inverse_transfer_function( elif reconstruction_algorithm == "TV": raise NotImplementedError + elif reconstruction_algorithm in ("RL", "RLGC"): + raise NotImplementedError("RL/RLGC reconstruction is only implemented for 3D fluorescence") + # Unpad if z_padding != 0: f_real = f_real[:, z_padding:-z_padding] @@ -457,7 +460,7 @@ def reconstruct( numerical_aperture_detection: Union[float, Tensor] = 1.2, invert_phase_contrast: bool = False, absorption_ratio: float = 0.0, - reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov", + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = "Tikhonov", regularization_strength: float = 1e-3, TV_rho_strength: float = 1e-3, TV_iterations: int = 10, diff --git a/waveorder/rlgc.py b/waveorder/rlgc.py new file mode 100644 index 00000000..65b91116 --- /dev/null +++ b/waveorder/rlgc.py @@ -0,0 +1,225 @@ +"""Richardson-Lucy (RL) and Gradient-Consensus (RLGC) inference core. + +This module provides an operator-agnostic, PyTorch implementation of +Richardson-Lucy deconvolution and Andrew G. York's "Gradient Consensus" +variant. It works with any linear forward operator ``H`` and its adjoint +``H_T``, so it is not specific to deconvolution: any Poisson-noisy linear +measurement (blurring, projection, binning, ...) can be inverted with it. + +Richardson-Lucy iteratively maximizes the Poisson log-likelihood of the +measurement. It sharpens well, but over many iterations it overfits the +noise, producing the characteristic "starry night" of spurious bright +speckles. Gradient Consensus resists this: at each iteration it splits the +photons into two random halves and only updates voxels where both halves +agree on the update direction, freezing voxels where the data disagrees +with itself. + +The scaled-gradient step, Poisson noise model, and log-likelihood are +adapted from Andrew G. York's Gradient Consensus demo, +`doi.org/10.5281/zenodo.10278918 `_. +""" + +from typing import Callable, Literal, Optional + +import torch +from torch import Tensor + +_EPS = 1e-12 + + +def clip(x: Tensor, eps: float = _EPS) -> Tensor: + """Clamp nonpositive entries up to a small positive number. + + FFT-based operators return small negative values from numerical error; + the Poisson model requires strictly positive rates, so we floor them. + + Parameters + ---------- + x : torch.Tensor + Input tensor. + eps : float, optional + Floor value, by default ``1e-12``. + + Returns + ------- + torch.Tensor + ``x`` with entries below ``eps`` replaced by ``eps``. + """ + return torch.clamp(x, min=eps) + + +def poisson_log_likelihood(expected_counts: Tensor, measured_counts: Tensor) -> Tensor: + """Poisson log-likelihood of a measurement given expected rates. + + Parameters + ---------- + expected_counts : torch.Tensor + Expected (mean) counts per pixel, i.e. ``H(estimate)``. + measured_counts : torch.Tensor + Measured counts per pixel. + + Returns + ------- + torch.Tensor + Scalar log-likelihood ``sum(m * log(e) - e - lgamma(1 + m))``. + """ + e = clip(expected_counts) + m = measured_counts + return torch.sum(m * torch.log(e) - e - torch.lgamma(1 + m)) + + +def _coinflip(counts: Tensor, probability: float, generator: Optional[torch.Generator]) -> Tensor: + """Binomially thin ``counts``, keeping each event with ``probability``. + + Equivalent to placing a beam splitter in the detection path: each photon + independently lands in the "heads" arm with the given probability. + """ + counts = torch.round(clip(counts, 0.0)) + probs = torch.full_like(counts, probability) + return torch.binomial(counts, probs, generator=generator) + + +def scaled_gradient_step( + estimate: Tensor, + measured: Tensor, + forward: Callable[[Tensor], Tensor], + transpose: Callable[[Tensor], Tensor], + transpose_ones: Tensor, + *, + method: Literal["RL", "RLGC"] = "RL", + background: float = 0.0, + generator: Optional[torch.Generator] = None, +) -> tuple[Tensor, Tensor]: + """One multiplicative RL (or RLGC) update of ``estimate``. + + The update is ``estimate + gradient * step_size`` where ``gradient`` is + the gradient of the Poisson log-likelihood and ``step_size`` is the + Richardson-Lucy step that guarantees the estimate stays nonnegative. + + Parameters + ---------- + estimate : torch.Tensor + Current object estimate. + measured : torch.Tensor + Measured photon counts. + forward : callable + Forward operator ``H`` (object space -> measurement space). + transpose : callable + Adjoint operator ``H_T`` (measurement space -> object space). + transpose_ones : torch.Tensor + Precomputed ``H_T(1)``, the Richardson-Lucy step-size normalizer. + method : {"RL", "RLGC"}, optional + ``"RL"`` maximizes the Poisson likelihood. ``"RLGC"`` additionally + freezes voxels where two random halves of the photons disagree on + the update direction over the ``H_T(H(.))`` crosstalk neighborhood. + By default ``"RL"``. + background : float, optional + Constant additive background (dark counts / offset) folded into the + expected rates, by default ``0.0``. + generator : torch.Generator, optional + RNG for the RLGC coin flip, for reproducibility. Unused for RL. + + Returns + ------- + updated_estimate : torch.Tensor + The updated estimate. + step_size : torch.Tensor + The per-voxel step size actually used (zeros mark frozen voxels for + RLGC); useful as a convergence diagnostic. + """ + rates = clip(forward(estimate) + background) + gradient = transpose(measured / rates - 1.0) + step_size = estimate / transpose_ones + if method == "RLGC": + heads = _coinflip(measured, 0.5, generator) + heads_gradient = transpose(heads / rates - 0.5) + tails_gradient = gradient - heads_gradient # faster than a second H_T call + # Crosstalk neighborhood H_T(H(.)) defines which voxels touch the + # same detector pixels; a nonpositive local dot product means the + # two photon halves (locally) disagree, so we freeze those voxels. + local_dot_product = transpose(forward(heads_gradient * tails_gradient)) + step_size = torch.where(local_dot_product <= 0, torch.zeros_like(step_size), step_size) + return estimate + gradient * step_size, step_size + + +def richardson_lucy( + measured: Tensor, + forward: Callable[[Tensor], Tensor], + transpose: Callable[[Tensor], Tensor], + *, + num_iterations: int, + method: Literal["RL", "RLGC"] = "RL", + background: float = 0.0, + stopping_tolerance: Optional[float] = None, + guess: Optional[Tensor] = None, + generator: Optional[torch.Generator] = None, +) -> Tensor: + """Reconstruct an object by iterating :func:`scaled_gradient_step`. + + Parameters + ---------- + measured : torch.Tensor + Measured photon counts. + forward : callable + Forward operator ``H``. + transpose : callable + Adjoint operator ``H_T``. + num_iterations : int + Maximum number of update steps. + method : {"RL", "RLGC"}, optional + Update rule, by default ``"RL"``. See :func:`scaled_gradient_step`. + background : float, optional + Constant additive background folded into the expected rates, by + default ``0.0``. + stopping_tolerance : float, optional + If set, stop early once the relative change of the estimate, + ``||new - old|| / ||old||``, falls below this value. RLGC also stops + automatically once every voxel is frozen. By default ``None`` (run + all iterations). + guess : torch.Tensor, optional + Initial estimate. Defaults to an array of ones, which is smoother + than the noisy measurement and avoids baking noise into the result. + generator : torch.Generator, optional + RNG for the RLGC coin flip, for reproducibility. + + Returns + ------- + torch.Tensor + The final object estimate, same shape as ``transpose(measured)``. + """ + if num_iterations < 1: + raise ValueError("num_iterations must be >= 1") + if method not in ("RL", "RLGC"): + raise ValueError(f"method must be 'RL' or 'RLGC', got {method!r}") + + transpose_ones = clip(transpose(torch.ones_like(measured))) + if guess is None: + estimate = torch.ones_like(transpose_ones) + else: + estimate = clip(guess.clone()) + + for _ in range(num_iterations): + updated, step_size = scaled_gradient_step( + estimate, + measured, + forward, + transpose, + transpose_ones, + method=method, + background=background, + generator=generator, + ) + updated = clip(updated) + if method == "RLGC" and torch.all(step_size == 0): + estimate = updated + break + if stopping_tolerance is not None: + change = torch.linalg.vector_norm(updated - estimate) + scale = torch.linalg.vector_norm(estimate) + estimate = updated + if scale > 0 and (change / scale) < stopping_tolerance: + break + else: + estimate = updated + + return estimate From 316a5cdeb604b3ec84593a5755d530672fcb085f Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 4 Aug 2026 10:20:06 -0700 Subject: [PATCH 02/11] Enforce nonnegative fluorescence PSF for RL/RLGC stability --- .../test_isotropic_fluorescent_thick_3d.py | 33 +++++++++++++++++++ tests/models/test_rlgc.py | 33 +++++++++++++++++++ .../models/isotropic_fluorescent_thick_3d.py | 29 +++++++++++++++- 3 files changed, 94 insertions(+), 1 deletion(-) diff --git a/tests/models/test_isotropic_fluorescent_thick_3d.py b/tests/models/test_isotropic_fluorescent_thick_3d.py index 4e1dca24..9d496709 100644 --- a/tests/models/test_isotropic_fluorescent_thick_3d.py +++ b/tests/models/test_isotropic_fluorescent_thick_3d.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import torch from waveorder import util @@ -6,6 +7,38 @@ from waveorder.models import isotropic_fluorescent_thick_3d +@pytest.mark.parametrize( + "yx_pixel_size, na, confocal_pinhole_diameter", + [ + (0.1, 1.2, None), + (0.325, 0.8, None), + (0.65, 0.45, None), + (0.65, 1.2, None), + (1.3, 1.2, None), + (0.325, 0.8, 0.5), # confocal + ], +) +def test_transfer_function_psf_nonnegative(yx_pixel_size, na, confocal_pinhole_diameter): + """The incoherent PSF must be nonnegative. + + Upsampling to Nyquist and then cropping the OTF in Fourier space rings the + PSF below zero (worst near sub-Nyquist sampling). This is unphysical and + breaks Richardson-Lucy, so ``calculate_transfer_function`` clips it away. + """ + otf = isotropic_fluorescent_thick_3d.calculate_transfer_function( + zyx_shape=(24, 64, 64), + yx_pixel_size=yx_pixel_size, + z_pixel_size=0.5, + wavelength_emission=0.515, + z_padding=0, + index_of_refraction_media=1.4, + numerical_aperture_detection=na, + confocal_pinhole_diameter=confocal_pinhole_diameter, + ) + psf = torch.real(torch.fft.ifftn(otf, dim=(-3, -2, -1))) + assert psf.min() >= -1e-6 * psf.max() + + def test_pinhole_aperture_otf_small_diameter(): """Test that pinhole OTF is broader for smaller diameter.""" yx_shape = (128, 128) diff --git a/tests/models/test_rlgc.py b/tests/models/test_rlgc.py index a489ca78..97bea4e3 100644 --- a/tests/models/test_rlgc.py +++ b/tests/models/test_rlgc.py @@ -187,6 +187,39 @@ def reconstruct(algorithm): assert float(gc[empty].max()) < 10.0 +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_rl_stable_on_coarse_sampling(algorithm): + """Coarse (sub-Nyquist) sampling is where the OTF crop used to leave + negative PSF lobes that can destabilize Richardson-Lucy. The forward PSF + must be nonnegative and RL/RLGC must stay finite and bounded.""" + zyx_shape = (20, 48, 48) + otf = thick.calculate_transfer_function( + zyx_shape, + yx_pixel_size=0.65, + z_pixel_size=0.65, + wavelength_emission=0.515, + z_padding=0, + index_of_refraction_media=1.4, + numerical_aperture_detection=0.8, + ) + psf = torch.real(torch.fft.ifftn(otf, dim=(-3, -2, -1))) + assert psf.min() >= -1e-6 * psf.max() + + obj = torch.zeros(zyx_shape) + for z, y, x in [(10, 16, 16), (10, 16, 32), (8, 30, 24)]: + obj[z, y, x] = 8000.0 + torch.manual_seed(0) + data = torch.poisson(thick.apply_transfer_function(obj, otf, z_padding=0, background=0).clamp(min=0)) + + torch.manual_seed(1) + recon = thick.apply_inverse_transfer_function( + data, otf, z_padding=0, reconstruction_algorithm=algorithm, rl_iterations=800 + ) + assert torch.all(torch.isfinite(recon)) + # No divergence: total recovered signal stays on the order of the input. + assert float(recon.sum()) < 10 * float(data.sum()) + + def test_stopping_tolerance_stops_early(): """A loose stopping tolerance should halt before the iteration cap and return a result close to the fully-iterated one.""" diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index 30b4b261..b3b88949 100644 --- a/waveorder/models/isotropic_fluorescent_thick_3d.py +++ b/waveorder/models/isotropic_fluorescent_thick_3d.py @@ -100,7 +100,34 @@ def calculate_transfer_function( confocal_pinhole_diameter, ) zyx_out_shape = (zyx_shape[0] + 2 * z_padding,) + zyx_shape[1:] - return sampling.nd_fourier_central_cuboid(optical_transfer_function, zyx_out_shape) + optical_transfer_function = sampling.nd_fourier_central_cuboid(optical_transfer_function, zyx_out_shape) + return _enforce_nonnegative_psf(optical_transfer_function) + + +def _enforce_nonnegative_psf(optical_transfer_function: Tensor) -> Tensor: + """Return an OTF whose real-space incoherent PSF is nonnegative. + + The intensity PSF is built as ``|field|**2`` and so is nonnegative, but + cropping the OTF to the working resolution (``nd_fourier_central_cuboid``, + an ideal Fourier-domain low-pass) makes the PSF ring below zero. A physical + fluorescence PSF cannot be negative, and Richardson-Lucy's convergence + guarantee requires a nonnegative forward operator, so we clip the (small, + sub-percent) negative lobes and rebuild the normalized OTF. + + Parameters + ---------- + optical_transfer_function : Tensor + 3D OTF, shape ``(Z, Y, X)``. + + Returns + ------- + Tensor + OTF whose inverse transform is nonnegative, normalized to unit peak. + """ + psf = torch.real(torch.fft.ifftn(optical_transfer_function, dim=(-3, -2, -1))) + psf = torch.clamp(psf, min=0) + otf = torch.fft.fftn(psf, dim=(-3, -2, -1)) + return otf / torch.clamp(torch.max(torch.abs(otf)), min=1e-12) def _calculate_pinhole_aperture_otf( From 1e82701720e930c789bc566b2bf83c959819733c Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Wed, 5 Aug 2026 16:58:00 -0700 Subject: [PATCH 03/11] feat: mismatched backprojectors! Signed-off-by: Sricharan Reddy Varra --- docs/examples/cli/configs/fluorescence_2d.yml | 5 + docs/examples/cli/configs/fluorescence_3d.yml | 5 + tests/models/test_backprojector.py | 358 ++++++++++++++++ waveorder/api/fluorescence.py | 27 ++ waveorder/backprojector.py | 402 ++++++++++++++++++ .../models/isotropic_fluorescent_thick_3d.py | 94 +++- .../models/isotropic_fluorescent_thin_3d.py | 15 + waveorder/rlgc.py | 10 +- 8 files changed, 909 insertions(+), 7 deletions(-) create mode 100644 tests/models/test_backprojector.py create mode 100644 waveorder/backprojector.py diff --git a/docs/examples/cli/configs/fluorescence_2d.yml b/docs/examples/cli/configs/fluorescence_2d.yml index b1839a77..7cb7ad2a 100644 --- a/docs/examples/cli/configs/fluorescence_2d.yml +++ b/docs/examples/cli/configs/fluorescence_2d.yml @@ -22,3 +22,8 @@ fluorescence: rl_iterations: 25 # maximum RL / RLGC iterations rl_background: 0.0 # constant background folded into the RL / RLGC Poisson forward model rl_stopping_tolerance: null # relative-change early-stop threshold for RL / RLGC (null = run all iterations) + rl_back_projector: matched # 'matched' is the matched transpose (classic RL); the unmatched 'gaussian'/'butterworth'/'wiener'/'wiener_butterworth' converge in far fewer iterations but are supported for 'RL' only, not 'RLGC' + rl_bp_alpha: null # Wiener regularization for the 'wiener'/'wiener_butterworth' back projectors (null = matched cutoff gain) + rl_bp_beta: null # cutoff gain for the 'butterworth'/'wiener_butterworth' back projectors (null = matched cutoff gain) + rl_bp_order: 8 # Butterworth order for the 'butterworth'/'wiener_butterworth' back projectors + rl_bp_resolution_mode: fwhm # cutoff-frequency rule for the back projector ('fwhm_over_sqrt2' suits iSIM) diff --git a/docs/examples/cli/configs/fluorescence_3d.yml b/docs/examples/cli/configs/fluorescence_3d.yml index 47498f94..519eb2b2 100644 --- a/docs/examples/cli/configs/fluorescence_3d.yml +++ b/docs/examples/cli/configs/fluorescence_3d.yml @@ -22,3 +22,8 @@ fluorescence: rl_iterations: 25 # maximum RL / RLGC iterations rl_background: 0.0 # constant background folded into the RL / RLGC Poisson forward model rl_stopping_tolerance: null # relative-change early-stop threshold for RL / RLGC (null = run all iterations) + rl_back_projector: matched # 'matched' is the matched transpose (classic RL); the unmatched 'gaussian'/'butterworth'/'wiener'/'wiener_butterworth' converge in far fewer iterations but are supported for 'RL' only, not 'RLGC' + rl_bp_alpha: null # Wiener regularization for the 'wiener'/'wiener_butterworth' back projectors (null = matched cutoff gain) + rl_bp_beta: null # cutoff gain for the 'butterworth'/'wiener_butterworth' back projectors (null = matched cutoff gain) + rl_bp_order: 8 # Butterworth order for the 'butterworth'/'wiener_butterworth' back projectors + rl_bp_resolution_mode: fwhm # cutoff-frequency rule for the back projector ('fwhm_over_sqrt2' suits iSIM) diff --git a/tests/models/test_backprojector.py b/tests/models/test_backprojector.py new file mode 100644 index 00000000..32e8946c --- /dev/null +++ b/tests/models/test_backprojector.py @@ -0,0 +1,358 @@ +"""Correctness tests for the unmatched back projectors. + +These check :mod:`waveorder.backprojector` and its wiring into 3D fluorescence +Richardson-Lucy: + +* the matched back projector still reproduces ``conj(OTF)`` exactly, +* the resolution limit inferred from the PSF FWHM lands on the real OTF band + edge, which is the assumption every unmatched choice rests on, +* Wiener-Butterworth flattens the spectral product across the passband, which + is the mechanism behind the speed-up, +* both beta conventions hit the cutoff gain they promise, +* Wiener-Butterworth reaches in one iteration what the matched back + projector needs tens of iterations for, +* unmatched back projectors are refused for RLGC and warn when over-iterated. +""" + +import warnings + +import pytest +import torch + +from waveorder import backprojector +from waveorder.api import fluorescence +from waveorder.backprojector import calculate_back_projector +from waveorder.models import isotropic_fluorescent_thick_3d as thick + +_OTF_KWARGS = dict( + yx_pixel_size=0.1, + z_pixel_size=0.3, + wavelength_emission=0.515, + z_padding=0, + index_of_refraction_media=1.4, + numerical_aperture_detection=1.2, +) + +# Guo et al. Table S2.1 lists 0.001-0.05 for both; these sit at the sharp end. +_FILTER_KWARGS = dict(alpha=0.001, beta=0.001, order=8) + +_UNMATCHED = ["gaussian", "butterworth", "wiener", "wiener_butterworth"] +_ALL_KINDS = ["matched"] + _UNMATCHED + + +@pytest.fixture(scope="module") +def otf(): + return thick.calculate_transfer_function((24, 64, 64), **_OTF_KWARGS) + + +def _cutoff_indices(otf): + fwhm = backprojector._psf_fwhm_zyx_px(otf) + return tuple(size / res for size, res in zip(otf.shape, fwhm)) + + +def _kx_profile(volume): + """Shifted |.| profile along kx through the DC row, for a (Z, Y, X) volume.""" + magnitude = torch.fft.fftshift(torch.abs(volume)) + return magnitude[magnitude.shape[0] // 2, magnitude.shape[1] // 2, :] + + +@pytest.mark.parametrize("kind", _ALL_KINDS) +def test_shape_dtype_and_device_are_preserved(otf, kind): + back_projector = calculate_back_projector(otf, kind, **_FILTER_KWARGS) + assert back_projector.shape == otf.shape + assert back_projector.dtype == otf.dtype + assert back_projector.device == otf.device + + +def test_matched_is_exactly_the_conjugate_otf(otf): + """The matched choice must be an independent, materialized conjugate.""" + back_projector = calculate_back_projector(otf, "matched") + assert torch.equal(back_projector, torch.conj(otf)) + assert not back_projector.is_conj() + assert back_projector.untyped_storage().data_ptr() != otf.untyped_storage().data_ptr() + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS unavailable") +def test_projectors_support_mps(otf): + mps_otf = otf.to("mps") + for kind in _ALL_KINDS: + expected = calculate_back_projector(otf, kind, **_FILTER_KWARGS) + actual = calculate_back_projector(mps_otf, kind, **_FILTER_KWARGS).cpu() + assert torch.allclose(actual, expected, rtol=2e-5, atol=2e-6) + + +@pytest.mark.parametrize("kind", ["matched", "gaussian", "butterworth"]) +def test_unit_dc_gain(otf, kind): + """These kernels sum to one, so they neither brighten nor dim the estimate.""" + back_projector = calculate_back_projector(otf, kind, **_FILTER_KWARGS) + assert float(torch.abs(back_projector[0, 0, 0])) == pytest.approx(1.0, abs=1e-5) + + +@pytest.mark.parametrize("kind", ["wiener", "wiener_butterworth"]) +def test_wiener_dc_gain_is_set_by_alpha(otf, kind): + """The Wiener term scales DC by 1/(1+alpha). + + Richardson-Lucy divides this constant back out through ``H_T(1)``, so it is + harmless, but it should stay predictable rather than drift. + """ + alpha = _FILTER_KWARGS["alpha"] + back_projector = calculate_back_projector(otf, kind, **_FILTER_KWARGS) + assert float(torch.abs(back_projector[0, 0, 0])) == pytest.approx(1.0 / (1.0 + alpha), abs=1e-5) + + +def test_fwhm_cutoff_lands_on_the_otf_band_edge(otf): + """The FWHM heuristic must actually locate the resolution limit. + + Every unmatched back projector places its Butterworth transition at this + inferred cutoff, so if the heuristic were far from the true band edge the + filters would either apodize inside the passband or amplify pure noise. + """ + profile = _kx_profile(otf) + center = profile.shape[0] // 2 + supported = [i for i in range(center) if float(profile[center + i]) > 1e-3] + assert max(supported) == pytest.approx(_cutoff_indices(otf)[2], rel=0.15) + + +def test_wiener_butterworth_flattens_the_spectral_product(otf): + """The mechanism behind the speed-up: a flatter ``|DFT(f) DFT(b)|``. + + Flatness is measured only where the OTF has real support. Outside it the + product is zero whatever the back projector, and the missing cone would + otherwise dominate the statistic. + """ + support = torch.abs(otf) > 0.05 + + def coefficient_of_variation(kind): + product = torch.abs(otf * calculate_back_projector(otf, kind, **_FILTER_KWARGS))[support] + return float(product.std() / product.mean()) + + matched = coefficient_of_variation("matched") + wiener_butterworth = coefficient_of_variation("wiener_butterworth") + assert wiener_butterworth < matched / 10 + # Every unmatched choice should improve on the matched one. + for kind in _UNMATCHED: + assert coefficient_of_variation(kind) < matched + + +def test_butterworth_passes_beta_at_the_cutoff(otf): + """``beta`` is defined as the gain at the resolution limit (Eq. 23).""" + beta = 0.01 + back_projector = calculate_back_projector(otf, "butterworth", beta=beta, order=8) + gain = backprojector._mean_gain_at_cutoff(_kx_profile(back_projector), _cutoff_indices(otf)[2]) + assert gain == pytest.approx(beta, rel=0.1) + + +def test_beta_conventions_differ_by_the_wiener_cutoff_gain(otf): + """Guo et al.'s text and reference code calibrate ``beta`` differently. + + Under ``"paper"`` the Wiener-Butterworth gain at the cutoff is ``beta`` + exactly; under ``"reference"`` it is ``beta * sqrt(beta_w)``. Both are the + same filter family, so this pins down which one we build. + """ + cutoff = _cutoff_indices(otf) + wiener = calculate_back_projector(otf, "wiener", alpha=_FILTER_KWARGS["alpha"]) + wiener_cutoff_gain = backprojector._wiener_cutoff_gain(wiener, cutoff) + # The Wiener term amplifies near the cutoff, which is why the two + # conventions never coincide in practice. + assert wiener_cutoff_gain > 2.0 + + def gain(convention): + back_projector = calculate_back_projector( + otf, "wiener_butterworth", beta_convention=convention, **_FILTER_KWARGS + ) + return backprojector._mean_gain_at_cutoff(_kx_profile(back_projector), cutoff[2]) + + beta = _FILTER_KWARGS["beta"] + assert gain("paper") == pytest.approx(beta, rel=0.15) + assert gain("reference") == pytest.approx(beta * wiener_cutoff_gain**0.5, rel=0.15) + + +def test_wiener_butterworth_factors_into_its_two_halves(otf): + """WB is the elementwise product of the Wiener and Butterworth terms.""" + wiener = calculate_back_projector(otf, "wiener", alpha=_FILTER_KWARGS["alpha"]) + wiener_butterworth = calculate_back_projector(otf, "wiener_butterworth", **_FILTER_KWARGS) + support = torch.abs(wiener) > 1e-6 + ratio = torch.abs(wiener_butterworth)[support] / torch.abs(wiener)[support] + # The quotient is the Butterworth mask: unity at DC, never amplifying. + assert float(torch.abs(wiener_butterworth[0, 0, 0]) / torch.abs(wiener[0, 0, 0])) == pytest.approx(1.0, abs=1e-4) + assert float(ratio.max()) <= 1.0 + 1e-4 + + +@pytest.mark.parametrize("kind", ["butterworth", "wiener_butterworth"]) +def test_apodized_kernels_have_negative_lobes(otf, kind): + """Butterworth apodization rings in real space (Supplementary Fig. 5). + + These negative lobes are expected and must survive: clipping them would + turn the filter back into something matched-like. Richardson-Lucy handles + them by clipping the *estimate* each iteration instead. + """ + back_projector = calculate_back_projector(otf, kind, **_FILTER_KWARGS) + kernel = torch.real(torch.fft.ifftn(back_projector)) + assert float(kernel.min()) < -1e-6 * float(kernel.max()) + + +def test_gaussian_takes_no_parameters(otf): + """The Gaussian back projector is fixed by the PSF FWHM alone.""" + baseline = calculate_back_projector(otf, "gaussian") + for kwargs in ({"alpha": 0.5}, {"beta": 0.5}, {"order": 2}, {"resolution_mode": "fwhm_over_sqrt2"}): + assert torch.equal(calculate_back_projector(otf, "gaussian", **kwargs), baseline) + + +def test_invalid_arguments_are_rejected(otf): + for invalid_projector in ("nonsense", "traditional"): + with pytest.raises(ValueError, match="back_projector"): + calculate_back_projector(otf, invalid_projector) + with pytest.raises(ValueError, match="requires resolution_zyx_px"): + calculate_back_projector(otf, "butterworth", resolution_mode="manual") + with pytest.raises(ValueError, match="only valid with resolution_mode"): + calculate_back_projector(otf, "butterworth", resolution_zyx_px=(2.0, 2.0, 2.0)) + with pytest.raises(ValueError, match="3 entries"): + calculate_back_projector(otf, "butterworth", resolution_mode="manual", resolution_zyx_px=(2.0, 2.0)) + with pytest.raises(ValueError, match="must be positive"): + calculate_back_projector(otf, "butterworth", resolution_mode="manual", resolution_zyx_px=(2.0, 2.0, 0.0)) + with pytest.raises(ValueError, match="order"): + calculate_back_projector(otf, "butterworth", order=0) + with pytest.raises(ValueError, match="beta <= 1"): + calculate_back_projector(otf, "butterworth", beta=2.0) + with pytest.raises(ValueError, match="must be complex"): + calculate_back_projector(torch.real(otf), "butterworth") + with pytest.raises(ValueError, match="must be 3D"): + calculate_back_projector(otf[0], "butterworth") + with pytest.raises(ValueError, match="alpha must be positive"): + calculate_back_projector( + torch.zeros_like(otf), + "wiener", + resolution_mode="manual", + resolution_zyx_px=(2.0, 2.0, 2.0), + ) + + +def test_undersampled_psf_reports_an_actionable_error(): + """A PSF that never falls to half maximum cannot yield a FWHM. + + The caller can recover by supplying the resolution explicitly, so the + error says so rather than silently propagating a NaN cutoff. + """ + # A Gaussian far wider than the array never crosses half maximum in z. + z, y, x = 8, 32, 32 + indices = [torch.fft.fftfreq(n) * n for n in (z, y, x)] + exponent = (indices[0].reshape(-1, 1, 1) / 100.0) ** 2 + (indices[1].reshape(1, -1, 1) / 3.0) ** 2 + exponent = exponent + (indices[2].reshape(1, 1, -1) / 3.0) ** 2 + otf = torch.fft.fftn(torch.exp(-0.5 * exponent).to(torch.complex64)) + + with pytest.raises(ValueError, match="undersampled"): + calculate_back_projector(otf, "wiener_butterworth", **_FILTER_KWARGS) + # The documented escape hatch works. + calculate_back_projector( + otf, "wiener_butterworth", resolution_mode="manual", resolution_zyx_px=(4.0, 3.0, 3.0), **_FILTER_KWARGS + ) + + +def _bead_concentration(volume, beads, half=2): + """Fraction of nonnegative energy inside small windows around ``beads``.""" + clamped = volume.clamp(min=0) + total = float(clamped.sum()) + local = sum( + float(clamped[z - half : z + half + 1, y - half : y + half + 1, x - half : x + half + 1].sum()) + for z, y, x in beads + ) + return local / total + + +def test_one_wiener_butterworth_iteration_beats_many_matched_ones(otf): + """The headline claim: an unmatched back projector collapses the iteration count. + + Energy concentration is the metric rather than bead FWHM, because the + matched back projector's effective kernel is a narrow cusp sitting on a + broad halo. Its FWHM looks respectable long before the halo clears, so FWHM + understates how much work is left; concentration does not. + """ + torch.manual_seed(0) + beads = [(12, 20, 20), (12, 20, 44), (12, 44, 20), (12, 44, 44)] + obj = torch.zeros(24, 64, 64) + for bead in beads: + obj[bead] = 5000.0 + data = torch.poisson(torch.clamp(thick.apply_transfer_function(obj, otf, 0, background=2), min=0)) + + def reconstruct(kind, iterations): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return thick.apply_inverse_transfer_function( + data, + otf, + 0, + reconstruction_algorithm="RL", + rl_iterations=iterations, + rl_back_projector=kind, + rl_bp_alpha=_FILTER_KWARGS["alpha"], + rl_bp_beta=_FILTER_KWARGS["beta"], + rl_bp_order=_FILTER_KWARGS["order"], + ) + + raw = _bead_concentration(data, beads) + one_iteration = _bead_concentration(reconstruct("wiener_butterworth", 1), beads) + forty_iterations = _bead_concentration(reconstruct("matched", 40), beads) + + assert one_iteration > 4 * raw + assert one_iteration > 0.8 * forty_iterations + + +def test_rlgc_refuses_unmatched_back_projectors(otf): + """RLGC's consensus test needs a true adjoint to be a valid inner product.""" + data = torch.rand(24, 64, 64) + for kind in _UNMATCHED: + with pytest.raises(NotImplementedError, match="RLGC"): + thick.apply_inverse_transfer_function(data, otf, 0, reconstruction_algorithm="RLGC", rl_back_projector=kind) + # The matched default is still allowed. + thick.apply_inverse_transfer_function(data, otf, 0, reconstruction_algorithm="RLGC", rl_iterations=1) + + +def test_over_iterating_an_unmatched_back_projector_warns(otf): + """Past a few iterations these degrade rather than converge.""" + data = torch.rand(24, 64, 64) + with pytest.warns(UserWarning, match="rl_iterations"): + thick.apply_inverse_transfer_function( + data, otf, 0, reconstruction_algorithm="RL", rl_iterations=25, rl_back_projector="wiener_butterworth" + ) + with warnings.catch_warnings(): + warnings.simplefilter("error") + thick.apply_inverse_transfer_function( + data, otf, 0, reconstruction_algorithm="RL", rl_iterations=25, rl_back_projector="matched" + ) + + +def test_default_settings_leave_existing_configs_unchanged(): + """Adding these knobs must not alter any reconstruction already in use.""" + settings = fluorescence.ApplyInverseSettings() + assert settings.rl_back_projector == "matched" + assert settings.rl_bp_alpha is None + assert settings.rl_bp_beta is None + + +def test_settings_round_trip(): + settings = fluorescence.ApplyInverseSettings( + reconstruction_algorithm="RL", + rl_iterations=1, + rl_back_projector="wiener_butterworth", + rl_bp_alpha=0.001, + rl_bp_beta=0.001, + rl_bp_order=10, + rl_bp_resolution_mode="fwhm_over_sqrt2", + ) + dumped = settings.model_dump() + assert dumped["rl_back_projector"] == "wiener_butterworth" + assert dumped["rl_bp_order"] == 10 + assert dumped["rl_bp_resolution_mode"] == "fwhm_over_sqrt2" + # The dump is splatted straight into the model function, so the keys must match. + zyx_shape = (16, 32, 32) + thick.apply_inverse_transfer_function( + torch.rand(*zyx_shape), thick.calculate_transfer_function(zyx_shape, **_OTF_KWARGS), 0, **dumped + ) + + +def test_rejects_invalid_settings(): + for invalid_projector in ("nonsense", "traditional"): + with pytest.raises(ValueError): + fluorescence.ApplyInverseSettings(rl_back_projector=invalid_projector) + with pytest.raises(ValueError): + fluorescence.ApplyInverseSettings(rl_bp_resolution_mode="manual") diff --git a/waveorder/api/fluorescence.py b/waveorder/api/fluorescence.py index c8de976e..b6a1b790 100644 --- a/waveorder/api/fluorescence.py +++ b/waveorder/api/fluorescence.py @@ -24,6 +24,7 @@ _to_tensor, _wrap_output_tensor, ) +from waveorder.backprojector import BackProjectorType from waveorder.device import resolve_device from waveorder.models import ( isotropic_fluorescent_thick_3d, @@ -83,6 +84,32 @@ class ApplyInverseSettings(FourierApplyInverseSettings): default=None, description="relative-change early-stop threshold for RL / RLGC (null = run all iterations)", ) + rl_back_projector: BackProjectorType = Field( + default="matched", + description=( + "'matched' is the matched transpose (classic RL); the unmatched " + "'gaussian'/'butterworth'/'wiener'/'wiener_butterworth' converge in far fewer " + "iterations but are supported for 'RL' only, not 'RLGC'" + ), + ) + rl_bp_alpha: Optional[PositiveFloat] = Field( + default=None, + description="Wiener regularization for the 'wiener'/'wiener_butterworth' back projectors " + "(null = matched cutoff gain)", + ) + rl_bp_beta: Optional[PositiveFloat] = Field( + default=None, + description="cutoff gain for the 'butterworth'/'wiener_butterworth' back projectors " + "(null = matched cutoff gain)", + ) + rl_bp_order: PositiveInt = Field( + default=8, + description="Butterworth order for the 'butterworth'/'wiener_butterworth' back projectors", + ) + rl_bp_resolution_mode: Literal["fwhm", "fwhm_over_sqrt2"] = Field( + default="fwhm", + description="cutoff-frequency rule for the back projector ('fwhm_over_sqrt2' suits iSIM)", + ) class Settings(MyBaseModel): diff --git a/waveorder/backprojector.py b/waveorder/backprojector.py new file mode 100644 index 00000000..a0ba2a45 --- /dev/null +++ b/waveorder/backprojector.py @@ -0,0 +1,402 @@ +"""Unmatched back projectors that accelerate Richardson-Lucy deconvolution. + +Richardson-Lucy traditionally uses a back projector ``b`` "matched" to the +forward projector ``f``, i.e. its transpose, which in Fourier space is +``conj(OTF)``. The back projector does not have to be the transpose, though. +Convergence is governed by the eigenvalue spectrum of the operator product, +which for a shift-invariant convolution is just ``DFT(f) * DFT(b)`` evaluated +per spatial frequency: a mode whose product is close to one converges in a +single iteration, while a mode with a small product needs roughly its +reciprocal in iterations. The matched choice gives a product of ``|OTF|**2``, +which spans orders of magnitude between DC and the resolution limit, so the +iteration count ends up set by the slowest, highest-frequency mode. + +Choosing ``b`` to flatten that product across the passband is therefore a +preconditioner, and it is what lets Richardson-Lucy reach a resolution-limited +result in one iteration instead of ten or more. This module builds the family +of such back projectors described in Guo et al. 2020, Supplementary Note 2 +(`doi.org/10.1038/s41587-020-0560-x `_), +following the authors' reference implementation ``BackProjector.m`` in +`eguomin/regDeconProject `_. + +Every kind except ``"gaussian"`` factors into an inversion term times an +apodization term: + +===================== ==================== ========================== +inversion apodization: none apodization: Butterworth +===================== ==================== ========================== +``conj(OTF)`` ``"matched"`` -- +Wiener ``"wiener"`` ``"wiener_butterworth"`` +``1`` (Dirac delta) (noise, unusable) ``"butterworth"`` +===================== ==================== ========================== + +``"gaussian"`` stands apart: it is designed in real space as a Gaussian whose +FWHM matches the PSF, has no free parameters, and only ever attenuates. The +Wiener term, by contrast, actively amplifies near the resolution limit, which +is why it flattens the spectral product far more effectively. + +Because these back projectors are not adjoints, they invalidate the usual +Richardson-Lucy convergence guarantee, and over-iterating with them introduces +artifacts. Guo et al. recommend a single iteration as a rule of thumb. +""" + +from __future__ import annotations + +import math +from typing import Literal, Optional + +import torch +from torch import Tensor + +_EPS = 1e-12 + +BackProjectorType = Literal[ + "matched", + "gaussian", + "butterworth", + "wiener", + "wiener_butterworth", +] +ResolutionMode = Literal["fwhm", "fwhm_over_sqrt2", "manual"] +BetaConvention = Literal["reference", "paper"] + +_BACK_PROJECTOR_TYPES = ( + "matched", + "gaussian", + "butterworth", + "wiener", + "wiener_butterworth", +) +_RESOLUTION_MODES = ("fwhm", "fwhm_over_sqrt2", "manual") +_BETA_CONVENTIONS = ("reference", "paper") + + +def calculate_back_projector( + optical_transfer_function: Tensor, + back_projector: BackProjectorType = "matched", + *, + alpha: Optional[float] = None, + beta: Optional[float] = None, + order: int = 8, + resolution_mode: ResolutionMode = "fwhm", + resolution_zyx_px: Optional[tuple[float, float, float]] = None, + beta_convention: BetaConvention = "reference", +) -> Tensor: + """Build a back projector in Fourier space from a forward-projector OTF. + + The returned tensor is a drop-in replacement for ``conj(OTF)`` in the + Richardson-Lucy back-projection step: it uses the same FFT convention as + the input (DC at index zero) and the same shape, device and dtype, so the + adjoint step stays ``ifftn(fftn(y) * back_projector)``. + + Parameters + ---------- + optical_transfer_function : Tensor + Forward-projector OTF, complex, shape ``(Z, Y, X)``, DC at index zero. + Every kind except ``"matched"`` assumes a unit-peak OTF and + normalizes internally if needed. + back_projector : {"matched", "gaussian", "butterworth", "wiener", \ +"wiener_butterworth"}, optional + Which back projector to build, by default ``"matched"`` (the + matched transpose, i.e. plain Richardson-Lucy). + alpha : float, optional + Wiener regularization, preventing division by a vanishing OTF. Read by + ``"wiener"`` and ``"wiener_butterworth"``. ``None`` (default) + substitutes the matched back projector's mean cutoff gain, which is + what ``alpha=1`` means in the reference implementation. Guo et al. + report good results in 0.001-0.05; the reference defaults are smaller. + beta : float, optional + Cutoff gain, the spectral amplitude passed at the resolution limit. + Read by ``"butterworth"`` and ``"wiener_butterworth"``. ``None`` + (default) substitutes the matched back projector's mean cutoff + gain. Guo et al. use 0.001-0.05 (Table S2.1). + order : int, optional + Butterworth filter order, setting the steepness of the transition at + the cutoff, by default 8. Read by ``"butterworth"`` and + ``"wiener_butterworth"``. This is coupled to iteration count: Guo et + al. pair ``order`` 8-10 with a single iteration for single- and + dual-view microscopes, but drop to 5 (needing 2-5 iterations) for + quad-view and reflective geometries, which ring more readily. + resolution_mode : {"fwhm", "fwhm_over_sqrt2", "manual"}, optional + How to set the resolution limit that defines the cutoff frequencies. + ``"fwhm"`` (default) uses the measured PSF FWHM; + ``"fwhm_over_sqrt2"`` uses FWHM / sqrt(2), appropriate for iSIM; + ``"manual"`` uses ``resolution_zyx_px``. Ignored by ``"matched"`` + and ``"gaussian"``, which always match the PSF FWHM. + resolution_zyx_px : tuple of float, optional + Resolution limit per axis **in pixels**, required by and only valid + with ``resolution_mode="manual"``. Callers holding a physical + resolution should divide by their pixel size first. + beta_convention : {"reference", "paper"}, optional + How ``beta`` calibrates the Wiener-Butterworth transition, by default + ``"reference"``. See Notes. Ignored by every other kind. + + Returns + ------- + Tensor + Complex back projector, same shape, device and dtype as the input OTF. + + Notes + ----- + Guo et al.'s text and their reference code disagree on how ``beta`` maps + onto the Butterworth transition width for the Wiener-Butterworth filter. + Both are members of one family, ``eps**2 = beta_w**p / beta**2 - 1``, where + ``beta_w`` is the Wiener term's own gain at the lateral cutoff: the paper's + Eq. 27 is ``p = 2`` and the reference code is ``p = 1``. They coincide only + when ``beta_w == 1``, which never happens in practice because the Wiener + term amplifies near the cutoff, making ``beta_w`` of order ten. + + The two are exactly interconvertible. Under ``"paper"`` the filter's actual + gain at the cutoff is ``beta``, so ``beta`` means literally what it says; + under ``"reference"`` it is ``beta * sqrt(beta_w)``. This module implements + the paper's formula and, for ``"reference"``, first rescales ``beta`` by + ``sqrt(beta_w)`` to reproduce the reference code exactly. ``"reference"`` + is the default so that the ``beta`` values published in Table S2.1 produce + the published results. + + References + ---------- + Guo, M. et al. Rapid image deconvolution and multiview fusion for optical + microscopy. *Nat. Biotechnol.* **38**, 1337-1346 (2020), Supplementary + Note 2. + """ + if back_projector not in _BACK_PROJECTOR_TYPES: + raise ValueError(f"back_projector must be one of {_BACK_PROJECTOR_TYPES}, got {back_projector!r}") + if resolution_mode not in _RESOLUTION_MODES: + raise ValueError(f"resolution_mode must be one of {_RESOLUTION_MODES}, got {resolution_mode!r}") + if beta_convention not in _BETA_CONVENTIONS: + raise ValueError(f"beta_convention must be one of {_BETA_CONVENTIONS}, got {beta_convention!r}") + if optical_transfer_function.ndim != 3: + raise ValueError( + f"optical_transfer_function must be 3D (Z, Y, X), got shape {tuple(optical_transfer_function.shape)}" + ) + if not optical_transfer_function.is_complex(): + raise ValueError(f"optical_transfer_function must be complex, got dtype {optical_transfer_function.dtype}") + + # The transpose is the true adjoint whatever the OTF normalization, so it + # short-circuits before any of the filter-design machinery below. + if back_projector == "matched": + return torch.conj_physical(optical_transfer_function) + + if resolution_mode == "manual": + if resolution_zyx_px is None: + raise ValueError("resolution_mode='manual' requires resolution_zyx_px") + if len(resolution_zyx_px) != 3: + raise ValueError(f"resolution_zyx_px must have 3 entries (Z, Y, X), got {len(resolution_zyx_px)}") + if any(r <= 0 for r in resolution_zyx_px): + raise ValueError(f"resolution_zyx_px entries must be positive, got {resolution_zyx_px}") + elif resolution_zyx_px is not None: + raise ValueError(f"resolution_zyx_px is only valid with resolution_mode='manual', got {resolution_mode!r}") + if order < 1: + raise ValueError(f"order must be >= 1, got {order}") + + shape = tuple(optical_transfer_function.shape) + device = optical_transfer_function.device + real_dtype = optical_transfer_function.real.dtype + + # The Gaussian back projector has no free parameters and, per the reference + # implementation, matches the PSF FWHM whatever resolution_mode says. + if back_projector == "gaussian": + fwhm_zyx_px = _psf_fwhm_zyx_px(optical_transfer_function) + gaussian = _gaussian_back_projector(shape, fwhm_zyx_px, device, real_dtype) + return gaussian.to(optical_transfer_function.dtype) + + if resolution_mode == "manual": + resolution = tuple(float(r) for r in resolution_zyx_px) + else: + fwhm_zyx_px = _psf_fwhm_zyx_px(optical_transfer_function) + divisor = 1.0 if resolution_mode == "fwhm" else math.sqrt(2.0) + resolution = tuple(f / divisor for f in fwhm_zyx_px) + + # Cutoff as a signed frequency index, matching the reference: the + # Fourier-domain pixel size is 1/S, so the cutoff sits at S / resolution. + cutoff_indices = tuple(size / res for size, res in zip(shape, resolution)) + + alpha_value = None if alpha is None else float(alpha) + beta_value = None if beta is None else float(beta) + if back_projector != "butterworth" or alpha_value is None or beta_value is None: + normalized_otf = optical_transfer_function / torch.clamp( + torch.max(torch.abs(optical_transfer_function)), min=_EPS + ) + normalized_magnitude = torch.abs(normalized_otf) + if alpha_value is None or beta_value is None: + matched_cutoff_gain = _matched_cutoff_gain(normalized_magnitude, cutoff_indices) + alpha_value = matched_cutoff_gain if alpha_value is None else alpha_value + beta_value = matched_cutoff_gain if beta_value is None else beta_value + + if alpha_value <= 0: + raise ValueError(f"alpha must be positive, got {alpha_value}") + if beta_value <= 0: + raise ValueError(f"beta must be positive, got {beta_value}") + + if back_projector == "butterworth": + # beta = 1 / sqrt(1 + eps**2), so eps**2 = 1 / beta**2 - 1. + if beta_value > 1.0: + raise ValueError(f"butterworth requires beta <= 1 (it is a gain at the cutoff), got {beta_value}") + epsilon_squared = 1.0 / beta_value**2 - 1.0 + mask = _butterworth_mask(shape, cutoff_indices, epsilon_squared, order, device, real_dtype) + return mask.to(optical_transfer_function.dtype) + + wiener = torch.conj(normalized_otf) / (normalized_magnitude.square() + alpha_value) + if back_projector == "wiener": + return wiener.to(optical_transfer_function.dtype) + + # Wiener-Butterworth. See Notes on the two beta conventions. + wiener_cutoff_gain = _wiener_cutoff_gain(wiener, cutoff_indices) + if beta_convention == "reference": + effective_beta = beta_value * math.sqrt(wiener_cutoff_gain) + else: + effective_beta = beta_value + if effective_beta > wiener_cutoff_gain: + raise ValueError( + f"wiener_butterworth requires an effective beta <= the Wiener cutoff gain " + f"({wiener_cutoff_gain:.4g}), got {effective_beta:.4g}. Lower beta or raise alpha." + ) + epsilon_squared = (wiener_cutoff_gain / effective_beta) ** 2 - 1.0 + mask = _butterworth_mask(shape, cutoff_indices, epsilon_squared, order, device, real_dtype) + return (wiener * mask).to(optical_transfer_function.dtype) + + +def _signed_frequency_indices(size: int, device: torch.device, dtype: torch.dtype) -> Tensor: + """Signed frequency index per FFT bin, ``[0, 1, ..., -2, -1]``, unshifted.""" + return torch.fft.fftfreq(size, device=device, dtype=dtype) * size + + +def _broadcast_along(values: Tensor, axis: int) -> Tensor: + """Reshape a 1D tensor so it broadcasts along ``axis`` of a 3D volume.""" + return values.reshape([-1 if a == axis else 1 for a in range(3)]) + + +def _ellipsoidal_radius_squared( + shape: tuple[int, int, int], + scales: tuple[float, float, float], + device: torch.device, + dtype: torch.dtype, +) -> Tensor: + """Squared radius on an unshifted ellipsoidal frequency grid.""" + radius_squared = torch.zeros(shape, device=device, dtype=dtype) + for axis, (size, scale) in enumerate(zip(shape, scales)): + indices = _signed_frequency_indices(size, device, dtype) / scale + radius_squared.add_(_broadcast_along(indices.square(), axis)) + return radius_squared + + +def _butterworth_mask( + shape: tuple[int, int, int], + cutoff_indices: tuple[float, float, float], + epsilon_squared: float, + order: int, + device: torch.device, + dtype: torch.dtype, +) -> Tensor: + """Butterworth low-pass over an ellipsoidal cutoff surface, unshifted. + + Evaluated in float64 where supported because ``radius**(2 * order)`` can + overflow float32. MPS uses an equivalent log-space float32 calculation. + """ + work_dtype = dtype if device.type == "mps" else torch.float64 + radius_squared = _ellipsoidal_radius_squared(shape, cutoff_indices, device, work_dtype) + if work_dtype == torch.float64: + mask = radius_squared.pow_(order).mul_(epsilon_squared).add_(1.0).sqrt_().reciprocal_() + elif epsilon_squared == 0: + mask = torch.ones_like(radius_squared) + else: + log_attenuation = radius_squared.log_().mul_(order).add_(math.log(epsilon_squared)) + mask = torch.nn.functional.softplus(log_attenuation).mul_(-0.5).exp_() + return mask.to(dtype) + + +def _gaussian_back_projector( + shape: tuple[int, int, int], + fwhm_zyx_px: tuple[float, float, float], + device: torch.device, + dtype: torch.dtype, +) -> Tensor: + """OTF of a unit-sum Gaussian kernel whose FWHM matches the PSF.""" + sigmas = tuple(fwhm / (2.0 * math.sqrt(2.0 * math.log(2.0))) for fwhm in fwhm_zyx_px) + work_dtype = dtype if device.type == "mps" else torch.float64 + exponent = _ellipsoidal_radius_squared(shape, sigmas, device, work_dtype) + kernel = torch.exp(-0.5 * exponent) + kernel.div_(torch.clamp(kernel.sum(), min=_EPS)) + return torch.fft.fftn(kernel.to(dtype), dim=(-3, -2, -1)) + + +def _psf_fwhm_zyx_px(optical_transfer_function: Tensor) -> tuple[float, float, float]: + """Measure the PSF full width at half maximum per axis, in pixels.""" + psf = torch.fft.fftshift(torch.real(torch.fft.ifftn(optical_transfer_function, dim=(-3, -2, -1)))) + peak_z, peak_y, peak_x = (int(index) for index in torch.unravel_index(torch.argmax(psf), psf.shape)) + return ( + _fwhm_1d(psf[:, peak_y, peak_x], "z"), + _fwhm_1d(psf[peak_z, :, peak_x], "y"), + _fwhm_1d(psf[peak_z, peak_y, :], "x"), + ) + + +def _fwhm_1d(profile: Tensor, axis_name: str) -> float: + """Full width at half maximum of a peaked 1D profile, in pixels. + + Both half-maximum crossings are located by linear interpolation between + the bracketing samples. A profile that never falls below half maximum on + one side is undersampled relative to the PSF, which the caller cannot + recover from, so this raises rather than returning a sentinel. + """ + peak_index = int(torch.argmax(profile)) + peak_value = float(profile[peak_index]) + if not peak_value > 0: + raise ValueError(f"cannot measure PSF FWHM along {axis_name}: the profile peak is not positive") + normalized = profile.cpu().to(torch.float64) / peak_value + + below_left = torch.nonzero(normalized[: peak_index + 1] < 0.5).flatten() + below_right = torch.nonzero(normalized[peak_index:] < 0.5).flatten() + if below_left.numel() == 0 or below_right.numel() == 0: + raise ValueError( + f"cannot measure PSF FWHM along {axis_name}: the profile never falls below half maximum, so " + f"the PSF is undersampled along this axis. Pass resolution_mode='manual' with an explicit " + f"resolution_zyx_px, or reconstruct on a finer grid." + ) + + def _interpolate(low_index: int, high_index: int) -> float: + low_value = float(normalized[low_index]) + high_value = float(normalized[high_index]) + if high_value == low_value: + raise ValueError(f"cannot measure PSF FWHM along {axis_name}: the profile is flat at half maximum") + return low_index + (0.5 - low_value) / (high_value - low_value) * (high_index - low_index) + + left_index = int(below_left[-1]) + right_index = peak_index + int(below_right[0]) + return _interpolate(right_index - 1, right_index) - _interpolate(left_index, left_index + 1) + + +def _mean_gain_at_cutoff(profile: Tensor, cutoff: float) -> float: + """Average a shifted 1D profile at the two cutoff frequencies.""" + size = profile.shape[0] + center = size // 2 + low = max(int(round(center - cutoff)), 0) + high = min(int(round(center + cutoff)), size - 1) + return float((profile[low] + profile[high]) / 2) + + +def _matched_cutoff_gain(normalized_magnitude: Tensor, cutoff_indices: tuple[float, float, float]) -> float: + """Mean cutoff gain of the matched back projector (Eq. 28). + + Per axis, the OTF magnitude is maximum-projected onto that axis and + sampled at both cutoff frequencies; the three per-axis gains are averaged. + This is the value substituted when ``alpha`` or ``beta`` is left unset. + """ + magnitude = torch.fft.fftshift(normalized_magnitude) + gains = [] + for axis, cutoff in enumerate(cutoff_indices): + other_axes = tuple(a for a in range(3) if a != axis) + gains.append(_mean_gain_at_cutoff(torch.amax(magnitude, dim=other_axes), cutoff)) + return sum(gains) / 3.0 + + +def _wiener_cutoff_gain(wiener: Tensor, cutoff_indices: tuple[float, float, float]) -> float: + """Gain of the Wiener term at the lateral (X) cutoff. + + Unlike :func:`_matched_cutoff_gain` this reads the central Z slice + rather than a maximum projection, matching the reference implementation. + """ + magnitude = torch.fft.fftshift(torch.abs(wiener)) + central_slice = magnitude[magnitude.shape[0] // 2] # (Y, X) + return _mean_gain_at_cutoff(torch.amax(central_slice, dim=0), cutoff_indices[2]) diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index b3b88949..c8967fe0 100644 --- a/waveorder/models/isotropic_fluorescent_thick_3d.py +++ b/waveorder/models/isotropic_fluorescent_thick_3d.py @@ -1,11 +1,13 @@ +import warnings from typing import Literal import numpy as np import torch from torch import Tensor -from waveorder import optics, rlgc, sampling, util +from waveorder import backprojector, optics, rlgc, sampling, util from waveorder._pixel_size import YXPixelSize +from waveorder.backprojector import BackProjectorType from waveorder.reconstruct import tikhonov_regularized_inverse_filter from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer @@ -272,6 +274,11 @@ def apply_inverse_transfer_function( rl_iterations: int = 25, rl_background: float = 0.0, rl_stopping_tolerance: float | None = None, + rl_back_projector: BackProjectorType = "matched", + rl_bp_alpha: float | None = None, + rl_bp_beta: float | None = None, + rl_bp_order: int = 8, + rl_bp_resolution_mode: Literal["fwhm", "fwhm_over_sqrt2"] = "fwhm", ) -> Tensor: """Reconstructs fluorescence density from defocus data. @@ -302,6 +309,25 @@ def apply_inverse_transfer_function( rl_stopping_tolerance : float, optional If set, RL / RLGC stop early once the relative change of the estimate falls below this value, by default None (run all iterations) + rl_back_projector : str, optional + Back projector for RL, by default "matched" (the matched transpose, + i.e. classic Richardson-Lucy). The unmatched alternatives "gaussian", + "butterworth", "wiener" and "wiener_butterworth" flatten the spectral + product and so converge in far fewer iterations; see + :mod:`waveorder.backprojector`. Unmatched choices are RL-only, and one + iteration is a good rule of thumb for them. + rl_bp_alpha : float, optional + Wiener regularization for the "wiener"/"wiener_butterworth" back + projectors, by default None (use the matched cutoff gain) + rl_bp_beta : float, optional + Cutoff gain for the "butterworth"/"wiener_butterworth" back projectors, + by default None (use the matched cutoff gain) + rl_bp_order : int, optional + Butterworth order for the "butterworth"/"wiener_butterworth" back + projectors, by default 8 + rl_bp_resolution_mode : str, optional + How the back projector sets its cutoff frequency, by default "fwhm". + Use "fwhm_over_sqrt2" for iSIM. Returns ------- @@ -330,11 +356,46 @@ def apply_inverse_transfer_function( # The OTF (shared, shape (Z,Y,X)) broadcasts over the batch axis. otf = optical_transfer_function + # RLGC reads the sign of transpose(forward(.)) to ask whether the two + # photon halves agree. Only a true adjoint makes that question + # meaningful: an unmatched back projector's negative lobes flip the + # sign on their own, freezing good voxels. + if reconstruction_algorithm == "RLGC" and rl_back_projector != "matched": + raise NotImplementedError( + f"rl_back_projector={rl_back_projector!r} is only supported with " + f"reconstruction_algorithm='RL'; RLGC requires the matched " + f"back projector for its gradient-consensus test." + ) + + # An unmatched back projector abandons Richardson-Lucy's fixed point at + # the maximum-likelihood solution, so past a few iterations the estimate + # degrades instead of settling. Guo et al. recommend a single iteration, + # or up to five at low filter orders. + if rl_back_projector != "matched" and rl_iterations > 5: + warnings.warn( + f"rl_iterations={rl_iterations} with rl_back_projector={rl_back_projector!r}: " + f"unmatched back projectors reach a resolution-limited result in 1-5 iterations " + f"and introduce artifacts beyond that. Consider rl_iterations=1.", + UserWarning, + stacklevel=2, + ) + + back_projector_otf = backprojector.calculate_back_projector( + otf, + rl_back_projector, + alpha=rl_bp_alpha, + beta=rl_bp_beta, + order=rl_bp_order, + resolution_mode=rl_bp_resolution_mode, + ) + def forward(x: Tensor) -> Tensor: return torch.real(torch.fft.ifftn(torch.fft.fftn(x, dim=(-3, -2, -1)) * otf, dim=(-3, -2, -1))) def transpose(y: Tensor) -> Tensor: - return torch.real(torch.fft.ifftn(torch.fft.fftn(y, dim=(-3, -2, -1)) * torch.conj(otf), dim=(-3, -2, -1))) + return torch.real( + torch.fft.ifftn(torch.fft.fftn(y, dim=(-3, -2, -1)) * back_projector_otf, dim=(-3, -2, -1)) + ) f_real = rlgc.richardson_lucy( torch.clamp(zyx_padded, min=0.0), @@ -375,6 +436,11 @@ def reconstruct( rl_iterations: int = 25, rl_background: float = 0.0, rl_stopping_tolerance: float | None = None, + rl_back_projector: BackProjectorType = "matched", + rl_bp_alpha: float | None = None, + rl_bp_beta: float | None = None, + rl_bp_order: int = 8, + rl_bp_resolution_mode: Literal["fwhm", "fwhm_over_sqrt2"] = "fwhm", ) -> Tensor: """Reconstruct 3D fluorescence density from a defocus stack. @@ -411,6 +477,25 @@ def reconstruct( Constant background folded into the RL / RLGC forward model, by default 0.0 rl_stopping_tolerance : float, optional Relative-change early-stop threshold for RL / RLGC, by default None + rl_back_projector : str, optional + Back projector for RL, by default "matched" (the matched transpose, + i.e. classic Richardson-Lucy). The unmatched alternatives "gaussian", + "butterworth", "wiener" and "wiener_butterworth" flatten the spectral + product and so converge in far fewer iterations; see + :mod:`waveorder.backprojector`. Unmatched choices are RL-only, and one + iteration is a good rule of thumb for them. + rl_bp_alpha : float, optional + Wiener regularization for the "wiener"/"wiener_butterworth" back + projectors, by default None (use the matched cutoff gain) + rl_bp_beta : float, optional + Cutoff gain for the "butterworth"/"wiener_butterworth" back projectors, + by default None (use the matched cutoff gain) + rl_bp_order : int, optional + Butterworth order for the "butterworth"/"wiener_butterworth" back + projectors, by default 8 + rl_bp_resolution_mode : str, optional + How the back projector sets its cutoff frequency, by default "fwhm". + Use "fwhm_over_sqrt2" for iSIM. Returns ------- @@ -440,4 +525,9 @@ def reconstruct( rl_iterations=rl_iterations, rl_background=rl_background, rl_stopping_tolerance=rl_stopping_tolerance, + rl_back_projector=rl_back_projector, + rl_bp_alpha=rl_bp_alpha, + rl_bp_beta=rl_bp_beta, + rl_bp_order=rl_bp_order, + rl_bp_resolution_mode=rl_bp_resolution_mode, ) diff --git a/waveorder/models/isotropic_fluorescent_thin_3d.py b/waveorder/models/isotropic_fluorescent_thin_3d.py index 8bbb53c3..567abc4b 100644 --- a/waveorder/models/isotropic_fluorescent_thin_3d.py +++ b/waveorder/models/isotropic_fluorescent_thin_3d.py @@ -279,6 +279,11 @@ def apply_inverse_transfer_function( rl_iterations: int = 25, rl_background: float = 0.0, rl_stopping_tolerance: float | None = None, + rl_back_projector: str = "matched", + rl_bp_alpha: float | None = None, + rl_bp_beta: float | None = None, + rl_bp_order: int = 8, + rl_bp_resolution_mode: Literal["fwhm", "fwhm_over_sqrt2"] = "fwhm", ) -> Tensor: """Reconstruct fluorescence density from zyx_data and singular system. @@ -303,6 +308,16 @@ def apply_inverse_transfer_function( Constant background for the RL / RLGC forward model (3D only), by default 0.0 rl_stopping_tolerance : float, optional Relative-change early-stop threshold for RL / RLGC (3D only), by default None + rl_back_projector : str, optional + Back projector for RL (3D only), by default "matched" + rl_bp_alpha : float, optional + Wiener regularization for the RL back projector (3D only), by default None + rl_bp_beta : float, optional + Cutoff gain for the RL back projector (3D only), by default None + rl_bp_order : int, optional + Butterworth order for the RL back projector (3D only), by default 8 + rl_bp_resolution_mode : str, optional + Cutoff-frequency rule for the RL back projector (3D only), by default "fwhm" Returns ------- diff --git a/waveorder/rlgc.py b/waveorder/rlgc.py index 65b91116..c3e91cb4 100644 --- a/waveorder/rlgc.py +++ b/waveorder/rlgc.py @@ -74,8 +74,8 @@ def _coinflip(counts: Tensor, probability: float, generator: Optional[torch.Gene Equivalent to placing a beam splitter in the detection path: each photon independently lands in the "heads" arm with the given probability. """ - counts = torch.round(clip(counts, 0.0)) - probs = torch.full_like(counts, probability) + counts = counts.clamp_min(0.0).round_() + probs = counts.new_tensor(probability).expand_as(counts) return torch.binomial(counts, probs, generator=generator) @@ -138,7 +138,7 @@ def scaled_gradient_step( # same detector pixels; a nonpositive local dot product means the # two photon halves (locally) disagree, so we freeze those voxels. local_dot_product = transpose(forward(heads_gradient * tails_gradient)) - step_size = torch.where(local_dot_product <= 0, torch.zeros_like(step_size), step_size) + step_size = torch.where(local_dot_product <= 0, 0.0, step_size) return estimate + gradient * step_size, step_size @@ -196,7 +196,7 @@ def richardson_lucy( if guess is None: estimate = torch.ones_like(transpose_ones) else: - estimate = clip(guess.clone()) + estimate = clip(guess) for _ in range(num_iterations): updated, step_size = scaled_gradient_step( @@ -210,7 +210,7 @@ def richardson_lucy( generator=generator, ) updated = clip(updated) - if method == "RLGC" and torch.all(step_size == 0): + if method == "RLGC" and torch.count_nonzero(step_size) == 0: estimate = updated break if stopping_tolerance is not None: From 39939976f7e2c7eeb68e338f801756505771e3f5 Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Thu, 6 Aug 2026 16:01:49 -0700 Subject: [PATCH 04/11] fix(backprojector): scale default alpha as gain squared Eq. 28 returns a spectral amplitude at the cutoff, which is beta's scale. alpha is added to |OTF|**2, so substituting the same value leaves it ~200x too large. On a confocal OTF with |H(cutoff)| = 0.0106 the Wiener spectral product at the cutoff was 0.0053 instead of 0.171, i.e. ~187 iterations to converge the finest frequencies rather than ~6 -- defeating the point of an unmatched back projector. --- waveorder/backprojector.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/waveorder/backprojector.py b/waveorder/backprojector.py index a0ba2a45..e424ba61 100644 --- a/waveorder/backprojector.py +++ b/waveorder/backprojector.py @@ -102,9 +102,11 @@ def calculate_back_projector( alpha : float, optional Wiener regularization, preventing division by a vanishing OTF. Read by ``"wiener"`` and ``"wiener_butterworth"``. ``None`` (default) - substitutes the matched back projector's mean cutoff gain, which is - what ``alpha=1`` means in the reference implementation. Guo et al. - report good results in 0.001-0.05; the reference defaults are smaller. + substitutes the SQUARE of the matched back projector's mean cutoff gain, + because alpha is added to ``|OTF|**2`` and so lives on the scale of a + squared amplitude, not of the gain itself (Eq. 28 returns the gain, and + is the right substitution for ``beta`` only). Guo et al. report good + results in 0.001-0.05; the reference defaults are smaller. beta : float, optional Cutoff gain, the spectral amplitude passed at the resolution limit. Read by ``"butterworth"`` and ``"wiener_butterworth"``. ``None`` @@ -221,7 +223,9 @@ def calculate_back_projector( normalized_magnitude = torch.abs(normalized_otf) if alpha_value is None or beta_value is None: matched_cutoff_gain = _matched_cutoff_gain(normalized_magnitude, cutoff_indices) - alpha_value = matched_cutoff_gain if alpha_value is None else alpha_value + # Eq. 28 returns a gain, which is beta's scale. alpha is added to + # |OTF|**2, so it needs that gain squared. + alpha_value = matched_cutoff_gain**2 if alpha_value is None else alpha_value beta_value = matched_cutoff_gain if beta_value is None else beta_value if alpha_value <= 0: From 9a96e5acac8a5afe959180b25dece8cec9ad6a4a Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Thu, 6 Aug 2026 16:01:49 -0700 Subject: [PATCH 05/11] perf(fluorescence): accept a prebuilt back projector calculate_back_projector depends only on the OTF and the rl_bp_* knobs, all fixed for a run, but sat inside apply_inverse_transfer_function and so rebuilt on every call. Callers reconstructing many tiles can now build it once. Output is bit-identical; matched is unaffected. --- .../models/isotropic_fluorescent_thick_3d.py | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index c8967fe0..875e3ecc 100644 --- a/waveorder/models/isotropic_fluorescent_thick_3d.py +++ b/waveorder/models/isotropic_fluorescent_thick_3d.py @@ -279,6 +279,7 @@ def apply_inverse_transfer_function( rl_bp_beta: float | None = None, rl_bp_order: int = 8, rl_bp_resolution_mode: Literal["fwhm", "fwhm_over_sqrt2"] = "fwhm", + back_projector_otf: Tensor | None = None, ) -> Tensor: """Reconstructs fluorescence density from defocus data. @@ -328,6 +329,11 @@ def apply_inverse_transfer_function( rl_bp_resolution_mode : str, optional How the back projector sets its cutoff frequency, by default "fwhm". Use "fwhm_over_sqrt2" for iSIM. + back_projector_otf : Tensor, optional + Prebuilt back projector, skipping the rl_bp_* construction. Building it + costs a few seconds on a large OTF, so callers reconstructing many tiles + should build it once with :func:`waveorder.backprojector.calculate_back_projector` + and pass it here. By default None (build it on every call). Returns ------- @@ -380,14 +386,18 @@ def apply_inverse_transfer_function( stacklevel=2, ) - back_projector_otf = backprojector.calculate_back_projector( - otf, - rl_back_projector, - alpha=rl_bp_alpha, - beta=rl_bp_beta, - order=rl_bp_order, - resolution_mode=rl_bp_resolution_mode, - ) + # Depends only on the OTF and the rl_bp_* knobs, all fixed for a run, so a + # caller that reconstructs many tiles should build it once and pass it in; + # otherwise it is rebuilt on every call. + if back_projector_otf is None: + back_projector_otf = backprojector.calculate_back_projector( + otf, + rl_back_projector, + alpha=rl_bp_alpha, + beta=rl_bp_beta, + order=rl_bp_order, + resolution_mode=rl_bp_resolution_mode, + ) def forward(x: Tensor) -> Tensor: return torch.real(torch.fft.ifftn(torch.fft.fftn(x, dim=(-3, -2, -1)) * otf, dim=(-3, -2, -1))) @@ -496,6 +506,11 @@ def reconstruct( rl_bp_resolution_mode : str, optional How the back projector sets its cutoff frequency, by default "fwhm". Use "fwhm_over_sqrt2" for iSIM. + back_projector_otf : Tensor, optional + Prebuilt back projector, skipping the rl_bp_* construction. Building it + costs a few seconds on a large OTF, so callers reconstructing many tiles + should build it once with :func:`waveorder.backprojector.calculate_back_projector` + and pass it here. By default None (build it on every call). Returns ------- From 0524fb18c548d867ca9bd01f15b1306a6459555f Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Fri, 7 Aug 2026 13:17:10 -0700 Subject: [PATCH 06/11] updated docstrings --- waveorder/backprojector.py | 169 ++++++++++++++++++++----------------- 1 file changed, 92 insertions(+), 77 deletions(-) diff --git a/waveorder/backprojector.py b/waveorder/backprojector.py index e424ba61..ecd8de02 100644 --- a/waveorder/backprojector.py +++ b/waveorder/backprojector.py @@ -1,46 +1,46 @@ """Unmatched back projectors that accelerate Richardson-Lucy deconvolution. -Richardson-Lucy traditionally uses a back projector ``b`` "matched" to the -forward projector ``f``, i.e. its transpose, which in Fourier space is -``conj(OTF)``. The back projector does not have to be the transpose, though. -Convergence is governed by the eigenvalue spectrum of the operator product, -which for a shift-invariant convolution is just ``DFT(f) * DFT(b)`` evaluated -per spatial frequency: a mode whose product is close to one converges in a -single iteration, while a mode with a small product needs roughly its -reciprocal in iterations. The matched choice gives a product of ``|OTF|**2``, -which spans orders of magnitude between DC and the resolution limit, so the -iteration count ends up set by the slowest, highest-frequency mode. - -Choosing ``b`` to flatten that product across the passband is therefore a -preconditioner, and it is what lets Richardson-Lucy reach a resolution-limited -result in one iteration instead of ten or more. This module builds the family -of such back projectors described in Guo et al. 2020, Supplementary Note 2 -(`doi.org/10.1038/s41587-020-0560-x `_), -following the authors' reference implementation ``BackProjector.m`` in -`eguomin/regDeconProject `_. +Richardson-Lucy usually back-projects with the transpose of the forward +projector, ``conj(OTF)``. It does not have to. What sets the convergence rate +is the **spectral product** ``|DFT(f) * DFT(b)|``, read per spatial frequency: + +- product near 1 -> that frequency is recovered in one iteration +- product of 0.01 -> it needs roughly 100 -Every kind except ``"gaussian"`` factors into an inversion term times an -apodization term: +The transpose gives ``|OTF|**2``, which spans orders of magnitude between DC and +the resolution limit, so the iteration count is set by the slowest, highest +frequency. Choosing ``b`` to flatten that product instead is a preconditioner, +and it is what turns "ten or more iterations" into one. + +Every kind except ``"gaussian"`` is an inversion term times an apodization term: ===================== ==================== ========================== inversion apodization: none apodization: Butterworth ===================== ==================== ========================== -``conj(OTF)`` ``"matched"`` -- +``conj(OTF)`` ``"matched"`` -- Wiener ``"wiener"`` ``"wiener_butterworth"`` ``1`` (Dirac delta) (noise, unusable) ``"butterworth"`` ===================== ==================== ========================== -``"gaussian"`` stands apart: it is designed in real space as a Gaussian whose -FWHM matches the PSF, has no free parameters, and only ever attenuates. The -Wiener term, by contrast, actively amplifies near the resolution limit, which -is why it flattens the spectral product far more effectively. +The inversion term (``alpha``) decides how flat the product is; the apodization +term (``beta``, ``order``) decides how hard everything past the resolution limit +is suppressed. Only the Wiener term amplifies, so only the kinds containing it +flatten the product appreciably -- which is why ``"wiener_butterworth"``, with +both, is the one Guo et al. recommend. + +``"gaussian"`` stands apart: designed in real space to match the PSF FWHM, no +free parameters, and it only ever attenuates. + +Caveat: an unmatched ``b`` is not an adjoint, so Richardson-Lucy's convergence +guarantee no longer holds and over-iterating introduces artifacts. Guo et al. +suggest one iteration as a rule of thumb (up to five at low ``order``). -Because these back projectors are not adjoints, they invalidate the usual -Richardson-Lucy convergence guarantee, and over-iterating with them introduces -artifacts. Guo et al. recommend a single iteration as a rule of thumb. +Reference: Guo et al. 2020, Supplementary Note 2 +(`doi.org/10.1038/s41587-020-0560-x `_), +following ``BackProjector.m`` in +`eguomin/regDeconProject `_. """ -from __future__ import annotations import math from typing import Literal, Optional @@ -84,51 +84,69 @@ def calculate_back_projector( ) -> Tensor: """Build a back projector in Fourier space from a forward-projector OTF. - The returned tensor is a drop-in replacement for ``conj(OTF)`` in the - Richardson-Lucy back-projection step: it uses the same FFT convention as - the input (DC at index zero) and the same shape, device and dtype, so the - adjoint step stays ``ifftn(fftn(y) * back_projector)``. + Drop-in replacement for ``conj(OTF)`` in the Richardson-Lucy back-projection + step: same FFT convention (DC at index zero), same shape, device and dtype, + so the adjoint stays ``ifftn(fftn(y) * back_projector)``. + + Which knobs each kind reads, and what turning them does: + + ==================== ===== ==== ===== ========================================== + kind alpha beta order effect + ==================== ===== ==== ===== ========================================== + ``matched`` -- -- -- plain Richardson-Lucy, no free parameters + ``gaussian`` -- -- -- attenuates only; the mildest option + ``butterworth`` -- x x suppresses past the cutoff, no inversion + ``wiener`` x -- -- inverts, but amplifies noise unchecked + ``wiener_butterworth`` x x x inverts and suppresses; the recommended one + ==================== ===== ==== ===== ========================================== + + Lower ``alpha`` -> flatter spectral product -> fewer iterations, more noise. + Lower ``beta`` / higher ``order`` -> harder suppression past the resolution + limit, at the cost of real-space ringing. Parameters ---------- optical_transfer_function : Tensor Forward-projector OTF, complex, shape ``(Z, Y, X)``, DC at index zero. - Every kind except ``"matched"`` assumes a unit-peak OTF and - normalizes internally if needed. + Every kind except ``"matched"`` assumes a unit-peak OTF and normalizes + internally if needed. back_projector : {"matched", "gaussian", "butterworth", "wiener", \ "wiener_butterworth"}, optional - Which back projector to build, by default ``"matched"`` (the - matched transpose, i.e. plain Richardson-Lucy). + Which one to build, by default ``"matched"``. alpha : float, optional - Wiener regularization, preventing division by a vanishing OTF. Read by - ``"wiener"`` and ``"wiener_butterworth"``. ``None`` (default) - substitutes the SQUARE of the matched back projector's mean cutoff gain, - because alpha is added to ``|OTF|**2`` and so lives on the scale of a - squared amplitude, not of the gain itself (Eq. 28 returns the gain, and - is the right substitution for ``beta`` only). Guo et al. report good - results in 0.001-0.05; the reference defaults are smaller. + Wiener regularization: the term added to ``|OTF|**2`` that keeps the + inversion from dividing by a vanishing OTF. It sets **how hard the + deconvolution inverts**, since the spectral product is + ``|OTF|**2 / (|OTF|**2 + alpha)`` -- flat (so one iteration) only when + ``alpha`` is near ``|OTF(cutoff)|**2``, and increasingly damped above + that. Guo et al. use 0.001-0.05. + + ``None`` (default) substitutes the mean cutoff gain (Eq. 28) SQUARED -- + that gain is ``beta``'s scale, and ``alpha`` is added to a squared + magnitude, so using it unsquared leaves ``alpha`` orders of magnitude + too large. beta : float, optional - Cutoff gain, the spectral amplitude passed at the resolution limit. - Read by ``"butterworth"`` and ``"wiener_butterworth"``. ``None`` - (default) substitutes the matched back projector's mean cutoff - gain. Guo et al. use 0.001-0.05 (Table S2.1). + Cutoff gain: the spectral amplitude still passed **at** the resolution + limit, so smaller means sharper suppression beyond it. Enters as + ``eps**2``; see Notes. Guo et al. use 0.001-0.05 (Table S2.1). ``None`` + (default) substitutes the mean cutoff gain (Eq. 28). order : int, optional - Butterworth filter order, setting the steepness of the transition at - the cutoff, by default 8. Read by ``"butterworth"`` and - ``"wiener_butterworth"``. This is coupled to iteration count: Guo et - al. pair ``order`` 8-10 with a single iteration for single- and - dual-view microscopes, but drop to 5 (needing 2-5 iterations) for - quad-view and reflective geometries, which ring more readily. + Butterworth order, the exponent that sets **how steep the transition + is** at the cutoff, by default 8. Higher is flatter in the passband but + closer to a brick wall, which rings in real space; lower is gentler but + gives up amplitude near the cutoff, so more iterations are needed. Guo + et al. pair 8-10 with a single iteration for single- and dual-view + microscopes, dropping to 5 (and 2-5 iterations) for quad-view and + reflective geometries, which ring more readily. resolution_mode : {"fwhm", "fwhm_over_sqrt2", "manual"}, optional - How to set the resolution limit that defines the cutoff frequencies. - ``"fwhm"`` (default) uses the measured PSF FWHM; - ``"fwhm_over_sqrt2"`` uses FWHM / sqrt(2), appropriate for iSIM; - ``"manual"`` uses ``resolution_zyx_px``. Ignored by ``"matched"`` - and ``"gaussian"``, which always match the PSF FWHM. + Where the cutoff frequency sits. ``"fwhm"`` (default) uses the measured + PSF FWHM; ``"fwhm_over_sqrt2"`` puts the cutoff a factor sqrt(2) higher, + for resolution-doubling instruments such as iSIM; ``"manual"`` takes + ``resolution_zyx_px``. Ignored by ``"matched"`` and ``"gaussian"``. resolution_zyx_px : tuple of float, optional - Resolution limit per axis **in pixels**, required by and only valid - with ``resolution_mode="manual"``. Callers holding a physical - resolution should divide by their pixel size first. + Resolution limit per axis **in pixels**, required by and only valid with + ``resolution_mode="manual"``. Divide a physical resolution by the pixel + size first. beta_convention : {"reference", "paper"}, optional How ``beta`` calibrates the Wiener-Butterworth transition, by default ``"reference"``. See Notes. Ignored by every other kind. @@ -140,21 +158,18 @@ def calculate_back_projector( Notes ----- - Guo et al.'s text and their reference code disagree on how ``beta`` maps - onto the Butterworth transition width for the Wiener-Butterworth filter. - Both are members of one family, ``eps**2 = beta_w**p / beta**2 - 1``, where - ``beta_w`` is the Wiener term's own gain at the lateral cutoff: the paper's - Eq. 27 is ``p = 2`` and the reference code is ``p = 1``. They coincide only - when ``beta_w == 1``, which never happens in practice because the Wiener - term amplifies near the cutoff, making ``beta_w`` of order ten. - - The two are exactly interconvertible. Under ``"paper"`` the filter's actual - gain at the cutoff is ``beta``, so ``beta`` means literally what it says; - under ``"reference"`` it is ``beta * sqrt(beta_w)``. This module implements - the paper's formula and, for ``"reference"``, first rescales ``beta`` by - ``sqrt(beta_w)`` to reproduce the reference code exactly. ``"reference"`` - is the default so that the ``beta`` values published in Table S2.1 produce - the published results. + Guo et al.'s text and their reference code disagree on how ``beta`` maps to + the Butterworth transition width. Both belong to one family, + ``eps**2 = beta_w**p / beta**2 - 1``, with ``beta_w`` the Wiener term's own + gain at the lateral cutoff: Eq. 27 has ``p = 2``, the reference code + ``p = 1``. They agree only at ``beta_w == 1``, which never occurs -- the + Wiener term amplifies near the cutoff, putting ``beta_w`` around ten. + + So the filter's actual gain at the cutoff is ``beta`` under ``"paper"`` and + ``beta * sqrt(beta_w)`` under ``"reference"``. This module implements the + paper's formula, rescaling ``beta`` by ``sqrt(beta_w)`` for ``"reference"``, + which is the default so published Table S2.1 values reproduce published + results. References ---------- From 308bf7613f81c31dca5d6af46a306d00d98ff473 Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Fri, 7 Aug 2026 15:16:26 -0700 Subject: [PATCH 07/11] style: format fix Signed-off-by: Sricharan Reddy Varra --- waveorder/backprojector.py | 1 - 1 file changed, 1 deletion(-) diff --git a/waveorder/backprojector.py b/waveorder/backprojector.py index ecd8de02..371b855d 100644 --- a/waveorder/backprojector.py +++ b/waveorder/backprojector.py @@ -41,7 +41,6 @@ `eguomin/regDeconProject `_. """ - import math from typing import Literal, Optional From f3ce5f04e8d02258d39b75569669328f7ba9d34c Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Tue, 11 Aug 2026 17:18:32 -0700 Subject: [PATCH 08/11] refactor(fluorescence): nest the RL knobs under an `rl` block The RL parameters sat flat alongside the Tikhonov and TV ones, so every generated config carried ten rl_* lines whether or not the algorithm read them, and the list would grow with each new back projector. They now live in an optional RLSettings block that follows the algorithm: filled with defaults for RL/RLGC, dropped with a warning otherwise. A Tikhonov config carries one `rl: null` line instead of ten. Dropping rather than rejecting a stray block keeps the napari plugin working, since it builds widgets from every field and so submits one whatever the algorithm. Two knobs leave the user-facing schema, both still reachable at the model level and both unchanged in behaviour: - rl_bp_order stays at 8. Measured across 2-16 on anisotropic confocal data, 8 sits where the passband has saturated and ringing has not. - rl_bp_resolution_mode stays at 'fwhm'. It declares an instrument class rather than tuning anything; 'fwhm_over_sqrt2' is for resolution-doubling optics such as iSIM. to_model_kwargs() is the seam between the nested config and the flat model signatures. Every apply_inverse call site uses it, including phase and birefringence, so a nested block added there later is picked up without touching the callers. Also corrects the shared reconstruction_algorithm description, which still advertised only Tikhonov and TV after RL/RLGC were added. That string regenerates into the phase and birefringence example configs. Refs #573 --- .../configs/birefringence-and-phase_3d.yml | 2 +- docs/examples/cli/configs/fluorescence_2d.yml | 9 +- docs/examples/cli/configs/fluorescence_3d.yml | 9 +- docs/examples/cli/configs/phase_2d.yml | 2 +- docs/examples/cli/configs/phase_3d.yml | 2 +- tests/models/test_backprojector.py | 49 +++++++---- tests/models/test_rlgc.py | 14 ++-- waveorder/api/_settings.py | 12 ++- waveorder/api/birefringence_and_phase.py | 6 +- waveorder/api/fluorescence.py | 82 ++++++++++++------- waveorder/api/phase.py | 4 +- 11 files changed, 113 insertions(+), 78 deletions(-) diff --git a/docs/examples/cli/configs/birefringence-and-phase_3d.yml b/docs/examples/cli/configs/birefringence-and-phase_3d.yml index 6950d22e..1525de21 100644 --- a/docs/examples/cli/configs/birefringence-and-phase_3d.yml +++ b/docs/examples/cli/configs/birefringence-and-phase_3d.yml @@ -28,7 +28,7 @@ phase: numerical_aperture_illumination: 0.9 # (optimizable) condenser numerical aperture invert_phase_contrast: false # invert contrast for positive/negative phase apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization + reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution (3D fluorescence only) regularization_strength: 0.001 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization diff --git a/docs/examples/cli/configs/fluorescence_2d.yml b/docs/examples/cli/configs/fluorescence_2d.yml index 7cb7ad2a..1e5e9b29 100644 --- a/docs/examples/cli/configs/fluorescence_2d.yml +++ b/docs/examples/cli/configs/fluorescence_2d.yml @@ -19,11 +19,4 @@ fluorescence: regularization_strength: 0.01 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization - rl_iterations: 25 # maximum RL / RLGC iterations - rl_background: 0.0 # constant background folded into the RL / RLGC Poisson forward model - rl_stopping_tolerance: null # relative-change early-stop threshold for RL / RLGC (null = run all iterations) - rl_back_projector: matched # 'matched' is the matched transpose (classic RL); the unmatched 'gaussian'/'butterworth'/'wiener'/'wiener_butterworth' converge in far fewer iterations but are supported for 'RL' only, not 'RLGC' - rl_bp_alpha: null # Wiener regularization for the 'wiener'/'wiener_butterworth' back projectors (null = matched cutoff gain) - rl_bp_beta: null # cutoff gain for the 'butterworth'/'wiener_butterworth' back projectors (null = matched cutoff gain) - rl_bp_order: 8 # Butterworth order for the 'butterworth'/'wiener_butterworth' back projectors - rl_bp_resolution_mode: fwhm # cutoff-frequency rule for the back projector ('fwhm_over_sqrt2' suits iSIM) + rl: null # Richardson-Lucy knobs; only valid with reconstruction_algorithm 'RL'/'RLGC', and filled with defaults if omitted there diff --git a/docs/examples/cli/configs/fluorescence_3d.yml b/docs/examples/cli/configs/fluorescence_3d.yml index 519eb2b2..62abc724 100644 --- a/docs/examples/cli/configs/fluorescence_3d.yml +++ b/docs/examples/cli/configs/fluorescence_3d.yml @@ -19,11 +19,4 @@ fluorescence: regularization_strength: 0.001 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization - rl_iterations: 25 # maximum RL / RLGC iterations - rl_background: 0.0 # constant background folded into the RL / RLGC Poisson forward model - rl_stopping_tolerance: null # relative-change early-stop threshold for RL / RLGC (null = run all iterations) - rl_back_projector: matched # 'matched' is the matched transpose (classic RL); the unmatched 'gaussian'/'butterworth'/'wiener'/'wiener_butterworth' converge in far fewer iterations but are supported for 'RL' only, not 'RLGC' - rl_bp_alpha: null # Wiener regularization for the 'wiener'/'wiener_butterworth' back projectors (null = matched cutoff gain) - rl_bp_beta: null # cutoff gain for the 'butterworth'/'wiener_butterworth' back projectors (null = matched cutoff gain) - rl_bp_order: 8 # Butterworth order for the 'butterworth'/'wiener_butterworth' back projectors - rl_bp_resolution_mode: fwhm # cutoff-frequency rule for the back projector ('fwhm_over_sqrt2' suits iSIM) + rl: null # Richardson-Lucy knobs; only valid with reconstruction_algorithm 'RL'/'RLGC', and filled with defaults if omitted there diff --git a/docs/examples/cli/configs/phase_2d.yml b/docs/examples/cli/configs/phase_2d.yml index 435b40a9..f87793d6 100644 --- a/docs/examples/cli/configs/phase_2d.yml +++ b/docs/examples/cli/configs/phase_2d.yml @@ -16,7 +16,7 @@ phase: numerical_aperture_illumination: 0.9 # (optimizable) condenser numerical aperture invert_phase_contrast: false # invert contrast for positive/negative phase apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization + reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution (3D fluorescence only) regularization_strength: 0.01 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization diff --git a/docs/examples/cli/configs/phase_3d.yml b/docs/examples/cli/configs/phase_3d.yml index 28a7ce3b..02f23e33 100644 --- a/docs/examples/cli/configs/phase_3d.yml +++ b/docs/examples/cli/configs/phase_3d.yml @@ -16,7 +16,7 @@ phase: numerical_aperture_illumination: 0.9 # (optimizable) condenser numerical aperture invert_phase_contrast: false # invert contrast for positive/negative phase apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization + reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution (3D fluorescence only) regularization_strength: 0.001 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization diff --git a/tests/models/test_backprojector.py b/tests/models/test_backprojector.py index 32e8946c..faef28bd 100644 --- a/tests/models/test_backprojector.py +++ b/tests/models/test_backprojector.py @@ -324,35 +324,50 @@ def test_over_iterating_an_unmatched_back_projector_warns(otf): def test_default_settings_leave_existing_configs_unchanged(): """Adding these knobs must not alter any reconstruction already in use.""" settings = fluorescence.ApplyInverseSettings() - assert settings.rl_back_projector == "matched" - assert settings.rl_bp_alpha is None - assert settings.rl_bp_beta is None + assert settings.rl is None + assert "rl_back_projector" not in settings.to_model_kwargs() + + +def test_rl_block_is_defaulted_in_and_rejected_out(): + """The block follows the algorithm: filled for RL/RLGC, refused otherwise.""" + settings = fluorescence.ApplyInverseSettings(reconstruction_algorithm="RL") + assert settings.rl.back_projector == "matched" + assert settings.rl.bp_alpha is None + + with pytest.warns(UserWarning, match="ignoring 'rl' settings"): + dropped = fluorescence.ApplyInverseSettings(reconstruction_algorithm="Tikhonov", rl={"iterations": 5}) + assert dropped.rl is None def test_settings_round_trip(): settings = fluorescence.ApplyInverseSettings( reconstruction_algorithm="RL", - rl_iterations=1, - rl_back_projector="wiener_butterworth", - rl_bp_alpha=0.001, - rl_bp_beta=0.001, - rl_bp_order=10, - rl_bp_resolution_mode="fwhm_over_sqrt2", + rl={ + "iterations": 1, + "back_projector": "wiener_butterworth", + "bp_alpha": 0.001, + "bp_beta": 0.001, + }, ) - dumped = settings.model_dump() - assert dumped["rl_back_projector"] == "wiener_butterworth" - assert dumped["rl_bp_order"] == 10 - assert dumped["rl_bp_resolution_mode"] == "fwhm_over_sqrt2" - # The dump is splatted straight into the model function, so the keys must match. + assert fluorescence.ApplyInverseSettings(**settings.model_dump()) == settings + + kwargs = settings.to_model_kwargs() + assert kwargs["rl_back_projector"] == "wiener_butterworth" + assert kwargs["rl_bp_alpha"] == 0.001 + assert "rl" not in kwargs + # These kwargs are splatted straight into the model function, so the keys must match. zyx_shape = (16, 32, 32) thick.apply_inverse_transfer_function( - torch.rand(*zyx_shape), thick.calculate_transfer_function(zyx_shape, **_OTF_KWARGS), 0, **dumped + torch.rand(*zyx_shape), thick.calculate_transfer_function(zyx_shape, **_OTF_KWARGS), 0, **kwargs ) def test_rejects_invalid_settings(): for invalid_projector in ("nonsense", "traditional"): with pytest.raises(ValueError): - fluorescence.ApplyInverseSettings(rl_back_projector=invalid_projector) + fluorescence.RLSettings(back_projector=invalid_projector) + # order and resolution_mode are model-level only; the config must not accept them. + with pytest.raises(ValueError): + fluorescence.RLSettings(bp_order=10) with pytest.raises(ValueError): - fluorescence.ApplyInverseSettings(rl_bp_resolution_mode="manual") + fluorescence.RLSettings(bp_resolution_mode="fwhm_over_sqrt2") diff --git a/tests/models/test_rlgc.py b/tests/models/test_rlgc.py index 97bea4e3..da6a8eed 100644 --- a/tests/models/test_rlgc.py +++ b/tests/models/test_rlgc.py @@ -288,13 +288,11 @@ def test_fluorescence_config_accepts_rl(algorithm): settings = fluorescence.Settings( apply_inverse={ "reconstruction_algorithm": algorithm, - "rl_iterations": 15, - "rl_background": 3.0, - "rl_stopping_tolerance": 1e-3, + "rl": {"iterations": 15, "background": 3.0, "stopping_tolerance": 1e-3}, } ) - dump = settings.apply_inverse.model_dump() - assert dump["reconstruction_algorithm"] == algorithm - assert dump["rl_iterations"] == 15 - assert dump["rl_background"] == 3.0 - assert dump["rl_stopping_tolerance"] == 1e-3 + kwargs = settings.apply_inverse.to_model_kwargs() + assert kwargs["reconstruction_algorithm"] == algorithm + assert kwargs["rl_iterations"] == 15 + assert kwargs["rl_background"] == 3.0 + assert kwargs["rl_stopping_tolerance"] == 1e-3 diff --git a/waveorder/api/_settings.py b/waveorder/api/_settings.py index 925783a5..d6fe788d 100644 --- a/waveorder/api/_settings.py +++ b/waveorder/api/_settings.py @@ -128,8 +128,18 @@ class FourierApplyInverseSettings(MyBaseModel): # (see waveorder.api.fluorescence); other modalities raise NotImplementedError. reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = Field( default="Tikhonov", - description="'Tikhonov' or 'TV' regularization", + description="'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution " + "(3D fluorescence only)", ) regularization_strength: NonNegativeFloat = Field(default=1e-3, description="strength of regularization") TV_rho_strength: PositiveFloat = Field(default=1e-3, description="ADMM rho parameter for TV regularization") TV_iterations: NonNegativeInt = Field(default=1, description="ADMM iterations for TV regularization") + + def to_model_kwargs(self) -> dict: + """Flatten to the keyword arguments of ``apply_inverse_transfer_function``. + + The config groups related knobs into blocks so a YAML only carries the + ones its algorithm reads; the model functions take one flat signature. + This is the seam between the two. + """ + return self.model_dump() diff --git a/waveorder/api/birefringence_and_phase.py b/waveorder/api/birefringence_and_phase.py index 7da9a5fa..74acab80 100644 --- a/waveorder/api/birefringence_and_phase.py +++ b/waveorder/api/birefringence_and_phase.py @@ -406,7 +406,7 @@ def apply_inverse_transfer_function( ) = isotropic_thin_3d.apply_inverse_transfer_function( brightfield_3d, _to_singular_system(transfer_function, "vector_singular_system"), - **settings_phase.apply_inverse.model_dump(), + **settings_phase.apply_inverse.to_model_kwargs(), ) retardance = radians_to_nanometers(reconstructed_parameters_2d[0], wavelength) @@ -430,7 +430,7 @@ def apply_inverse_transfer_function( _to_tensor(transfer_function, "real_potential_transfer_function"), _to_tensor(transfer_function, "imaginary_potential_transfer_function"), z_padding=settings_phase.transfer_function.z_padding, - **settings_phase.apply_inverse.model_dump(), + **settings_phase.apply_inverse.to_model_kwargs(), ) retardance = radians_to_nanometers(reconstructed_parameters_3d[0], wavelength) @@ -444,7 +444,7 @@ def apply_inverse_transfer_function( szyx_data=stokes, singular_system=_to_singular_system(transfer_function, "vector_singular_system"), intensity_to_stokes_matrix=None, - **settings_phase.apply_inverse.model_dump(), + **settings_phase.apply_inverse.to_model_kwargs(), ) new_ret = (joint_recon_params[1] ** 2 + joint_recon_params[2] ** 2) ** (0.5) diff --git a/waveorder/api/fluorescence.py b/waveorder/api/fluorescence.py index b6a1b790..acccbb90 100644 --- a/waveorder/api/fluorescence.py +++ b/waveorder/api/fluorescence.py @@ -63,28 +63,19 @@ def warn_wavelength_consistency(self): return self -class ApplyInverseSettings(FourierApplyInverseSettings): - """Fluorescence inverse settings. - - Extends the shared Fourier (Tikhonov / TV) settings with the iterative - Richardson-Lucy ("RL") and Gradient-Consensus ("RLGC") options, which are - available for 3D fluorescence reconstruction only. - """ +class RLSettings(MyBaseModel): + """Richardson-Lucy knobs, read only when ``reconstruction_algorithm`` is 'RL' or 'RLGC'.""" - reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = Field( - default="Tikhonov", - description="'Tikhonov'/'TV' filters or 'RL'/'RLGC' iterative deconvolution", - ) - rl_iterations: PositiveInt = Field(default=25, description="maximum RL / RLGC iterations") - rl_background: NonNegativeFloat = Field( + iterations: PositiveInt = Field(default=25, description="maximum RL / RLGC iterations") + background: NonNegativeFloat = Field( default=0.0, description="constant background folded into the RL / RLGC Poisson forward model", ) - rl_stopping_tolerance: Optional[NonNegativeFloat] = Field( + stopping_tolerance: Optional[NonNegativeFloat] = Field( default=None, - description="relative-change early-stop threshold for RL / RLGC (null = run all iterations)", + description="relative-change early-stop threshold (null = run all iterations)", ) - rl_back_projector: BackProjectorType = Field( + back_projector: BackProjectorType = Field( default="matched", description=( "'matched' is the matched transpose (classic RL); the unmatched " @@ -92,25 +83,60 @@ class ApplyInverseSettings(FourierApplyInverseSettings): "iterations but are supported for 'RL' only, not 'RLGC'" ), ) - rl_bp_alpha: Optional[PositiveFloat] = Field( + bp_alpha: Optional[PositiveFloat] = Field( default=None, description="Wiener regularization for the 'wiener'/'wiener_butterworth' back projectors " - "(null = matched cutoff gain)", + "(null = matched cutoff gain squared); lower inverts harder, converging in fewer " + "iterations but amplifying noise", ) - rl_bp_beta: Optional[PositiveFloat] = Field( + bp_beta: Optional[PositiveFloat] = Field( default=None, description="cutoff gain for the 'butterworth'/'wiener_butterworth' back projectors " - "(null = matched cutoff gain)", + "(null = matched cutoff gain); lower suppresses harder past the resolution limit", ) - rl_bp_order: PositiveInt = Field( - default=8, - description="Butterworth order for the 'butterworth'/'wiener_butterworth' back projectors", + + +class ApplyInverseSettings(FourierApplyInverseSettings): + """Fluorescence inverse settings. + + Extends the shared Fourier (Tikhonov / TV) settings with the iterative + Richardson-Lucy ("RL") and Gradient-Consensus ("RLGC") options, which are + available for 3D fluorescence reconstruction only. + """ + + reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = Field( + default="Tikhonov", + description="'Tikhonov'/'TV' filters or 'RL'/'RLGC' iterative deconvolution", ) - rl_bp_resolution_mode: Literal["fwhm", "fwhm_over_sqrt2"] = Field( - default="fwhm", - description="cutoff-frequency rule for the back projector ('fwhm_over_sqrt2' suits iSIM)", + rl: Optional[RLSettings] = Field( + default=None, + description="Richardson-Lucy knobs; only valid with reconstruction_algorithm 'RL'/'RLGC', " + "and filled with defaults if omitted there", ) + @model_validator(mode="after") + def _rl_block_matches_algorithm(self): + """Keep the block and the algorithm in step, so a config only carries what it reads.""" + if self.reconstruction_algorithm in ("RL", "RLGC"): + if self.rl is None: + self.rl = RLSettings() + elif self.rl is not None: + # Dropped rather than rejected: the napari plugin builds its widgets from + # every field and so always submits a block, whatever the algorithm. + warnings.warn( + f"ignoring 'rl' settings: reconstruction_algorithm is " + f"{self.reconstruction_algorithm!r}, not 'RL' or 'RLGC'", + UserWarning, + ) + self.rl = None + return self + + def to_model_kwargs(self) -> dict: + kwargs = self.model_dump(exclude={"rl"}) + if self.rl is not None: + kwargs.update({f"rl_{name}": value for name, value in self.rl.model_dump().items()}) + return kwargs + class Settings(MyBaseModel): transfer_function: TransferFunctionSettings = TransferFunctionSettings() @@ -359,7 +385,7 @@ def apply_inverse_transfer_function( output = isotropic_fluorescent_thin_3d.apply_inverse_transfer_function( zyx_tensor, (U.to(device), S.to(device), Vh.to(device)), - **settings.apply_inverse.model_dump(), + **settings.apply_inverse.to_model_kwargs(), ) # [fluo, 3] elif recon_dim == 3: @@ -367,7 +393,7 @@ def apply_inverse_transfer_function( zyx_tensor, _to_tensor(transfer_function, "optical_transfer_function").to(device), settings.transfer_function.z_padding, - **settings.apply_inverse.model_dump(), + **settings.apply_inverse.to_model_kwargs(), ) # Wrap output tensor(s) back into xr.DataArray(s) diff --git a/waveorder/api/phase.py b/waveorder/api/phase.py index ff3588fb..4f005e33 100644 --- a/waveorder/api/phase.py +++ b/waveorder/api/phase.py @@ -334,7 +334,7 @@ def apply_inverse_transfer_function( _, output = isotropic_thin_3d.apply_inverse_transfer_function( zyx_tensor, (U.to(device), S.to(device), Vh.to(device)), - **settings.apply_inverse.model_dump(), + **settings.apply_inverse.to_model_kwargs(), ) # [phase only, 3] elif recon_dim == 3: @@ -343,7 +343,7 @@ def apply_inverse_transfer_function( _to_tensor(transfer_function, "real_potential_transfer_function").to(device), _to_tensor(transfer_function, "imaginary_potential_transfer_function").to(device), z_padding=settings.transfer_function.z_padding, - **settings.apply_inverse.model_dump(), + **settings.apply_inverse.to_model_kwargs(), ) # Wrap output tensor(s) back into xr.DataArray(s) From c8ae510542c839c32292a9d28a48d5984c48d841 Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Tue, 11 Aug 2026 17:24:51 -0700 Subject: [PATCH 09/11] revert(settings): keep the shared algorithm description as it was The previous commit widened the shared reconstruction_algorithm description to mention RL/RLGC. That field is on FourierApplyInverseSettings, so the text regenerated into the phase and birefringence example configs -- modalities that raise NotImplementedError for RL/RLGC, where the original 'Tikhonov' or 'TV' wording was already correct. Fluorescence overrides the field with its own description, so its configs were never affected either way. --- docs/examples/cli/configs/birefringence-and-phase_3d.yml | 2 +- docs/examples/cli/configs/phase_2d.yml | 2 +- docs/examples/cli/configs/phase_3d.yml | 2 +- waveorder/api/_settings.py | 3 +-- 4 files changed, 4 insertions(+), 5 deletions(-) diff --git a/docs/examples/cli/configs/birefringence-and-phase_3d.yml b/docs/examples/cli/configs/birefringence-and-phase_3d.yml index 1525de21..6950d22e 100644 --- a/docs/examples/cli/configs/birefringence-and-phase_3d.yml +++ b/docs/examples/cli/configs/birefringence-and-phase_3d.yml @@ -28,7 +28,7 @@ phase: numerical_aperture_illumination: 0.9 # (optimizable) condenser numerical aperture invert_phase_contrast: false # invert contrast for positive/negative phase apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution (3D fluorescence only) + reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization regularization_strength: 0.001 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization diff --git a/docs/examples/cli/configs/phase_2d.yml b/docs/examples/cli/configs/phase_2d.yml index f87793d6..435b40a9 100644 --- a/docs/examples/cli/configs/phase_2d.yml +++ b/docs/examples/cli/configs/phase_2d.yml @@ -16,7 +16,7 @@ phase: numerical_aperture_illumination: 0.9 # (optimizable) condenser numerical aperture invert_phase_contrast: false # invert contrast for positive/negative phase apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution (3D fluorescence only) + reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization regularization_strength: 0.01 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization diff --git a/docs/examples/cli/configs/phase_3d.yml b/docs/examples/cli/configs/phase_3d.yml index 02f23e33..28a7ce3b 100644 --- a/docs/examples/cli/configs/phase_3d.yml +++ b/docs/examples/cli/configs/phase_3d.yml @@ -16,7 +16,7 @@ phase: numerical_aperture_illumination: 0.9 # (optimizable) condenser numerical aperture invert_phase_contrast: false # invert contrast for positive/negative phase apply_inverse: - reconstruction_algorithm: Tikhonov # 'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution (3D fluorescence only) + reconstruction_algorithm: Tikhonov # 'Tikhonov' or 'TV' regularization regularization_strength: 0.001 # strength of regularization TV_rho_strength: 0.001 # ADMM rho parameter for TV regularization TV_iterations: 1 # ADMM iterations for TV regularization diff --git a/waveorder/api/_settings.py b/waveorder/api/_settings.py index d6fe788d..c4227be0 100644 --- a/waveorder/api/_settings.py +++ b/waveorder/api/_settings.py @@ -128,8 +128,7 @@ class FourierApplyInverseSettings(MyBaseModel): # (see waveorder.api.fluorescence); other modalities raise NotImplementedError. reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = Field( default="Tikhonov", - description="'Tikhonov'/'TV' regularization, or 'RL'/'RLGC' iterative deconvolution " - "(3D fluorescence only)", + description="'Tikhonov' or 'TV' regularization", ) regularization_strength: NonNegativeFloat = Field(default=1e-3, description="strength of regularization") TV_rho_strength: PositiveFloat = Field(default=1e-3, description="ADMM rho parameter for TV regularization") From 319b4c86b54ffde89c77ce7a1f02474afbc051d7 Mon Sep 17 00:00:00 2001 From: talonchandler Date: Tue, 18 Aug 2026 10:35:26 -0700 Subject: [PATCH 10/11] fix(settings): reject unsupported RL/RLGC configs at parse time Widening the shared FourierApplyInverseSettings to accept "RL"/"RLGC" let a phase or birefringence config validate and then fail deep in the reconstruction, after the transfer function had been computed. Keep the shared model to the Fourier filters and let fluorescence widen it in its own subclass. Two more pairings are now caught while parsing rather than mid-run: 2D fluorescence with RL/RLGC (checked on ReconstructionSettings, the only model that sees reconstruction_dimension), and RLGC with an unmatched back projector. The model-level NotImplementedError guards stay for direct callers. --- tests/models/test_rlgc.py | 51 +++++++++++++++++++++++++++++++---- waveorder/api/_settings.py | 7 ++--- waveorder/api/fluorescence.py | 10 +++++++ waveorder/cli/settings.py | 14 ++++++++++ 4 files changed, 74 insertions(+), 8 deletions(-) diff --git a/tests/models/test_rlgc.py b/tests/models/test_rlgc.py index da6a8eed..cce66976 100644 --- a/tests/models/test_rlgc.py +++ b/tests/models/test_rlgc.py @@ -12,9 +12,11 @@ import pytest import torch +from pydantic import ValidationError from waveorder import rlgc from waveorder.api import fluorescence, phase +from waveorder.cli.settings import ReconstructionSettings from waveorder.models import ( isotropic_fluorescent_thick_3d as thick, ) @@ -275,11 +277,50 @@ def test_phase_2d_not_implemented_for_rl(algorithm): @pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) -def test_phase_config_accepts_rl_request(algorithm): - """RL/RLGC are valid config values everywhere (so the request reaches the - model), even though only fluorescence implements them.""" - settings = phase.Settings(apply_inverse={"reconstruction_algorithm": algorithm}) - assert settings.apply_inverse.reconstruction_algorithm == algorithm +def test_phase_config_rejects_rl_request(algorithm): + """Only fluorescence implements RL/RLGC, so a phase config naming one is + rejected while parsing rather than mid-reconstruction.""" + with pytest.raises(ValidationError): + phase.Settings(apply_inverse={"reconstruction_algorithm": algorithm}) + + +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_fluorescence_2d_config_rejects_rl(algorithm): + """RL/RLGC need thick (3D) fluorescence, and that pairing is caught at parse time.""" + with pytest.raises(ValidationError, match="reconstruction_dimension 3"): + ReconstructionSettings( + input_channel_names=["GFP"], + reconstruction_dimension=2, + fluorescence={"apply_inverse": {"reconstruction_algorithm": algorithm}}, + ) + + +@pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) +def test_fluorescence_3d_config_accepts_rl(algorithm): + """The supported pairing still parses.""" + settings = ReconstructionSettings( + input_channel_names=["GFP"], + reconstruction_dimension=3, + fluorescence={"apply_inverse": {"reconstruction_algorithm": algorithm}}, + ) + assert settings.fluorescence.apply_inverse.reconstruction_algorithm == algorithm + + +@pytest.mark.parametrize("back_projector", ["gaussian", "butterworth", "wiener", "wiener_butterworth"]) +def test_rlgc_config_rejects_unmatched_back_projector(back_projector): + """RLGC's consensus test needs a true adjoint, so the combination is refused at parse time.""" + with pytest.raises(ValidationError, match="matched"): + fluorescence.ApplyInverseSettings( + reconstruction_algorithm="RLGC", + rl={"back_projector": back_projector}, + ) + + +@pytest.mark.parametrize("back_projector", ["matched", "gaussian", "wiener_butterworth"]) +def test_rl_config_accepts_any_back_projector(back_projector): + """RL supports the unmatched back projectors.""" + settings = fluorescence.ApplyInverseSettings(reconstruction_algorithm="RL", rl={"back_projector": back_projector}) + assert settings.rl.back_projector == back_projector @pytest.mark.parametrize("algorithm", ["RL", "RLGC"]) diff --git a/waveorder/api/_settings.py b/waveorder/api/_settings.py index c4227be0..6b5954e2 100644 --- a/waveorder/api/_settings.py +++ b/waveorder/api/_settings.py @@ -124,9 +124,10 @@ def resolve_floats(self): class FourierApplyInverseSettings(MyBaseModel): - # "RL"/"RLGC" are accepted here but only implemented for 3D fluorescence - # (see waveorder.api.fluorescence); other modalities raise NotImplementedError. - reconstruction_algorithm: Literal["Tikhonov", "TV", "RL", "RLGC"] = Field( + # Only the Fourier filters live here, so a phase or birefringence config that + # asks for "RL"/"RLGC" is rejected while parsing rather than deep in the + # reconstruction. Fluorescence widens this in its own ApplyInverseSettings. + reconstruction_algorithm: Literal["Tikhonov", "TV"] = Field( default="Tikhonov", description="'Tikhonov' or 'TV' regularization", ) diff --git a/waveorder/api/fluorescence.py b/waveorder/api/fluorescence.py index acccbb90..7b536588 100644 --- a/waveorder/api/fluorescence.py +++ b/waveorder/api/fluorescence.py @@ -129,6 +129,16 @@ def _rl_block_matches_algorithm(self): UserWarning, ) self.rl = None + + # RLGC reads the sign of transpose(forward(.)) for its consensus test, which + # only means anything for a true adjoint. Catch it here so a config fails + # while parsing rather than after the transfer function has been computed. + if self.reconstruction_algorithm == "RLGC" and self.rl.back_projector != "matched": + raise ValueError( + f"reconstruction_algorithm 'RLGC' requires rl.back_projector 'matched', " + f"got {self.rl.back_projector!r}. The unmatched back projectors are " + f"supported for 'RL' only." + ) return self def to_model_kwargs(self) -> dict: diff --git a/waveorder/cli/settings.py b/waveorder/cli/settings.py index f02e6b5f..0fc6ce76 100644 --- a/waveorder/cli/settings.py +++ b/waveorder/cli/settings.py @@ -90,6 +90,20 @@ def validate_reconstruction_types(self): f"{num_channel_names} channels names provided. Please provide a single channel for fluorescence/phase reconstructions." ) + # RL/RLGC are implemented for thick (3D) fluorescence only. reconstruction_dimension + # lives here rather than on the fluorescence block, so this is the only place the + # pairing can be checked while parsing instead of mid-reconstruction. + if ( + self.fluorescence is not None + and self.reconstruction_dimension == 2 + and self.fluorescence.apply_inverse.reconstruction_algorithm in ("RL", "RLGC") + ): + raise ValueError( + f"reconstruction_algorithm " + f"{self.fluorescence.apply_inverse.reconstruction_algorithm!r} requires " + f"reconstruction_dimension 3; it is not implemented for thin (2D) fluorescence." + ) + return self @property From e7f54627b6734b974368a5b47832b6cd086ca692 Mon Sep 17 00:00:00 2001 From: talonchandler Date: Tue, 18 Aug 2026 10:39:40 -0700 Subject: [PATCH 11/11] docs(fluorescence): drop back_projector_otf from reconstruct's parameters reconstruct() computes the transfer function itself and has no back_projector_otf parameter, so documenting one promised a knob that does not exist. Point callers who want to reuse a back projector at the calculate/apply pair instead, which is where the parameter lives. --- waveorder/models/isotropic_fluorescent_thick_3d.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index 875e3ecc..30604351 100644 --- a/waveorder/models/isotropic_fluorescent_thick_3d.py +++ b/waveorder/models/isotropic_fluorescent_thick_3d.py @@ -506,16 +506,20 @@ def reconstruct( rl_bp_resolution_mode : str, optional How the back projector sets its cutoff frequency, by default "fwhm". Use "fwhm_over_sqrt2" for iSIM. - back_projector_otf : Tensor, optional - Prebuilt back projector, skipping the rl_bp_* construction. Building it - costs a few seconds on a large OTF, so callers reconstructing many tiles - should build it once with :func:`waveorder.backprojector.calculate_back_projector` - and pass it here. By default None (build it on every call). Returns ------- Tensor Fluorescence density, shape ``(Z, Y, X)`` or ``(B, Z, Y, X)`` + + Notes + ----- + This recomputes the transfer function on every call, so it does not take a + prebuilt back projector. Callers reconstructing many tiles with RL should + call :func:`calculate_transfer_function` and + :func:`apply_inverse_transfer_function` directly, building the back projector + once with :func:`waveorder.backprojector.calculate_back_projector` and + passing it as ``back_projector_otf``. """ # Use last 3 dims as zyx_shape for TF computation zyx_shape = zyx_data.shape[-3:]