Add optika.zernikes and ZernikeSag

docs/refs.bib

@@ -236,3 +236,16 @@ @article{Ramanathan2020
   url = {https://link.aps.org/doi/10.1103/PhysRevD.102.063026}
 }
 
+@article{Noll1976,
+  title = {Zernike polynomials and atmospheric turbulence},
+  author = {Noll, Robert J.},
+  journal = {J. Opt. Soc. Am.},
+  volume = {66},
+  number = {3},
+  pages = {207--211},
+  year = {1976},
+  month = {Mar},
+  publisher = {Optica Publishing Group},
+  doi = {10.1364/JOSA.66.000207},
+  url = {https://opg.optica.org/abstract.cfm?URI=josa-66-3-207}
+}

optika/__init__.py

@@ -5,6 +5,7 @@
 from ._caching import memory
 from . import mixins
 from ._util import shape, direction, angles
+from . import zernikes
 from . import vectors
 from . import targets
 from . import rays
@@ -27,6 +28,7 @@
     "shape",
     "direction",
     "angles",
+    "zernikes",
     "vectors",
     "targets",
     "rays",

optika/_tests/test_zernikes.py

@@ -0,0 +1,129 @@
+import pytest
+import numpy as np
+import named_arrays as na
+import optika
+
+_position_random = na.Cartesian2dVectorArray(
+    x=na.random.uniform(-0.7, 0.7, shape_random=dict(s=101), seed=42),
+    y=na.random.uniform(-0.7, 0.7, shape_random=dict(s=101), seed=43),
+)
+
+_position_origin = na.Cartesian2dVectorArray(
+    x=na.ScalarArray(np.array([0.0, 0.5]), axes="s"),
+    y=na.ScalarArray(np.array([0.0, -0.5]), axes="s"),
+)
+
+
+@pytest.mark.parametrize(
+    argnames="j,n,m",
+    argvalues=[
+        (1, 0, 0),
+        (2, 1, 1),
+        (3, 1, -1),
+        (4, 2, 0),
+        (5, 2, -2),
+        (6, 2, 2),
+        (7, 3, -1),
+        (8, 3, 1),
+        (9, 3, -3),
+        (10, 3, 3),
+        (11, 4, 0),
+        (12, 4, 2),
+        (13, 4, -2),
+        (14, 4, 4),
+        (15, 4, -4),
+    ],
+)
+def test_noll(j: int, n: int, m: int):
+    assert optika.zernikes.noll(j) == (n, m)
+
+
+def test_noll_invalid():
+    with pytest.raises(ValueError):
+        optika.zernikes.noll(0)
+
+
+@pytest.mark.parametrize(
+    argnames="position",
+    argvalues=[
+        _position_random,
+        _position_origin,
+    ],
+)
+class TestClosedForms:
+
+    def test_piston(self, position: na.AbstractCartesian2dVectorArray):
+        result = optika.zernikes.zernike(position, 1)
+        assert np.allclose(result, 1 + 0 * position.x)
+
+    def test_tilt_x(self, position: na.AbstractCartesian2dVectorArray):
+        result = optika.zernikes.zernike(position, 2)
+        assert np.allclose(result, 2 * position.x)
+
+    def test_tilt_y(self, position: na.AbstractCartesian2dVectorArray):
+        result = optika.zernikes.zernike(position, 3)
+        assert np.allclose(result, 2 * position.y)
+
+    def test_defocus(self, position: na.AbstractCartesian2dVectorArray):
+        result = optika.zernikes.zernike(position, 4)
+        rho2 = np.square(position.length)
+        assert np.allclose(result, np.sqrt(3) * (2 * rho2 - 1))
+
+    def test_spherical(self, position: na.AbstractCartesian2dVectorArray):
+        result = optika.zernikes.zernike(position, 11)
+        rho2 = np.square(position.length)
+        expected = np.sqrt(5) * (6 * np.square(rho2) - 6 * rho2 + 1)
+        assert np.allclose(result, expected)
+
+
+def test_orthonormality():
+    position = na.Cartesian2dVectorStratifiedRandomSpace(
+        start=-1,
+        stop=1,
+        axis=na.Cartesian2dVectorArray("x", "y"),
+        num=512,
+        seed=42,
+    ).explicit
+
+    where = position.length <= 1
+    num_inside = where.sum()
+
+    js = range(1, 12)
+    basis = {j: optika.zernikes.zernike(position, j) for j in js}
+
+    for j1 in js:
+        for j2 in js:
+            if j2 < j1:
+                continue
+            inner = (basis[j1] * basis[j2]).sum(where=where) / num_inside
+            expected = 1 if j1 == j2 else 0
+            assert np.abs(inner - expected) < 0.03, (j1, j2)
+
+
+@pytest.mark.parametrize("j", range(1, 16))
+@pytest.mark.parametrize(
+    argnames="position",
+    argvalues=[
+        _position_random,
+        _position_origin,
+    ],
+)
+def test_zernike_gradient(j: int, position: na.AbstractCartesian2dVectorArray):
+    result = optika.zernikes.zernike_gradient(position, j)
+
+    assert isinstance(result, na.AbstractCartesian2dVectorArray)
+
+    h = 1e-6
+    dx = na.Cartesian2dVectorArray(h, 0)
+    dy = na.Cartesian2dVectorArray(0, h)
+
+    derivative_x = optika.zernikes.zernike(position + dx, j)
+    derivative_x = derivative_x - optika.zernikes.zernike(position - dx, j)
+    derivative_x = derivative_x / (2 * h)
+
+    derivative_y = optika.zernikes.zernike(position + dy, j)
+    derivative_y = derivative_y - optika.zernikes.zernike(position - dy, j)
+    derivative_y = derivative_y / (2 * h)
+
+    assert np.allclose(result.x, derivative_x, atol=1e-5)
+    assert np.allclose(result.y, derivative_y, atol=1e-5)

optika/apertures/_apertures.py

@@ -618,29 +618,49 @@ def __call__(
 
         return mask
 
+    def _bound_center_half(
+        self,
+    ) -> tuple[na.Cartesian3dVectorArray, na.Cartesian3dVectorArray]:
+        """
+        The center and per-component half-extent of the axis-aligned bounding
+        box, computed analytically so the bound is exact even when
+        :attr:`transformation` rotates the ellipse.
+
+        A point on the ellipse boundary is
+        :math:`p(t) = c + a\\,\\hat{e}_a \\cos t + b\\,\\hat{e}_b \\sin t`,
+        where :math:`c` is the center and :math:`\\hat{e}_a, \\hat{e}_b` are the
+        images of the local axes under the transformation.  The extent along
+        any world component is then
+        :math:`\\sqrt{(a\\,\\hat{e}_a)^2 + (b\\,\\hat{e}_b)^2}`, since
+        :math:`\\max_t (A \\cos t + B \\sin t) = \\sqrt{A^2 + B^2}`.
+        """
+        radius = self.radius
+        zero = na.Cartesian3dVectorArray() * radius.x
+        axis_a = na.Cartesian3dVectorArray(x=radius.x, y=0 * radius.x, z=0 * radius.x)
+        axis_b = na.Cartesian3dVectorArray(x=0 * radius.y, y=radius.y, z=0 * radius.y)
+
+        center = zero
+        if self.transformation is not None:
+            center = self.transformation(zero)
+            axis_a = self.transformation(axis_a) - center
+            axis_b = self.transformation(axis_b) - center
+
+        half = na.Cartesian3dVectorArray(
+            x=np.sqrt(np.square(axis_a.x) + np.square(axis_b.x)),
+            y=np.sqrt(np.square(axis_a.y) + np.square(axis_b.y)),
+            z=np.sqrt(np.square(axis_a.z) + np.square(axis_b.z)),
+        )
+        return center, half
+
     @property
     def bound_lower(self) -> na.Cartesian3dVectorArray:
-        unit = na.unit(self.radius)
-        result = na.Cartesian3dVectorArray()
-        if unit is not None:
-            result = result * unit
-        if self.transformation is not None:
-            result = self.transformation(result)
-        result.x = result.x - self.radius.x
-        result.y = result.y - self.radius.y
-        return result
+        center, half = self._bound_center_half()
+        return center - half
 
     @property
     def bound_upper(self) -> na.Cartesian3dVectorArray:
-        unit = na.unit(self.radius)
-        result = na.Cartesian3dVectorArray()
-        if unit is not None:
-            result = result * unit
-        if self.transformation is not None:
-            result = self.transformation(result)
-        result.x = result.x + self.radius.x
-        result.y = result.y + self.radius.y
-        return result
+        center, half = self._bound_center_half()
+        return center + half
 
     @property
     def vertices(self) -> None:
@@ -717,11 +737,17 @@ def __call__(
 
     @property
     def bound_lower(self) -> na.AbstractCartesian3dVectorArray:
-        return self.vertices.min(axis="vertex")
+        vertices = self.vertices
+        if self.transformation is not None:
+            vertices = self.transformation(vertices)
+        return vertices.min(axis="vertex")
 
     @property
     def bound_upper(self) -> na.AbstractCartesian3dVectorArray:
-        return self.vertices.max(axis="vertex")
+        vertices = self.vertices
+        if self.transformation is not None:
+            vertices = self.transformation(vertices)
+        return vertices.max(axis="vertex")
 
     def wire(self, num: None | int = None) -> na.Cartesian3dVectorArray:
         if num is None:
@@ -874,24 +900,24 @@ def __call__(
         self,
         position: na.AbstractCartesian3dVectorArray,
     ) -> na.AbstractScalar:
-        bound_lower = self.bound_lower
-        bound_upper = self.bound_upper
+        half_width = na.asanyarray(
+            self.half_width,
+            like=na.Cartesian2dVectorArray(),
+        )
         active = self.active
         inverted = self.inverted
         if self.transformation is not None:
             position = self.transformation.inverse(position)
+        position = position.xy
 
-        shape = na.shape_broadcasted(
-            bound_lower, bound_upper, active, inverted, position
-        )
+        shape = na.shape_broadcasted(half_width, active, inverted, position)
 
-        bound_lower = na.broadcast_to(bound_lower, shape)
-        bound_upper = na.broadcast_to(bound_upper, shape)
+        half_width = na.broadcast_to(half_width, shape)
         active = na.broadcast_to(active, shape)
         inverted = na.broadcast_to(inverted, shape)
         position = na.broadcast_to(position, shape)
 
-        mask = (bound_lower <= position) & (position <= bound_upper)
+        mask = (-half_width <= position) & (position <= half_width)
         mask = mask.x & mask.y
 
         mask[inverted] = ~mask[inverted]