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

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):
@@ -133,7 +132,7 @@ def generate_pupil(frr, NA, lamb_in):
                     numerical aperture of the pupil function (normalized by the refractive index of the immersion media)
 
         lamb_in : float
-                    wavelength of the light (inside the immersion media)
+                    wavelength of the light in free space
                     in units of length (inverse of frr's units)
 
     Returns
@@ -225,6 +224,103 @@ def gen_sector_Pupil(fxx, fyy, NA, lamb_in, sector_angle, rotation_angle):
     return Pupil_sector
 
 
+def rotation_matrix(nu_z, nu_y, nu_x, wavelength):
+    nu_perp_squared = nu_x**2 + nu_y**2
+    nu_zz = wavelength * nu_z - 1
+
+    R_xx = (wavelength * nu_x**2 * nu_z + nu_y**2) / nu_perp_squared
+    R_yy = (wavelength * nu_y**2 * nu_z + nu_x**2) / nu_perp_squared
+    R_xy = nu_x * nu_y * nu_zz / nu_perp_squared
+
+    row0 = torch.stack((-wavelength * nu_y, -wavelength * nu_x), dim=0)
+    row1 = torch.stack((R_yy, R_xy), dim=0)
+    row2 = torch.stack((R_xy, R_xx), dim=0)
+
+    out = torch.stack((row0, row1, row2), dim=0)
+
+    # KLUDGE: fix the DC term manually, avoiding nan
+    out[..., 0, 0] = torch.tensor([[0, 0], [1, 0], [0, 1]])[..., None]
+
+    return torch.nan_to_num(out, nan=0.0)
+
+
+def generate_vector_source_defocus_pupil(
+    x_frequencies,
+    y_frequencies,
+    z_position_list,
+    defocus_pupil,
+    input_jones,
+    ill_pupil,
+    wavelength,
+):
+    ill_pupil_3d = torch.einsum(
+        "zyx,yx->zyx", torch.fft.fft(defocus_pupil, dim=0), ill_pupil
+    ).abs()  # make this real
+
+    freq_shape = z_position_list.shape + x_frequencies.shape
+
+    y_broadcast = torch.broadcast_to(y_frequencies[None, :, :], freq_shape)
+    x_broadcast = torch.broadcast_to(x_frequencies[None, :, :], freq_shape)
+    z_broadcast = np.sqrt(wavelength ** (-2) - x_broadcast**2 - y_broadcast**2)
+
+    # Calculate rotation matrix
+    rotations = rotation_matrix(
+        z_broadcast, y_broadcast, x_broadcast, wavelength
+    ).type(torch.complex64)
+
+    # TEMPORARY SIMPLIFY ROTATIONS "TURN OFF ROTATIONS"
+    # 3x2 IDENTITY MATRIX
+    rotations = torch.zeros_like(rotations)
+    rotations[1, 0, ...] = 1
+    rotations[2, 1, ...] = 1
+
+    # Main calculation in the frequency domain
+    source_pupil = torch.einsum(
+        "ijzyx,j,zyx->izyx", rotations, input_jones, ill_pupil_3d
+    )
+
+    # Convert back to defocus pupil
+    source_defocus_pupil = torch.fft.ifft(source_pupil, dim=-3)
+
+    return source_defocus_pupil
+
+
+def generate_vector_detection_defocus_pupil(
+    x_frequencies,
+    y_frequencies,
+    z_position_list,
+    det_defocus_pupil,
+    det_pupil,
+    wavelength,
+):
+    # TODO: refactor redundancy with illumination pupil
+    det_pupil_3d = torch.einsum(
+        "zyx,yx->zyx", torch.fft.ifft(det_defocus_pupil, dim=0), det_pupil
+    )
+
+    # Calculate zyx_frequency grid (inelegant)
+    z_frequencies = torch.fft.ifft(z_position_list)
+    freq_shape = z_frequencies.shape + x_frequencies.shape
+    z_broadcast = torch.broadcast_to(z_frequencies[:, None, None], freq_shape)
+    y_broadcast = torch.broadcast_to(y_frequencies[None, :, :], freq_shape)
+    x_broadcast = torch.broadcast_to(x_frequencies[None, :, :], freq_shape)
+
+    # Calculate rotation matrix
+    rotations = rotation_matrix(
+        z_broadcast, y_broadcast, x_broadcast, wavelength
+    ).type(torch.complex64)
+
+    # Main calculation in the frequency domain
+    vector_detection_pupil = torch.einsum(
+        "jizyx,zyx->ijzyx", rotations, det_pupil_3d
+    )
+
+    # Convert back to defocus pupil
+    detection_defocus_pupil = torch.fft.fft(vector_detection_pupil, dim=-3)
+
+    return detection_defocus_pupil
+
+
 def Source_subsample(Source_cont, NAx_coord, NAy_coord, subsampled_NA=0.1):
     """
 
@@ -300,7 +396,7 @@ def generate_propagation_kernel(
         wavelength      : float
                         wavelength of the light in the immersion media
 
-        z_position_list : torch.tensor or list
+        z_position_list : torch.tensor
                         1D array of defocused z positions with the size of (Z)
 
     Returns
@@ -310,15 +406,16 @@ def generate_propagation_kernel(
 
     """
 
-    oblique_factor = (
-        (1 - wavelength**2 * radial_frequencies**2) * pupil_support
-    ) ** (1 / 2) / wavelength
+    oblique_factor = ((1 - wavelength**2 * radial_frequencies**2)) ** (
+        1 / 2
+    ) / wavelength
+    oblique_factor = torch.nan_to_num(oblique_factor, nan=0.0)
 
     propagation_kernel = pupil_support[None, :, :] * torch.exp(
         1j
         * 2
         * np.pi
-        * torch.tensor(z_position_list)[:, None, None]
+        * z_position_list[:, None, None]
         * oblique_factor[None, :, :]
     )
 
@@ -326,7 +423,11 @@ def generate_propagation_kernel(
 
 
 def generate_greens_function_z(
-    radial_frequencies, pupil_support, wavelength_illumination, z_position_list
+    radial_frequencies,
+    pupil_support,
+    wavelength_illumination,
+    z_position_list,
+    axially_even=True,
 ):
     """
 
@@ -343,9 +444,14 @@ def generate_greens_function_z(
         wavelength_illumination       : float
                         wavelength of the light in the immersion media
 
-        z_position_list      : torch.tensor or list
+        z_position_list : torch.tensor
                         1D array of defocused z position with the size of (Z,)
 
+        axially_even    : bool
+                        For backwards compatibility with legacy phase reconstruction.
+                        Ideally the legacy phase reconstruction should be unified with
+                        the new reconstructions, and this parameter should be removed.
+
     Returns
     -------
         greens_function_z       : torch.tensor
@@ -358,47 +464,97 @@ def generate_greens_function_z(
         * pupil_support
     ) ** (1 / 2) / wavelength_illumination
 
+    if axially_even:
+        z_positions = torch.abs(z_position_list[:, None, None])
+    else:
+        z_positions = z_position_list[:, None, None]
+
     greens_function_z = (
         -1j
         / 4
         / np.pi
         * pupil_support[None, :, :]
-        * torch.exp(
-            1j
-            * 2
-            * np.pi
-            * torch.tensor(z_position_list)[:, None, None]
-            * oblique_factor[None, :, :]
-        )
+        * torch.exp(1j * 2 * np.pi * z_positions * oblique_factor[None, :, :])
         / (oblique_factor[None, :, :] + 1e-15)
     )
 
     return greens_function_z
 
 
-def gen_dyadic_Greens_tensor_z(fxx, fyy, G_fun_z, Pupil_support, lambda_in):
+def generate_defocus_greens_tensor(
+    fxx, fyy, G_fun_z, Pupil_support, lambda_in
+):
     """
 
     generate forward dyadic Green's function in u_x, u_y, z space
 
     Parameters
     ----------
-        fxx           : numpy.ndarray
+        fxx           : tensor.Tensor
                         x component of 2D spatial frequency array with the size of (Ny, Nx)
 
-        fyy           : numpy.ndarray
+        fyy           : tensor.Tensor
                         y component of 2D spatial frequency array with the size of (Ny, Nx)
 
-        G_fun_z       : numpy.ndarray
-                        forward Green's function in u_x, u_y, z space with size of (Ny, Nx, Nz)
+        G_fun_z       : tensor.Tensor
+                        forward Green's function in u_x, u_y, z space with size of (Nz, Ny, Nx)
 
-        Pupil_support : numpy.ndarray
+        Pupil_support : tensor.Tensor
                         the array that defines the support of the pupil function with the size of (Ny, Nx)
 
         lambda_in     : float
                         wavelength of the light in the immersion media
 
     Returns
+    -------
+        G_tensor_z    : tensor.Tensor
+                        forward dyadic Green's function in u_x, u_y, z space with the size of (3, 3, Nz, Ny, Nx)
+    """
+
+    fr = (fxx**2 + fyy**2) ** (1 / 2)
+    oblique_factor = ((1 - lambda_in**2 * fr**2) * Pupil_support) ** (
+        1 / 2
+    ) / lambda_in
+
+    diff_filter = torch.zeros((3,) + G_fun_z.shape, dtype=torch.complex64)
+    diff_filter[0] = (1j * 2 * np.pi * oblique_factor)[None, ...]
+    diff_filter[1] = (1j * 2 * np.pi * fyy * Pupil_support)[None, ...]
+    diff_filter[2] = (1j * 2 * np.pi * fxx * Pupil_support)[None, ...]
+
+    G_tensor_z = torch.zeros((3, 3) + G_fun_z.shape, dtype=torch.complex64)
+
+    for i in range(3):
+        for j in range(3):
+            G_tensor_z[i, j] = (
+                G_fun_z
+                * diff_filter[i]
+                * diff_filter[j]
+                / (2 * np.pi / lambda_in) ** 2
+            )
+            if i == j:
+                G_tensor_z[i, i] += G_fun_z
+
+    return G_tensor_z
+
+
+def gen_dyadic_Greens_tensor_z(fxx, fyy, G_fun_z, Pupil_support, lambda_in):
+    """
+    keeping for backwards compatibility
+
+    generate forward dyadic Green's function in u_x, u_y, z space
+    Parameters
+    ----------
+        fxx           : numpy.ndarray
+                        x component of 2D spatial frequency array with the size of (Ny, Nx)
+        fyy           : numpy.ndarray
+                        y component of 2D spatial frequency array with the size of (Ny, Nx)
+        G_fun_z       : numpy.ndarray
+                        forward Green's function in u_x, u_y, z space with size of (Ny, Nx, Nz)
+        Pupil_support : numpy.ndarray
+                        the array that defines the support of the pupil function with the size of (Ny, Nx)
+        lambda_in     : float
+                        wavelength of the light in the immersion media
+    Returns
     -------
         G_tensor_z    : numpy.ndarray
                         forward dyadic Green's function in u_x, u_y, z space with the size of (3, 3, Ny, Nx, Nz)
@@ -427,7 +583,6 @@ def gen_dyadic_Greens_tensor_z(fxx, fyy, G_fun_z, Pupil_support, lambda_in):
             )
             if i == j:
                 G_tensor_z[i, i] += G_fun_z
-
     return G_tensor_z
 
 
@@ -561,6 +716,58 @@ def gen_dyadic_Greens_tensor(G_real, ps, psz, lambda_in, space="real"):
         )
 
 
