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

waveorder/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.2.0rc0'
-__version_tuple__ = version_tuple = (2, 2, 0)
+__version__ = version = '2.2.1b0'
+__version_tuple__ = version_tuple = (2, 2, 1)

waveorder/background_estimator.py

@@ -1,12 +1,12 @@
 """Estimate flat field images"""
 
-import numpy as np
 import itertools
 
+import numpy as np
 
 """
 
-This script is adopted from 
+This script is adopted from
 
 https://github.com/mehta-lab/reconstruct-order
 

waveorder/correction.py

@@ -2,7 +2,7 @@
 
 import torch
 import torch.nn.functional as F
-from torch import Tensor, Size
+from torch import Size, Tensor
 
 
 def _sample_block_medians(image: Tensor, block_size) -> Tensor:

waveorder/filter.py

@@ -0,0 +1,206 @@
+import itertools
+
+import torch
+
+
+def apply_filter_bank(
+    io_filter_bank: torch.Tensor,
+    i_input_array: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Applies a filter bank to an input array.
+
+    io_filter_bank.shape must be smaller or equal to i_input_array.shape in all
+    dimensions. When io_filter_bank is smaller, it is effectively "stretched"
+    to apply the filter.
+
+    io_filter_bank is in "wrapped" format, i.e., the zero frequency is the
+    zeroth element.
+
+    i_input_array and io_filter_bank must have inverse sample spacing, i.e.,
+    is input_array contains samples spaced by dx, then io_filter_bank must
+    have extent 1/dx. Note that there is no need for io_filter_bank to have
+    sample spacing 1/(n*dx) because io_filter_bank will be stretched.
+
+    Parameters
+    ----------
+    io_filter_bank : torch.Tensor
+        The filter bank to be applied in the frequency domain.
+        The spatial extent of io_filter_bank must be 1/dx, where dx is the
+        sample spacing of i_input_array.
+
+        Leading dimensions are the input and output dimensions.
+        io_filter_bank.shape[:2] == (num_input_channels, num_output_channels)
+
+        Trailing dimensions are spatial frequency dimensions.
+        io_filter_bank.shape[2:] == (Z', Y', X') or (Y', X')
+
+    i_input_array : torch.Tensor
+        The real-valued input array with sample spacing dx to be filtered.
+
+        Leading dimension is the input dimension, matching the filter bank.
+        i_input_array.shape[0] == i
+
+        Trailing dimensions are spatial dimensions.
+        i_input_array.shape[1:] == (Z, Y, X) or (Y, X)
+
+    Returns
+    -------
+    torch.Tensor
+        The filtered real-valued output array with shape
+        (num_output_channels, Z, Y, X) or (num_output_channels, Y, X).
+
+    """
+
+    # Ensure all dimensions of transfer_function are smaller than or equal to input_array
+    if any(
+        t > i
+        for t, i in zip(io_filter_bank.shape[2:], i_input_array.shape[1:])
+    ):
+        raise ValueError(
+            "All spatial dimensions of io_filter_bank must be <= i_input_array."
+        )
+
+    # Ensure the number of spatial dimensions match
+    if io_filter_bank.ndim - i_input_array.ndim != 1:
+        raise ValueError(
+            "io_filter_bank and i_input_array must have the same number of spatial dimensions."
+        )
+
+    # Ensure the input dimensions match
+    if io_filter_bank.shape[0] != i_input_array.shape[0]:
+        raise ValueError(
+            "io_filter_bank.shape[0] and i_input_array.shape[0] must be the same."
+        )
+
+    num_input_channels, num_output_channels = io_filter_bank.shape[:2]
+    spatial_dims = io_filter_bank.shape[2:]
+
+    # Pad input_array until each dimension is divisible by transfer_function
+    pad_sizes = [
+        (0, (t - (i % t)) % t)
+        for t, i in zip(
+            io_filter_bank.shape[2:][::-1], i_input_array.shape[1:][::-1]
+        )
+    ]
+    flat_pad_sizes = list(itertools.chain(*pad_sizes))
+    padded_input_array = torch.nn.functional.pad(i_input_array, flat_pad_sizes)
+
+    # Apply the transfer function in the frequency domain
+    fft_dims = [d for d in range(1, i_input_array.ndim)]
+    padded_input_spectrum = torch.fft.fftn(padded_input_array, dim=fft_dims)
+
+    # Matrix-vector multiplication over f
+    # If this is a bottleneck, consider extending `stretched_multiply` to
+    # a `stretched_matrix_multiply` that uses an call like
+    # torch.einsum('io..., i... -> o...', io_filter_bank, padded_input_spectrum)
+    #
+    # Further optimization is likely with a combination of
+    # torch.baddbmm, torch.pixel_shuffle, torch.pixel_unshuffle.
+    padded_output_spectrum = torch.zeros(
+        (num_output_channels,) + spatial_dims,
+        dtype=padded_input_spectrum.dtype,
+        device=padded_input_spectrum.device,
+    )
+    for input_channel_idx in range(num_input_channels):
+        for output_channel_idx in range(num_output_channels):
+            padded_output_spectrum[output_channel_idx] += stretched_multiply(
+                io_filter_bank[input_channel_idx, output_channel_idx],
+                padded_input_spectrum[input_channel_idx],
+            )
+
+    # Cast to real, ignoring imaginary part
+    padded_result = torch.real(
+        torch.fft.ifftn(padded_output_spectrum, dim=fft_dims)
+    )
+
+    # Remove padding and return
+    slices = tuple(slice(0, i) for i in i_input_array.shape)
+    return padded_result[slices]
+
+
+def stretched_multiply(
+    small_array: torch.Tensor, large_array: torch.Tensor
+) -> torch.Tensor:
+    """
+    Effectively "stretches" small_array onto large_array before multiplying.
+
+    Each dimension of large_array must be divisible by each dimension of small_array.
+
+    Instead of upsampling small_array, this function uses a "block element-wise"
+    multiplication by breaking the large_array into blocks before element-wise
+    multiplication with the small_array.
+
+    For example, a `stretched_multiply` of a 3x3 array by a 99x99 array will
+    divide the 99x99 array into 33x33 blocks
+    [[33x33, 33x33, 33x33],
+     [33x33, 33x33, 33x33],
+     [33x33, 33x33, 33x33]]
+     and multiply each block by the corresponding element in the 3x3 array.
+
+    Returns an array with the same shape as large_array.
+
+    Works for arbitrary dimensions.
+
+    Parameters
+    ----------
+    small_array : torch.Tensor
+        A smaller array whose elements will be "stretched" onto blocks in the large array.
+    large_array : torch.Tensor
+        A larger array that will be divided into blocks and multiplied by the small array.
+
+    Returns
+    -------
+    torch.Tensor
+        Resulting tensor with shape matching large_array.
+
+    Example
+    -------
+    small_array = torch.tensor([[1, 2],
+                                [3, 4]])
+
+    large_array = torch.tensor([[1,   2,  3,  4],
+                                [5,   6,  7,  8],
+                                [9,  10, 11, 12],
+                                [13, 14, 15, 16]])
+
+    stretched_multiply(small_array, large_array) returns
+
+    [[  1,  2,  6,  8],
+     [  5,  6, 14, 16],
+     [ 27, 30, 44, 48],
+     [ 39, 42, 60, 64]]
+    """
+
+    # Ensure each dimension of large_array is divisible by each dimension of small_array
+    if any(l % s != 0 for s, l in zip(small_array.shape, large_array.shape)):
+        raise ValueError(
+            "Each dimension of large_array must be divisible by each dimension of small_array"
+        )
+
+    # Ensure the number of dimensions match
+    if small_array.ndim != large_array.ndim:
+        raise ValueError(
+            "small_array and large_array must have the same number of dimensions"
+        )
+
+    # Get shapes
+    s_shape = small_array.shape
+    l_shape = large_array.shape
+
+    # Reshape both array into blocks
+    block_shape = tuple(p // s for p, s in zip(l_shape, s_shape))
+    new_large_shape = tuple(itertools.chain(*zip(s_shape, block_shape)))
+    new_small_shape = tuple(
+        itertools.chain(*zip(s_shape, small_array.ndim * (1,)))
+    )
+    reshaped_large_array = large_array.reshape(new_large_shape)
+    reshaped_small_array = small_array.reshape(new_small_shape)
+
+    # Multiply the reshaped arrays
+    reshaped_result = reshaped_large_array * reshaped_small_array
+
+    # Reshape the result back to the large array shape
+    result = reshaped_result.reshape(l_shape)
+
+    return result

waveorder/focus.py

@@ -1,9 +1,11 @@
-from scipy.signal import peak_widths
+import warnings
 from typing import Literal, Optional
-from waveorder import util
+
 import matplotlib.pyplot as plt
 import numpy as np
-import warnings
+from scipy.signal import peak_widths
+
+from waveorder import util
 
 
 def focus_from_transverse_band(

waveorder/models/inplane_oriented_thick_pol3d.py

@@ -7,7 +7,9 @@ from torch import Tensor
 from waveorder import correction, stokes, util
 
 
-def generate_test_phantom(yx_shape):
+def generate_test_phantom(
+    yx_shape: Tuple[int, int],
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
     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 +19,15 @@ def generate_test_phantom(yx_shape):
 
 
 def calculate_transfer_function(
-    swing,
-    scheme,
-):
+    swing: float,
+    scheme: str,
+) -> Tensor:
     return stokes.calculate_intensity_to_stokes_matrix(swing, scheme=scheme)
 
 
-def visualize_transfer_function(viewer, intensity_to_stokes_matrix):
+def visualize_transfer_function(
+    viewer, intensity_to_stokes_matrix: Tensor
+) -> None:
     viewer.add_image(
         intensity_to_stokes_matrix.cpu().numpy(),
         name="Intensity to stokes matrix",
@@ -31,12 +35,12 @@ def visualize_transfer_function(viewer, intensity_to_stokes_matrix):
 
 
 def apply_transfer_function(
-    retardance,
-    orientation,
-    transmittance,
-    depolarization,
-    intensity_to_stokes_matrix,
-):
+    retardance: Tensor,
+    orientation: Tensor,
+    transmittance: Tensor,
+    depolarization: Tensor,
+    intensity_to_stokes_matrix: Tensor,
+) -> Tensor:
     stokes_params = stokes.stokes_after_adr(
         retardance, orientation, transmittance, depolarization
     )
@@ -59,7 +63,7 @@ def apply_inverse_transfer_function(
     project_stokes_to_2d: bool = False,
     flip_orientation: bool = False,
     rotate_orientation: bool = False,
-) -> Tuple[Tensor]:
+) -> Tuple[Tensor, Tensor, Tensor, 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/inplane_oriented_thick_pol3d_vector.py

@@ -0,0 +1,320 @@
+from typing import Literal
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.nn.functional import avg_pool3d
+
+from waveorder import optics, sampling, stokes, util
+from waveorder.filter import apply_filter_bank
+from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
+
+
+def generate_test_phantom(zyx_shape: tuple[int, int, int]) -> torch.Tensor:
+    # Simulate
+    yx_star, yx_theta, _ = util.generate_star_target(
+        yx_shape=zyx_shape[1:],
+        blur_px=1,
+        margin=50,
+    )
+    c00 = yx_star
+    c2_2 = -torch.sin(2 * yx_theta) * yx_star  # torch.zeros_like(c00)
+    c22 = -torch.cos(2 * yx_theta) * yx_star  # torch.zeros_like(c00)  #
+
+    # Put in a center slices of a 3D object
+    center_slice_object = torch.stack((c00, c2_2, c22), dim=0)
+    object = torch.zeros((3,) + zyx_shape)
+    object[:, zyx_shape[0] // 2, ...] = center_slice_object
+    return object
+
+
+def calculate_transfer_function(
+    swing: float,
+    scheme: str,
+    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,
+    fourier_oversample_factor: int = 1,
+) -> tuple[
+    torch.Tensor, torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+]:
+    if z_padding != 0:
+        raise NotImplementedError("Padding not implemented for this model")
+
+    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))
+
+    print("YX factor:", yx_factor)
+    print("Z factor:", z_factor)
+
+    tf_calculation_shape = (
+        zyx_shape[0] * z_factor * fourier_oversample_factor,
+        int(np.ceil(zyx_shape[1] * yx_factor * fourier_oversample_factor)),
+        int(np.ceil(zyx_shape[2] * yx_factor * fourier_oversample_factor)),
+    )
+
+    (
+        sfZYX_transfer_function,
+        intensity_to_stokes_matrix,
+    ) = _calculate_wrap_unsafe_transfer_function(
+        swing,
+        scheme,
+        tf_calculation_shape,
+        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,
+    )
+
+    # avg_pool3d does not support complex numbers
+    pooled_sfZYX_transfer_function_real = avg_pool3d(
+        sfZYX_transfer_function.real, (fourier_oversample_factor,) * 3
+    )
+    pooled_sfZYX_transfer_function_imag = avg_pool3d(
+        sfZYX_transfer_function.imag, (fourier_oversample_factor,) * 3
+    )
+    pooled_sfZYX_transfer_function = (
+        pooled_sfZYX_transfer_function_real
+        + 1j * pooled_sfZYX_transfer_function_imag
+    )
+
+    # Crop to original size
+    sfzyx_out_shape = (
+        pooled_sfZYX_transfer_function.shape[0],
+        pooled_sfZYX_transfer_function.shape[1],
+        zyx_shape[0] + 2 * z_padding,
+    ) + zyx_shape[1:]
+
+    cropped = sampling.nd_fourier_central_cuboid(
+        pooled_sfZYX_transfer_function, sfzyx_out_shape
+    )
+
+    # Compute singular system on cropped and downsampled
+    singular_system = calculate_singular_system(cropped)
+
+    return (
+        cropped,
+        intensity_to_stokes_matrix,
+        singular_system,
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    swing,
+    scheme,
+    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,
+):
+    print("Computing transfer function")
+    intensity_to_stokes_matrix = stokes.calculate_intensity_to_stokes_matrix(
+        swing, scheme=scheme
+    )
+
+    input_jones = torch.tensor([0.0 - 1.0j, 1.0 + 0j])  # circular
+    # input_jones = torch.tensor([0 + 0j, 1 + 0j]) # linear
+
+    # Calculate frequencies
+    y_frequencies, x_frequencies = util.generate_frequencies(
+        zyx_shape[1:], yx_pixel_size
+    )
+    radial_frequencies = torch.sqrt(x_frequencies**2 + y_frequencies**2)
+
+    z_total = zyx_shape[0] + 2 * z_padding
+    z_position_list = torch.fft.ifftshift(
+        (torch.arange(z_total) - z_total // 2) * z_pixel_size
+    )
+    if (
+        not invert_phase_contrast
+    ):  # opposite sign of direct phase reconstruction
+        z_position_list = torch.flip(z_position_list, dims=(0,))
+    z_frequencies = torch.fft.fftfreq(z_total, d=z_pixel_size)
+
+    # 2D pupils
+    print("\tCalculating pupils...")
+    ill_pupil = optics.generate_pupil(
+        radial_frequencies,
+        numerical_aperture_illumination,
+        wavelength_illumination,
+    )
+    det_pupil = optics.generate_pupil(
+        radial_frequencies,
+        numerical_aperture_detection,
+        wavelength_illumination,
+    )
+    pupil = optics.generate_pupil(
+        radial_frequencies,
+        index_of_refraction_media,  # largest possible NA
+        wavelength_illumination,
+    )
+
+    # Defocus pupils
+    defocus_pupil = optics.generate_propagation_kernel(
+        radial_frequencies,
+        pupil,
+        wavelength_illumination / index_of_refraction_media,
+        z_position_list,
+    )
+
+    # Calculate vector defocus pupils
+    S = optics.generate_vector_source_defocus_pupil(
+        x_frequencies,
+        y_frequencies,
+        z_position_list,
+        defocus_pupil,
+        input_jones,
+        ill_pupil,
+        wavelength_illumination / index_of_refraction_media,
+    )
+
+    # Simplified scalar pupil
+    P = optics.generate_propagation_kernel(
+        radial_frequencies,
+        det_pupil,
+        wavelength_illumination / index_of_refraction_media,
+        z_position_list,
+    )
+
+    P_3D = torch.abs(torch.fft.ifft(P, dim=-3)).type(torch.complex64)
+    S_3D = torch.fft.ifft(S, dim=-3)
+
+    print("\tCalculating greens tensor spectrum...")
+    G_3D = optics.generate_greens_tensor_spectrum(
+        zyx_shape=(z_total, zyx_shape[1], zyx_shape[2]),
+        zyx_pixel_size=(z_pixel_size, yx_pixel_size, yx_pixel_size),
+        wavelength=wavelength_illumination / index_of_refraction_media,
+    )
+
+    # Main part
+    PG_3D = torch.einsum("zyx,ipzyx->ipzyx", P_3D, G_3D)
+    PS_3D = torch.einsum("zyx,jzyx,kzyx->jkzyx", P_3D, S_3D, torch.conj(S_3D))
+
+    del P_3D, G_3D, S_3D
+
+    print("\tComputing pg and ps...")
+    pg = torch.fft.fftn(PG_3D, dim=(-3, -2, -1))
+    ps = torch.fft.fftn(PS_3D, dim=(-3, -2, -1))
+
+    del PG_3D, PS_3D
+
+    print("\tComputing H1 and H2...")
+    H1 = torch.fft.ifftn(
+        torch.einsum("ipzyx,jkzyx->ijpkzyx", pg, torch.conj(ps)),
+        dim=(-3, -2, -1),
+    )
+
+    H2 = torch.fft.ifftn(
+        torch.einsum("ikzyx,jpzyx->ijpkzyx", ps, torch.conj(pg)),
+        dim=(-3, -2, -1),
+    )
+
+    H_re = H1[1:, 1:] + H2[1:, 1:]  # drop data-side z components
+    # H_im = 1j * (H1 - H2) # ignore absorptive terms
+
+    del H1, H2
+
+    H_re /= torch.amax(torch.abs(H_re))
+
+    s = util.pauli()[[0, 1, 2, 3]]  # select s0, s1, and s2
+    Y = util.gellmann()[[0, 4, 8]]
+    # select phase f00 and transverse linear isotropic terms 2-2, and f22
+
+    print("\tComputing final transfer function...")
+    sfZYX_transfer_function = torch.einsum(
+        "sik,ikpjzyx,lpj->slzyx", s, H_re, Y
+    )
+    return (
+        sfZYX_transfer_function,
+        intensity_to_stokes_matrix,
+    )
+
+
+def calculate_singular_system(sfZYX_transfer_function):
+    # Compute regularized inverse filter
+    print("Computing SVD")
+    ZYXsf_transfer_function = sfZYX_transfer_function.permute(2, 3, 4, 0, 1)
+    U, S, Vh = torch.linalg.svd(ZYXsf_transfer_function, full_matrices=False)
+    singular_system = (
+        U.permute(3, 4, 0, 1, 2),
+        S.permute(3, 0, 1, 2),
+        Vh.permute(3, 4, 0, 1, 2),
+    )
+    return singular_system
+
+
+def visualize_transfer_function(
+    viewer: "napari.Viewer",
+    sfZYX_transfer_function: torch.Tensor,
+    zyx_scale: tuple[float, float, float],
+) -> None:
+    add_transfer_function_to_viewer(
+        viewer,
+        sfZYX_transfer_function,
+        zyx_scale=zyx_scale,
+        layer_name="Transfer Function",
+        complex_rgb=True,
+        clim_factor=0.5,
+    )
+
+
+def apply_transfer_function(
+    fzyx_object: torch.Tensor,
+    sfZYX_transfer_function: torch.Tensor,
+    intensity_to_stokes_matrix: torch.Tensor,  # TODO use this to simulate intensities
+) -> torch.Tensor:
+    fZYX_object = torch.fft.fftn(fzyx_object, dim=(1, 2, 3))
+    sZYX_data = torch.einsum(
+        "fzyx,sfzyx->szyx", fZYX_object, sfZYX_transfer_function
+    )
+    szyx_data = torch.fft.ifftn(sZYX_data, dim=(1, 2, 3))
+
+    return 50 * szyx_data  # + 0.1 * torch.randn(szyx_data.shape)
+
+
+def apply_inverse_transfer_function(
+    szyx_data: Tensor,
+    singular_system: tuple[Tensor],
+    intensity_to_stokes_matrix: Tensor,
+    reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov",
+    regularization_strength: float = 1e-3,
+    TV_rho_strength: float = 1e-3,
+    TV_iterations: int = 10,
+):
+    # Key computation
+    print("Computing inverse filter")
+    U, S, Vh = singular_system
+    S_reg = S / (S**2 + regularization_strength)
+    sfzyx_inverse_filter = torch.einsum(
+        "sjzyx,jzyx,jfzyx->sfzyx", U, S_reg, Vh
+    )
+
+    fzyx_recon = apply_filter_bank(sfzyx_inverse_filter, szyx_data)
+
+    return fzyx_recon