Match parameters to simplify recOrder-waveorder interface (#131)

* add axial_flip to `isotropic_thin_3d`

* `illumination` -> `emission` for fluorescence

* simplify parameters for usage with recOrder

* fix test

tests/models/test_isotropic_fluorescent_thick_3d.py

@@ -11,7 +11,7 @@ def test_calculate_transfer_function(axial_flip):
             zyx_shape=(20, 100, 101),
             yx_pixel_size=6.5 / 40,
             z_pixel_size=2,
-            wavelength_illumination=0.5,
+            wavelength_emission=0.5,
             z_padding=z_padding,
             index_of_refraction_media=1.0,
             numerical_aperture_detection=0.55,
@@ -34,8 +34,8 @@ def test_apply_inverse_transfer_function():
             zyx_data,
             optical_transfer_function,
             z_padding,
-            method="Tikhonov",
-            reg_re=1e-3,
+            reconstruction_algorithm="Tikhonov",
+            regularization_strength=1e-3,
         )
     )
     assert result_tikhonov.shape == (10, 5, 5)

tests/models/test_isotropic_thin_3d.py

@@ -1,15 +1,19 @@
+import pytest
+import torch
 from waveorder.models import isotropic_thin_3d
 
 
-def test_calculate_transfer_function():
+@pytest.mark.parametrize("axial_flip", (True, False))
+def test_calculate_transfer_function(axial_flip):
     Hu, Hp = isotropic_thin_3d.calculate_transfer_function(
         yx_shape=(100, 101),
         yx_pixel_size=6.5 / 40,
-        z_position_list=[-1, 0, 1],
+        z_position_list=torch.tensor([-1, 0, 1]),
         wavelength_illumination=0.5,
         index_of_refraction_media=1.0,
         numerical_aperture_illumination=0.4,
         numerical_aperture_detection=0.55,
+        axial_flip=axial_flip,
     )
 
     assert Hu.shape == (3, 100, 101)

waveorder/models/inplane_oriented_thick_pol3d.py

