waveorder handles all rotational quantities (including retardance) …

…in radians (#149)

* return retardance in radians and document

* document fluorescence

* fluorescence docs improvments

* isotroic_thin_3d documentation

* phase 3d documentation

* minor synchronization

* improved documentation of polarization reconstruction

examples/models/inplane_oriented_thick_pol3D.py

@@ -1,6 +1,5 @@
 import napari
-import numpy as np
-from waveorder import util
+
 from waveorder.models import inplane_oriented_thick_pol3d
 
 # Parameters
@@ -46,13 +45,13 @@
 arrays = [
     (inplane_oriented_parameters_recon[3], "Depolarization - recon"),
     (inplane_oriented_parameters_recon[2], "Transmittance - recon"),
-    (inplane_oriented_parameters_recon[1], "Orientation - recon"),
-    (inplane_oriented_parameters_recon[0], "Retardance - recon"),
+    (inplane_oriented_parameters_recon[1], "Orientation (rad) - recon"),
+    (inplane_oriented_parameters_recon[0], "Retardance (rad) - recon"),
     (czyx_data, "Data"),
     (inplane_oriented_parameters[3], "Depolarization"),
     (inplane_oriented_parameters[2], "Transmittance"),
-    (inplane_oriented_parameters[1], "Orientation"),
-    (inplane_oriented_parameters[0], "Retardance"),
+    (inplane_oriented_parameters[1], "Orientation (rad)"),
+    (inplane_oriented_parameters[0], "Retardance (rad)"),
 ]
 
 for array in arrays:

waveorder/models/inplane_oriented_thick_pol3d.py

@@ -1,5 +1,8 @@
+from typing import Optional, Tuple
+
 import numpy as np
 import torch
+from torch import Tensor
 
 from waveorder import background_estimator, stokes, util
 
@@ -49,38 +52,86 @@ def apply_transfer_function(
 
 
 def apply_inverse_transfer_function(
-    czyx_data,
-    intensity_to_stokes_matrix,
-    wavelength_illumination=0.5,  # TOOD: MOVE THIS PARAM TO OTF? (leaky param)
-    cyx_no_sample_data=None,  # if not None, use this data for background correction
-    project_stokes_to_2d=False,
-    remove_estimated_background=False,  # if True estimate background from czyx_data and remove it
-    flip_orientation=False,
-    rotate_orientation=False,
-):
+    czyx_data: Tensor,
+    intensity_to_stokes_matrix: Tensor,
+    cyx_no_sample_data: Optional[Tensor] = None,
+    remove_estimated_background: bool = False,
+    project_stokes_to_2d: bool = False,
+    flip_orientation: bool = False,
+    rotate_orientation: bool = False,
+) -> Tuple[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.
+
+    Parameters
+    ----------
+    czyx_data : Tensor
+        4D raw data, first dimension is the polarization dimension, remaining
+        dimensions are spatial
+    intensity_to_stokes_matrix : Tensor
+        Forward model, see calculate_transfer_function above
+    cyx_no_sample_data : Tensor, optional
+        3D raw background data, by default None
+        First dimension is the polarization dimension, remaining dimensions are spatial.
+        cyx shape must match in this parameter and czxy_data
+        If provided, this background will be removed.
+        If None, no background will be removed.
+    remove_estimated_background : bool, optional
+        Estimate a background from the data and remove it, by default False
+    project_stokes_to_2d : bool, optional
+        Project stokes to 2D for SNR improvement in thin samples, by default False
+    flip_orientation : bool, optional
+        Flip the reconstructed orientation about the x axis, by default False
+    rotate_orientation : bool, optional
+        Add 90 degrees to the reconstructed orientation, by default False
+
+    Notes
+    -----
+    cyx_no_sample_data and remove_estimated_background provide background correction options
+
+    flip_orientation and rotate_orientation modify the reconstructed orientation.
+    We recommend using these parameters when a test target with a known orientation
+    is available.
+
+    Returns
+    -------
+    Tuple[Tensor]
+        zyx_retardance (radians)
+        zyx_orientation (radians)
+        zyx_transmittance (unitless)
+        zyx_depolarization (unitless)
+    """
     data_stokes = stokes.mmul(intensity_to_stokes_matrix, czyx_data)
 
-    # "Measured" background correction
+    # Apply a "Measured" background correction
     if cyx_no_sample_data is None:
         background_corrected_stokes = data_stokes
     else:
+        # Find the no-sample Stokes parameters from the background data
         measured_no_sample_stokes = stokes.mmul(
             intensity_to_stokes_matrix, cyx_no_sample_data
         )
+        # Estimate the attenuating, depolarizing, retarder's inverse Mueller
+        # matrix that caused this background data
         inverse_background_mueller = stokes.mueller_from_stokes(
             *measured_no_sample_stokes, model="adr", direction="inverse"
         )
+        # Apply this background-correction Mueller matrix to the data to remove
+        # the background contribution
         background_corrected_stokes = stokes.mmul(
             inverse_background_mueller, data_stokes
         )
 
-    # "Estimated" background correction
+    # 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(
                 background_corrected_stokes[stokes_index], dim=0
             )
+            # Estimate the background and subtract
             background_corrected_stokes[
                 stokes_index
             ] -= estimator.get_background(
@@ -94,16 +145,17 @@ def apply_inverse_transfer_function(
             background_corrected_stokes, dim=1
         )[:, None, ...]
 
+    # Estimate an attenuating, depolarizing, retarder's parameters,
+    # i.e. (retardance, orientation, transmittance, depolarization)
+    # from the background-corrected Stokes values
     adr_parameters = stokes.estimate_adr_from_stokes(
         *background_corrected_stokes
     )
 
-    # Return retardance in distance units (matching wavelength_illumination)
-    retardance = adr_parameters[0] * wavelength_illumination / (2 * np.pi)
-
     # Apply orientation transformations
     orientation = stokes.apply_orientation_offset(
         adr_parameters[1], rotate=rotate_orientation, flip=flip_orientation
     )
 
-    return retardance, orientation, adr_parameters[2], adr_parameters[3]
+    # Return (retardance, orientation, transmittance, depolarization)
+    return adr_parameters[0], orientation, adr_parameters[2], adr_parameters[3]

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -1,4 +1,7 @@
+from typing import Literal
+
 import torch
+from torch import Tensor
 
 from waveorder import optics, util
 
@@ -33,7 +36,7 @@ def calculate_transfer_function(
     z_position_list = torch.fft.ifftshift(
         (torch.arange(z_total) - z_total // 2) * z_pixel_size
     )
-    
+
     det_pupil = optics.generate_pupil(
         radial_frequencies,
         numerical_aperture_detection,
@@ -101,14 +104,49 @@ def apply_transfer_function(zyx_object, optical_transfer_function, z_padding):
 
 
 def apply_inverse_transfer_function(
-    zyx_data,
-    optical_transfer_function,
-    z_padding,
-    reconstruction_algorithm="Tikhonov",
-    regularization_strength=1e-3,
-    TV_rho_strength=1e-3,
-    TV_iterations=10,
+    zyx_data: Tensor,
+    optical_transfer_function: Tensor,
+    z_padding: int,
+    reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov",
+    regularization_strength: float = 1e-3,
+    TV_rho_strength: float = 1e-3,
+    TV_iterations: int = 10,
 ):
+    """Reconstructs fluorescence density from zyx_data and
+    an optical_transfer_function, providing options for z padding and
+    reconstruction algorithms.
+
+    Parameters
+    ----------
+    zyx_data : Tensor
+        3D raw data, fluorescence defocus stack
+    optical_transfer_function : Tensor
+        3D optical transfer function, see calculate_transfer_function above
+    z_padding : int
+        Padding for axial dimension. Use zero for defocus stacks that
+        extend ~3 PSF widths beyond the sample. Pad by ~3 PSF widths otherwise.
+    reconstruction_algorithm : str, optional
+        "Tikhonov" or "TV", by default "Tikhonov"
+        "TV" is not implemented.
+    regularization_strength : float, optional
+        regularization parameter, by default 1e-3
+    TV_rho_strength : _type_, optional
+        TV-specific regularization parameter, by default 1e-3
+        "TV" is not implemented.
+    TV_iterations : int, optional
+        TV-specific number of iterations, by default 10
+        "TV" is not implemented.
+
+    Returns
+    -------
+    Tensor
+        zyx_fluorescence_density (fluorophores per volumes)
+
+    Raises
+    ------
+    NotImplementedError
+        TV is not implemented
+    """
     # Handle padding
     zyx_padded = util.pad_zyx_along_z(zyx_data, z_padding)
 

waveorder/models/isotropic_thin_3d.py

@@ -1,5 +1,7 @@
-import numpy as np
+from typing import Literal, Tuple
+
 import torch
+from torch import Tensor
 
 from waveorder import optics, util
 
@@ -144,16 +146,55 @@ def apply_transfer_function(
 
 
 def apply_inverse_transfer_function(
-    zyx_data,
-    absorption_2d_to_3d_transfer_function,
-    phase_2d_to_3d_transfer_function,
-    reconstruction_algorithm="Tikhonov",
-    regularization_strength=1e-6,
-    reg_p=1e-6,  # TODO: use this parameter
-    TV_rho_strength=1e-3,
-    TV_iterations=10,
-    bg_filter=True,
-):
+    zyx_data: Tensor,
+    absorption_2d_to_3d_transfer_function: Tensor,
+    phase_2d_to_3d_transfer_function: Tensor,
+    reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov",
+    regularization_strength: float = 1e-6,
+    reg_p: float = 1e-6,  # TODO: use this parameter
+    TV_rho_strength: float = 1e-3,
+    TV_iterations: int = 10,
+    bg_filter: bool = True,
+) -> Tuple[Tensor]:
+    """Reconstructs absorption and phase from zyx_data and a pair of
+    3D-to-2D transfer functions named absorption_2d_to_3d_transfer_function and
+    phase_2d_to_3d_transfer_function, providing options for reconstruction
+    algorithms.
+
+    Parameters
+    ----------
+    zyx_data : Tensor
+        3D raw data, label-free defocus stack
+    absorption_2d_to_3d_transfer_function : Tensor
+        3D-to-2D absorption transfer function, see calculate_transfer_function above
+    phase_2d_to_3d_transfer_function : Tensor
+        3D-to-2D phase transfer function, see calculate_transfer_function above
+    reconstruction_algorithm : Literal["Tikhonov", "TV"], optional
+        "Tikhonov" or "TV", by default "Tikhonov"
+        "TV" is not implemented.
+    regularization_strength : float, optional
+        regularization parameter, by default 1e-6
+    reg_p : float, optional
+        TV-specific phase regularization parameter, by default 1e-6
+        "TV" is not implemented.
+    TV_iterations : int, optional
+        TV-specific number of iterations, by default 10
+        "TV" is not implemented.
+    bg_filter : bool, optional
+        option for slow-varying 2D background normalization with uniform filter
+        by default True
+
+    Returns
+    -------
+    Tuple[Tensor]
+        yx_absorption (unitless)
+        yx_phase (radians)
+
+    Raises
+    ------
+    NotImplementedError
+        TV is not implemented
+    """
     zyx_data_normalized = util.inten_normalization(
         zyx_data, bg_filter=bg_filter
     )

waveorder/models/phase_thick_3d.py

@@ -1,5 +1,8 @@
+from typing import Literal
+
 import numpy as np
 import torch
+from torch import Tensor
 
 from waveorder import optics, util
 from waveorder.models import isotropic_fluorescent_thick_3d
@@ -127,18 +130,67 @@ def apply_transfer_function(
 
 
 def apply_inverse_transfer_function(
-    zyx_data,
-    real_potential_transfer_function,
-    imaginary_potential_transfer_function,
-    z_padding,
-    z_pixel_size,  # TODO: MOVE THIS PARAM TO OTF? (leaky param)
-    wavelength_illumination,  # TOOD: MOVE THIS PARAM TO OTF? (leaky param)
-    absorption_ratio=0.0,
-    reconstruction_algorithm="Tikhonov",
-    regularization_strength=1e-3,
-    TV_rho_strength=1e-3,
-    TV_iterations=10,
+    zyx_data: Tensor,
+    real_potential_transfer_function: Tensor,
+    imaginary_potential_transfer_function: Tensor,
+    z_padding: int,
+    z_pixel_size: float,  # TODO: MOVE THIS PARAM TO OTF? (leaky param)
+    wavelength_illumination: float,  # TOOD: MOVE THIS PARAM TO OTF? (leaky param)
+    absorption_ratio: float = 0.0,
+    reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov",
+    regularization_strength: float = 1e-3,
+    TV_rho_strength: float = 1e-3,
+    TV_iterations: int = 10,
 ):
+    """Reconstructs 3D phase from labelfree defocus zyx_data and a pair of
+    complex 3D transfer functions real_potential_transfer_function and
+    imag_potential_transfer_function, providing options for reconstruction
+    algorithms.
+
+    Parameters
+    ----------
+    zyx_data : Tensor
+        3D raw data, label-free defocus stack
+    real_potential_transfer_function : Tensor
+        Real potential transfer function, see calculate_transfer_function abov
+    imaginary_potential_transfer_function : Tensor
+        Imaginary potential transfer function, see calculate_transfer_function abov
+    z_padding : int
+        Padding for axial dimension. Use zero for defocus stacks that
+        extend ~3 PSF widths beyond the sample. Pad by ~3 PSF widths otherwise.
+    z_pixel_size : float
+        spacing between axial samples in sample space
+        units must be consistent with wavelength_illumination
+        TODO: move this leaky parameter to calculate_transfer_function
+    wavelength_illumination : float,
+        illumination wavelength
+        units must be consistent with z_pixel_size
+        TODO: move this leaky parameter to calculate_transfer_function
+    absorption_ratio : float, optional,
+        Absorption-to-phase ratio in the sample.
+        Use default 0 for purely phase objects.
+    reconstruction_algorithm : str, optional
+        "Tikhonov" or "TV", by default "Tikhonov"
+        "TV" is not implemented.
+    regularization_strength : float, optional
+        regularization parameter, by default 1e-3
+    TV_rho_strength : _type_, optional
+        TV-specific regularization parameter, by default 1e-3
+        "TV" is not implemented.
+    TV_iterations : int, optional
+        TV-specific number of iterations, by default 10
+        "TV" is not implemented.
+
+    Returns
+    -------
+    Tensor
+        zyx_phase (radians)
+
+    Raises
+    ------
+    NotImplementedError
+        TV is not implemented
+    """
     # Handle padding
     zyx_padded = util.pad_zyx_along_z(zyx_data, z_padding)