From 783b916da53c014cf2acde767ba090f2deed40c8 Mon Sep 17 00:00:00 2001
From: Talon Chandler <talonchandler@gmail.com>
Date: Mon, 9 Sep 2024 17:27:48 -0700
Subject: [PATCH] 3d phase wrap safety

---
 waveorder/models/phase_thick_3d.py | 55 +++++++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/waveorder/models/phase_thick_3d.py b/waveorder/models/phase_thick_3d.py
index ac29d35..868b15a 100644
--- a/waveorder/models/phase_thick_3d.py
+++ b/waveorder/models/phase_thick_3d.py
@@ -4,7 +4,7 @@
 import torch
 from torch import Tensor
 
-from waveorder import optics, util
+from waveorder import optics, sampling, util
 from waveorder.models import isotropic_fluorescent_thick_3d
 
 
@@ -40,6 +40,59 @@ 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,
+    )
+    axial_nyquist = sampling.axial_nyquist(
+        wavelength_illumination,
+        numerical_aperture_detection,
+        index_of_refraction_media,
+    )
+
+    yx_factor = int(np.ceil(yx_pixel_size / transverse_nyquist))
+    z_factor = int(np.ceil(z_pixel_size / axial_nyquist))
+
+    real_potential_transfer_function, imag_potential_transfer_function = (
+        _calculate_wrap_unsafe_transfer_function(
+            (
+                zyx_shape[0] * z_factor,
+                zyx_shape[1] * yx_factor,
+                zyx_shape[2] * yx_factor,
+            ),
+            yx_pixel_size / yx_factor,
+            z_pixel_size / z_factor,
+            wavelength_illumination,
+            z_padding,
+            index_of_refraction_media,
+            numerical_aperture_illumination,
+            numerical_aperture_detection,
+            invert_phase_contrast=invert_phase_contrast,
+        )
+    )
+
+    return (
+        sampling.nd_fourier_central_cuboid(
+            real_potential_transfer_function, zyx_shape
+        ),
+        sampling.nd_fourier_central_cuboid(
+            imag_potential_transfer_function, zyx_shape
+        ),
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    zyx_shape,
+    yx_pixel_size,
+    z_pixel_size,
+    wavelength_illumination,
+    z_padding,
+    index_of_refraction_media,
+    numerical_aperture_illumination,
+    numerical_aperture_detection,
+    invert_phase_contrast=False,
 ):
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size