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/sags/__init__.py

@@ -7,6 +7,7 @@
 from ._conic import AbstractConicSag, ConicSag
 from ._parabolic import ParabolicSag
 from ._toroidal import ToroidalSag
+from ._zernike import ZernikeSag
 
 __all__ = [
     "AbstractSag",
@@ -18,4 +19,5 @@
     "ConicSag",
     "ParabolicSag",
     "ToroidalSag",
+    "ZernikeSag",
 ]

optika/sags/_tests/_zernike_test.py

@@ -0,0 +1,120 @@
+import pytest
+import numpy as np
+import astropy.units as u
+import named_arrays as na
+import optika
+from ._abc_test import AbstractTestAbstractSag
+
+
+@pytest.mark.parametrize(
+    argnames="a",
+    argvalues=[
+        optika.sags.ZernikeSag(
+            coefficients=[0, 0, 0, 200e-6] * u.mm,
+            radius=50 * u.mm,
+        ),
+        optika.sags.ZernikeSag(
+            base=optika.sags.SphericalSag(radius=500 * u.mm),
+            coefficients=[100, 50, 50, 200, 30, 30, 10, 10, 5, 5, 20] * u.nm,
+            radius=50 * u.mm,
+        ),
+        optika.sags.ZernikeSag(
+            base=optika.sags.ParabolicSag(focal_length=1000 * u.mm),
+            coefficients=na.ScalarArray(
+                ndarray=[0, 0, 0, 100, 0, 50] * u.nm,
+                axes="zernike",
+            ),
+            radius=25 * u.mm,
+            transformation=na.transformations.Cartesian3dTranslation(
+                z=10 * u.mm,
+            ),
+        ),
+    ],
+)
+class TestZernikeSag(
+    AbstractTestAbstractSag,
+):
+    pass
+
+
+def test_coefficients_invalid_axis():
+    sag = optika.sags.ZernikeSag(
+        coefficients=na.ScalarArray([0, 0, 0, 1] * u.um, axes="not_zernike"),
+        radius=50 * u.mm,
+    )
+    with pytest.raises(ValueError):
+        sag(na.Cartesian3dVectorArray() * u.mm)
+
+
+def test_defocus_closed_form():
+    """
+    A pure-defocus ZernikeSag over a flat profile should equal the analytic
+    form of the Noll-normalized defocus polynomial.
+    """
+    c = 100 * u.nm
+    radius = 50 * u.mm
+
+    sag = optika.sags.ZernikeSag(
+        coefficients=[0 * u.nm, 0 * u.nm, 0 * u.nm, c],
+        radius=radius,
+    )
+
+    position = na.Cartesian3dVectorArray(
+        x=na.linspace(-50, 50, axis="x", num=11) * u.mm,
+        y=na.linspace(-50, 50, axis="y", num=11) * u.mm,
+        z=0 * u.mm,
+    )
+
+    rho2 = np.square(position.xy.length / radius)
+    expected = c * np.sqrt(3) * (2 * rho2 - 1)
+
+    assert np.allclose(sag(position), expected)
+
+
+def test_defocus_focal_shift():
+    """
+    Adding a Zernike defocus term to a paraboloid yields another exact
+    paraboloid with a focal length given by
+
+    .. math::
+
+        \\frac{1}{f'} = \\frac{1}{f} + \\frac{8 \\sqrt{3} c}{R^2},
+
+    so collimated rays should still focus perfectly, at the shifted focus.
+    """
+    f = 1000 * u.mm
+    radius = 25 * u.mm
+    c = 1 * u.um
+
+    sag = optika.sags.ZernikeSag(
+        base=optika.sags.ParabolicSag(focal_length=f),
+        coefficients=[0 * u.um, 0 * u.um, 0 * u.um, c],
+        radius=radius,
+    )
+
+    position = na.Cartesian3dVectorArray(
+        x=na.linspace(-radius, radius, axis="rx", num=11),
+        y=na.linspace(-radius, radius, axis="ry", num=11),
+        z=2 * f,
+    )
+    rays = optika.rays.RayVectorArray(
+        wavelength=500 * u.nm,
+        position=position,
+        direction=na.Cartesian3dVectorArray(0, 0, -1),
+    )
+
+    rays = sag.intercept(rays)
+    normal = sag.normal(rays.position)
+    direction = rays.direction
+    direction = direction - 2 * (direction @ normal) * normal
+
+    f_perturbed = 1 / (1 / f + 8 * np.sqrt(3) * c / np.square(radius))
+    z_focus = f_perturbed - np.sqrt(3) * c
+
+    def rms_spot(z: u.Quantity) -> u.Quantity:
+        t = (z - rays.position.z) / direction.z
+        spot = rays.position + direction * t
+        return np.sqrt(np.mean(np.square(spot.x) + np.square(spot.y)))
+
+    assert rms_spot(z_focus) < 100 * u.nm
+    assert rms_spot(f) > 50 * u.um