waveorder 2.2.0rc0__py3-none-any.whl → 2.2.1__py3-none-any.whl

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -1,17 +1,21 @@
 from typing import Literal
 
+import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, util
+from waveorder import optics, sampling, util
+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
 
 
 def generate_test_phantom(
-    zyx_shape,
-    yx_pixel_size,
-    z_pixel_size,
-    sphere_radius,
-):
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    sphere_radius: float,
+) -> Tensor:
     sphere, _, _ = util.generate_sphere_target(
         zyx_shape, yx_pixel_size, z_pixel_size, sphere_radius
     )
@@ -20,14 +24,56 @@ def generate_test_phantom(
 
 
 def calculate_transfer_function(
-    zyx_shape,
-    yx_pixel_size,
-    z_pixel_size,
-    wavelength_emission,
-    z_padding,
-    index_of_refraction_media,
-    numerical_aperture_detection,
-):
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    wavelength_emission: float,
+    z_padding: int,
+    index_of_refraction_media: float,
+    numerical_aperture_detection: float,
+) -> Tensor:
+    transverse_nyquist = sampling.transverse_nyquist(
+        wavelength_emission,
+        numerical_aperture_detection,  # ill = det for fluorescence
+        numerical_aperture_detection,
+    )
+    axial_nyquist = sampling.axial_nyquist(
+        wavelength_emission,
+        numerical_aperture_detection,
+        index_of_refraction_media,
+    )
+
+    yx_factor = int(np.ceil(yx_pixel_size / transverse_nyquist))
+    z_factor = int(np.ceil(z_pixel_size / axial_nyquist))
+
+    optical_transfer_function = _calculate_wrap_unsafe_transfer_function(
+        (
+            zyx_shape[0] * z_factor,
+            zyx_shape[1] * yx_factor,
+            zyx_shape[2] * yx_factor,
+        ),
+        yx_pixel_size / yx_factor,
+        z_pixel_size / z_factor,
+        wavelength_emission,
+        z_padding,
+        index_of_refraction_media,
+        numerical_aperture_detection,
+    )
+    zyx_out_shape = (zyx_shape[0] + 2 * z_padding,) + zyx_shape[1:]
+    return sampling.nd_fourier_central_cuboid(
+        optical_transfer_function, zyx_out_shape
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    wavelength_emission: float,
+    z_padding: int,
+    index_of_refraction_media: float,
+    numerical_aperture_detection: float,
+) -> Tensor:
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size
     )
