From 614340f6165758c70d09a699501edac3e185b7e3 Mon Sep 17 00:00:00 2001
From: Talon Chandler <talonchandler@gmail.com>
Date: Mon, 15 Apr 2024 15:56:00 -0700
Subject: [PATCH] test voxel-size invariance

---
 tests/models/test_phase_thick_3d.py | 81 ++++++++++++++++++++++++++++-
 1 file changed, 80 insertions(+), 1 deletion(-)

diff --git a/tests/models/test_phase_thick_3d.py b/tests/models/test_phase_thick_3d.py
index 224c96c..d60c7fa 100644
--- a/tests/models/test_phase_thick_3d.py
+++ b/tests/models/test_phase_thick_3d.py
@@ -1,5 +1,5 @@
 import pytest
-
+import numpy as np
 from waveorder.models import phase_thick_3d
 
 
@@ -20,3 +20,82 @@ def test_calculate_transfer_function(invert_phase_contrast):
 
     assert H_re.shape == (20 + 2 * z_padding, 100, 101)
     assert H_im.shape == (20 + 2 * z_padding, 100, 101)
+
+
+# Helper function for testing reconstruction invariances
+def simulate_phase_recon(
+    z_pixel_size_um=0.1,
+    yx_pixel_size_um=6.5 / 63,
+):
+
+    z_fov_um = 50
+    yx_fov_um = 50
+
+    n_z = np.int32(z_fov_um / z_pixel_size_um)
+    n_yx = np.int32(yx_fov_um / yx_pixel_size_um)
+
+    # Parameters
+    # all lengths must use consistent units e.g. um
+    simulation_arguments = {
+        "zyx_shape": (n_z, n_yx, n_yx),
+        "yx_pixel_size": yx_pixel_size_um,
+        "z_pixel_size": z_pixel_size_um,
+        "index_of_refraction_media": 1.3,
+    }
+    phantom_arguments = {
+        "index_of_refraction_sample": 1.40,
+        "sphere_radius": 5,
+    }
+    transfer_function_arguments = {
+        "z_padding": 0,
+        "wavelength_illumination": 0.532,
+        "numerical_aperture_illumination": 0.9,
+        "numerical_aperture_detection": 1.3,
+    }
+
+    # Create a phantom
+    zyx_phase = phase_thick_3d.generate_test_phantom(
+        **simulation_arguments, **phantom_arguments
+    )
+
+    # Calculate transfer function
+    (
+        real_potential_transfer_function,
+        imag_potential_transfer_function,
+    ) = phase_thick_3d.calculate_transfer_function(
+        **simulation_arguments, **transfer_function_arguments
+    )
+
+    # Simulate
+    zyx_data = phase_thick_3d.apply_transfer_function(
+        zyx_phase,
+        real_potential_transfer_function,
+        transfer_function_arguments["z_padding"],
+        brightness=1000,
+    )
+
+    # Reconstruct
+    zyx_recon = phase_thick_3d.apply_inverse_transfer_function(
+        zyx_data,
+        real_potential_transfer_function,
+        imag_potential_transfer_function,
+        transfer_function_arguments["z_padding"],
+        regularization_strength=1e-3,
+    )
+
+    Z, Y, X = zyx_phase.shape
+    recon_center = zyx_recon[Z // 2, Y // 2, X // 2].numpy()
+
+    return recon_center
+
+
+def test_phase_invariance():
+    recon = simulate_phase_recon()
+
+    # test z pixel size invariance
+    recon1 = simulate_phase_recon(z_pixel_size_um=0.3)
+    assert np.abs((recon1 - recon) / recon) < 0.02
+
+    # test yx pixel size invariance
+    recon2 = simulate_phase_recon(yx_pixel_size_um=0.7 * 6.5 / 63)
+    assert np.abs((recon2 - recon) / recon) < 0.02
\ No newline at end of file