waveorder 2.2.0__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.0'
-__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: Tuple[int, int]) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+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)
@@ -23,7 +25,9 @@ def calculate_transfer_function(
     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: Tensor
+) -> None:
     viewer.add_image(
         intensity_to_stokes_matrix.cpu().numpy(),
         name="Intensity to stokes matrix",

waveorder/models/inplane_oriented_thick_pol3d_vector.py

@@ -1,10 +1,12 @@
-import torch
-import numpy as np
+from typing import Literal
 
+import numpy as np
+import torch
 from torch import Tensor
-from typing import Literal
-from torch.nn.functional import avg_pool3d, interpolate
+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
 
 
@@ -39,7 +41,6 @@ def calculate_transfer_function(
     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]
 ]:
@@ -65,38 +66,25 @@ def calculate_transfer_function(
 
     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
-            )
-        ),
+        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,
-        )
+    (
+        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
@@ -123,25 +111,12 @@ def calculate_transfer_function(
     )
 
     # 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)
+    singular_system = calculate_singular_system(cropped)
 
     return (
-        full_cropped,
+        cropped,
         intensity_to_stokes_matrix,
-        (full_U, full_S, full_Vh),
+        singular_system,
     )
 
 
@@ -332,20 +307,14 @@ def apply_inverse_transfer_function(
     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(
+    sfzyx_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
-    )
+    fzyx_recon = apply_filter_bank(sfzyx_inverse_filter, szyx_data)
 
-    return torch.real(torch.fft.ifftn(fZYX_reconstructed, dim=(1, 2, 3)))
+    return fzyx_recon

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -5,6 +5,8 @@ import torch
 from torch import Tensor
 
 from waveorder import optics, sampling, util
+from waveorder.filter import apply_filter_bank
+from waveorder.reconstruct import tikhonov_regularized_inverse_filter
 from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
 
 
@@ -30,7 +32,6 @@ def calculate_transfer_function(
     index_of_refraction_media: float,
     numerical_aperture_detection: float,
 ) -> Tensor:
-
     transverse_nyquist = sampling.transverse_nyquist(
         wavelength_emission,
         numerical_aperture_detection,  # ill = det for fluorescence
@@ -108,7 +109,11 @@ 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:
+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),
@@ -118,7 +123,10 @@ def visualize_transfer_function(viewer, optical_transfer_function: Tensor, zyx_s
 
 
 def apply_transfer_function(
-    zyx_object: Tensor, optical_transfer_function: Tensor, z_padding: int, background: int = 10
+    zyx_object: Tensor,
+    optical_transfer_function: Tensor,
+    z_padding: int,
+    background: int = 10,
 ) -> Tensor:
     """Simulate imaging by applying a transfer function
 
@@ -205,12 +213,15 @@ def apply_inverse_transfer_function(
 
     # Reconstruct
     if reconstruction_algorithm == "Tikhonov":
-        f_real = util.single_variable_tikhonov_deconvolution_3D(
-            zyx_padded,
-            optical_transfer_function,
-            reg_re=regularization_strength,
+        inverse_filter = tikhonov_regularized_inverse_filter(
+            optical_transfer_function, regularization_strength
         )
 
+        # [None]s and [0] are for applying a 1x1 "bank" of filters.
+        # For further uniformity, consider returning (1, Z, Y, X)
+        f_real = apply_filter_bank(
+            inverse_filter[None, None], zyx_padded[None]
+        )[0]
     elif reconstruction_algorithm == "TV":
         raise NotImplementedError
         f_real = util.single_variable_admm_tv_deconvolution_3D(

waveorder/models/isotropic_thin_3d.py

@@ -6,6 +6,7 @@ from torch import Tensor
 
 from waveorder import optics, sampling, util
 
+
 def generate_test_phantom(
     yx_shape: Tuple[int, int],
     yx_pixel_size: float,
@@ -50,20 +51,21 @@ def calculate_transfer_function(
     )
     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,
+        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(
@@ -151,9 +153,9 @@ def 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 
+    """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 
+    util.add_transfer_function_to_viewer. If more 3D-to-2D transfer functions
     are added, consider refactoring.
     """
     arrays = [
@@ -286,6 +288,7 @@ def apply_inverse_transfer_function(
     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
+    # TODO Reformulate to use filter.apply_filter_bank
     AHA = [
         torch.sum(torch.abs(absorption_2d_to_3d_transfer_function) ** 2, dim=0)
         + regularization_strength,

waveorder/models/phase_thick_3d.py

@@ -5,7 +5,9 @@ import torch
 from torch import Tensor
 
 from waveorder import optics, sampling, util
+from waveorder.filter import apply_filter_bank
 from waveorder.models import isotropic_fluorescent_thick_3d
+from waveorder.reconstruct import tikhonov_regularized_inverse_filter
 from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
 
 
@@ -56,22 +58,23 @@ def calculate_transfer_function(
     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,
-        )
+    (
+        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:]
@@ -167,7 +170,10 @@ def visualize_transfer_function(
 
 
 def apply_transfer_function(
-    zyx_object: np.ndarray, real_potential_transfer_function: np.ndarray, z_padding: int, brightness: float
+    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
@@ -249,10 +255,14 @@ def apply_inverse_transfer_function(
 
     # Reconstruct
     if reconstruction_algorithm == "Tikhonov":
-        f_real = util.single_variable_tikhonov_deconvolution_3D(
-            zyx, effective_transfer_function, reg_re=regularization_strength
+        inverse_filter = tikhonov_regularized_inverse_filter(
+            effective_transfer_function, regularization_strength
         )
 
+        # [None]s and [0] are for applying a 1x1 "bank" of filters.
+        # For further uniformity, consider returning (1, Z, Y, X)
+        f_real = apply_filter_bank(inverse_filter[None, None], zyx[None])[0]
+
     elif reconstruction_algorithm == "TV":
         raise NotImplementedError
         f_real = util.single_variable_admm_tv_deconvolution_3D(

waveorder/optics.py

@@ -1,9 +1,8 @@
+import itertools
+
 import numpy as np
 import torch
-import matplotlib.pyplot as plt
-import gc
-import itertools
-from numpy.fft import fft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
+from numpy.fft import fft2, fftn, fftshift, ifft2, ifftn, ifftshift
 
 
 def Jones_sample(Ein, t, sa):
@@ -272,7 +271,7 @@ def generate_vector_source_defocus_pupil(
     # TEMPORARY SIMPLIFY ROTATIONS "TURN OFF ROTATIONS"
     # 3x2 IDENTITY MATRIX
     rotations = torch.zeros_like(rotations)
-    rotations[1, 0, ...] = 1 
+    rotations[1, 0, ...] = 1
     rotations[2, 1, ...] = 1
 
     # Main calculation in the frequency domain
@@ -718,10 +717,10 @@ def gen_dyadic_Greens_tensor(G_real, ps, psz, lambda_in, space="real"):
 
 
 def generate_greens_tensor_spectrum(
-        zyx_shape, 
-        zyx_pixel_size,
-        wavelength,
-    ):
+    zyx_shape,
+    zyx_pixel_size,
+    wavelength,
+):
     """
     Parameters
     ----------
@@ -733,14 +732,12 @@ def generate_greens_tensor_spectrum(
     Returns
     -------
     torch.tensor
-        Green's tensor spectrum 
+        Green's tensor spectrum
     """
     Z, Y, X = zyx_shape
     dZ, dY, dX = zyx_pixel_size
 
-    z_step = torch.fft.ifftshift(
-        (torch.arange(Z) - Z // 2) * dZ
-    )
+    z_step = torch.fft.ifftshift((torch.arange(Z) - Z // 2) * dZ)
     y_step = torch.fft.ifftshift((torch.arange(Y) - Y // 2) * dY)
     x_step = torch.fft.ifftshift((torch.arange(X) - X // 2) * dX)
 
@@ -769,7 +766,7 @@ def generate_greens_tensor_spectrum(
     G_3D /= torch.amax(torch.abs(G_3D))
 
     return G_3D
-    
+
 
 def compute_weak_object_transfer_function_2d(
     illumination_pupil, detection_pupil

waveorder/reconstruct.py

@@ -0,0 +1,28 @@
+import torch
+
+
+def tikhonov_regularized_inverse_filter(
+    forward_filter: torch.Tensor, regularization_strength: float
+):
+    """Compute the Tikhonov regularized inverse filter from a forward filter.
+
+    Parameters
+    ----------
+    forward_filter : torch.Tensor
+        The forward filter tensor.
+    regularization_strength : float
+        The strength of the regularization term.
+    Returns
+    -------
+    torch.Tensor
+        The Tikhonov regularized inverse filter.
+    """
+
+    if forward_filter.ndim == 3:
+        forward_filter_conj = torch.conj(forward_filter)
+        return forward_filter_conj / (
+            (forward_filter_conj * forward_filter) + regularization_strength
+        )
+    else:
+        # TC TODO INTEGRATE THE 5D FILTER BANK CASE
+        raise NotImplementedError("Only 3D tensors are supported.")

waveorder/stokes.py

@@ -2,7 +2,7 @@
 Overview
 --------
 
-This module collects Stokes- and Mueller-related calculations. 
+This module collects Stokes- and Mueller-related calculations.
 
 The functions are roughly organized into groups:
 
@@ -29,8 +29,8 @@ y = mmul(A, x)
 Usage
 -----
 
-All functions are intended to be used with torch.Tensors with Stokes- or 
-Mueller-indices as the first axes. 
+All functions are intended to be used with torch.Tensors with Stokes- or
+Mueller-indices as the first axes.
 
 For example, the following usage modes of stokes_after_adr are valid:
 
@@ -46,6 +46,7 @@ For example, the following usage modes of stokes_after_adr are valid:
 >>> stokes_after_adr(*adr_params) # * expands along the first axis
 
 """
+
 import numpy as np
 import torch
 

waveorder/util.py

@@ -1,15 +1,14 @@
+import re
+import time
+from collections import namedtuple
+
 import numpy as np
-import matplotlib.pyplot as plt
 import pywt
-import time
 import torch
-
-from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
+from numpy.fft import fft, fft2, fftn, fftshift, ifft, ifftn, ifftshift
 from scipy.ndimage import uniform_filter
-from collections import namedtuple
-from .optics import scattering_potential_tensor_to_3D_orientation_PN
 
-import re
+from .optics import scattering_potential_tensor_to_3D_orientation_PN
 
 numbers = re.compile(r"(\d+)")
 
@@ -2289,5 +2288,6 @@ def gellmann():
             [[e, 0, 0], [0, d, 0], [0, 0, d]],
             [[0, 0, c], [0, 0, 0], [c, 0, 0]],
             [[0, 0, 0], [0, -c, 0], [0, 0, c]],  #
-        ], dtype=torch.complex64
-    )
+        ],
+        dtype=torch.complex64,
+    )

