From 995c90df8a4b80d3168e9a75d98e75910a44bcd5 Mon Sep 17 00:00:00 2001
From: Talon Chandler <talonchandler@gmail.com>
Date: Fri, 27 Sep 2024 10:00:08 -0700
Subject: [PATCH] wrap-safe vector transfer function

---
 .../inplane_oriented_thick_pol3d_vector.py    | 59 ++++++++++++++++++-
 1 file changed, 57 insertions(+), 2 deletions(-)

diff --git a/waveorder/models/inplane_oriented_thick_pol3d_vector.py b/waveorder/models/inplane_oriented_thick_pol3d_vector.py
index 874afac..56da40f 100644
--- a/waveorder/models/inplane_oriented_thick_pol3d_vector.py
+++ b/waveorder/models/inplane_oriented_thick_pol3d_vector.py
@@ -2,7 +2,7 @@
 import numpy as np
 
 from torch import Tensor
-from waveorder import optics, stokes, util
+from waveorder import optics, sampling, stokes, util
 from waveorder.visuals.napari_visuals import add_transfer_function_to_viewer
 
 
@@ -37,6 +37,61 @@ 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))
+
+    sfZYX_transfer_function, intensity_to_stokes_matrix = (
+        _calculate_wrap_unsafe_transfer_function(
+            swing,
+            scheme,
+            (
+                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,
+        )
+    )
+    sfzyx_out_shape = (3, 3, zyx_shape[0] + 2 * z_padding,) + zyx_shape[1:]
+    return (
+        sampling.nd_fourier_central_cuboid(
+            sfZYX_transfer_function, sfzyx_out_shape
+        ),
+        intensity_to_stokes_matrix,
+    )
+
+
+def _calculate_wrap_unsafe_transfer_function(
+    swing,
+    scheme,
+    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,
 ):
     intensity_to_stokes_matrix = stokes.calculate_intensity_to_stokes_matrix(
         swing, scheme=scheme
@@ -108,7 +163,6 @@ def calculate_transfer_function(
         z_position_list,
     )
 
-
     G = optics.generate_defocus_greens_tensor(
         x_frequencies,
         y_frequencies,
@@ -181,6 +235,7 @@ def visualize_transfer_function(viewer, sfZYX_transfer_function, zyx_scale):
         layer_name="Transfer Function",
     )
 
+
 def apply_transfer_function(
     fzyx_object,
     sfZYX_transfer_function,