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

waveorder/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.2.0'
+__version__ = version = '2.2.0rc0'
 __version_tuple__ = version_tuple = (2, 2, 0)

waveorder/models/inplane_oriented_thick_pol3d.py

@@ -7,7 +7,7 @@ from torch import Tensor
 from waveorder import correction, stokes, util
 
 
-def generate_test_phantom(yx_shape: Tuple[int, int]) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+def generate_test_phantom(yx_shape):
     star, theta, _ = util.generate_star_target(yx_shape, blur_px=0.1)
     retardance = 0.25 * star
     orientation = (theta % np.pi) * (star > 1e-3)
@@ -17,13 +17,13 @@ def generate_test_phantom(yx_shape: Tuple[int, int]) -> Tuple[Tensor, Tensor, Te
 
 
 def calculate_transfer_function(
-    swing: float,
-    scheme: str,
-) -> Tensor:
+    swing,
+    scheme,
+):
     return stokes.calculate_intensity_to_stokes_matrix(swing, scheme=scheme)
 
 
-def visualize_transfer_function(viewer, intensity_to_stokes_matrix: Tensor) -> None:
+def visualize_transfer_function(viewer, intensity_to_stokes_matrix):
     viewer.add_image(
         intensity_to_stokes_matrix.cpu().numpy(),
         name="Intensity to stokes matrix",
@@ -31,12 +31,12 @@ def visualize_transfer_function(viewer, intensity_to_stokes_matrix: Tensor) -> N
 
 
 def apply_transfer_function(
-    retardance: Tensor,
-    orientation: Tensor,
-    transmittance: Tensor,
-    depolarization: Tensor,
-    intensity_to_stokes_matrix: Tensor,
-) -> Tensor:
+    retardance,
+    orientation,
+    transmittance,
+    depolarization,
+    intensity_to_stokes_matrix,
+):
     stokes_params = stokes.stokes_after_adr(
         retardance, orientation, transmittance, depolarization
     )
@@ -59,7 +59,7 @@ def apply_inverse_transfer_function(
     project_stokes_to_2d: bool = False,
     flip_orientation: bool = False,
     rotate_orientation: bool = False,
-) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+) -> Tuple[Tensor]:
     """Reconstructs retardance, orientation, transmittance, and depolarization
     from czyx_data and an intensity_to_stokes_matrix, providing options for
     background correction, projection, and orientation transformations.

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -1,19 +1,17 @@
 from typing import Literal
 
-import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, sampling, util
-from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
+from waveorder import optics, util
 
 
 def generate_test_phantom(
-    zyx_shape: tuple[int, int, int],
-    yx_pixel_size: float,
-    z_pixel_size: float,
-    sphere_radius: float,
-) -> Tensor:
+    zyx_shape,
+    yx_pixel_size,
+    z_pixel_size,
+    sphere_radius,
+):
     sphere, _, _ = util.generate_sphere_target(
         zyx_shape, yx_pixel_size, z_pixel_size, sphere_radius
     )
@@ -22,57 +20,14 @@ def generate_test_phantom(
 
 
 def calculate_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:
-
-    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:
+    zyx_shape,
+    yx_pixel_size,
+    z_pixel_size,
+    wavelength_emission,
+    z_padding,
+    index_of_refraction_media,
+    numerical_aperture_detection,
+):
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size
     )
@@ -108,18 +63,27 @@ def _calculate_wrap_unsafe_transfer_function(
     return optical_transfer_function
 
 
-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 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 apply_transfer_function(
-    zyx_object: Tensor, optical_transfer_function: Tensor, z_padding: int, background: int = 10
-) -> Tensor:
+    zyx_object, optical_transfer_function, z_padding, background=10
+):
     """Simulate imaging by applying a transfer function
 
     Parameters
