waveorder 0.2.2rc0__py3-none-any.whl

waveorder/models/inplane_oriented_thick_pol3d.py

@@ -0,0 +1,159 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from waveorder import correction, stokes, util
+
+
+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)
+    transmittance = 0.9 * torch.ones_like(retardance)
+    depolarization = 0.9 * torch.ones_like(retardance)
+    return retardance, orientation, transmittance, depolarization
+
+
+def calculate_transfer_function(
+    swing,
+    scheme,
+):
+    return stokes.calculate_intensity_to_stokes_matrix(swing, scheme=scheme)
+
+
+def visualize_transfer_function(viewer, intensity_to_stokes_matrix):
+    viewer.add_image(
+        intensity_to_stokes_matrix.cpu().numpy(),
+        name="Intensity to stokes matrix",
+    )
+
+
+def apply_transfer_function(
+    retardance,
+    orientation,
+    transmittance,
+    depolarization,
+    intensity_to_stokes_matrix,
+):
+    stokes_params = stokes.stokes_after_adr(
+        retardance, orientation, transmittance, depolarization
+    )
+    stokes_to_intensity_matrix = torch.linalg.pinv(intensity_to_stokes_matrix)
+
+    cyx_intensities = stokes.mmul(
+        stokes_to_intensity_matrix, torch.stack(stokes_params)
+    )
+
+    # Return in czyx shape
+    # TODO: make this simulation more realistic with defocussed data
+    return cyx_intensities[:, None, ...] + 0.1
+
+
+def apply_inverse_transfer_function(
+    czyx_data: Tensor,
+    intensity_to_stokes_matrix: Tensor,
+    cyx_no_sample_data: Optional[Tensor] = None,
+    remove_estimated_background: bool = False,
+    project_stokes_to_2d: bool = False,
+    flip_orientation: bool = False,
+    rotate_orientation: bool = False,
+) -> 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.
+
+    Parameters
+    ----------
+    czyx_data : Tensor
+        4D raw data, first dimension is the polarization dimension, remaining
+        dimensions are spatial
+    intensity_to_stokes_matrix : Tensor
+        Forward model, see calculate_transfer_function above
+    cyx_no_sample_data : Tensor, optional
+        3D raw background data, by default None
+        First dimension is the polarization dimension, remaining dimensions are spatial.
+        cyx shape must match in this parameter and czxy_data
+        If provided, this background will be removed.
+        If None, no background will be removed.
+    remove_estimated_background : bool, optional
+        Estimate a background from the data and remove it, by default False
+    project_stokes_to_2d : bool, optional
+        Project stokes to 2D for SNR improvement in thin samples, by default False
+    flip_orientation : bool, optional
+        Flip the reconstructed orientation about the x axis, by default False
+    rotate_orientation : bool, optional
+        Add 90 degrees to the reconstructed orientation, by default False
+
+    Notes
+    -----
+    cyx_no_sample_data and remove_estimated_background provide background correction options
+
+    flip_orientation and rotate_orientation modify the reconstructed orientation.
+    We recommend using these parameters when a test target with a known orientation
+    is available.
+
+    Returns
+    -------
+    Tuple[Tensor]
+        zyx_retardance (radians)
+        zyx_orientation (radians)
+        zyx_transmittance (unitless)
+        zyx_depolarization (unitless)
+    """
+    data_stokes = stokes.mmul(intensity_to_stokes_matrix, czyx_data)
+
+    # Apply a "Measured" background correction
+    if cyx_no_sample_data is None:
+        background_corrected_stokes = data_stokes
+    else:
+        # Find the no-sample Stokes parameters from the background data
+        measured_no_sample_stokes = stokes.mmul(
+            intensity_to_stokes_matrix, cyx_no_sample_data
+        )
+        # Estimate the attenuating, depolarizing, retarder's inverse Mueller
+        # matrix that caused this background data
+        inverse_background_mueller = stokes.mueller_from_stokes(
+            *measured_no_sample_stokes, model="adr", direction="inverse"
+        )
+        # Apply this background-correction Mueller matrix to the data to remove
+        # the background contribution
+        background_corrected_stokes = stokes.mmul(
+            inverse_background_mueller, data_stokes
+        )
+
+    # Apply an "Estimated" background correction
+    if remove_estimated_background:
+        for stokes_index in range(background_corrected_stokes.shape[0]):
+            # Project to 2D
+            z_projection = torch.mean(
+                background_corrected_stokes[stokes_index], dim=0
+            )
+            # Estimate the background and subtract
+            background_corrected_stokes[
+                stokes_index
+            ] -= correction.estimate_background(
+                z_projection, order=2, block_size=32
+            )
+
+    # Project to 2D (typically for SNR reasons)
+    if project_stokes_to_2d:
+        background_corrected_stokes = torch.mean(
+            background_corrected_stokes, dim=1
+        )[:, None, ...]
+
+    # Estimate an attenuating, depolarizing, retarder's parameters,
+    # i.e. (retardance, orientation, transmittance, depolarization)
+    # from the background-corrected Stokes values
+    adr_parameters = stokes.estimate_adr_from_stokes(
+        *background_corrected_stokes
+    )
+
+    # Apply orientation transformations
+    orientation = stokes.apply_orientation_offset(
+        adr_parameters[1], rotate=rotate_orientation, flip=flip_orientation
+    )
+
+    # Return (retardance, orientation, transmittance, depolarization)
+    return adr_parameters[0], orientation, adr_parameters[2], adr_parameters[3]

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -0,0 +1,192 @@
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from waveorder import optics, util
+
+
+def generate_test_phantom(
+    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
+    )
+
+    return sphere
+
+
+def calculate_transfer_function(
+    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
+    )
+
+    z_total = zyx_shape[0] + 2 * z_padding
+    z_position_list = torch.fft.ifftshift(
+        (torch.arange(z_total) - z_total // 2) * z_pixel_size
+    )
+
+    det_pupil = optics.generate_pupil(
+        radial_frequencies,
+        numerical_aperture_detection,
+        wavelength_emission,
+    )
+
+    propagation_kernel = optics.generate_propagation_kernel(
+        radial_frequencies,
+        det_pupil,
+        wavelength_emission / index_of_refraction_media,
+        z_position_list,
+    )
+
+    point_spread_function = (
+        torch.abs(torch.fft.ifft2(propagation_kernel, dim=(1, 2))) ** 2
+    )
+    optical_transfer_function = torch.fft.fftn(
+        point_spread_function, dim=(0, 1, 2)
+    )
+    optical_transfer_function /= torch.max(
+        torch.abs(optical_transfer_function)
+    )  # normalize
+
+    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 apply_transfer_function(
+    zyx_object, optical_transfer_function, z_padding, background=10
+):
+    """Simulate imaging by applying a transfer function
+
+    Parameters
+    ----------
+    zyx_object : torch.Tensor
+    optical_transfer_function : torch.Tensor
+    z_padding : int
+    background : int, optional
+        constant background counts added to each voxel, by default 10
+
+    Returns
+    -------
+    Simulated data : torch.Tensor
+        
+    """
+    if (
+        zyx_object.shape[0] + 2 * z_padding
+        != optical_transfer_function.shape[0]
+    ):
+        raise ValueError(
+            "Please check padding: ZYX_obj.shape[0] + 2 * Z_pad != H_re.shape[0]"
+        )
+    if z_padding > 0:
+        optical_transfer_function = optical_transfer_function[
+            z_padding:-z_padding
+        ]
+
+    # Very simple simulation, consider adding noise and bkg knobs
+    zyx_obj_hat = torch.fft.fftn(zyx_object)
+    zyx_data = zyx_obj_hat * optical_transfer_function
+    data = torch.real(torch.fft.ifftn(zyx_data))
+
+    data += background  # Add a direct background
+    return data
+
+
+def apply_inverse_transfer_function(
+    zyx_data: Tensor,
+    optical_transfer_function: Tensor,
+    z_padding: int,
+    reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov",
+    regularization_strength: float = 1e-3,
+    TV_rho_strength: float = 1e-3,
+    TV_iterations: int = 10,
+):
+    """Reconstructs fluorescence density from zyx_data and
+    an optical_transfer_function, providing options for z padding and
+    reconstruction algorithms.
+
+    Parameters
+    ----------
+    zyx_data : Tensor
+        3D raw data, fluorescence defocus stack
+    optical_transfer_function : Tensor
+        3D optical transfer function, see calculate_transfer_function above
+    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.
+    reconstruction_algorithm : str, optional
+        "Tikhonov" or "TV", by default "Tikhonov"
+        "TV" is not implemented.
+    regularization_strength : float, optional
+        regularization parameter, by default 1e-3
+    TV_rho_strength : _type_, optional
+        TV-specific regularization parameter, by default 1e-3
+        "TV" is not implemented.
+    TV_iterations : int, optional
+        TV-specific number of iterations, by default 10
+        "TV" is not implemented.
+
+    Returns
+    -------
+    Tensor
+        zyx_fluorescence_density (fluorophores per volumes)
+
+    Raises
+    ------
+    NotImplementedError
+        TV is not implemented
+    """
+    # Handle padding
+    zyx_padded = util.pad_zyx_along_z(zyx_data, z_padding)
+
+    # Reconstruct
+    if reconstruction_algorithm == "Tikhonov":
+        f_real = util.single_variable_tikhonov_deconvolution_3D(
+            zyx_padded,
+            optical_transfer_function,
+            reg_re=regularization_strength,
+        )
+
+    elif reconstruction_algorithm == "TV":
+        raise NotImplementedError
+        f_real = util.single_variable_admm_tv_deconvolution_3D(
+            zyx_padded,
+            optical_transfer_function,
+            reg_re=regularization_strength,
+            rho=TV_rho_strength,
+            itr=TV_iterations,
+        )
+
+    # Unpad
+    if z_padding != 0:
+        f_real = f_real[z_padding:-z_padding]
+
+    return f_real

waveorder/models/isotropic_thin_3d.py

@@ -0,0 +1,281 @@
+from typing import Literal, Tuple
+
+import torch
+from torch import Tensor
+
+from waveorder import optics, util
+
+
+def generate_test_phantom(
+    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,
+        z_pixel_size=1.0,
+        radius=sphere_radius,
+        blur_size=2 * yx_pixel_size,
+    )
+    yx_phase = (
+        sphere[1]
+        * (index_of_refraction_sample - index_of_refraction_media)
+        * 0.1
+        / wavelength_illumination
+    )  # phase in radians
+
+    yx_absorption = 0.02 * sphere[1]
+
+    return yx_absorption, yx_phase
+
+
+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,
+):
+    if invert_phase_contrast:
+        z_position_list = torch.flip(torch.tensor(z_position_list), dims=(0,))
+
+    radial_frequencies = util.generate_radial_frequencies(
+        yx_shape, yx_pixel_size
+    )
+
+    illumination_pupil = optics.generate_pupil(
+        radial_frequencies,
+        numerical_aperture_illumination,
+        wavelength_illumination,
+    )
+    detection_pupil = optics.generate_pupil(
+        radial_frequencies,
+        numerical_aperture_detection,
+        wavelength_illumination,
+    )
+    propagation_kernel = optics.generate_propagation_kernel(
+        radial_frequencies,
+        detection_pupil,
+        wavelength_illumination / index_of_refraction_media,
+        torch.tensor(z_position_list),
+    )
+
+    zyx_shape = (len(z_position_list),) + tuple(yx_shape)
+    absorption_2d_to_3d_transfer_function = torch.zeros(
+        zyx_shape, dtype=torch.complex64
+    )
+    phase_2d_to_3d_transfer_function = torch.zeros(
+        zyx_shape, dtype=torch.complex64
+    )
+    for z in range(len(z_position_list)):
+        (
+            absorption_2d_to_3d_transfer_function[z],
+            phase_2d_to_3d_transfer_function[z],
+        ) = optics.compute_weak_object_transfer_function_2d(
+            illumination_pupil, detection_pupil * propagation_kernel[z]
+        )
+
+    return (
+        absorption_2d_to_3d_transfer_function,
+        phase_2d_to_3d_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
+    arrays = [
+        (torch.imag(absorption_2d_to_3d_transfer_function), "Im(absorb TF)"),
+        (torch.real(absorption_2d_to_3d_transfer_function), "Re(absorb TF)"),
+        (torch.imag(phase_2d_to_3d_transfer_function), "Im(phase TF)"),
+        (torch.real(phase_2d_to_3d_transfer_function), "Re(phase TF)"),
+    ]
+
+    for array in arrays:
+        lim = 0.5 * torch.max(torch.abs(array[0]))
+        viewer.add_image(
+            torch.fft.ifftshift(array[0], dim=(1, 2)).cpu().numpy(),
+            name=array[1],
+            colormap="bwr",
+            contrast_limits=(-lim, lim),
+            scale=(1, 1, 1),
+        )
+    viewer.dims.order = (0, 1, 2)
+
+
+def visualize_point_spread_function(
+    viewer,
+    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"),
+    ]
+
+    for array in arrays:
+        lim = 0.5 * torch.max(torch.abs(array[0]))
+        viewer.add_image(
+            torch.fft.ifftshift(array[0], dim=(1, 2)).cpu().numpy(),
+            name=array[1],
+            colormap="bwr",
+            contrast_limits=(-lim, lim),
+            scale=(1, 1, 1),
+        )
+    viewer.dims.order = (0, 1, 2)
+
+
+def apply_transfer_function(
+    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
+    yx_absorption_hat = torch.fft.fftn(yx_absorption)
+    zyx_absorption_data_hat = yx_absorption_hat[None, ...] * torch.real(
+        absorption_2d_to_3d_transfer_function
+    )
+    zyx_absorption_data = torch.real(
+        torch.fft.ifftn(zyx_absorption_data_hat, dim=(1, 2))
+    )
+
+    # simulate phase object
+    yx_phase_hat = torch.fft.fftn(yx_phase)
+    zyx_phase_data_hat = yx_phase_hat[None, ...] * torch.real(
+        phase_2d_to_3d_transfer_function
+    )
+    zyx_phase_data = torch.real(
+        torch.fft.ifftn(zyx_phase_data_hat, dim=(1, 2))
+    )
+
+    # sum and add background
+    data = zyx_absorption_data + zyx_phase_data
+    data += 10  # Add a direct background
+    return data
+
+
+def apply_inverse_transfer_function(
+    zyx_data: Tensor,
+    absorption_2d_to_3d_transfer_function: Tensor,
+    phase_2d_to_3d_transfer_function: Tensor,
+    reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov",
+    regularization_strength: float = 1e-6,
+    reg_p: float = 1e-6,  # TODO: use this parameter
+    TV_rho_strength: float = 1e-3,
+    TV_iterations: int = 10,
+    bg_filter: bool = True,
+) -> 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
+    algorithms.
+
+    Parameters
+    ----------
+    zyx_data : Tensor
+        3D raw data, label-free defocus stack
+    absorption_2d_to_3d_transfer_function : Tensor
+        3D-to-2D absorption transfer function, see calculate_transfer_function above
+    phase_2d_to_3d_transfer_function : Tensor
+        3D-to-2D phase transfer function, see calculate_transfer_function above
+    reconstruction_algorithm : Literal["Tikhonov", "TV"], optional
+        "Tikhonov" or "TV", by default "Tikhonov"
+        "TV" is not implemented.
+    regularization_strength : float, optional
+        regularization parameter, by default 1e-6
+    reg_p : float, optional
+        TV-specific phase regularization parameter, by default 1e-6
+        "TV" is not implemented.
+    TV_iterations : int, optional
+        TV-specific number of iterations, by default 10
+        "TV" is not implemented.
+    bg_filter : bool, optional
+        option for slow-varying 2D background normalization with uniform filter
+        by default True
+
+    Returns
+    -------
+    Tuple[Tensor]
+        yx_absorption (unitless)
+        yx_phase (radians)
+
+    Raises
+    ------
+    NotImplementedError
+        TV is not implemented
+    """
+    zyx_data_normalized = util.inten_normalization(
+        zyx_data, bg_filter=bg_filter
+    )
+
+    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
+    AHA = [
+        torch.sum(torch.abs(absorption_2d_to_3d_transfer_function) ** 2, dim=0)
+        + regularization_strength,
+        torch.sum(
+            torch.conj(absorption_2d_to_3d_transfer_function)
+            * phase_2d_to_3d_transfer_function,
+            dim=0,
+        ),
+        torch.sum(
+            torch.conj(
+                phase_2d_to_3d_transfer_function,
+            )
+            * absorption_2d_to_3d_transfer_function,
+            dim=0,
+        ),
+        torch.sum(
+            torch.abs(
+                phase_2d_to_3d_transfer_function,
+            )
+            ** 2,
+            dim=0,
+        )
+        + reg_p,
+    ]
+
+    b_vec = [
+        torch.sum(
+            torch.conj(absorption_2d_to_3d_transfer_function) * zyx_data_hat,
+            dim=0,
+        ),
+        torch.sum(
+            torch.conj(
+                phase_2d_to_3d_transfer_function,
+            )
+            * zyx_data_hat,
+            dim=0,
+        ),
+    ]
+
+    # Deconvolution with Tikhonov regularization
+    if reconstruction_algorithm == "Tikhonov":
+        absorption, phase = util.dual_variable_tikhonov_deconvolution_2d(
+            AHA, b_vec
+        )
+
+    # ADMM deconvolution with anisotropic TV regularization
+    elif reconstruction_algorithm == "TV":
+        raise NotImplementedError
+        absorption, phase = util.dual_variable_admm_tv_deconv_2d(
+            AHA, b_vec, rho=TV_rho_strength, itr=TV_iterations
+        )
+
+    phase -= torch.mean(phase)
+
+    return absorption, phase