diff --git a/tests/api_tests/test_reconstruct.py b/tests/api_tests/test_reconstruct.py index b21f81ec..cec15684 100644 --- a/tests/api_tests/test_reconstruct.py +++ b/tests/api_tests/test_reconstruct.py @@ -258,7 +258,7 @@ def settings_with(rolloff): assert spectrum[ny, :].max() < 1e-3 * hard_spectrum[ny, :].max() -def test_fluorescence_apodization_rolloff_parses(make_czyx): +def test_fluorescence_apodization_rolloff(make_czyx): settings = fluorescence.Settings( transfer_function=fluorescence.TransferFunctionSettings( yx_pixel_size=6.5 / 20, @@ -269,7 +269,12 @@ def test_fluorescence_apodization_rolloff_parses(make_czyx): ), apply_inverse=fluorescence.ApplyInverseSettings(apodization_rolloff=0.25), ) - assert settings.apply_inverse.to_model_kwargs()["apodization_rolloff"] == 0.25 + data = make_czyx() + apod = fluorescence.reconstruct(data, recon_dim=3, settings=settings) + settings.apply_inverse.apodization_rolloff = 0.0 + hard = fluorescence.reconstruct(data, recon_dim=3, settings=settings) - result = fluorescence.reconstruct(make_czyx(), recon_dim=3, settings=settings) - assert np.all(np.isfinite(result.values)) + spectrum = np.abs(np.fft.fft2(apod.values[0, apod.sizes["z"] // 2])) + hard_spectrum = np.abs(np.fft.fft2(hard.values[0, hard.sizes["z"] // 2])) + ny = spectrum.shape[-2] // 2 + assert spectrum[ny, :].max() < 1e-3 * hard_spectrum[ny, :].max() diff --git a/tests/optim/test_autoreg.py b/tests/optim/test_autoreg.py index df2fc833..e3c47d78 100644 --- a/tests/optim/test_autoreg.py +++ b/tests/optim/test_autoreg.py @@ -16,6 +16,7 @@ from waveorder.models import isotropic_fluorescent_thick_3d, phase_thick_3d from waveorder.optim import autoreg from waveorder.optim.autoreg import AutoRegularizationSettings +from waveorder.reconstruct import tikhonov_regularized_inverse_filter PHASE_TF_KWARGS = dict( yx_pixel_size=0.15, @@ -55,7 +56,7 @@ def _sweep_one(data, tf, contrast, z_padding, strength, apodization_rolloff=0.0) spectrum = torch.fft.fftn(measurement, dim=(-3, -2, -1)) recon = torch.real( torch.fft.ifftn( - spectrum * autoreg._inverse_filter(tf, strength, apodization_rolloff), + spectrum * tikhonov_regularized_inverse_filter(tf, strength, apodization_rolloff=apodization_rolloff), dim=(-3, -2, -1), ) ) diff --git a/tests/test_reconstruct.py b/tests/test_reconstruct.py new file mode 100644 index 00000000..b97c0530 --- /dev/null +++ b/tests/test_reconstruct.py @@ -0,0 +1,59 @@ +import pytest +import torch + +from waveorder.reconstruct import tikhonov_regularized_inverse_filter + + +def test_inverse_filter_without_apodization(): + forward = torch.tensor([1 + 2j, 2 - 1j, 0j]) + expected = torch.tensor([1 - 2j, 2 + 1j, 0j]) / 5.25 + + for kwargs in [{}, {"apodization_rolloff": 0.0}]: + actual = tikhonov_regularized_inverse_filter(forward, 0.25, **kwargs) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(forward, torch.tensor([1 + 2j, 2 - 1j, 0j])) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.complex64, torch.complex128]) +@pytest.mark.parametrize( + "rolloff,y_window,x_window", + [ + (0.5, [1, 1, 1, 0.5, 0, 0.5, 1, 1], [1, 1, 1, 1, 0.75, 0.25, 0, 0.25, 0.75, 1, 1, 1]), + ( + 1.0, + [1, 0.654508497, 0.095491503, 0.095491503, 0.654508497], + [1, 0.654508497, 0.095491503, 0.095491503, 0.654508497], + ), + ], +) +def test_inverse_filter_apodizes_only_transverse_axes(dtype, rolloff, y_window, x_window): + value = 1 + 2j if dtype.is_complex else 2.0 + inverse_value = (1 - 2j) / 5.25 if dtype.is_complex else 2.0 / 4.25 + forward = torch.full((2, 3, len(y_window), len(x_window)), value, dtype=dtype) + expected_window = torch.tensor(y_window)[:, None] * torch.tensor(x_window) + expected = (inverse_value * expected_window.to(dtype)).expand_as(forward) + + actual = tikhonov_regularized_inverse_filter(forward, 0.25, apodization_rolloff=rolloff) + + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) + assert actual.dtype == forward.dtype + assert actual.device == forward.device + torch.testing.assert_close(forward, torch.full_like(forward, value)) + + +def test_svd_apodization_preserves_inverse_and_gradients(): + generator = torch.Generator().manual_seed(584) + U = torch.randn(2, 2, 2, 8, 12, dtype=torch.complex64, generator=generator) + S = (torch.rand(2, 2, 8, 12, generator=generator) + 0.1).requires_grad_() + Vh = torch.randn(2, 2, 4, 8, 12, dtype=torch.complex64, generator=generator) + window = torch.tensor([1, 1, 1, 0.5, 0, 0.5, 1, 1])[:, None] * torch.tensor( + [1, 1, 1, 1, 0.75, 0.25, 0, 0.25, 0.75, 1, 1, 1] + ) + expected = torch.einsum("bsj...,bj...,bjf...->bfs...", U, S / (S**2 + 1e-3), Vh) * window + S_reg = tikhonov_regularized_inverse_filter(S, 1e-3, apodization_rolloff=0.5) + actual = torch.einsum("bsj...,bj...,bjf...->bfs...", U, S_reg, Vh) + + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5) + expected_grad = torch.autograd.grad(expected.abs().square().mean(), S, retain_graph=True)[0] + actual_grad = torch.autograd.grad(actual.abs().square().mean(), S)[0] + torch.testing.assert_close(actual_grad, expected_grad, rtol=2e-5, atol=2e-5) diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 40226ddd..43474221 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -1,6 +1,6 @@ import torch -from waveorder.sampling import nd_fourier_central_cuboid, raised_cosine_window +from waveorder.sampling import apply_transverse_apodization, nd_fourier_central_cuboid, raised_cosine_window def test_nd_fourier_central_cuboid(): @@ -19,3 +19,17 @@ def test_raised_cosine_window(): assert torch.allclose(window[1:4], window.flip(0)[:3]) # monotone roll-off assert torch.all(window[1:5] <= window[:4]) + + +def test_transverse_apodization_preserves_input_and_gradients(): + source = torch.full((2, 3, 8, 8), 1 + 2j, requires_grad=True) + assert apply_transverse_apodization(source, 0.0) is source + + axis_window = torch.tensor([1, 1, 1, 0.5, 0, 0.5, 1, 1]) + window = (axis_window[:, None] * axis_window).to(source.dtype).expand_as(source) + result = apply_transverse_apodization(source, 0.5) + + torch.testing.assert_close(result, (1 + 2j) * window) + torch.testing.assert_close(source, torch.full_like(source, 1 + 2j)) + result.real.sum().backward() + torch.testing.assert_close(source.grad, window) diff --git a/waveorder/models/inplane_oriented_thick_pol3d_vector.py b/waveorder/models/inplane_oriented_thick_pol3d_vector.py index 11146cb8..06be332f 100644 --- a/waveorder/models/inplane_oriented_thick_pol3d_vector.py +++ b/waveorder/models/inplane_oriented_thick_pol3d_vector.py @@ -8,6 +8,7 @@ from waveorder import optics, sampling, stokes, util from waveorder._pixel_size import YXPixelSize from waveorder.filter import apply_filter_bank +from waveorder.reconstruct import tikhonov_regularized_inverse_filter from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer @@ -287,19 +288,9 @@ def apply_inverse_transfer_function( ): # Key computation U, S, Vh = singular_system - S_reg = S / (S**2 + regularization_strength) + S_reg = tikhonov_regularized_inverse_filter(S, regularization_strength, apodization_rolloff=apodization_rolloff) sfzyx_inverse_filter = torch.einsum("sjzyx,jzyx,jfzyx->sfzyx", U, S_reg, Vh) - if apodization_rolloff > 0: - # Raised-cosine roll-off of the inverse filter at the transverse - # Nyquist edge; suppresses Nyquist-rate checkerboard artifacts. - window = sampling.raised_cosine_window( - sfzyx_inverse_filter.shape[-2], apodization_rolloff, device=sfzyx_inverse_filter.device - )[:, None] * sampling.raised_cosine_window( - sfzyx_inverse_filter.shape[-1], apodization_rolloff, device=sfzyx_inverse_filter.device - ) - sfzyx_inverse_filter = sfzyx_inverse_filter * window - fzyx_recon = apply_filter_bank(sfzyx_inverse_filter, szyx_data) return fzyx_recon diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index 00217ccf..d1d1d9a3 100644 --- a/waveorder/models/isotropic_fluorescent_thick_3d.py +++ b/waveorder/models/isotropic_fluorescent_thick_3d.py @@ -358,15 +358,9 @@ def apply_inverse_transfer_function( # Reconstruct if reconstruction_algorithm == "Tikhonov": - inverse_filter = tikhonov_regularized_inverse_filter(optical_transfer_function, regularization_strength) - - if apodization_rolloff > 0: - window = sampling.raised_cosine_window( - inverse_filter.shape[-2], apodization_rolloff, device=inverse_filter.device - )[:, None] * sampling.raised_cosine_window( - inverse_filter.shape[-1], apodization_rolloff, device=inverse_filter.device - ) - inverse_filter = inverse_filter * window + inverse_filter = tikhonov_regularized_inverse_filter( + optical_transfer_function, regularization_strength, apodization_rolloff=apodization_rolloff + ) # Batched FFT multiply: inverse_filter (Z,Y,X) broadcasts over B zyx_fft = torch.fft.fftn(zyx_padded, dim=(-3, -2, -1)) diff --git a/waveorder/models/isotropic_fluorescent_thin_3d.py b/waveorder/models/isotropic_fluorescent_thin_3d.py index d6a4b538..a5a76ae7 100644 --- a/waveorder/models/isotropic_fluorescent_thin_3d.py +++ b/waveorder/models/isotropic_fluorescent_thin_3d.py @@ -9,6 +9,7 @@ from waveorder import optics, sampling, util from waveorder._pixel_size import YXPixelSize from waveorder.filter import apply_filter_bank +from waveorder.reconstruct import tikhonov_regularized_inverse_filter def generate_test_phantom( @@ -344,17 +345,9 @@ def apply_inverse_transfer_function( if reconstruction_algorithm == "Tikhonov": U, S, Vh = singular_system - S_reg = S / (S**2 + regularization_strength) + S_reg = tikhonov_regularized_inverse_filter(S, regularization_strength, apodization_rolloff=apodization_rolloff) sfyx_inverse_filter = torch.einsum("sj...,j...,jf...->fs...", U, S_reg, Vh) - if apodization_rolloff > 0: - window = sampling.raised_cosine_window( - sfyx_inverse_filter.shape[-2], apodization_rolloff, device=sfyx_inverse_filter.device - )[:, None] * sampling.raised_cosine_window( - sfyx_inverse_filter.shape[-1], apodization_rolloff, device=sfyx_inverse_filter.device - ) - sfyx_inverse_filter = sfyx_inverse_filter * window - results = [] for b in range(zyx_data.shape[0]): results.append(apply_filter_bank(sfyx_inverse_filter, zyx_data[b])[0]) diff --git a/waveorder/models/isotropic_thin_3d.py b/waveorder/models/isotropic_thin_3d.py index 0ecb5f89..0f958413 100644 --- a/waveorder/models/isotropic_thin_3d.py +++ b/waveorder/models/isotropic_thin_3d.py @@ -10,6 +10,7 @@ from waveorder import optics, sampling, util from waveorder._pixel_size import YXPixelSize from waveorder.filter import apply_filter_bank +from waveorder.reconstruct import tikhonov_regularized_inverse_filter def generate_test_phantom( @@ -426,31 +427,20 @@ def apply_inverse_transfer_function( U, S, Vh = singular_system batched_ss = S.ndim == 4 # (B, 2, Vy, Vx) - if apodization_rolloff > 0: - window = sampling.raised_cosine_window(S.shape[-2], apodization_rolloff, device=S.device)[ - :, None - ] * sampling.raised_cosine_window(S.shape[-1], apodization_rolloff, device=S.device) - else: - window = None + S_reg = tikhonov_regularized_inverse_filter(S, regularization_strength, apodization_rolloff=apodization_rolloff) if not batched_ss: # Shared singular system: compute inverse filter once - S_reg = S / (S**2 + regularization_strength) sfyx_inverse_filter = torch.einsum("sj...,j...,jf...->fs...", U, S_reg, Vh) - if window is not None: - sfyx_inverse_filter = sfyx_inverse_filter * window results = [] for b in range(zyx.shape[0]): results.append(apply_filter_bank(sfyx_inverse_filter, zyx[b])) output = torch.stack(results, dim=0) # (B, 2, Y, X) else: # Per-tile singular system: compute inverse filter per tile - S_reg = S / (S**2 + regularization_strength) results = [] for b in range(zyx.shape[0]): filt_b = torch.einsum("sj...,j...,jf...->fs...", U[b], S_reg[b], Vh[b]) - if window is not None: - filt_b = filt_b * window results.append(apply_filter_bank(filt_b, zyx[b])) output = torch.stack(results, dim=0) # (B, 2, Y, X) diff --git a/waveorder/models/phase_thick_3d.py b/waveorder/models/phase_thick_3d.py index 848ccec6..03bc8d39 100644 --- a/waveorder/models/phase_thick_3d.py +++ b/waveorder/models/phase_thick_3d.py @@ -436,15 +436,9 @@ def apply_inverse_transfer_function( # Reconstruct if reconstruction_algorithm == "Tikhonov": - inverse_filter = tikhonov_regularized_inverse_filter(effective_transfer_function, regularization_strength) - - if apodization_rolloff > 0: - window = sampling.raised_cosine_window( - inverse_filter.shape[-2], apodization_rolloff, device=inverse_filter.device - )[:, None] * sampling.raised_cosine_window( - inverse_filter.shape[-1], apodization_rolloff, device=inverse_filter.device - ) - inverse_filter = inverse_filter * window + inverse_filter = tikhonov_regularized_inverse_filter( + effective_transfer_function, regularization_strength, apodization_rolloff=apodization_rolloff + ) # Batched FFT multiply: inverse_filter (Z,Y,X) broadcasts over B zyx_fft = torch.fft.fftn(zyx, dim=(-3, -2, -1)) diff --git a/waveorder/optim/autoreg.py b/waveorder/optim/autoreg.py index ef9d4414..4c987ebb 100644 --- a/waveorder/optim/autoreg.py +++ b/waveorder/optim/autoreg.py @@ -51,7 +51,7 @@ from scipy.signal import find_peaks from torch import Tensor -from waveorder import sampling, util +from waveorder import util from waveorder.reconstruct import tikhonov_regularized_inverse_filter # Fixed rather than exposed as settings: these are implementation choices of the @@ -590,24 +590,6 @@ def _window_totals(integral): return int(y_starts[flat // len(x_starts)]), int(x_starts[flat % len(x_starts)]), size -def _inverse_filter(transfer_function: Tensor, strength: float, apodization_rolloff: float) -> Tensor: - """Build the inverse filter the 3D models build, apodization included. - - Kept in step with ``phase_thick_3d`` and ``isotropic_fluorescent_thick_3d``: - the sweep has to score the reconstruction that will actually be produced, and - ``apodization_rolloff`` reshapes the filter after regularization. - """ - inverse_filter = tikhonov_regularized_inverse_filter(transfer_function, strength) - if apodization_rolloff > 0: - window = sampling.raised_cosine_window( - inverse_filter.shape[-2], apodization_rolloff, device=inverse_filter.device - )[:, None] * sampling.raised_cosine_window( - inverse_filter.shape[-1], apodization_rolloff, device=inverse_filter.device - ) - inverse_filter = inverse_filter * window - return inverse_filter - - # --- the sweep --- @@ -722,7 +704,10 @@ def select_regularization( def _reconstruct(strength: float) -> Tensor: return torch.real( torch.fft.ifftn( - measurement_fft * _inverse_filter(transfer_function, strength, apodization_rolloff), + measurement_fft + * tikhonov_regularized_inverse_filter( + transfer_function, strength, apodization_rolloff=apodization_rolloff + ), dim=(-3, -2, -1), ) ) diff --git a/waveorder/reconstruct.py b/waveorder/reconstruct.py index 9686170e..db867309 100644 --- a/waveorder/reconstruct.py +++ b/waveorder/reconstruct.py @@ -1,15 +1,29 @@ import torch +from waveorder import sampling -def tikhonov_regularized_inverse_filter(forward_filter: torch.Tensor, regularization_strength: float): + +def tikhonov_regularized_inverse_filter( + forward_filter: torch.Tensor, + regularization_strength: float, + *, + apodization_rolloff: float = 0.0, +) -> torch.Tensor: """Compute the Tikhonov regularized inverse filter from a forward filter. Parameters ---------- forward_filter : torch.Tensor - The forward filter tensor. + The forward filter tensor, or real singular values. With apodization + enabled, the final two axes must be Y and X in unshifted FFT order. regularization_strength : float The strength of the regularization term. + apodization_rolloff : float, optional + Raised-cosine roll-off fraction between 0 and 1, by default 0.0 + (disabled). Applied after regularization along Y and X only; + leading batch, channel, and Z axes are preserved. + For an SVD inverse, regularized singular values may be apodized + before contraction because the window is shared across modes. Returns ------- torch.Tensor @@ -17,4 +31,5 @@ def tikhonov_regularized_inverse_filter(forward_filter: torch.Tensor, regulariza """ forward_filter_conj = torch.conj(forward_filter) - return forward_filter_conj / ((forward_filter_conj * forward_filter) + regularization_strength) + inverse_filter = forward_filter_conj / ((forward_filter_conj * forward_filter) + regularization_strength) + return sampling.apply_transverse_apodization(inverse_filter, apodization_rolloff) diff --git a/waveorder/sampling.py b/waveorder/sampling.py index 5a74dded..aebd86b7 100644 --- a/waveorder/sampling.py +++ b/waveorder/sampling.py @@ -96,6 +96,21 @@ def raised_cosine_window(size, rolloff, device=None): return window +def apply_transverse_apodization(fourier_filter: torch.Tensor, rolloff: float) -> torch.Tensor: + """Apply a raised-cosine window to the final Y/X axes in unshifted FFT order. + + ``rolloff`` is a fraction between 0 and 1. Zero returns the input tensor + unchanged, without allocation. Positive rolloff returns a new tensor, + preserving leading axes and leaving the input untouched. + """ + if rolloff > 0: + window = raised_cosine_window(fourier_filter.shape[-2], rolloff, device=fourier_filter.device)[ + :, None + ] * raised_cosine_window(fourier_filter.shape[-1], rolloff, device=fourier_filter.device) + return fourier_filter * window + return fourier_filter + + def nd_fourier_central_cuboid(source, target_shape): """Central cuboid of an N-D Fourier transform.