@@ -50,7 +50,7 @@ def apply_transfer_function(
 def apply_inverse_transfer_function(
     czyx_data,
     intensity_to_stokes_matrix,
-    wavelength_illumination,  # TOOD: MOVE THIS PARAM TO OTF? (leaky param)
+    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

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -19,7 +19,7 @@ def calculate_transfer_function(
     zyx_shape,
     yx_pixel_size,
     z_pixel_size,
-    wavelength_illumination,
+    wavelength_emission,
     z_padding,
     index_of_refraction_media,
     numerical_aperture_detection,
@@ -39,13 +39,13 @@ def calculate_transfer_function(
     det_pupil = optics.generate_pupil(
         radial_frequencies,
         numerical_aperture_detection,
-        wavelength_illumination,
+        wavelength_emission,
     )
 
     propagation_kernel = optics.generate_propagation_kernel(
         radial_frequencies,
         det_pupil,
-        wavelength_illumination / index_of_refraction_media,
+        wavelength_emission / index_of_refraction_media,
         z_position_list,
     )
 
@@ -106,28 +106,30 @@ def apply_inverse_transfer_function(
     zyx_data,
     optical_transfer_function,
     z_padding,
-    method="Tikhonov",
-    reg_re=1e-3,
-    rho=1e-3,
-    itr=10,
+    reconstruction_algorithm="Tikhonov",
+    regularization_strength=1e-3,
+    TV_rho_strength=1e-3,
+    TV_iterations=10,
 ):
     # Handle padding
     zyx_padded = util.pad_zyx_along_z(zyx_data, z_padding)
 
     # Reconstruct
-    if method == "Tikhonov":
+    if reconstruction_algorithm == "Tikhonov":
         f_real = util.single_variable_tikhonov_deconvolution_3D(
-            zyx_padded, optical_transfer_function, reg_re=reg_re
+            zyx_padded,
+            optical_transfer_function,
+            reg_re=regularization_strength,
         )
 
-    elif method == "TV":
+    elif reconstruction_algorithm == "TV":
         raise NotImplementedError
         f_real = util.single_variable_admm_tv_deconvolution_3D(
             zyx_padded,
             optical_transfer_function,
-            reg_re=reg_re,
-            rho=rho,
-            itr=itr,
+            reg_re=regularization_strength,
+            rho=TV_rho_strength,
+            itr=TV_iterations,
         )
 
     # Unpad

waveorder/models/isotropic_thin_3d.py

@@ -38,7 +38,11 @@ def calculate_transfer_function(
     index_of_refraction_media,
     numerical_aperture_illumination,
     numerical_aperture_detection,
+    axial_flip=False,
 ):
+    if axial_flip:
+        z_position_list = torch.flip(z_position_list, dims=(0,))
+
     radial_frequencies = util.generate_radial_frequencies(
         yx_shape, yx_pixel_size
     )
@@ -142,11 +146,11 @@ def apply_inverse_transfer_function(
     zyx_data,
     absorption_2d_to_3d_transfer_function,
     phase_2d_to_3d_transfer_function,
-    method="Tikhonov",
-    reg_u=1e-6,
-    reg_p=1e-6,
-    rho=1e-3,
-    itr=10,
+    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_normalized = util.inten_normalization(
@@ -158,7 +162,7 @@ def apply_inverse_transfer_function(
     # TODO AHA and b_vec calculations should be moved into tikhonov/tv calculations
     AHA = [
         torch.sum(torch.abs(absorption_2d_to_3d_transfer_function) ** 2, dim=0)
-        + reg_u,
+        + regularization_strength,
         torch.sum(
             torch.conj(absorption_2d_to_3d_transfer_function)
             * phase_2d_to_3d_transfer_function,
@@ -196,16 +200,16 @@ def apply_inverse_transfer_function(
     ]
 
     # Deconvolution with Tikhonov regularization
-    if method == "Tikhonov":
+    if reconstruction_algorithm == "Tikhonov":
         absorption, phase = util.dual_variable_tikhonov_deconvolution_2d(
             AHA, b_vec
         )
 
     # ADMM deconvolution with anisotropic TV regularization
-    elif method == "TV":
+    elif reconstruction_algorithm == "TV":
         raise NotImplementedError
         absorption, phase = util.dual_variable_admm_tv_deconv_2d(
-            AHA, b_vec, rho=rho, itr=itr
+            AHA, b_vec, rho=TV_rho_strength, itr=TV_iterations
         )
 
     phase -= torch.mean(phase)

waveorder/models/phase_thick_3d.py

@@ -133,10 +133,10 @@ def apply_inverse_transfer_function(
     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,
-    method="Tikhonov",
-    reg_re=1e-3,
-    rho=1e-3,
-    itr=10,
+    reconstruction_algorithm="Tikhonov",
+    regularization_strength=1e-3,
+    TV_rho_strength=1e-3,
+    TV_iterations=10,
 ):
     # Handle padding
     zyx_padded = util.pad_zyx_along_z(zyx_data, z_padding)
@@ -151,15 +151,19 @@ def apply_inverse_transfer_function(
     )
 
     # Reconstruct
-    if method == "Tikhonov":
+    if reconstruction_algorithm == "Tikhonov":
         f_real = util.single_variable_tikhonov_deconvolution_3D(
-            zyx, effective_transfer_function, reg_re=reg_re
+            zyx, effective_transfer_function, reg_re=regularization_strength
         )
 
-    elif method == "TV":
+    elif reconstruction_algorithm == "TV":
         raise NotImplementedError
         f_real = util.single_variable_admm_tv_deconvolution_3D(
-            zyx, effective_transfer_function, reg_re=reg_re, rho=rho, itr=itr
+            zyx,
+            effective_transfer_function,
+            reg_re=regularization_strength,
+            rho=TV_rho_strength,
+            itr=TV_iterations,
         )
 
     # Unpad