Consolidate inverse-filter apodization

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()

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

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)

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)

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

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))

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])

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)
 

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))

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

waveorder/reconstruct.py

@@ -1,20 +1,35 @@
 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
         The Tikhonov regularized inverse filter.
     """
 
     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)

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.