diff --git a/examples/models/inplane_oriented_thick_pol3D.py b/examples/models/inplane_oriented_thick_pol3D.py index cea1f791..a1235c22 100644 --- a/examples/models/inplane_oriented_thick_pol3D.py +++ b/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: diff --git a/waveorder/models/inplane_oriented_thick_pol3d.py b/waveorder/models/inplane_oriented_thick_pol3d.py index 8aceef2f..5d3ed46e 100644 --- a/waveorder/models/inplane_oriented_thick_pol3d.py +++ b/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] diff --git a/waveorder/models/isotropic_fluorescent_thick_3d.py b/waveorder/models/isotropic_fluorescent_thick_3d.py index b0ef5c37..58234852 100644 --- a/waveorder/models/isotropic_fluorescent_thick_3d.py +++ b/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) diff --git a/waveorder/models/isotropic_thin_3d.py b/waveorder/models/isotropic_thin_3d.py index f74fba80..b581e13f 100644 --- a/waveorder/models/isotropic_thin_3d.py +++ b/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 ) diff --git a/waveorder/models/phase_thick_3d.py b/waveorder/models/phase_thick_3d.py index 77883a3d..39555b50 100644 --- a/waveorder/models/phase_thick_3d.py +++ b/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)