diff --git a/docs/refs.bib b/docs/refs.bib index a70bae3..efea534 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 3f2766f..67ab71b 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 0000000..ca644df --- /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/sags/__init__.py b/optika/sags/__init__.py index 4241548..d0c4d81 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 0000000..bff7532 --- /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 0000000..0e6c5d6 --- /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 0000000..ade45cf --- /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), + )