waveorder/visuals/jupyter_visuals.py

@@ -1,22 +1,16 @@
-import numpy as np
-import matplotlib.pyplot as plt
-import ipywidgets as widgets
-import os
 import io
+import os
+
+import ipywidgets as widgets
+import matplotlib.pyplot as plt
+import numpy as np
+from ipywidgets import HBox, Image, Layout, interact
+from matplotlib.colors import Normalize, hsv_to_rgb
+from numpy.typing import NDArray
 from PIL import Image as PImage
-from ipywidgets import (
-    Image,
-    Layout,
-    interact,
-    HBox,
-)
-from matplotlib.colors import hsv_to_rgb
-from matplotlib.colors import Normalize
 from scipy.ndimage import uniform_filter
 from scipy.stats import binned_statistic_2d
 
-from numpy.typing import NDArray
-
 
 def im_bit_convert(im, bit=16, norm=False, limit=[]):
     im = im.astype(
@@ -172,7 +166,7 @@ def image_stack_viewer_fast(
         else:
             raise ValueError('origin can only be either "upper" or "lower"')
 
-    im_wgt = Image(        
+    im_wgt = Image(
         value=im_dict[0],
         layout=Layout(height=str(size[0]) + "px", width=str(size[1]) + "px"),
     )
@@ -1046,15 +1040,17 @@ def plotVectorField(
     # plot vector field representaiton of the orientation map
 
     # Compute U, V such that they are as long as line-length when anisotropy = 1.
-    U, V = anisotropy * linelength * np.cos(
-        2 * orientation
-    ), anisotropy * linelength * np.sin(2 * orientation)
+    U, V = (
+        anisotropy * linelength * np.cos(2 * orientation),
+        anisotropy * linelength * np.sin(2 * orientation),
+    )
     USmooth = uniform_filter(U, (window, window))  # plot smoothed vector field
     VSmooth = uniform_filter(V, (window, window))  # plot smoothed vector field
     azimuthSmooth = 0.5 * np.arctan2(VSmooth, USmooth)
     RSmooth = np.sqrt(USmooth**2 + VSmooth**2)
-    USmooth, VSmooth = RSmooth * np.cos(azimuthSmooth), RSmooth * np.sin(
-        azimuthSmooth
+    USmooth, VSmooth = (
+        RSmooth * np.cos(azimuthSmooth),
+        RSmooth * np.sin(azimuthSmooth),
     )
 
     nY, nX = img.shape
@@ -1639,8 +1635,9 @@ def plot3DVectorField(
     VSmooth = uniform_filter(V, (window, window))  # plot smoothed vector field
     azimuthSmooth = 0.5 * np.arctan2(VSmooth, USmooth)
     RSmooth = np.sqrt(USmooth**2 + VSmooth**2)
-    USmooth, VSmooth = RSmooth * np.cos(azimuthSmooth), RSmooth * np.sin(
-        azimuthSmooth
+    USmooth, VSmooth = (
+        RSmooth * np.cos(azimuthSmooth),
+        RSmooth * np.sin(azimuthSmooth),
     )
 
     nY, nX = img.shape

waveorder/visuals/matplotlib_visuals.py

@@ -23,7 +23,7 @@ def plot_5d_ortho(
     Plot 5D multi-channel data in a grid or ortho-slice views.
 
     Input data is a 6D array with (row, column, channels, Z, Y, X) dimensions.
-    
+
     `color_funcs` permits different RGB color maps for each row and column.
 
     Parameters
@@ -32,7 +32,7 @@ def plot_5d_ortho(
         5D array with shape (R, C, Ch, Z, Y, X) containing the data to plot.
         [r]ows and [c]olumns form a grid
         [C]hannels contain multiple color channels
-        [ZYX] contain 3D volumes. 
+        [ZYX] contain 3D volumes.
     filename : str
         Path to save the output plot.
     voxel_size : tuple[float, float, float]
@@ -40,8 +40,8 @@ def plot_5d_ortho(
     zyx_slice : tuple[int, int, int]
         Indices of the ortho-slices to plot in (Z, Y, X) indices.
     color_funcs : list[list[callable]]
-        A list of lists of callables, one for each element of the plot grid, 
-        with len(color_funcs) == R and len(colors_funcs[0] == C). 
+        A list of lists of callables, one for each element of the plot grid,
+        with len(color_funcs) == R and len(colors_funcs[0] == C).
         Each callable accepts [C]hannel arguments and returns RGB color values,
         enabling different RGB color maps for each member of the grid.
     row_labels : list[str], optional

waveorder/visuals/napari_visuals.py

@@ -1,9 +1,8 @@
-from waveorder.visuals.utils import complex_tensor_to_rgb
-from typing import TYPE_CHECKING
-
 import numpy as np
 import torch
 
+from waveorder.visuals.utils import complex_tensor_to_rgb
+
 
 def add_transfer_function_to_viewer(
     viewer: "napari.Viewer",

waveorder/visuals/utils.py

@@ -1,11 +1,10 @@
-import numpy as np
 import matplotlib.colors as mcolors
+import numpy as np
 
 
 # Main function to convert a complex-valued torch tensor to RGB numpy array
 # with red at +1, green at +i, blue at -1, and purple at -i
 def complex_tensor_to_rgb(array, saturate_clim_fraction=1.0):
-
     # Calculate magnitude and phase for the entire array
     magnitude = np.abs(array)
     phase = np.angle(array)

waveorder/waveorder_reconstructor.py

@@ -1,15 +1,15 @@
-import numpy as np
-import matplotlib.pyplot as plt
 import itertools
 import time
-import os
 import warnings
-from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
+
+import matplotlib.pyplot as plt
+import numpy as np
 from IPython import display
-from scipy.ndimage import uniform_filter
-from .util import *
-from .optics import *
+from numpy.fft import fft2, fftn, fftshift, ifft, ifft2, ifftn, ifftshift
+
 from .background_estimator import *
+from .optics import *
+from .util import *
 
 
 def intensity_mapping(img_stack):
@@ -1474,8 +1474,9 @@ class waveorder_microscopy:
                 torch.tensor(z.astype("complex64").transpose((2, 1, 0))),
                 torch.tensor(self.psz),
             )
-            return H_re.numpy().transpose((1, 2, 0)), H_im.numpy().transpose(
-                (1, 2, 0)
+            return (
+                H_re.numpy().transpose((1, 2, 0)),
+                H_im.numpy().transpose((1, 2, 0)),
             )
 
         for i in range(self.N_pattern):

waveorder/waveorder_simulator.py

@@ -1,13 +1,13 @@
-import numpy as np
-import matplotlib.pyplot as plt
 import itertools
 import time
-import os
-import torch
-from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
 from concurrent.futures import ProcessPoolExecutor
-from .util import *
+
+import numpy as np
+import torch
+from numpy.fft import fft2, fftn, fftshift, ifft2, ifftn, ifftshift
+
 from .optics import *
+from .util import *
 
 
 def Jones_PC_forward_model(

{waveorder-2.2.0.dist-info → waveorder-2.2.1b0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: waveorder
-Version: 2.2.0
+Version: 2.2.1b0
 Summary: Wave-optical simulations and deconvolution of optical properties
 Author-email: CZ Biohub SF <compmicro@czbiohub.org>
 Maintainer-email: Talon Chandler <talon.chandler@czbiohub.org>, Shalin Mehta <shalin.mehta@czbiohub.org>
@@ -62,6 +62,7 @@ Requires-Dist: torch>=2.4.1
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: black==25.1.0; extra == "dev"
 Provides-Extra: examples
 Requires-Dist: napari[all]; extra == "examples"
 Requires-Dist: jupyter; extra == "examples"

waveorder-2.2.1b0.dist-info/RECORD

@@ -0,0 +1,27 @@
+waveorder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+waveorder/_version.py,sha256=GmZKo0FXy1i2hmm_Tns9LgX3RbtJXrUCizylCrNMoI4,513
+waveorder/background_estimator.py,sha256=GFy3N0qqp5M6JVZBbIUvTYhSin8Njg7zA2WN69pKLAE,12398
+waveorder/correction.py,sha256=uAWDKXq-FYwi1obxxWq0A-suNVf1cvqUnPsDC-LIlsM,3460
+waveorder/filter.py,sha256=BSNYTMYy-HrqtupcXMm9OzWPMlX_NF7KSNjeN0-0ghQ,7350
+waveorder/focus.py,sha256=ELI-bhtoodSFL7mKlXt0S8HqCc02gJMr7gCl3kYW5bE,6314
+waveorder/optics.py,sha256=RikR3kkPFoVUS86jiGQZey7Jtz5Ga-xKWnXsexo3Okc,44242
+waveorder/reconstruct.py,sha256=nF00-uxYF67QVcNClp1VIB31ElUfbjO3RNxbxc0ELH4,832
+waveorder/sampling.py,sha256=OAqlfjEemX6tQV2a2S8X0wQx9JCSykcvVxvKr1CY3H0,2521
+waveorder/stokes.py,sha256=G_hcqS-GbslVnCWmS13y77K1YZ8cvwIl-DDKyC84LSI,15184
+waveorder/util.py,sha256=K5wdR_da1MYSVHUSp0O-m_FzMKRX60_KXCv3IXtCgKA,71243
+waveorder/waveorder_reconstructor.py,sha256=-MluWmnZnCZm7Xu-8V8QWX9ma0_5oAZO2Xlyrg1bRes,152072
+waveorder/waveorder_simulator.py,sha256=uRRX_wcWzJzlVcfToLpIlh4e8Xt9NjTvdonyGEf2Z1c,45805
+waveorder/models/inplane_oriented_thick_pol3d.py,sha256=bx5yViNz7wY5BBXUea1Tw0RhYsEzErENRnAbgpa34C0,5992
+waveorder/models/inplane_oriented_thick_pol3d_vector.py,sha256=tAHkIV02CZ1DbjvikKON2bvkrRfR6jEiUenDhrTtdtI,9754
+waveorder/models/isotropic_fluorescent_thick_3d.py,sha256=mrai8u3aQ9TxI_TZZXFXZWei6g_ZcxkLxjlUDv1rHJA,7050
+waveorder/models/isotropic_thin_3d.py,sha256=ADn1py_-bkXKD9_Uw4av1xE-X-Y6Wyp3uSdaA9C7JPU,10956
+waveorder/models/phase_thick_3d.py,sha256=sOHphv6_SfarLxrhj4bcZwM-5bdjXS4-j6RCreOBTmA,8654
+waveorder/visuals/jupyter_visuals.py,sha256=6kxICjEtP1qc1EuETc_NJ6Y4A7nVaC-bP3wl_CQNPfg,58096
+waveorder/visuals/matplotlib_visuals.py,sha256=v1zi0ZlXEV5CcpNzTWL0vDJ2Md2-RSHnc3rAB61wimg,10915
+waveorder/visuals/napari_visuals.py,sha256=jgyKRlWqJv1PvCRWfpgqRUv6rwIsI4AXQGLL5IftPnM,2352
+waveorder/visuals/utils.py,sha256=QC5WSc2yzPMjk66IjA39iNFgO_0It6evma201hH8Lg4,1001
+waveorder-2.2.1b0.dist-info/LICENSE,sha256=auz4oGH1A-xZtoiR2zuXIk-Hii4v9aGgFVBqn7nfpms,1509
+waveorder-2.2.1b0.dist-info/METADATA,sha256=ZFf0LKaV4-RlN6AkT9K6_BzLqqG72i9wRPQDc9bscms,9724
+waveorder-2.2.1b0.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+waveorder-2.2.1b0.dist-info/top_level.txt,sha256=i3zReXiiMTnyPk93W7aEz_oEfsLnfR_Kzl7PW7kUslA,10
+waveorder-2.2.1b0.dist-info/RECORD,,

{waveorder-2.2.0.dist-info → waveorder-2.2.1b0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (76.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

waveorder-2.2.0.dist-info/RECORD

@@ -1,25 +0,0 @@
-waveorder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-waveorder/_version.py,sha256=pbL_Q6fDSZl5UbKP04ZFdzrJpd1PO1gH7IwJCwLV7mk,411
-waveorder/background_estimator.py,sha256=gCIO6-232H0CGH4o6gnqW9KSYGOrXf5E9nD67WeF304,12399
-waveorder/correction.py,sha256=N0Ic6mqw3U7mqow4dKTOkNx2QYOLwedGNH7HiKV-M6s,3460
-waveorder/focus.py,sha256=4mg84Fe4V-oFplsuaU_VQU1_TEDoEfPggIAv6Is2dE4,6312
-waveorder/optics.py,sha256=Z0N9IN5AJ563eOISlFVJLS5K7lyfSjOcHVkzcLSKP1E,44325
-waveorder/sampling.py,sha256=OAqlfjEemX6tQV2a2S8X0wQx9JCSykcvVxvKr1CY3H0,2521
-waveorder/stokes.py,sha256=Wk9ZimzICIZLh1CkB0kQSCSBLeugkDeydwXTPd-M-po,15186
-waveorder/util.py,sha256=goT2OD6Zej4v3RT1tu_XAQtDWi94HiDlL4pJk2itt6s,71272
-waveorder/waveorder_reconstructor.py,sha256=SSSru4TOLB1VUOWMLHzMSboMbazgfslFXrjOpnwmqFk,152107
-waveorder/waveorder_simulator.py,sha256=_HCmDZkACUGzgwnaI-q0PjsL1gRE55IQuaWw-wtAjCU,45856
-waveorder/models/inplane_oriented_thick_pol3d.py,sha256=jEpMcAZ6tIg9Kg-lHpbH0vBAphSGZGUAqLyDc0hv_bs,5979
-waveorder/models/inplane_oriented_thick_pol3d_vector.py,sha256=09-Qu6Ka3S2GmkcmGIyVbQCYP1I_S9O1HvDBZDQ7AlQ,10817
-waveorder/models/isotropic_fluorescent_thick_3d.py,sha256=TuE5QlScy6pdCZyrWJifJ6UJMQT15lv8HhkNa_cQfFs,6713
-waveorder/models/isotropic_thin_3d.py,sha256=c-1eKknNWKxDiMNUi0wBf6YeXhUQiUR-zguH8M7En_k,10941
-waveorder/models/phase_thick_3d.py,sha256=2gq5pd6BPxDmkqnf9bvbOfraD2CeGBr0GU1s9cYBTms,8374
-waveorder/visuals/jupyter_visuals.py,sha256=w-vlMtfyl3I1ACNfYIW4fbS9TIMAVittNj3GbjlRYz4,58121
-waveorder/visuals/matplotlib_visuals.py,sha256=e-4LrHPFU--j3gbUoZrO8WHpDIYNZDFu8vqBZyhyiG4,10922
-waveorder/visuals/napari_visuals.py,sha256=gI420pgOzun4Elx__txdk1eEBcTIBB6Gpln-6n8Wo1k,2385
-waveorder/visuals/utils.py,sha256=6vdZmpvFGHwSwxeV8vCKWQ0MBOrDokSIJhjdBtJLHeM,1002
-waveorder-2.2.0.dist-info/LICENSE,sha256=auz4oGH1A-xZtoiR2zuXIk-Hii4v9aGgFVBqn7nfpms,1509
-waveorder-2.2.0.dist-info/METADATA,sha256=D7tFPrufseWBEz4-kPpR6pejV_GtA91ll7SsO4u3qas,9677
-waveorder-2.2.0.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-waveorder-2.2.0.dist-info/top_level.txt,sha256=i3zReXiiMTnyPk93W7aEz_oEfsLnfR_Kzl7PW7kUslA,10
-waveorder-2.2.0.dist-info/RECORD,,