waveorder 2.0.0rc3__py3-none-any.whl → 2.2.0__py3-none-any.whl

waveorder/_version.py

@@ -1,4 +1,16 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-__version__ = version = '2.0.0rc3'
-__version_tuple__ = version_tuple = (2, 0, 0)
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '2.2.0'
+__version_tuple__ = version_tuple = (2, 2, 0)

waveorder/correction.py

@@ -0,0 +1,107 @@
+"""Background correction methods"""
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, Size
+
+
+def _sample_block_medians(image: Tensor, block_size) -> Tensor:
+    """
+    Sample densely tiled square blocks from a 2D image and return their medians.
+    Incomplete blocks (overhangs) will be ignored.
+
+    Parameters
+    ----------
+    image : Tensor
+        2D image
+    block_size : int, optional
+        Width and height of the blocks
+
+    Returns
+    -------
+    Tensor
+        Median intensity values for each block, flattened
+    """
+    if not image.dtype.is_floating_point:
+        image.to(torch.float)
+    blocks = F.unfold(image[None, None], block_size, stride=block_size)[0]
+    return blocks.median(0)[0]
+
+
+def _grid_coordinates(image: Tensor, block_size: int) -> Tensor:
+    """Build image coordinates from the center points of square blocks"""
+    coords = torch.meshgrid(
+        [
+            torch.arange(
+                0 + block_size / 2,
+                boundary - block_size / 2 + 1,
+                block_size,
+                device=image.device,
+            )
+            for boundary in image.shape
+        ]
+    )
+    return torch.stack(coords, dim=-1).reshape(-1, 2)
+
+
+def _fit_2d_polynomial_surface(
+    coords: Tensor, values: Tensor, order: int, surface_shape: Size
+) -> Tensor:
+    """Fit a 2D polynomial to a set of coordinates and their values,
+    and return the surface evaluated at every point."""
+    n_coeffs = int((order + 1) * (order + 2) / 2)
+    if n_coeffs >= len(values):
+        raise ValueError(
+            f"Cannot fit a {order} degree 2D polynomial "
+            f"with {len(values)} sampled values"
+        )
+    orders = torch.arange(order + 1, device=coords.device)
+    order_pairs = torch.stack(torch.meshgrid(orders, orders), -1)
+    order_pairs = order_pairs[order_pairs.sum(-1) <= order].reshape(-1, 2)
+    terms = torch.stack(
+        [coords[:, 0] ** i * coords[:, 1] ** j for i, j in order_pairs], -1
+    )
+    # use "gels" driver for precision and GPU consistency
+    coeffs = torch.linalg.lstsq(terms, values, driver="gels").solution
+    dense_coords = torch.meshgrid(
+        [
+            torch.arange(s, dtype=values.dtype, device=values.device)
+            for s in surface_shape
+        ]
+    )
+    dense_terms = torch.stack(
+        [dense_coords[0] ** i * dense_coords[1] ** j for i, j in order_pairs],
+        -1,
+    )
+    return torch.matmul(dense_terms, coeffs)
+
+
+def estimate_background(image: Tensor, order: int = 2, block_size: int = 32):
+    """
+    Combine sampling and polynomial surface fit for background estimation.
+    To background correct an image, divide it by the background.
+
+    Parameters
+    ----------
+    image : Tensor
+        2D image
+    order : int, optional
+        Order of polynomial, by default 2
+    block_size : int, optional
+        Width and height of the blocks, by default 32
+
+    Returns
+    -------
+    Tensor
+        Background image
+    """
+    if image.ndim != 2:
+        raise ValueError(f"Image must be 2D, got shape {image.shape}")
+    height, width = image.shape
+    if block_size > width:
+        raise ValueError("Block size larger than image height")
+    if block_size > height:
+        raise ValueError("Block size larger than image width")
+    medians = _sample_block_medians(image, block_size)
+    coords = _grid_coordinates(image, block_size)
+    return _fit_2d_polynomial_surface(coords, medians, order, image.shape)

waveorder/focus.py

@@ -3,6 +3,7 @@ from typing import Literal, Optional
 from waveorder import util
 import matplotlib.pyplot as plt
 import numpy as np