@@ -133,7 +97,7 @@ def apply_transfer_function(
     Returns
     -------
     Simulated data : torch.Tensor
-
+        
     """
     if (
         zyx_object.shape[0] + 2 * z_padding
@@ -164,7 +128,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.

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, sampling, util
+from waveorder import optics, util
+
 
 def generate_test_phantom(
-    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]:
+    yx_shape,
+    yx_pixel_size,
+    wavelength_illumination,
+    index_of_refraction_media,
+    index_of_refraction_sample,
+    sphere_radius,
+):
     sphere, _, _ = util.generate_sphere_target(
         (3,) + yx_shape,
         yx_pixel_size,
@@ -34,73 +34,15 @@ def generate_test_phantom(
 
 
 def calculate_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]:
-    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]:
+    yx_shape,
+    yx_pixel_size,
+    z_position_list,
+    wavelength_illumination,
+    index_of_refraction_media,
+    numerical_aperture_illumination,
+    numerical_aperture_detection,
+    invert_phase_contrast=False,
+):
     if invert_phase_contrast:
         z_position_list = torch.flip(torch.tensor(z_position_list), dims=(0,))
 
@@ -148,14 +90,10 @@ def _calculate_wrap_unsafe_transfer_function(
 
 def visualize_transfer_function(
     viewer,
-    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.
-    """
+    absorption_2d_to_3d_transfer_function,
+    phase_2d_to_3d_transfer_function,
+):
+    # TODO: consider generalizing w/ phase_thick_3d.visualize_transfer_function
     arrays = [
         (torch.imag(absorption_2d_to_3d_transfer_function), "Im(absorb TF)"),
         (torch.real(absorption_2d_to_3d_transfer_function), "Re(absorb TF)"),
@@ -172,14 +110,14 @@ def visualize_transfer_function(
             contrast_limits=(-lim, lim),
             scale=(1, 1, 1),
         )
-    viewer.dims.order = (2, 0, 1)
+    viewer.dims.order = (0, 1, 2)
 
 
 def visualize_point_spread_function(
     viewer,
-    absorption_2d_to_3d_transfer_function: Tensor,
-    phase_2d_to_3d_transfer_function: Tensor,
-) -> None:
+    absorption_2d_to_3d_transfer_function,
+    phase_2d_to_3d_transfer_function,
+):
     arrays = [
         (torch.fft.ifftn(absorption_2d_to_3d_transfer_function), "absorb PSF"),
         (torch.fft.ifftn(phase_2d_to_3d_transfer_function), "phase PSF"),
@@ -198,11 +136,11 @@ def visualize_point_spread_function(
 
 
 def apply_transfer_function(
-    yx_absorption: Tensor,
-    yx_phase: Tensor,
-    phase_2d_to_3d_transfer_function: Tensor,
-    absorption_2d_to_3d_transfer_function: Tensor,
-) -> Tensor:
+    yx_absorption,
+    yx_phase,
+    phase_2d_to_3d_transfer_function,
+    absorption_2d_to_3d_transfer_function,
+):
     # Very simple simulation, consider adding noise and bkg knobs
 
     # simulate absorbing object
@@ -239,7 +177,7 @@ def apply_inverse_transfer_function(
     TV_rho_strength: float = 1e-3,
     TV_iterations: int = 10,
     bg_filter: bool = True,
-) -> Tuple[Tensor, Tensor]:
+) -> Tuple[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,18 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, sampling, util
+from waveorder import optics, 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: 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:
+    zyx_shape,
+    yx_pixel_size,
+    z_pixel_size,
+    index_of_refraction_media,
+    index_of_refraction_sample,
+    sphere_radius,
+):
     sphere, _, _ = util.generate_sphere_target(
         zyx_shape,
         yx_pixel_size,
@@ -32,70 +31,16 @@ def generate_test_phantom(
 
 
 def calculate_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]:
-    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]:
+    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,
+):
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size
     )
@@ -127,7 +72,6 @@ def _calculate_wrap_unsafe_transfer_function(
         det_pupil,
         wavelength_illumination / index_of_refraction_media,
         z_position_list,
-        axially_even=False,
     )
 
     (
@@ -147,28 +91,33 @@ def _calculate_wrap_unsafe_transfer_function(
 
 def visualize_transfer_function(
     viewer,
-    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",
-    )
+    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)
 
 
 def apply_transfer_function(
-    zyx_object: np.ndarray, real_potential_transfer_function: np.ndarray, z_padding: int, brightness: float
-) -> np.ndarray:
+    zyx_object, real_potential_transfer_function, z_padding, brightness
+):
     # This simplified forward model only handles phase, so it resuses the fluorescence forward model
     # TODO: extend to absorption
     return (
@@ -193,7 +142,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