diff --git a/docs/refs.bib b/docs/refs.bib index a70bae3e..efea5344 100644 --- a/docs/refs.bib +++ b/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} +} diff --git a/optika/__init__.py b/optika/__init__.py index 3f2766fd..67ab71bf 100644 --- a/optika/__init__.py +++ b/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", diff --git a/optika/_tests/test_zernikes.py b/optika/_tests/test_zernikes.py new file mode 100644 index 00000000..ca644df8 --- /dev/null +++ b/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) diff --git a/optika/apertures/_apertures.py b/optika/apertures/_apertures.py index 1db224b4..9854d9a4 100644 --- a/optika/apertures/_apertures.py +++ b/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] diff --git a/optika/apertures/_apertures_test.py b/optika/apertures/_apertures_test.py index 01291007..57e70f80 100644 --- a/optika/apertures/_apertures_test.py +++ b/optika/apertures/_apertures_test.py @@ -25,6 +25,12 @@ ] +def _nominal(x: na.AbstractScalar) -> na.AbstractScalar: + if isinstance(x, na.AbstractUncertainScalarArray): + return x.nominal + return x + + class AbstractTestAbstractAperture( test_mixins.AbstractTestDxfWritable, test_mixins.AbstractTestPrintable, @@ -57,6 +63,17 @@ def test__call__(self, a: optika.apertures.AbstractAperture): assert np.all(result[~na.as_named_array(a.active)]) + # the centroid of the wire must be inside an active, non-inverted + # aperture (every aperture here is star-shaped about its centroid) + if (a.active is True) and (a.inverted is False): + centroid = a.wire().mean("wire") + centroid = na.Cartesian3dVectorArray( + x=_nominal(centroid.x), + y=_nominal(centroid.y), + z=_nominal(centroid.z), + ) + assert np.all(_nominal(na.as_named_array(a(centroid)))) + @pytest.mark.parametrize("rays", optika.rays._tests.test_ray_vectors.rays) def test_clip_rays( self, @@ -75,12 +92,28 @@ def test_clip_rays( assert np.mean(result.unvignetted) <= np.mean(rays.unvignetted) def test_bound_lower(self, a: optika.apertures.AbstractAperture): - assert isinstance(a.bound_lower, na.AbstractCartesian3dVectorArray) + result = a.bound_lower + assert isinstance(result, na.AbstractCartesian3dVectorArray) + # compare nominal values only, since uncertain parameters may be + # redrawn between independent evaluations of the geometry + wire = a.wire() + for component in ("x", "y"): + bound = _nominal(getattr(result, component)) + edge = _nominal(getattr(wire, component).min("wire")) + tolerance = 1e-9 * (np.abs(bound) + np.abs(edge)) + assert np.all(bound <= edge + tolerance) def test_bound_upper(self, a: optika.apertures.AbstractAperture): - assert isinstance(a.bound_upper, na.AbstractCartesian3dVectorArray) - assert np.all(a.bound_upper.x != a.bound_lower.x) - assert np.all(a.bound_upper.y != a.bound_lower.y) + result = a.bound_upper + assert isinstance(result, na.AbstractCartesian3dVectorArray) + assert np.all(result.x != a.bound_lower.x) + assert np.all(result.y != a.bound_lower.y) + wire = a.wire() + for component in ("x", "y"): + bound = _nominal(getattr(result, component)) + edge = _nominal(getattr(wire, component).max("wire")) + tolerance = 1e-9 * (np.abs(bound) + np.abs(edge)) + assert np.all(bound >= edge - tolerance) def test_vertices(self, a: optika.apertures.AbstractAperture): if a.vertices is not None: @@ -400,3 +433,67 @@ class TestIsoscelesTrapezoidalAperture( AbstractTestAbstractIsoscelesTrapezoidalAperture, ): pass + + +def test_rectangular_aperture_decentered(): + """ + A rectangular aperture decentered by its internal transformation must + accept points inside the translated rectangle and reject points inside + the untranslated one. + + Regression test: ``RectangularAperture.__call__`` used to express the + rectangle in terms of :attr:`bound_lower`/:attr:`bound_upper` while + also inverse-transforming the position, applying the internal + transformation twice. + """ + half_width = 5 * u.mm + decenter = 12 * u.mm + a = optika.apertures.RectangularAperture( + half_width=half_width, + transformation=na.transformations.Cartesian3dTranslation(x=decenter), + ) + point_inside = na.Cartesian3dVectorArray(12, 0, 0) * u.mm + point_outside = na.Cartesian3dVectorArray(0, 0, 0) * u.mm + assert np.all(na.as_named_array(a(point_inside))) + assert not np.any(na.as_named_array(a(point_outside))) + + +def test_elliptical_aperture_bounds_analytic(): + """ + The bounding box of a rotated, decentered elliptical aperture is computed + analytically and must match the true extent of the ellipse. + + Regression test: the bounds used to be sampled from a coarse + :meth:`wire`, which underestimates the extent between samples whenever the + transformation rotates the ellipse off the coordinate axes. + """ + a = optika.apertures.EllipticalAperture( + radius=na.Cartesian2dVectorArray(100, 50) * u.mm, + transformation=( + na.transformations.Cartesian3dRotationZ(30 * u.deg) + @ na.transformations.Cartesian3dTranslation(x=5 * u.mm, y=-2 * u.mm) + ), + ) + + bound_lower = a.bound_lower + bound_upper = a.bound_upper + + # a densely-sampled wire approaches the true extent from the inside + wire = a.wire(num=100001) + for component in ("x", "y"): + lower = getattr(bound_lower, component) + upper = getattr(bound_upper, component) + edge_lower = getattr(wire, component).min("wire") + edge_upper = getattr(wire, component).max("wire") + scale = np.abs(upper - lower) + + # the analytic bound encloses the dense wire ... + assert np.all(lower <= edge_lower + 1e-9 * scale) + assert np.all(upper >= edge_upper - 1e-9 * scale) + # ... and is tight: it does not overshoot it + assert np.all(np.abs(lower - edge_lower) < 1e-6 * scale) + assert np.all(np.abs(upper - edge_upper) < 1e-6 * scale) + + # an in-plane rotation leaves the aperture flat in z + assert np.all(bound_lower.z == 0 * u.mm) + assert np.all(bound_upper.z == 0 * u.mm) diff --git a/optika/sags/__init__.py b/optika/sags/__init__.py index 4241548f..d0c4d818 100644 --- a/optika/sags/__init__.py +++ b/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", ] diff --git a/optika/sags/_tests/_zernike_test.py b/optika/sags/_tests/_zernike_test.py new file mode 100644 index 00000000..bff7532b --- /dev/null +++ b/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 diff --git a/optika/sags/_zernike.py b/optika/sags/_zernike.py new file mode 100644 index 00000000..0e6c5d67 --- /dev/null +++ b/optika/sags/_zernike.py @@ -0,0 +1,174 @@ +import dataclasses +import numpy as np +import astropy.units as u +import named_arrays as na +import optika +from ._abc import AbstractSag +from ._flat import NoSag + +__all__ = [ + "ZernikeSag", +] + + +@dataclasses.dataclass(eq=False, repr=False) +class ZernikeSag( + AbstractSag, +): + r""" + A sag profile consisting of a base profile plus a sum of Zernike + polynomials. + + This is useful for representing measured or modeled figure errors of an + optical surface, since the perturbation modifies the actual shape of the + surface, it is seen consistently by both geometric raytraces and + physical-optics calculations. + + Examples + -------- + Plot a slice through a parabolic sag profile with a large coma + perturbation. + + .. jupyter-execute:: + + import numpy as np + import matplotlib.pyplot as plt + import astropy.units as u + import astropy.visualization + import named_arrays as na + import optika + + sag = optika.sags.ZernikeSag( + base=optika.sags.ParabolicSag(focal_length=500 * u.mm), + coefficients=[0, 0, 0, 0, 0, 0, 0, 1] * u.mm, + radius=50 * u.mm, + ) + + # the coma term (Noll index 8) varies along x, so slice along x + position = na.Cartesian3dVectorArray( + x=na.linspace(-50, 50, axis="x", num=101) * u.mm, + y=0 * u.mm, + z=0 * u.mm, + ) + + z = sag(position) + z_base = sag.base(position) + + with astropy.visualization.quantity_support(): + plt.figure() + na.plt.plot(position.x, z, axis="x", label="perturbed") + na.plt.plot(position.x, z_base, axis="x", label="base") + plt.legend() + """ + + base: None | AbstractSag = None + """ + The base sag profile to be perturbed. + If :obj:`None` (the default), a flat profile, :class:`optika.sags.NoSag`, + is used. + """ + + coefficients: u.Quantity | na.AbstractScalar = 0 * u.mm + """ + The magnitudes of the Zernike polynomial terms along the logical axis + `axis`, where element :math:`i` corresponds to Noll index :math:`j = i + 1`. + If given as a bare array or scalar, it is interpreted as being along + `axis`. + """ + + radius: u.Quantity | na.AbstractScalar = 1 * u.mm + """ + The radius of the unit disk on which the Zernike polynomials are defined, + in the same coordinate system as the evaluation points. + """ + + axis: str = "zernike" + """The logical axis of `coefficients` indexing the Noll terms.""" + + def __post_init__(self): + if self.base is None: + self.base = NoSag() + + @property + def _coefficients_normalized(self) -> na.AbstractScalar: + """The coefficients as a named array guaranteed to contain `axis`.""" + result = self.coefficients + if not isinstance(result, na.AbstractArray): + result = np.atleast_1d(u.Quantity(result)) + result = na.ScalarArray(result, axes=(self.axis,)) + if self.axis not in result.shape: + raise ValueError( + f"`coefficients` must vary along `axis`, {self.axis!r}, " + f"got an array with shape {result.shape}." + ) + return result + + @property + def shape(self) -> dict[str, int]: + shape_coefficients = dict(self._coefficients_normalized.shape) + shape_coefficients.pop(self.axis, None) + return na.broadcast_shapes( + optika.shape(self.base), + shape_coefficients, + optika.shape(self.radius), + optika.shape(self.transformation), + optika.shape(self.parameters_slope_error), + optika.shape(self.parameters_roughness), + optika.shape(self.parameters_microroughness), + ) + + def __call__( + self, + position: na.AbstractCartesian3dVectorArray, + ) -> na.AbstractScalar: + + if self.transformation is not None: + position = self.transformation.inverse(position) + + result = self.base(position) + + coefficients = self._coefficients_normalized + position_normalized = position.xy / self.radius + + for i in range(coefficients.shape[self.axis]): + c = coefficients[{self.axis: i}] + result = result + c * optika.zernikes.zernike( + position=position_normalized, + j=i + 1, + ) + + return result + + def normal( + self, + position: na.AbstractCartesian3dVectorArray, + ) -> na.AbstractCartesian3dVectorArray: + + if self.transformation is not None: + position = self.transformation.inverse(position) + + normal_base = self.base.normal(position) + + gradient_x = normal_base.x / -normal_base.z + gradient_y = normal_base.y / -normal_base.z + + coefficients = self._coefficients_normalized + radius = self.radius + position_normalized = position.xy / radius + + for i in range(coefficients.shape[self.axis]): + c = coefficients[{self.axis: i}] + gradient = optika.zernikes.zernike_gradient( + position=position_normalized, + j=i + 1, + ) + gradient_x = gradient_x + c * gradient.x / radius + gradient_y = gradient_y + c * gradient.y / radius + + norm = np.sqrt(np.square(gradient_x) + np.square(gradient_y) + 1) + + return na.Cartesian3dVectorArray( + x=gradient_x / norm, + y=gradient_y / norm, + z=-1 / norm, + ) diff --git a/optika/zernikes.py b/optika/zernikes.py new file mode 100644 index 00000000..ade45cf2 --- /dev/null +++ b/optika/zernikes.py @@ -0,0 +1,243 @@ +r""" +Zernike polynomials, useful for describing wavefront and figure errors. + +The Zernike polynomials form an orthonormal basis on the unit disk and are +indexed here using Noll's convention :cite:p:`Noll1976`, where each polynomial +is normalized to have unit RMS over the unit disk. +""" + +import math +import numpy as np +import named_arrays as na + +__all__ = [ + "noll", + "zernike", + "zernike_gradient", +] + + +def noll(j: int) -> tuple[int, int]: + r""" + Convert a Noll index :math:`j` into the corresponding Zernike quantum + numbers :math:`(n, m)`. + + Parameters + ---------- + j + The Noll index of the Zernike polynomial. + Must be greater than or equal to one. + + Examples + -------- + + Find the quantum numbers of the first four Zernike polynomials: + piston, :math:`x` tilt, :math:`y` tilt, and defocus. + + .. jupyter-execute:: + + import optika + + [optika.zernikes.noll(j) for j in (1, 2, 3, 4)] + + Notes + ----- + The radial degree :math:`n` and the signed azimuthal degree :math:`m` + follow Noll's ordering :cite:p:`Noll1976`: + even :math:`j` corresponds to the cosine polynomials (:math:`m > 0`) + and odd :math:`j` corresponds to the sine polynomials (:math:`m < 0`). + """ + if j < 1: + raise ValueError(f"Noll index must be greater than zero, got {j=}.") + n = 0 + k = j - 1 + while k > n: + n += 1 + k -= n + m = (-1) ** j * ((n % 2) + 2 * ((k + ((n + 1) % 2)) // 2)) + return n, m + + +def _coefficients_radial(n: int, m: int) -> list[tuple[float, int]]: + r""" + The coefficients and exponents of the radial Zernike polynomial + :math:`R_n^m(\rho) = \sum_k c_k \rho^{e_k}`. + + Parameters + ---------- + n + The radial degree of the Zernike polynomial. + m + The unsigned azimuthal degree of the Zernike polynomial. + """ + result = [] + for k in range((n - m) // 2 + 1): + c = (-1) ** k * math.factorial(n - k) + c = c / math.factorial(k) + c = c / math.factorial((n + m) // 2 - k) + c = c / math.factorial((n - m) // 2 - k) + result.append((c, n - 2 * k)) + return result + + +def zernike( + position: na.AbstractCartesian2dVectorArray, + j: int, +) -> na.AbstractScalar: + r""" + Evaluate the Zernike polynomial with the given Noll index at the given + points on the unit disk. + + Parameters + ---------- + position + The normalized, dimensionless points at which to evaluate the + polynomial. + Points satisfying :math:`|\text{position}| \leq 1` are inside + the unit disk. + j + The Noll index of the Zernike polynomial. + Must be greater than or equal to one. + + Examples + -------- + + Plot the defocus polynomial, :math:`Z_4`, on the unit disk. + + .. jupyter-execute:: + + import numpy as np + import matplotlib.pyplot as plt + import named_arrays as na + import optika + + position = na.Cartesian2dVectorLinearSpace( + start=-1, + stop=1, + axis=na.Cartesian2dVectorArray("x", "y"), + num=101, + ).explicit + + z4 = optika.zernikes.zernike(position, 4) + z4[position.length > 1] = np.nan + + fig, ax = plt.subplots(constrained_layout=True) + na.plt.pcolormesh(position, C=z4, ax=ax) + ax.set_aspect("equal"); + + Notes + ----- + The Zernike polynomial with quantum numbers :math:`(n, m)` is + + .. math:: + + Z_n^m(\rho, \phi) = \begin{cases} + \sqrt{n + 1} \, R_n^0(\rho), & m = 0 \\ + \sqrt{2 (n + 1)} \, R_n^m(\rho) \cos(m \phi), & m > 0 \\ + \sqrt{2 (n + 1)} \, R_n^{|m|}(\rho) \sin(|m| \phi), & m < 0, + \end{cases} + + where the radial polynomial is + + .. math:: + + R_n^m(\rho) = \sum_{k=0}^{(n - m) / 2} + \frac{(-1)^k (n - k)!}{k! \left( \frac{n + m}{2} - k \right)! + \left( \frac{n - m}{2} - k \right)!} \rho^{n - 2 k}. + + The normalization follows :cite:t:`Noll1976`, so that each polynomial has + unit RMS over the unit disk. + """ + n, m = noll(j) + + rho = position.length + phi = np.arctan2(position.y, position.x) + + radial = 0 * rho + for c, e in _coefficients_radial(n, abs(m)): + radial = radial + c * rho**e + + if m == 0: + return np.sqrt(n + 1) * radial + elif m > 0: + return np.sqrt(2 * (n + 1)) * radial * np.cos(m * phi) + else: + return np.sqrt(2 * (n + 1)) * radial * np.sin(-m * phi) + + +def zernike_gradient( + position: na.AbstractCartesian2dVectorArray, + j: int, +) -> na.Cartesian2dVectorArray: + r""" + Evaluate the gradient of the Zernike polynomial with the given Noll index + at the given points on the unit disk. + + Parameters + ---------- + position + The normalized, dimensionless points at which to evaluate the + gradient. + j + The Noll index of the Zernike polynomial. + Must be greater than or equal to one. + + Notes + ----- + The gradient is computed analytically using the chain rule in polar + coordinates, + + .. math:: + + \frac{\partial Z}{\partial x} + &= N \left[ R'(\rho) \, T(m \phi) \cos \phi + - \frac{R(\rho)}{\rho} \, T'(m \phi) \sin \phi \right] \\ + \frac{\partial Z}{\partial y} + &= N \left[ R'(\rho) \, T(m \phi) \sin \phi + + \frac{R(\rho)}{\rho} \, T'(m \phi) \cos \phi \right], + + where :math:`N` is the normalization constant, :math:`R` is the radial + polynomial, and :math:`T` is the azimuthal sinusoid. + Since the lowest-order term of :math:`R_n^m` is :math:`\rho^m`, + the quotient :math:`R(\rho) / \rho` is itself a polynomial whenever + :math:`m \geq 1`, and is evaluated as such to avoid dividing by zero at + the origin. + """ + n, m = noll(j) + mu = abs(m) + + rho = position.length + phi = np.arctan2(position.y, position.x) + cos_phi = np.cos(phi) + sin_phi = np.sin(phi) + + coefficients = _coefficients_radial(n, mu) + + d_radial = 0 * rho + for c, e in coefficients: + if e != 0: + d_radial = d_radial + c * e * rho ** (e - 1) + + if m == 0: + norm = np.sqrt(n + 1) + return na.Cartesian2dVectorArray( + x=norm * d_radial * cos_phi, + y=norm * d_radial * sin_phi, + ) + + radial_over_rho = 0 * rho + for c, e in coefficients: + radial_over_rho = radial_over_rho + c * rho ** (e - 1) + + norm = np.sqrt(2 * (n + 1)) + if m > 0: + t = np.cos(m * phi) + dt = -m * np.sin(m * phi) + else: + t = np.sin(mu * phi) + dt = mu * np.cos(mu * phi) + + return na.Cartesian2dVectorArray( + x=norm * (d_radial * t * cos_phi - radial_over_rho * dt * sin_phi), + y=norm * (d_radial * t * sin_phi + radial_over_rho * dt * cos_phi), + )