From 22a059b6fc323a7988540f61acc5975f9b77ade3 Mon Sep 17 00:00:00 2001
From: Talon Chandler <talonchandler@gmail.com>
Date: Mon, 9 Sep 2024 20:07:48 -0700
Subject: [PATCH] 2d phase wrap safety

---
 waveorder/models/isotropic_thin_3d.py | 61 ++++++++++++++++++++++++++-
 1 file changed, 60 insertions(+), 1 deletion(-)

diff --git a/waveorder/models/isotropic_thin_3d.py b/waveorder/models/isotropic_thin_3d.py
index c1bb96a..1af9ed8 100644
--- a/waveorder/models/isotropic_thin_3d.py
+++ b/waveorder/models/isotropic_thin_3d.py
@@ -1,9 +1,10 @@
 from typing import Literal, Tuple
 
+import numpy as np
 import torch
 from torch import Tensor
 
-from waveorder import optics, util
+from waveorder import optics, sampling, util
 
 
 def generate_test_phantom(
@@ -42,6 +43,64 @@ def calculate_transfer_function(
     numerical_aperture_illumination,
     numerical_aperture_detection,
     invert_phase_contrast=False,
+):
+    transverse_nyquist = sampling.transverse_nyquist(
+        wavelength_illumination,
+        numerical_aperture_illumination,
+        numerical_aperture_detection,
+    )
+    yx_factor = int(np.ceil(yx_pixel_size / transverse_nyquist))
+
+    absorption_2d_to_3d_transfer_function, phase_2d_to_3d_transfer_function = (
+        _calculate_wrap_unsafe_transfer_function(
+            (
+                yx_shape[0] * yx_factor,
+                yx_shape[1] * yx_factor,
+            ),
+            yx_pixel_size / yx_factor,
+            z_position_list,
+            wavelength_illumination,
+            index_of_refraction_media,
+            numerical_aperture_illumination,
+            numerical_aperture_detection,
+            invert_phase_contrast=invert_phase_contrast,
+        )
+    )
+
+    absorption_2d_to_3d_transfer_function_out = torch.zeros(
+        (len(z_position_list),) + tuple(yx_shape), dtype=torch.complex64
+    )
+    phase_2d_to_3d_transfer_function_out = torch.zeros(
+        (len(z_position_list),) + tuple(yx_shape), dtype=torch.complex64
+    )
+
+    for z in range(len(z_position_list)):
+        absorption_2d_to_3d_transfer_function_out[z] = (
+            sampling.nd_fourier_central_cuboid(
+                absorption_2d_to_3d_transfer_function[z], yx_shape
+            )
+        )
+        phase_2d_to_3d_transfer_function_out[z] = (
+            sampling.nd_fourier_central_cuboid(
+                phase_2d_to_3d_transfer_function[z], yx_shape
+            )
+        )
+
+    return (
+        absorption_2d_to_3d_transfer_function_out,
+        phase_2d_to_3d_transfer_function_out,
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    yx_shape,
+    yx_pixel_size,
+    z_position_list,
+    wavelength_illumination,
+    index_of_refraction_media,
+    numerical_aperture_illumination,
+    numerical_aperture_detection,
+    invert_phase_contrast=False,
 ):
     if invert_phase_contrast:
         z_position_list = torch.flip(torch.tensor(z_position_list), dims=(0,))