fix(phase): apply the generalized Snell's law in one step

Phase.ray_reaction refracted with plain Snell's law, which drops a ray
past the ordinary critical angle, before diffract added the phase
gradient. Yu et al. (Science 2011) show the gradient shifts the critical
angle: a ray is lost only when |n1 t1 + lambda/(2 pi) grad phi| >= n2.
So a phase on glass exiting into air lost steep rays the gradient would
transmit.

diffract now takes n1 and computes n2 t2 = n1 t1 + m lambda/(2 pi) grad phi
directly (Snell's law is the grad phi = 0 case), deciding validity only
on |t2|. Docstring cites Yu 2011 and Aieta et al. (Nano Lett. 2012) for
the three-dimensional form.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: Mihou

deeplens/phase_surface/base_phase.py

@@ -177,8 +177,8 @@ def ray_reaction(self, ray, n1, n2):
         """Trace a ray through the phase surface.
 
         Transforms the ray to local coordinates, intersects it with the plane,
-        applies refraction then diffraction, and transforms back to global
-        coordinates.
+        applies the generalized law of refraction (`diffract`), and transforms
+        back to global coordinates.
 
         Args:
             ray (Ray): Incident ray in global coordinates.
@@ -190,8 +190,7 @@ def ray_reaction(self, ray, n1, n2):
         """
         ray = self.to_local_coord(ray)
         ray = self.intersect(ray, n1)
-        ray = self.refract(ray, n1 / n2)
-        ray = self.diffract(ray, n2=n2)
+        ray = self.diffract(ray, n1, n2)
         ray = self.to_global_coord(ray)
         return ray
 
@@ -241,28 +240,32 @@ def intersect(self, ray, n=1.0):
 
         return ray
 
-    def diffract(self, ray, n2=1.0):
-        """Apply phase-surface diffraction to a ray.
+    def diffract(self, ray, n1=1.0, n2=1.0):
+        """Refract a ray through the phase surface by the generalized Snell's law.
 
         Two effects are applied:
 
         1. The phase $\\phi$ (in radians) adds to the optical path length as
-           $m\\,\\phi \\cdot \\lambda / (2\\pi)$, where $\\lambda$ is converted
-           from [µm] to [mm] internally.
-        2. The phase gradient bends the ray by the vector grating equation on
-           the plane (normal $z$),
-           $n_2 \\mathbf{t}_2 = n_1 \\mathbf{t}_1 + m\\,\\lambda / (2\\pi)\\,\\nabla\\phi$,
-           for the tangential direction components $\\mathbf{t} = (d_x, d_y)$.
-           Refraction is already applied, so the tangential components gain
-           $m\\,\\lambda / (2\\pi n_2)\\,\\nabla\\phi$ and $d_z$ is recomputed so
-           $|\\mathbf{d}| = 1$. An order with $|\\mathbf{t}_2| \\geq 1$ is
-           evanescent and the ray is marked invalid. The relation is the same in
-           both trace directions, so a backward trace retraces the forward path.
+           $m\\,\\phi \\cdot \\lambda / (2\\pi)$, where $\\lambda$ is the vacuum
+           wavelength, converted from [µm] to [mm] internally.
+        2. The phase gradient bends the ray by the generalized law of
+           refraction [2] in its three-dimensional form [3], on the plane
+           (normal $z$):
+           $n_2 \\mathbf{t}_2 = n_1 \\mathbf{t}_1 + m\\,\\lambda / (2\\pi)\\,\\nabla\\phi$
+           for the tangential direction components $\\mathbf{t} = (d_x, d_y)$,
+           with $d_z$ recomputed so $|\\mathbf{d}| = 1$. With $\\nabla\\phi = 0$
+           this is Snell's law, so no separate refraction is applied: the ray
+           is lost only when $|\\mathbf{t}_2| \\geq 1$, the critical angle the
+           phase gradient shifts [2], not at the plain-refraction one. The
+           relation is the same in both trace directions, so a backward trace
+           retraces the forward path.
 
         Args:
             ray (Ray): Ray with position, direction, and wavelength in [µm].
-            n2 (float, optional): Refractive index of the medium after the surface;
-                the deflection scales as $1/n_2$. Defaults to 1.0.
+            n1 (float, optional): Refractive index of the incident medium.
+                Defaults to 1.0.
+            n2 (float, optional): Refractive index of the transmission medium.
+                Defaults to 1.0.
 
         Returns:
             ray (Ray): Ray with updated direction `d`, `is_valid` (evanescent
@@ -277,7 +280,8 @@ def diffract(self, ray, n2=1.0):
 
         Reference:
             [1] https://support.zemax.com/hc/en-us/articles/1500005489061-How-diffractive-surfaces-are-modeled-in-OpticStudio
-            [2] Light propagation with phase discontinuities: generalized laws of reflection and refraction. Science 2011.
+            [2] N. Yu et al., Light propagation with phase discontinuities: generalized laws of reflection and refraction. Science 334, 333 (2011).
+            [3] F. Aieta et al., Out-of-plane reflection and refraction of light by anisotropic optical antenna metasurfaces with phase discontinuities. Nano Lett. 12, 1702 (2012).
         """
         valid = ray.is_valid > 0
         order = self.diffraction_order
@@ -290,11 +294,11 @@ def diffract(self, ray, n2=1.0):
             )
             ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)
 
-        # Step 2: vector grating equation (see docstring).
+        # Step 2: generalized Snell's law (see docstring).
         dphidx, dphidy = self.dphi_dxy(ray.o[..., 0], ray.o[..., 1])
-        scale = order * ray.wvln * 1e-3 / (2 * torch.pi * n2)
-        tx = ray.d[..., 0] + scale * dphidx
-        ty = ray.d[..., 1] + scale * dphidy
+        scale = order * ray.wvln * 1e-3 / (2 * torch.pi)
+        tx = (n1 * ray.d[..., 0] + scale * dphidx) / n2
+        ty = (n1 * ray.d[..., 1] + scale * dphidy) / n2
         # Decide validity on detached values. Invalid rays get a zero tangent so
         # the masked sqrt never sees a negative radicand (NaN gradient); the
         # clamp keeps its gradient finite as |t| -> 1.

test/test_phase_surfaces.py

@@ -261,6 +261,20 @@ def test_backward_trace_retraces_forward_path(self):
         back = s.ray_reaction(back, torch.tensor(1.5), torch.tensor(1.0))
         assert torch.allclose(back.d, -d_in, atol=1e-5)
 
+    def test_gradient_moves_the_critical_angle(self):
+        """Glass to air past the plain critical angle, rescued by the gradient.
+
+        Yu et al. 2011: TIR needs |n1 sin(theta1) + lambda/(2 pi) dphi/dy| > n2,
+        so n1 sin(theta1) = 1.05 with lambda/(2 pi) dphi/dy = -0.1 transmits at
+        n2 sin(theta2) = 0.95 instead of being lost to plain refraction.
+        """
+        alpha = -0.1 * 2 * math.pi / (WVLN * 1e-3) * 5.0
+        ray = _oblique_ray(math.degrees(math.asin(0.7)))
+        out = _grating(alpha).ray_reaction(ray, torch.tensor(1.5), torch.tensor(1.0))
+        assert out.is_valid.item() == 1
+        assert out.d[0, 1].item() == pytest.approx(0.95, abs=1e-5)
+        assert out.d[0, 2].item() == pytest.approx(math.sqrt(1 - 0.95**2), abs=1e-5)
+
     def test_evanescent_order_is_blocked_with_finite_gradient(self):
         """An order with |t| >= 1 is invalid and leaves no NaN gradient."""
         s = _grating(1e5)  # m lambda/(2 pi) |grad phi| ~ 1.75 > 1