Repository: MatthewFilipovich/torchoptics

Issue:
Add aberration profiles
- Zernike Polynomials
- Seidel aberrations

Discussion:

Solution patch:
diff --git a/torchoptics/profiles/__init__.py b/torchoptics/profiles/__init__.py
index d71d773..26bde53 100644
--- a/torchoptics/profiles/__init__.py
+++ b/torchoptics/profiles/__init__.py
@@ -8,3 +8,4 @@ from .lens import lens
 from .shapes import checkerboard, circle, rectangle, square, triangle
 from .spatial_coherence import gaussian_schell_model, schell_model
 from .special import airy, sinc
+from .zernike import zernike
diff --git a/torchoptics/profiles/zernike.py b/torchoptics/profiles/zernike.py
new file mode 100644
index 0000000..07bc4c9
--- /dev/null
+++ b/torchoptics/profiles/zernike.py
@@ -0,0 +1,91 @@
+"""This module defines functions to generate Zernike polynomial profiles."""
+
+import math
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from ..planar_geometry import PlanarGeometry
+from ..type_defs import Int, Scalar, Vector2
+from ..utils import initialize_tensor
+
+__all__ = ["zernike"]
+
+
+def zernike(
+    shape: Vector2,
+    n: Int,
+    m: Int,
+    radius: Scalar,
+    spacing: Optional[Vector2] = None,
+    offset: Optional[Vector2] = None,
+) -> Tensor:
+    r"""
+    Generates a Zernike polynomial profile over a 2D grid.
+
+    `Zernike polynomials <https://en.wikipedia.org/wiki/Zernike_polynomials>`_ are a set of orthogonal
+    polynomials defined on the unit disk, commonly used in optics for wavefront analysis and aberration
+    characterization. They are expressed in polar coordinates :math:`(\rho, \theta)` as:
+
+    .. math::
+        Z_n^m(\rho, \theta) = R_n^m(\rho) \exp(i m \theta)
+
+    where:
+
+        - :math:`n` is the radial order (non-negative integer),
+        - :math:`m` is the azimuthal index (integer satisfying :math:`|m| \leq n` and :math:`n - |m|` even),
+        - :math:`R_n^m(\rho)` is the radial polynomial,
+        - :math:`\rho` is the normalized radial coordinate, and
+        - :math:`\theta` is the angular coordinate.
+
+    The radial polynomial :math:`R_n^m(\rho)` is computed using:
+
+    .. 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-2k}
+
+    Args:
+        shape (Vector2): Number of grid points along the planar dimensions.
+        n (int): The radial order of the Zernike polynomial.
+        m (int): The azimuthal order of the Zernike polynomial.
+        radius (Scalar): The radius of the unit disk.
+        spacing (Optional[Vector2]): Distance between grid points along planar dimensions. Default: if
+            `None`, uses a global default.
+        offset (Optional[Vector2]): Center coordinates of the profile. Default: `(0, 0)`.
+
+    Returns:
+        Tensor: The generated Zernike polynomial profile.
+    """
+
+    n = initialize_tensor("n", n, is_scalar=True, is_integer=True, is_non_negative=True)
+    m = initialize_tensor("m", m, is_scalar=True, is_integer=True)
+    radius = initialize_tensor("radius", radius, is_scalar=True, is_positive=True)
+
+    if m.abs() > n:
+        raise ValueError("Azimuthal index m must satisfy |m| <= n.")
+    if (n - m.abs()) % 2 != 0:
+        raise ValueError("The difference n - |m| must be even for a valid Zernike polynomial.")
+
+    x, y = PlanarGeometry(shape, spacing=spacing, offset=offset).meshgrid()
+
+    rho = torch.sqrt(x**2 + y**2) / radius
+    theta = torch.atan2(y, x)
+
+    radial_poly = zernike_radial(n, m.abs(), rho)
+    angular_component = torch.cos(m * theta) if m >= 0 else torch.sin(m.abs() * theta)
+
+    # Apply the unit disk mask
+    mask = rho <= 1.0
+    zernike_profile = radial_poly * angular_component * mask
+
+    return zernike_profile
+
+
+def zernike_radial(n: Tensor, m: Tensor, rho: Tensor) -> Tensor:
+    """Computes the radial component of the Zernike polynomial."""
+    radial = torch.zeros_like(rho)
+    for k in range((n - m.abs()) // 2 + 1):
+        coeff = (-1) ** k * math.comb(n - k, k) * math.comb(n - 2 * k, (n - m.abs()) // 2 - k)
+        radial += coeff * rho ** (n - 2 * k)
+    return radial


Test patch:
diff --git a/tests/test_profiles.py b/tests/test_profiles.py
index 16a0f94..8bdcace 100644
--- a/tests/test_profiles.py
+++ b/tests/test_profiles.py
@@ -393,5 +393,35 @@ class TestBesselProfile(unittest.TestCase):
         self.assertTrue(torch.all(self.profile.abs() <= 1))
 
 
+class TestZernikeProfile(unittest.TestCase):
+    def setUp(self):
+        self.shape = (100, 100)
+        self.radius = 50.0
+        self.spacing = (1.0, 1.0)
+        self.offset = (0.0, 0.0)
+        self.n = 3
+        self.m = 1
+
+    def test_zernike_profile(self):
+        profile = zernike(
+            shape=self.shape,
+            n=self.n,
+            m=self.m,
+            radius=self.radius,
+            spacing=self.spacing,
+            offset=self.offset,
+        )
+        self.assertEqual(profile.shape, self.shape)
+        self.assertFalse(torch.is_complex(profile))
+        self.assertTrue(torch.all(profile >= -1))
+        self.assertTrue(torch.all(profile <= 1))
+
+    def test_invalid_zernike_parameters(self):
+        with self.assertRaises(ValueError):
+            zernike(shape=self.shape, n=2, m=3, radius=self.radius, spacing=self.spacing, offset=self.offset)
+        with self.assertRaises(ValueError):
+            zernike(shape=self.shape, n=3, m=2, radius=self.radius, spacing=self.spacing, offset=self.offset)
+
+
 if __name__ == "__main__":
     unittest.main()