From 005ea5486848a61c9e06a9c29e52a39827d42b9b Mon Sep 17 00:00:00 2001
From: Talon Chandler <talonchandler@gmail.com>
Date: Mon, 9 Sep 2024 15:15:35 -0700
Subject: [PATCH] correct phase recon regression, legacy recon assumes axially
 even green's function

---
 waveorder/models/phase_thick_3d.py |  1 +
 waveorder/optics.py                | 23 ++++++++++++++---------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/waveorder/models/phase_thick_3d.py b/waveorder/models/phase_thick_3d.py
index 2a2903e..5ecc373 100644
--- a/waveorder/models/phase_thick_3d.py
+++ b/waveorder/models/phase_thick_3d.py
@@ -72,6 +72,7 @@ def calculate_transfer_function(
         det_pupil,
         wavelength_illumination / index_of_refraction_media,
         z_position_list,
+        axially_even=False,
     )
 
     (
diff --git a/waveorder/optics.py b/waveorder/optics.py
index 5ec4eeb..65ac130 100644
--- a/waveorder/optics.py
+++ b/waveorder/optics.py
@@ -422,7 +422,8 @@ def generate_propagation_kernel(
 
 
 def generate_greens_function_z(
-    radial_frequencies, pupil_support, wavelength_illumination, z_position_list
+    radial_frequencies, pupil_support, wavelength_illumination, z_position_list, 
+    axially_even=True,
 ):
     """
 
@@ -439,9 +440,14 @@ def generate_greens_function_z(
         wavelength_illumination       : float
                         wavelength of the light in the immersion media
 
-        z_position_list      : torch.tensor or list
+        z_position_list : torch.tensor or list
                         1D array of defocused z position with the size of (Z,)
 
+        axially_even    : bool
+                        For backwards compatibility with legacy phase reconstruction.
+                        Ideally the legacy phase reconstruction should be unified with 
+                        the new reconstructions, and this parameter should be removed.
+
     Returns
     -------
         greens_function_z       : torch.tensor
@@ -454,18 +460,17 @@ def generate_greens_function_z(
         * pupil_support
     ) ** (1 / 2) / wavelength_illumination
 
+    if axially_even:
+        z_positions = torch.abs(torch.tensor(z_position_list)[:, None, None])
+    else:
+        z_positions = torch.tensor(z_position_list)[:, None, None]
+        
     greens_function_z = (
         -1j
         / 4
         / np.pi
         * pupil_support[None, :, :]
-        * torch.exp(
-            1j
-            * 2
-            * np.pi
-            * torch.abs(torch.tensor(z_position_list)[:, None, None])
-            * oblique_factor[None, :, :]
-        )
+        * torch.exp(1j * 2 * np.pi * z_positions * oblique_factor[None, :, :])
         / (oblique_factor[None, :, :] + 1e-15)
     )