Add option for axial flip of phase_thick_3d transfer function (#124)

* `illumination_wavelength` -> `wavelength_illumination` * add option for axial flip of transfer function * test axial flip --------- Co-authored-by: Ziwen Liu <[email protected]>
mehta-lab · Jun 13, 2023 · 1943167 · 1943167
1 parent 5cec51b
commit 1943167
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/tests/models/test_phase_thick_3d.py b/tests/models/test_phase_thick_3d.py
@@ -1,7 +1,9 @@
+import pytest
 from waveorder.models import phase_thick_3d
 
 
-def test_calculate_transfer_function():
+@pytest.mark.parametrize("axial_flip", (True, False))
+def test_calculate_transfer_function(axial_flip):
     z_padding = 5
     H_re, H_im = phase_thick_3d.calculate_transfer_function(
         zyx_shape=(20, 100, 101),
@@ -12,6 +14,7 @@ def test_calculate_transfer_function():
         index_of_refraction_media=1.0,
         numerical_aperture_illumination=0.45,
         numerical_aperture_detection=0.55,
+        axial_flip=axial_flip
     )
 
     assert H_re.shape == (20 + 2 * z_padding, 100, 101)

diff --git a/waveorder/models/phase_thick_3d.py b/waveorder/models/phase_thick_3d.py
@@ -38,6 +38,7 @@ def calculate_transfer_function(
     index_of_refraction_media,
     numerical_aperture_illumination,
     numerical_aperture_detection,
+    axial_flip=False,
 ):
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size
@@ -46,6 +47,8 @@ def calculate_transfer_function(
     z_position_list = torch.fft.ifftshift(
         (torch.arange(z_total) - z_total // 2) * z_pixel_size
     )
+    if axial_flip:
+        z_position_list = torch.flip(z_position_list, dims=(0,))
 
     ill_pupil = optics.generate_pupil(
         radial_frequencies,