+def generate_greens_tensor_spectrum(
+    zyx_shape,
+    zyx_pixel_size,
+    wavelength,
+):
+    """
+    Parameters
+    ----------
+    zyx_shape : tuple
+    zyx_pixel_size : tuple
+    wavelength : float
+        wavelength in medium
+
+    Returns
+    -------
+    torch.tensor
+        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)
+    y_step = torch.fft.ifftshift((torch.arange(Y) - Y // 2) * dY)
+    x_step = torch.fft.ifftshift((torch.arange(X) - X // 2) * dX)
+
+    zz = torch.broadcast_to(z_step[:, None, None], (Z, Y, X))
+    yy = torch.broadcast_to(y_step[None, :, None], (Z, Y, X))
+    xx = torch.broadcast_to(x_step[None, None, :], (Z, Y, X))
+
+    rr = torch.sqrt(xx**2 + yy**2 + zz**2)
+    rhat = torch.stack([zz, yy, xx], dim=0) / rr
+
+    scalar_g = torch.exp(1j * 2 * torch.pi * rr / wavelength) / (
+        4 * torch.pi * rr
+    )
+
+    eye = torch.zeros((3, 3, Z, Y, X))
+    eye[0, 0] = 1
+    eye[1, 1] = 1
+    eye[2, 2] = 1
+
+    Q = eye - torch.einsum("izyx,jzyx->ijzyx", rhat, rhat)
+    g_3d = Q * scalar_g
+    g_3d = torch.nan_to_num(g_3d)
+
+    G_3D = torch.fft.fftn(g_3d, dim=(-3, -2, -1))
+    G_3D = torch.imag(G_3D) * 1j
+    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/sampling.py

@@ -0,0 +1,94 @@
+import numpy as np
+import torch
+
+
+def transverse_nyquist(
+    wavelength_emission,
+    numerical_aperture_illumination,
+    numerical_aperture_detection,
+):
+    """Transverse Nyquist sample spacing in `wavelength_emission` units.
+
+    For widefield label-free imaging, the transverse Nyquist sample spacing is
+    lambda / (2 * (NA_ill + NA_det)).
+
+    Perhaps surprisingly, the transverse Nyquist sample spacing for widefield
+    fluorescence is lambda / (4 * NA), which is equivalent to the above formula
+    when NA_ill = NA_det.
+
+    Parameters
+    ----------
+    wavelength_emission : float
+        Output units match these units
+    numerical_aperture_illumination : float
+        For widefield fluorescence, set to numerical_aperture_detection
+    numerical_aperture_detection : float
+
+    Returns
+    -------
+    float
+        Transverse Nyquist sample spacing
+
+    """
+    return wavelength_emission / (
+        2 * (numerical_aperture_detection + numerical_aperture_illumination)
+    )
+
+
+def axial_nyquist(
+    wavelength_emission,
+    numerical_aperture_detection,
+    index_of_refraction_media,
+):
+    """Axial Nyquist sample spacing in `wavelength_emission` units.
+
+    For widefield microscopes, the axial Nyquist cutoff frequency is:
+
+    (n/lambda) - sqrt( (n/lambda)^2 - (NA_det/lambda)^2 ),
+
+    and the axial Nyquist sample spacing is 1 / (2 * cutoff_frequency).
+
+    Perhaps surprisingly, the axial Nyquist sample spacing is independent of
+    the illumination numerical aperture.
+
+    Parameters
+    ----------
+    wavelength_emission : float
+        Output units match these units
+    numerical_aperture_detection : float
+    index_of_refraction_media: float
+
+    Returns
+    -------
+    float
+        Axial Nyquist sample spacing
+
+    """
+    n_on_lambda = index_of_refraction_media / wavelength_emission
+    cutoff_frequency = n_on_lambda - np.sqrt(
+        n_on_lambda**2
+        - (numerical_aperture_detection / wavelength_emission) ** 2
+    )
+    return 1 / (2 * cutoff_frequency)
+
+
+def nd_fourier_central_cuboid(source, target_shape):
+    """Central cuboid of an N-D Fourier transform.
+
+    Parameters
+    ----------
+    source : torch.Tensor
+        Source tensor
+    target_shape : tuple of int
+
+    Returns
+    -------
+    torch.Tensor
+        Center cuboid in Fourier space
+
+    """
+    center_slices = tuple(
+        slice((s - o) // 2, (s - o) // 2 + o)
+        for s, o in zip(source.shape, target_shape)
+    )
+    return torch.fft.ifftshift(torch.fft.fftshift(source)[center_slices])

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+)")
 
@@ -331,12 +330,15 @@ def gen_coordinate(img_dim, ps):
     return (xx, yy, fxx, fyy)
 
 
-def generate_radial_frequencies(img_dim, ps):
+def generate_frequencies(img_dim, ps):
     fy = torch.fft.fftfreq(img_dim[0], ps)
     fx = torch.fft.fftfreq(img_dim[1], ps)
-
     fyy, fxx = torch.meshgrid(fy, fx, indexing="ij")
+    return fyy, fxx
+
 
+def generate_radial_frequencies(img_dim, ps):
+    fyy, fxx = generate_frequencies(img_dim, ps)
     return torch.sqrt(fyy**2 + fxx**2)
 
 
@@ -2239,3 +2241,53 @@ def orientation_3D_continuity_map(
     retardance_pr_avg /= np.max(retardance_pr_avg)
 
     return retardance_pr_avg
+
+
+def pauli():
+    # yx order
+    # trace-orthogonal normalization
+    # torch.einsum("kij,lji->kl", pauli(), pauli()) == torch.eye(4)
+
+    # intensity, x-y, +45-(-45), LCP-RCP
+    # yx
+    # yx
+    a = 2**-0.5
+    sigma = torch.tensor(
+        [
+            [[a, 0], [0, a]],
+            [[-a, 0], [0, a]],
+            [[0, a], [a, 0]],
+            [[0, 1j * a], [-1j * a, 0]],
+        ]
+    )
+    return sigma
+
+
+def gellmann():
+    # zyx order
+    # trace-orthogonal normalization
+    # torch.einsum("kij,lji->kl", gellmann(), gellmann()) == torch.eye(9)
+    #
+    # lexicographical order of the Gell-Mann matrices
+    # 00, 1-1, 10, 11, 2-2, 2-1, 20, 21, 22
+    #
+    # zyx
+    # zyx
+    a = 3**-0.5
+    c = 2**-0.5
+    d = -(6**-0.5)
+    e = 2 * (6**-0.5)
+    return torch.tensor(
+        [
+            [[a, 0, 0], [0, a, 0], [0, 0, a]],
+            [[0, 0, -c], [0, 0, 0], [c, 0, 0]],
+            [[0, 0, 0], [0, 0, -c], [0, c, 0]],
+            [[0, -c, 0], [c, 0, 0], [0, 0, 0]],
+            [[0, 0, 0], [0, 0, c], [0, c, 0]],  #
+            [[0, c, 0], [c, 0, 0], [0, 0, 0]],
+            [[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,
+    )

waveorder/{visual.py → visuals/jupyter_visuals.py}

@@ -1,26 +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,
-    interactive,
-    fixed,
-    interact_manual,
-    HBox,
-    VBox,
-)
-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(
@@ -1050,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
@@ -1643,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
@@ -1928,4 +1921,4 @@ def orientation_3D_hist(
                 if colorbar:
                     fig.colorbar(img, ax=ax[row_idx, col_idx])
 
-    return fig, ax
+    return fig, ax