+import warnings
 
 
 def focus_from_transverse_band(
@@ -60,10 +61,19 @@ def focus_from_transverse_band(
     >>> slice = focus_from_transverse_band(zyx_array, NA_det=0.55, lambda_ill=0.532, pixel_size=6.5/20)
     >>> in_focus_data = data[slice,:,:]
     """
-    minmaxfunc = _check_focus_inputs(
-        zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions, mode
+    minmaxfunc = _mode_to_minmaxfunc(mode)
+
+    _check_focus_inputs(
+        zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions
     )
 
+    # Check for single slice
+    if zyx_array.shape[0] == 1:
+        warnings.warn(
+            "The dataset only contained a single slice. Returning trivial slice index = 0."
+        )
+        return 0
+
     # Calculate coordinates
     _, Y, X = zyx_array.shape
     _, _, fxx, fyy = util.gen_coordinate((Y, X), pixel_size)
@@ -94,25 +104,35 @@ def focus_from_transverse_band(
     # Plot
     if plot_path is not None:
         _plot_focus_metric(
-            plot_path, midband_sum, peak_index, in_focus_index, peak_results, threshold_FWHM
+            plot_path,
+            midband_sum,
+            peak_index,
+            in_focus_index,
+            peak_results,
+            threshold_FWHM,
         )
 
     return in_focus_index
 
 
+def _mode_to_minmaxfunc(mode):
+    if mode == "min":
+        minmaxfunc = np.argmin
+    elif mode == "max":
+        minmaxfunc = np.argmax
+    else:
+        raise ValueError("mode must be either `min` or `max`")
+    return minmaxfunc
+
+
 def _check_focus_inputs(
-    zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions, mode
+    zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions
 ):
     N = len(zyx_array.shape)
     if N != 3:
         raise ValueError(
             f"{N}D array supplied. `focus_from_transverse_band` only accepts 3D arrays."
         )
-    if zyx_array.shape[0] == 1:
-        print(
-            "WARNING: The dataset only contained a single slice. Returning trivial slice index = 0."
-        )
-        return 0
 
     if NA_det < 0:
         raise ValueError("NA must be > 0")
@@ -121,7 +141,7 @@ def _check_focus_inputs(
     if pixel_size < 0:
         raise ValueError("pixel_size must be > 0")
     if not 0.4 < lambda_ill / pixel_size < 10:
-        print(
+        warnings.warn(
             f"WARNING: lambda_ill/pixel_size = {lambda_ill/pixel_size}."
             f"Did you use the same units?"
             f"Did you enter the pixel size in (demagnified) object-space units?"
@@ -134,17 +154,15 @@ def _check_focus_inputs(
         raise ValueError("midband_fractions[0] must be between 0 and 1")
     if not (0 <= midband_fractions[1] <= 1):
         raise ValueError("midband_fractions[1] must be between 0 and 1")
-    if mode == "min":
-        minmaxfunc = np.argmin
-    elif mode == "max":
-        minmaxfunc = np.argmax
-    else:
-        raise ValueError("mode must be either `min` or `max`")
-    return minmaxfunc
 
 
 def _plot_focus_metric(
-    plot_path, midband_sum, peak_index, in_focus_index, peak_results, threshold_FWHM
+    plot_path,
+    midband_sum,
+    peak_index,
+    in_focus_index,
+    peak_results,
+    threshold_FWHM,
 ):
     _, ax = plt.subplots(1, 1, figsize=(4, 4))
     ax.plot(midband_sum, "-k")

waveorder/models/inplane_oriented_thick_pol3d.py

@@ -4,10 +4,10 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import background_estimator, stokes, util
+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 +17,13 @@ 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 +31,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 +59,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.
@@ -125,7 +125,6 @@ def apply_inverse_transfer_function(
 
     # Apply an "Estimated" background correction
     if remove_estimated_background:
-        estimator = background_estimator.BackgroundEstimator2D()
         for stokes_index in range(background_corrected_stokes.shape[0]):
             # Project to 2D
             z_projection = torch.mean(
@@ -134,9 +133,8 @@ def apply_inverse_transfer_function(
             # Estimate the background and subtract
             background_corrected_stokes[
                 stokes_index
-            ] -= estimator.get_background(
-                z_projection,
-                normalize=False,
+            ] -= correction.estimate_background(
+                z_projection, order=2, block_size=32
             )
 
     # Project to 2D (typically for SNR reasons)

waveorder/models/inplane_oriented_thick_pol3d_vector.py

@@ -0,0 +1,351 @@
+import torch
+import numpy as np
+
+from torch import Tensor
+from typing import Literal
+from torch.nn.functional import avg_pool3d, interpolate
+from waveorder import optics, sampling, stokes, util
+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,
+    transverse_downsample_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
+                / transverse_downsample_factor
+            )
+        ),
+        int(
+            np.ceil(
+                zyx_shape[2]
+                * yx_factor
+                * fourier_oversample_factor
+                / transverse_downsample_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
+    U, S, Vh = calculate_singular_system(cropped)
+
+    # Interpolate to final size in YX
+    def complex_interpolate(
+        tensor: torch.Tensor, zyx_shape: tuple[int, int, int]
+    ) -> torch.Tensor:
+        interpolated_real = interpolate(tensor.real, size=zyx_shape)
+        interpolated_imag = interpolate(tensor.imag, size=zyx_shape)
+        return interpolated_real + 1j * interpolated_imag
+
+    full_cropped = complex_interpolate(cropped, zyx_shape)
+    full_U = complex_interpolate(U, zyx_shape)
+    full_S = interpolate(S[None], size=zyx_shape)[0]  # S is real
+    full_Vh = complex_interpolate(Vh, zyx_shape)
+
+    return (
+        full_cropped,
+        intensity_to_stokes_matrix,
+        (full_U, full_S, full_Vh),
+    )
+
+
+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,
+):
+    sZYX_data = torch.fft.fftn(szyx_data, dim=(1, 2, 3))
+
+    # Key computation
+    print("Computing inverse filter")
+    U, S, Vh = singular_system
+    S_reg = S / (S**2 + regularization_strength)
+
+    ZYXsf_inverse_filter = torch.einsum(
+        "sjzyx,jzyx,jfzyx->sfzyx", U, S_reg, Vh
+    )
+
+    # Apply inverse filter
+    fZYX_reconstructed = torch.einsum(
+        "szyx,sfzyx->fzyx", sZYX_data, ZYXsf_inverse_filter
+    )
+
+    return torch.real(torch.fft.ifftn(fZYX_reconstructed, dim=(1, 2, 3)))