@@ -63,27 +109,25 @@ def calculate_transfer_function(
     return optical_transfer_function
 
 
-def visualize_transfer_function(viewer, optical_transfer_function, zyx_scale):
-    arrays = [
-        (torch.imag(optical_transfer_function), "Im(OTF)"),
-        (torch.real(optical_transfer_function), "Re(OTF)"),
-    ]
-
-    for array in arrays:
-        lim = 0.1 * torch.max(torch.abs(array[0]))
-        viewer.add_image(
-            torch.fft.ifftshift(array[0]).cpu().numpy(),
-            name=array[1],
-            colormap="bwr",
-            contrast_limits=(-lim, lim),
-            scale=1 / zyx_scale,
-        )
-    viewer.dims.order = (0, 1, 2)
+def visualize_transfer_function(
+    viewer,
+    optical_transfer_function: Tensor,
+    zyx_scale: tuple[float, float, float],
+) -> None:
+    add_transfer_function_to_viewer(
+        viewer,
+        torch.real(optical_transfer_function),
+        zyx_scale,
+        clim_factor=0.05,
+    )
 
 
 def apply_transfer_function(
-    zyx_object, optical_transfer_function, z_padding, background=10
-):
+    zyx_object: Tensor,
+    optical_transfer_function: Tensor,
+    z_padding: int,
+    background: int = 10,
+) -> Tensor:
     """Simulate imaging by applying a transfer function
 
     Parameters
@@ -97,7 +141,7 @@ def apply_transfer_function(
     Returns
     -------
     Simulated data : torch.Tensor
-        
+
     """
     if (
         zyx_object.shape[0] + 2 * z_padding
@@ -128,7 +172,7 @@ def apply_inverse_transfer_function(
     regularization_strength: float = 1e-3,
     TV_rho_strength: float = 1e-3,
     TV_iterations: int = 10,
-):
+) -> Tensor:
     """Reconstructs fluorescence density from zyx_data and
     an optical_transfer_function, providing options for z padding and
     reconstruction algorithms.
@@ -169,12 +213,15 @@ def apply_inverse_transfer_function(
 
     # Reconstruct
     if reconstruction_algorithm == "Tikhonov":
-        f_real = util.single_variable_tikhonov_deconvolution_3D(
-            zyx_padded,
-            optical_transfer_function,
-            reg_re=regularization_strength,
+        inverse_filter = tikhonov_regularized_inverse_filter(
+            optical_transfer_function, regularization_strength
         )
 
+        # [None]s and [0] are for applying a 1x1 "bank" of filters.
+        # For further uniformity, consider returning (1, Z, Y, X)
+        f_real = apply_filter_bank(
+            inverse_filter[None, None], zyx_padded[None]
+        )[0]
     elif reconstruction_algorithm == "TV":
         raise NotImplementedError
         f_real = util.single_variable_admm_tv_deconvolution_3D(

waveorder/models/isotropic_thin_3d.py

@@ -1,19 +1,20 @@
 from typing import Literal, Tuple
 
+import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, util
+from waveorder import optics, sampling, util
 
 
 def generate_test_phantom(
-    yx_shape,
-    yx_pixel_size,
-    wavelength_illumination,
-    index_of_refraction_media,
-    index_of_refraction_sample,
-    sphere_radius,
-):
+    yx_shape: Tuple[int, int],
+    yx_pixel_size: float,
+    wavelength_illumination: float,
+    index_of_refraction_media: float,
+    index_of_refraction_sample: float,
+    sphere_radius: float,
+) -> Tuple[Tensor, Tensor]:
     sphere, _, _ = util.generate_sphere_target(
         (3,) + yx_shape,
         yx_pixel_size,
@@ -34,15 +35,74 @@ def generate_test_phantom(
 
 
 def calculate_transfer_function(
-    yx_shape,
-    yx_pixel_size,
-    z_position_list,
-    wavelength_illumination,
-    index_of_refraction_media,
-    numerical_aperture_illumination,
-    numerical_aperture_detection,
-    invert_phase_contrast=False,
-):
+    yx_shape: Tuple[int, int],
+    yx_pixel_size: float,
+    z_position_list: list,
+    wavelength_illumination: float,
+    index_of_refraction_media: float,
+    numerical_aperture_illumination: float,
+    numerical_aperture_detection: float,
+    invert_phase_contrast: bool = False,
+) -> Tuple[Tensor, Tensor]:
+    transverse_nyquist = sampling.transverse_nyquist(
+        wavelength_illumination,
+        numerical_aperture_illumination,
+        numerical_aperture_detection,
+    )
+    yx_factor = int(np.ceil(yx_pixel_size / transverse_nyquist))
+
+    (
+        absorption_2d_to_3d_transfer_function,
+        phase_2d_to_3d_transfer_function,
+    ) = _calculate_wrap_unsafe_transfer_function(
+        (
+            yx_shape[0] * yx_factor,
+            yx_shape[1] * yx_factor,
+        ),
+        yx_pixel_size / yx_factor,
+        z_position_list,
+        wavelength_illumination,
+        index_of_refraction_media,
+        numerical_aperture_illumination,
+        numerical_aperture_detection,
+        invert_phase_contrast=invert_phase_contrast,
+    )
+
+    absorption_2d_to_3d_transfer_function_out = torch.zeros(
+        (len(z_position_list),) + tuple(yx_shape), dtype=torch.complex64
+    )
+    phase_2d_to_3d_transfer_function_out = torch.zeros(
+        (len(z_position_list),) + tuple(yx_shape), dtype=torch.complex64
+    )
+
+    for z in range(len(z_position_list)):
+        absorption_2d_to_3d_transfer_function_out[z] = (
+            sampling.nd_fourier_central_cuboid(
+                absorption_2d_to_3d_transfer_function[z], yx_shape
+            )
+        )
+        phase_2d_to_3d_transfer_function_out[z] = (
+            sampling.nd_fourier_central_cuboid(
+                phase_2d_to_3d_transfer_function[z], yx_shape
+            )
+        )
+
+    return (
+        absorption_2d_to_3d_transfer_function_out,
+        phase_2d_to_3d_transfer_function_out,
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    yx_shape: Tuple[int, int],
+    yx_pixel_size: float,
+    z_position_list: list,
+    wavelength_illumination: float,
+    index_of_refraction_media: float,
+    numerical_aperture_illumination: float,
+    numerical_aperture_detection: float,
+    invert_phase_contrast: bool = False,
+) -> Tuple[Tensor, Tensor]:
     if invert_phase_contrast:
         z_position_list = torch.flip(torch.tensor(z_position_list), dims=(0,))
 
@@ -90,10 +150,14 @@ def calculate_transfer_function(
 
 def visualize_transfer_function(
     viewer,
-    absorption_2d_to_3d_transfer_function,
-    phase_2d_to_3d_transfer_function,
-):
-    # TODO: consider generalizing w/ phase_thick_3d.visualize_transfer_function
+    absorption_2d_to_3d_transfer_function: Tensor,
+    phase_2d_to_3d_transfer_function: Tensor,
+) -> None:
+    """Note: unlike other `visualize_transfer_function` calls, this transfer
+    function is a mixed 3D-to-2D transfer function, so it cannot reuse
+    util.add_transfer_function_to_viewer. If more 3D-to-2D transfer functions
+    are added, consider refactoring.
+    """
     arrays = [
         (torch.imag(absorption_2d_to_3d_transfer_function), "Im(absorb TF)"),
         (torch.real(absorption_2d_to_3d_transfer_function), "Re(absorb TF)"),
@@ -110,14 +174,14 @@ def visualize_transfer_function(
             contrast_limits=(-lim, lim),
             scale=(1, 1, 1),
         )
-    viewer.dims.order = (0, 1, 2)
+    viewer.dims.order = (2, 0, 1)
 
 
 def visualize_point_spread_function(
     viewer,
-    absorption_2d_to_3d_transfer_function,
-    phase_2d_to_3d_transfer_function,
-):
+    absorption_2d_to_3d_transfer_function: Tensor,
+    phase_2d_to_3d_transfer_function: Tensor,
+) -> None:
     arrays = [
         (torch.fft.ifftn(absorption_2d_to_3d_transfer_function), "absorb PSF"),
         (torch.fft.ifftn(phase_2d_to_3d_transfer_function), "phase PSF"),
@@ -136,11 +200,11 @@ def visualize_point_spread_function(
 
 
 def apply_transfer_function(
-    yx_absorption,
-    yx_phase,
-    phase_2d_to_3d_transfer_function,
-    absorption_2d_to_3d_transfer_function,
-):
+    yx_absorption: Tensor,
+    yx_phase: Tensor,
+    phase_2d_to_3d_transfer_function: Tensor,
+    absorption_2d_to_3d_transfer_function: Tensor,
+) -> Tensor:
     # Very simple simulation, consider adding noise and bkg knobs
 
     # simulate absorbing object
@@ -177,7 +241,7 @@ def apply_inverse_transfer_function(
     TV_rho_strength: float = 1e-3,
     TV_iterations: int = 10,
     bg_filter: bool = True,
-) -> Tuple[Tensor]:
+) -> Tuple[Tensor, Tensor]:
     """Reconstructs absorption and phase from zyx_data and a pair of
     3D-to-2D transfer functions named absorption_2d_to_3d_transfer_function and
     phase_2d_to_3d_transfer_function, providing options for reconstruction
