Fix Zernike tests for radius and k scaling

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}."