waveorder 2.1.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

waveorder/models/isotropic_thin_3d.py

@@ -1,19 +1,19 @@
 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 +34,73 @@ 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 +148,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 +172,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 +198,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 +239,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

waveorder/models/phase_thick_3d.py

@@ -4,19 +4,19 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, util
+from waveorder import optics, sampling, util
 from waveorder.models import isotropic_fluorescent_thick_3d
+from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
 
 
 def generate_test_phantom(
-    zyx_shape,
-    yx_pixel_size,
-    z_pixel_size,
-    wavelength_illumination,
-    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,
@@ -24,27 +24,78 @@ def generate_test_phantom(
         radius=sphere_radius,
         blur_size=2 * yx_pixel_size,
     )
-    zyx_phase = (
-        sphere
-        * (index_of_refraction_sample - index_of_refraction_media)
-        * z_pixel_size
-        / wavelength_illumination
-    )  # phase in radians
+    zyx_phase = sphere * (
+        index_of_refraction_sample - index_of_refraction_media
+    )  # refractive index increment
 
     return zyx_phase
 
 
 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
     )
@@ -76,6 +127,7 @@ def calculate_transfer_function(
         det_pupil,
         wavelength_illumination / index_of_refraction_media,
         z_position_list,
+        axially_even=False,
     )
 
     (
@@ -95,37 +147,39 @@ 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
-):
+    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 isotropic_fluorescent_thick_3d.apply_transfer_function(
-        zyx_object, real_potential_transfer_function, z_padding
+    return (
+        isotropic_fluorescent_thick_3d.apply_transfer_function(
+            zyx_object,
+            real_potential_transfer_function,
+            z_padding,
+            background=0,
+        )
+        * brightness
+        + brightness
     )
 
 
@@ -134,14 +188,12 @@ def apply_inverse_transfer_function(
     real_potential_transfer_function: Tensor,
     imaginary_potential_transfer_function: Tensor,
     z_padding: int,
-    z_pixel_size: float,  # TODO: MOVE THIS PARAM TO OTF? (leaky param)
-    wavelength_illumination: float,  # TOOD: MOVE THIS PARAM TO OTF? (leaky param)
     absorption_ratio: float = 0.0,
     reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov",
     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
@@ -158,14 +210,6 @@ 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.
-    z_pixel_size : float
-        spacing between axial samples in sample space
-        units must be consistent with wavelength_illumination
-        TODO: move this leaky parameter to calculate_transfer_function
-    wavelength_illumination : float,
-        illumination wavelength
-        units must be consistent with z_pixel_size
-        TODO: move this leaky parameter to calculate_transfer_function
     absorption_ratio : float, optional,
         Absorption-to-phase ratio in the sample.
         Use default 0 for purely phase objects.
@@ -223,4 +267,4 @@ def apply_inverse_transfer_function(
     if z_padding != 0:
         f_real = f_real[z_padding:-z_padding]
 
-    return f_real * z_pixel_size / 4 / np.pi * wavelength_illumination
+    return f_real