@@ -224,6 +288,7 @@ def apply_inverse_transfer_function(
     zyx_data_hat = torch.fft.fft2(zyx_data_normalized, dim=(1, 2))
 
     # TODO AHA and b_vec calculations should be moved into tikhonov/tv calculations
+    # TODO Reformulate to use filter.apply_filter_bank
     AHA = [
         torch.sum(torch.abs(absorption_2d_to_3d_transfer_function) ** 2, dim=0)
         + regularization_strength,

waveorder/models/phase_thick_3d.py

@@ -4,18 +4,21 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, util
+from waveorder import optics, sampling, util
+from waveorder.filter import apply_filter_bank
 from waveorder.models import isotropic_fluorescent_thick_3d
+from waveorder.reconstruct import tikhonov_regularized_inverse_filter
+from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
 
 
 def generate_test_phantom(
-    zyx_shape,
-    yx_pixel_size,
-    z_pixel_size,
-    index_of_refraction_media,
-    index_of_refraction_sample,
-    sphere_radius,
-):
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    index_of_refraction_media: float,
+    index_of_refraction_sample: float,
+    sphere_radius: float,
+) -> np.ndarray:
     sphere, _, _ = util.generate_sphere_target(
         zyx_shape,
         yx_pixel_size,
@@ -31,16 +34,71 @@ def generate_test_phantom(
 
 
 def calculate_transfer_function(
-    zyx_shape,
-    yx_pixel_size,
-    z_pixel_size,
-    wavelength_illumination,
-    z_padding,
-    index_of_refraction_media,
-    numerical_aperture_illumination,
-    numerical_aperture_detection,
-    invert_phase_contrast=False,
-):
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    wavelength_illumination: float,
+    z_padding: int,
+    index_of_refraction_media: float,
+    numerical_aperture_illumination: float,
+    numerical_aperture_detection: float,
+    invert_phase_contrast: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
+    transverse_nyquist = sampling.transverse_nyquist(
+        wavelength_illumination,
+        numerical_aperture_illumination,
+        numerical_aperture_detection,
+    )
+    axial_nyquist = sampling.axial_nyquist(
+        wavelength_illumination,
+        numerical_aperture_detection,
+        index_of_refraction_media,
+    )
+
+    yx_factor = int(np.ceil(yx_pixel_size / transverse_nyquist))
+    z_factor = int(np.ceil(z_pixel_size / axial_nyquist))
+
+    (
+        real_potential_transfer_function,
+        imag_potential_transfer_function,
+    ) = _calculate_wrap_unsafe_transfer_function(
+        (
+            zyx_shape[0] * z_factor,
+            zyx_shape[1] * yx_factor,
+            zyx_shape[2] * yx_factor,
+        ),
+        yx_pixel_size / yx_factor,
+        z_pixel_size / z_factor,
+        wavelength_illumination,
+        z_padding,
+        index_of_refraction_media,
+        numerical_aperture_illumination,
+        numerical_aperture_detection,
+        invert_phase_contrast=invert_phase_contrast,
+    )
+
+    zyx_out_shape = (zyx_shape[0] + 2 * z_padding,) + zyx_shape[1:]
+    return (
+        sampling.nd_fourier_central_cuboid(
+            real_potential_transfer_function, zyx_out_shape
+        ),
+        sampling.nd_fourier_central_cuboid(
+            imag_potential_transfer_function, zyx_out_shape
+        ),
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    wavelength_illumination: float,
+    z_padding: int,
+    index_of_refraction_media: float,
+    numerical_aperture_illumination: float,
+    numerical_aperture_detection: float,
+    invert_phase_contrast: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size
     )
@@ -72,6 +130,7 @@ def calculate_transfer_function(
         det_pupil,
         wavelength_illumination / index_of_refraction_media,
         z_position_list,
+        axially_even=False,
     )
 
     (
@@ -91,33 +150,31 @@ def calculate_transfer_function(
 
 def visualize_transfer_function(
     viewer,
-    real_potential_transfer_function,
-    imag_potential_transfer_function,
-    zyx_scale,
-):
-    # TODO: consider generalizing w/ phase2Dto3D.visualize_TF
-    arrays = [
-        (torch.real(imag_potential_transfer_function), "Re(imag pot. TF)"),
-        (torch.imag(imag_potential_transfer_function), "Im(imag pot. TF)"),
-        (torch.real(real_potential_transfer_function), "Re(real pot. TF)"),
-        (torch.imag(real_potential_transfer_function), "Im(real pot. TF)"),
-    ]
-
-    for array in arrays:
-        lim = 0.5 * torch.max(torch.abs(array[0]))
-        viewer.add_image(
-            torch.fft.ifftshift(array[0]).cpu().numpy(),
-            name=array[1],
-            colormap="bwr",
-            contrast_limits=(-lim, lim),
-            scale=1 / zyx_scale,
-        )
-    viewer.dims.order = (0, 1, 2)
+    real_potential_transfer_function: np.ndarray,
+    imag_potential_transfer_function: np.ndarray,
+    zyx_scale: tuple[float, float, float],
+) -> None:
+    add_transfer_function_to_viewer(
+        viewer,
+        imag_potential_transfer_function,
+        zyx_scale,
+        layer_name="Imag pot. TF",
+    )
+
+    add_transfer_function_to_viewer(
+        viewer,
+        real_potential_transfer_function,
+        zyx_scale,
+        layer_name="Real pot. TF",
+    )
 
 
 def apply_transfer_function(
-    zyx_object, real_potential_transfer_function, z_padding, brightness
-):
+    zyx_object: np.ndarray,
+    real_potential_transfer_function: np.ndarray,
+    z_padding: int,
+    brightness: float,
+) -> np.ndarray:
     # This simplified forward model only handles phase, so it resuses the fluorescence forward model
     # TODO: extend to absorption
     return (
@@ -142,7 +199,7 @@ def apply_inverse_transfer_function(
     regularization_strength: float = 1e-3,
     TV_rho_strength: float = 1e-3,
     TV_iterations: int = 10,
-):
+) -> Tensor:
     """Reconstructs 3D phase from labelfree defocus zyx_data and a pair of
     complex 3D transfer functions real_potential_transfer_function and
     imag_potential_transfer_function, providing options for reconstruction
@@ -198,10 +255,14 @@ def apply_inverse_transfer_function(
 
     # Reconstruct
     if reconstruction_algorithm == "Tikhonov":
-        f_real = util.single_variable_tikhonov_deconvolution_3D(
-            zyx, effective_transfer_function, reg_re=regularization_strength
+        inverse_filter = tikhonov_regularized_inverse_filter(
+            effective_transfer_function, regularization_strength
         )
 
+        # [None]s and [0] are for applying a 1x1 "bank" of filters.
+        # For further uniformity, consider returning (1, Z, Y, X)
+        f_real = apply_filter_bank(inverse_filter[None, None], zyx[None])[0]
+
     elif reconstruction_algorithm == "TV":
         raise NotImplementedError
         f_real = util.single_variable_admm_tv_deconvolution_3D(