diff --git a/docs/examples/cli/configs/fluorescence_2d.yml b/docs/examples/cli/configs/fluorescence_2d.yml index e01946ac..1e5e9b29 100644 --- a/docs/examples/cli/configs/fluorescence_2d.yml +++ b/docs/examples/cli/configs/fluorescence_2d.yml @@ -15,7 +15,8 @@ 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: 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 29803bfa..62abc724 100644 --- a/docs/examples/cli/configs/fluorescence_3d.yml +++ b/docs/examples/cli/configs/fluorescence_3d.yml @@ -15,7 +15,8 @@ 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: null # Richardson-Lucy knobs; only valid with reconstruction_algorithm 'RL'/'RLGC', and filled with defaults if omitted there diff --git a/tests/models/test_backprojector.py b/tests/models/test_backprojector.py new file mode 100644 index 00000000..faef28bd --- /dev/null +++ b/tests/models/test_backprojector.py @@ -0,0 +1,373 @@ +"""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 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, + "back_projector": "wiener_butterworth", + "bp_alpha": 0.001, + "bp_beta": 0.001, + }, + ) + 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, **kwargs + ) + + +def test_rejects_invalid_settings(): + for invalid_projector in ("nonsense", "traditional"): + with pytest.raises(ValueError): + 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.RLSettings(bp_resolution_mode="fwhm_over_sqrt2") 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 new file mode 100644 index 00000000..cce66976 --- /dev/null +++ b/tests/models/test_rlgc.py @@ -0,0 +1,339 @@ +"""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 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, +) +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 + + +@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.""" + 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_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"]) +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, "background": 3.0, "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 bb0376a3..6b5954e2 100644 --- a/waveorder/api/_settings.py +++ b/waveorder/api/_settings.py @@ -124,6 +124,9 @@ def resolve_floats(self): class FourierApplyInverseSettings(MyBaseModel): + # 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", @@ -131,3 +134,12 @@ class FourierApplyInverseSettings(MyBaseModel): 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 0c9b22b9..7b536588 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 ( @@ -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, @@ -62,7 +63,89 @@ def warn_wavelength_consistency(self): return self -ApplyInverseSettings = FourierApplyInverseSettings +class RLSettings(MyBaseModel): + """Richardson-Lucy knobs, read only when ``reconstruction_algorithm`` is 'RL' or 'RLGC'.""" + + 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", + ) + stopping_tolerance: Optional[NonNegativeFloat] = Field( + default=None, + description="relative-change early-stop threshold (null = run all iterations)", + ) + 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'" + ), + ) + bp_alpha: Optional[PositiveFloat] = Field( + default=None, + description="Wiener regularization for the 'wiener'/'wiener_butterworth' back projectors " + "(null = matched cutoff gain squared); lower inverts harder, converging in fewer " + "iterations but amplifying noise", + ) + bp_beta: Optional[PositiveFloat] = Field( + default=None, + description="cutoff gain for the 'butterworth'/'wiener_butterworth' back projectors " + "(null = matched cutoff gain); lower suppresses harder past the resolution limit", + ) + + +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: 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 + + # 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: + 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): @@ -312,7 +395,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: @@ -320,7 +403,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) diff --git a/waveorder/backprojector.py b/waveorder/backprojector.py new file mode 100644 index 00000000..371b855d --- /dev/null +++ b/waveorder/backprojector.py @@ -0,0 +1,420 @@ +"""Unmatched back projectors that accelerate Richardson-Lucy deconvolution. + +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 + +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"`` -- +Wiener ``"wiener"`` ``"wiener_butterworth"`` +``1`` (Dirac delta) (noise, unusable) ``"butterworth"`` +===================== ==================== ========================== + +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``). + +Reference: Guo et al. 2020, Supplementary Note 2 +(`doi.org/10.1038/s41587-020-0560-x `_), +following ``BackProjector.m`` in +`eguomin/regDeconProject `_. +""" + +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. + + 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. + back_projector : {"matched", "gaussian", "butterworth", "wiener", \ +"wiener_butterworth"}, optional + Which one to build, by default ``"matched"``. + alpha : float, optional + 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 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 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 + 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"``. 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. + + 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 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 + ---------- + 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) + # 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: + 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/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 diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index f251d6c0..30604351 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, 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 @@ -100,7 +102,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( @@ -238,10 +267,19 @@ 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, + 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", + back_projector_otf: Tensor | None = None, ) -> Tensor: """Reconstructs fluorescence density from defocus data. @@ -254,14 +292,48 @@ 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) + 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. + 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 ------- @@ -286,6 +358,68 @@ 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 + + # 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, + ) + + # 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))) + + def transpose(y: Tensor) -> Tensor: + 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), + 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 +439,18 @@ 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, + 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. @@ -330,19 +472,54 @@ 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 + 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 ------- 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:] @@ -364,4 +541,12 @@ 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, + 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 10e5cfeb..567abc4b 100644 --- a/waveorder/models/isotropic_fluorescent_thin_3d.py +++ b/waveorder/models/isotropic_fluorescent_thin_3d.py @@ -272,10 +272,18 @@ 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, + 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. @@ -285,20 +293,42 @@ 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 + 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 ------- 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..c3e91cb4 --- /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 = counts.clamp_min(0.0).round_() + probs = counts.new_tensor(probability).expand_as(counts) + 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, 0.0, 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) + + 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.count_nonzero(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