Optimized ParabolicSag.intercept() using numexpr. (#143)

Author: roytsmart

optika/_tests/test_sags.py

@@ -80,6 +80,10 @@ def test_intercept(
 
         assert np.allclose(a(result.position), result.position.z)
 
+        result_expected = optika.sags.AbstractSag.intercept(a, rays)
+
+        assert np.allclose(result, result_expected)
+
 
 @pytest.mark.parametrize(
     argnames="a",

optika/sags.py

@@ -727,6 +727,101 @@ def normal(
             "position / sqrt((_x / r)**2 + (_y / r)**2 + 1) / r",
         )
 
+    def intercept(
+        self,
+        rays: optika.rays.AbstractRayVectorArray,
+    ) -> optika.rays.RayVectorArray:
+        r"""
+        Compute the intercept of a ray with this parabolic surface.
+
+        Parameters
+        ----------
+        rays
+            The rays to find the intercept for.
+
+        Notes
+        -----
+
+        The equation for a paraboloid is
+
+        .. math::
+
+            z(x, y) = \frac{x^2 + y^2}{4 f},
+
+        where :math:`x`, :math:`y`, and :math:`z` are points on the paraboloid
+        and :math:`f` is the focal length of the paraboloid.
+        The equation for a line is
+
+        .. math::
+
+            \mathbf{x} = \mathbf{o} + d \mathbf{u},
+
+        where :math:`\mathbf{o}` is the starting point of the line,
+        :math:`\mathbf{u}` is a unit vector pointing in the direction of the line,
+        and :math:`d` is the distance from the origin to the intercept with
+        the paraboloid.
+        Combining these equations gives
+
+        .. math::
+
+            o_z + d u_z = \frac{(o_x + d u_x)^2 + (o_y + d u_y)^2}{4 f},
+
+        which can then be solved for :math:`d` using the quadratic equation,
+
+        .. math::
+
+            d = \frac{-o_x u_x - o_y u_y + 2 f u_z
+                      - \text{sgn}(f u_z) \sqrt{-o_y^2 u_x^2 - o_x^2 u_y^2 + 2 o_y u_y (o_x u_x - 2 f u_z)
+                              + 4 f (o_z (u_x^2 + u_y^2) - o_x u_x u_z + f u_z^2}
+                      }{u_x^2 + u_y^2}.
+
+        If the line is parallel to the :math:`z` axis, then the above equation
+        is singular and we need to solve the corresponding linear equation to find
+
+        .. math::
+
+            d = \frac{o_x^2 + o_y^2 - 4 f o_z}{4 f u_z}.
+        """
+
+        if self.transformation is not None:
+            rays = self.transformation.inverse(rays)
+
+        f = self.focal_length
+
+        unit = f.unit
+
+        f = f.value  # noqa: F841
+
+        o = rays.position.to(unit).value
+
+        u = rays.direction.value
+
+        ox = o.x  # noqa: F841
+        oy = o.y  # noqa: F841
+        oz = o.z  # noqa: F841
+
+        ux = u.x  # noqa: F841
+        uy = u.y  # noqa: F841
+        uz = u.z  # noqa: F841
+
+        position = na.numexpr.evaluate(
+            "o + u * where("
+            "   (ux**2 + uy**2) > 1e-10,"
+            "   (-ox * ux - oy * uy + 2 * f * uz - sign(f * uz) * sqrt("
+            "       -(oy * ux)**2 - (ox * uy)**2 + 2 * oy * uy * (ox * ux - 2 * f * uz)"
+            "       + 4 * f * (oz * (ux**2 + uy**2) - ox * ux * uz + f * uz**2)"
+            "   )) / (ux**2 + uy**2),"
+            "   (ox**2 + oy**2 - 4 * f * oz) / (4 * f * uz),"
+            ")"
+        )
+
+        result = rays.replace(position=position << unit)
+
+        if self.transformation is not None:
+            result = self.transformation(result)
+
+        return result
+
 
 @dataclasses.dataclass(eq=False, repr=False)
 class ToroidalSag(

pyproject.toml

@@ -15,7 +15,7 @@ classifiers = [
 ]
 dependencies = [
     "astropy!=6.1.5",
-    "named-arrays>=0.27.0",
+    "named-arrays>=0.27.1",
     "pymupdf",
     "joblib",
 ]