waveorder 2.1.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

waveorder/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.1.0'
-__version_tuple__ = version_tuple = (2, 1, 0)
+__version__ = version = '2.2.0'
+__version_tuple__ = version_tuple = (2, 2, 0)

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

@@ -7,7 +7,7 @@ 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 +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.

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)))

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -1,17 +1,19 @@
 from typing import Literal
 
+import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, util
+from waveorder import optics, sampling, util
+from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
 
 
 def generate_test_phantom(
-    zyx_shape,
-    yx_pixel_size,
-    z_pixel_size,
-    sphere_radius,
-):
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    sphere_radius: float,
+) -> Tensor:
     sphere, _, _ = util.generate_sphere_target(
         zyx_shape, yx_pixel_size, z_pixel_size, sphere_radius
     )
@@ -20,14 +22,57 @@ def generate_test_phantom(
 
 
 def calculate_transfer_function(
-    zyx_shape,
-    yx_pixel_size,
-    z_pixel_size,
-    wavelength_emission,
-    z_padding,
-    index_of_refraction_media,
-    numerical_aperture_detection,
-):
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    wavelength_emission: float,
+    z_padding: int,
+    index_of_refraction_media: float,
+    numerical_aperture_detection: float,
+) -> Tensor:
+
+    transverse_nyquist = sampling.transverse_nyquist(
+        wavelength_emission,
+        numerical_aperture_detection,  # ill = det for fluorescence
+        numerical_aperture_detection,
+    )
+    axial_nyquist = sampling.axial_nyquist(
+        wavelength_emission,
+        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))
+
+    optical_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_emission,
+        z_padding,
+        index_of_refraction_media,
+        numerical_aperture_detection,
+    )
+    zyx_out_shape = (zyx_shape[0] + 2 * z_padding,) + zyx_shape[1:]
+    return sampling.nd_fourier_central_cuboid(
+        optical_transfer_function, zyx_out_shape
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    zyx_shape: tuple[int, int, int],
+    yx_pixel_size: float,
+    z_pixel_size: float,
+    wavelength_emission: float,
+    z_padding: int,
+    index_of_refraction_media: float,
+    numerical_aperture_detection: float,
+) -> Tensor:
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size
     )
@@ -63,25 +108,33 @@ def calculate_transfer_function(
     return optical_transfer_function
 
 
-def visualize_transfer_function(viewer, optical_transfer_function, zyx_scale):
-    arrays = [
-        (torch.imag(optical_transfer_function), "Im(OTF)"),
-        (torch.real(optical_transfer_function), "Re(OTF)"),
-    ]
-
-    for array in arrays:
-        lim = 0.1 * torch.max(torch.abs(array[0]))
-        viewer.add_image(
-            torch.fft.ifftshift(array[0]).cpu().numpy(),
-            name=array[1],
-            colormap="bwr",
-            contrast_limits=(-lim, lim),
-            scale=1 / zyx_scale,
-        )
-    viewer.dims.order = (0, 1, 2)
+def 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),
+        zyx_scale,
+        clim_factor=0.05,
+    )
 
 
-def apply_transfer_function(zyx_object, optical_transfer_function, z_padding):
+def apply_transfer_function(
+    zyx_object: Tensor, optical_transfer_function: Tensor, z_padding: int, background: int = 10
+) -> Tensor:
+    """Simulate imaging by applying a transfer function
+
+    Parameters
+    ----------
+    zyx_object : torch.Tensor
+    optical_transfer_function : torch.Tensor
+    z_padding : int
+    background : int, optional
+        constant background counts added to each voxel, by default 10
+
+    Returns
+    -------
+    Simulated data : torch.Tensor
+
+    """
     if (
         zyx_object.shape[0] + 2 * z_padding
         != optical_transfer_function.shape[0]
@@ -99,7 +152,7 @@ def apply_transfer_function(zyx_object, optical_transfer_function, z_padding):
     zyx_data = zyx_obj_hat * optical_transfer_function
     data = torch.real(torch.fft.ifftn(zyx_data))
 
-    data += 10  # Add a direct background
+    data += background  # Add a direct background
     return data
 
 
@@ -111,7 +164,7 @@ def apply_inverse_transfer_function(
     regularization_strength: float = 1e-3,
     TV_rho_strength: float = 1e-3,
     TV_iterations: int = 10,
-):
+) -> Tensor:
     """Reconstructs fluorescence density from zyx_data and
     an optical_transfer_function, providing options for z padding and
     reconstruction algorithms.