From aedcd4b799858f954af2b0140061465e6700c7f1 Mon Sep 17 00:00:00 2001
From: Diptodip <dd@d2d.sh>
Date: Wed, 11 Dec 2024 14:28:22 -0500
Subject: [PATCH] Fix Zernike tests for radius and k scaling

---
 tests/test_zernikes.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_zernikes.py b/tests/test_zernikes.py
index 7f9b962..a108704 100644
--- a/tests/test_zernikes.py
+++ b/tests/test_zernikes.py
@@ -20,7 +20,7 @@ def test_first_ten_zernikes():
     grid = create_grid(size, spacing)
     # Normalize coordinates from -1 to 1 within radius R
     grid = grid_spatial_to_pupil(grid, f, NA, n)
-    rho = jnp.sum(grid**2, axis=0)  # radial coordinate
+    rho = jnp.sqrt(jnp.sum(grid**2, axis=0))  # radial coordinate
     mask = rho <= 1
     radius = rho * mask
     angle = jnp.arctan2(*grid) * mask  # angle coordinate
@@ -55,7 +55,7 @@ def test_first_ten_zernikes():
         )
         assert phase.shape == size
         assert jnp.allclose(
-            phase.squeeze(), expected[idx]
+            phase.squeeze(), 2 * jnp.pi * expected[idx] / wavelength
         ), f"Mismatch in Zernike polynomial {idx}."