Piecewise Even Asphere Surface

[goldengrape]

This piecewise evenAsphere surface is similar to the binary4 surface in Zemax, but it does not include the diffractive component. It can be used to describe complex multi-arc axisymmetric surfaces, such as the concave surface of an orthokeratology lens. Although, in theory, an evenAsphere surface with sufficiently many terms can represent any axisymmetric surface, the piecewise approach offers better interpretability.

I used "Design by Contract" programming, which requires import icontract, but this is not strictly necessary—removing @icontract terms will not affect the program’s execution. Design by Contract programming just enhances the robustness of the code.

# -*- coding: utf-8 -*-
"""
Defines a Piecewise Even Asphere geometry.

This geometry is composed of multiple concentric annular zones, where each zone
is defined by an independent even asphere equation. The implementation
ensures C0 continuity between adjacent zones, meaning the sag values are
equal at the boundaries.

"""

from __future__ import annotations
from typing import List, Optional

import icontract
import matplotlib.pyplot as plt
import numpy as np

import optiland.backend as be
from optiland.coordinate_system import CoordinateSystem
from optiland.geometries.newton_raphson import NewtonRaphsonGeometry


class PiecewiseEvenAsphere(NewtonRaphsonGeometry):
    """
    Represents a piecewise even asphere geometry.

    This surface consists of multiple concentric even asphere zones, with C0
    continuity guaranteed at the boundaries. The base sag for each zone `i`
    is defined by the formula:
    z_i(r) = r^2 / (R_i * (1 + sqrt(1 - (1 + k_i) * r^2 / R_i^2))) + sum(C_ij * r^(2j))

    The final surface sag Z(r) is defined piecewise:
    Z(r) = z_i(r) + offset_i,  for semi_diameter_{i-1} <= r < semi_diameter_i

    The `offset_i` is a constant calculated to ensure C0 continuity.

    Args:
        coordinate_system (CoordinateSystem): The coordinate system of the geometry.
        radii (list[float]): A list of curvature radii (R_i) for each zone.
        semi_diameters (list[float]): A list of outer semi-diameters for each zone.
                                      Must be sorted in ascending order.
        conics (list[float], optional): A list of conic constants (k_i) for each zone.
                                        If None, defaults to a list of zeros.
        coefficients (list[list[float]], optional): A list of aspheric coefficient
                                                    lists (C_ij) for each zone.
                                                    If None, defaults to a list of
                                                    empty lists.
        tol (float, optional): Tolerance for the Newton-Raphson iteration.
                               Defaults to 1e-10.
        max_iter (int, optional): Maximum iterations for Newton-Raphson.
                                  Defaults to 100.
    """

    @icontract.require(
        lambda radii, semi_diameters: len(radii) == len(semi_diameters),
        "The 'radii' and 'semi_diameters' lists must have the same length."
    )
    @icontract.require(
        lambda radii, conics: conics is None or len(radii) == len(conics),
        "If provided, the 'conics' list must have the same length as 'radii'."
    )
    @icontract.require(
        lambda radii, coefficients: coefficients is None or len(radii) == len(coefficients),
        "If provided, the 'coefficients' list must have the same length as 'radii'."
    )
    @icontract.require(
        lambda semi_diameters: len(semi_diameters) > 0,
        "At least one zone must be defined."
    )
    @icontract.require(
        lambda semi_diameters: all(d > 0 for d in semi_diameters),
        "All semi-diameter values must be positive."
    )
    @icontract.require(
        lambda semi_diameters: all(
            semi_diameters[i] < semi_diameters[i+1]
            for i in range(len(semi_diameters) - 1)
        ),
        "The semi-diameters list must be strictly sorted in ascending order."
    )
    def __init__(
        self,
        coordinate_system: CoordinateSystem,
        radii: List[float],
        semi_diameters: List[float],
        conics: Optional[List[float]] = None,
        coefficients: Optional[List[List[float]]] = None,
        tol: float = 1e-10,
        max_iter: int = 100,
    ):
        num_zones = len(radii)

        # Handle optional arguments by providing sensible defaults.
        if conics is None:
            processed_conics = [0.0] * num_zones
        else:
            processed_conics = conics

        if coefficients is None:
            processed_coefficients = [[] for _ in range(num_zones)]
        else:
            processed_coefficients = coefficients

        # The parent class requires a base radius of curvature; we use the first zone's.
        super().__init__(coordinate_system, radii[0], processed_conics[0], tol, max_iter)

        self.radii = radii
        self.conics = processed_conics
        self.coefficients = processed_coefficients
        self.semi_diameters = semi_diameters
        self.num_zones = num_zones
        
        self.is_symmetric = True
        self.order = 2

        self._offsets = self._calculate_offsets()

    def _calculate_offsets(self) -> list[float]:
        """Calculates the sag offset for each zone to ensure C0 continuity."""
        offsets = [0.0] * self.num_zones
        for i in range(1, self.num_zones):
            boundary_r2 = self.semi_diameters[i - 1] ** 2
            sag_inner = (
                self._zone_sag_formula(boundary_r2, self.radii[i - 1], self.conics[i - 1], self.coefficients[i - 1])
                + offsets[i - 1]
            )
            sag_outer = self._zone_sag_formula(boundary_r2, self.radii[i], self.conics[i], self.coefficients[i])
            offsets[i] = sag_inner - sag_outer
        return offsets

    @staticmethod
    def _zone_sag_formula(r2: be.ndarray, R: float, k: float, C: list[float]) -> be.ndarray:
        """Calculates the sag of a single even asphere."""
        sqrt_arg = 1 - (1 + k) * r2 / R**2
        z = r2 / (R * (1 + be.sqrt(be.maximum(0, sqrt_arg))))
        for i, Ci in enumerate(C):
            z += Ci * r2 ** (i + 1)
        return z

    @staticmethod
    def _zone_derivative_r_div_r(r2: be.ndarray, R: float, k: float, C: list[float]) -> be.ndarray:
        """Calculates the derivative of sag w.r.t. r, divided by r (df/dr / r)."""
        sqrt_arg = 1 - (1 + k) * r2 / R**2
        denom = R * be.sqrt(be.maximum(1e-12, sqrt_arg))
        dfdr_div_r = 1 / denom
        for i, Ci in enumerate(C):
            dfdr_div_r += 2 * (i + 1) * Ci * r2**i
        return dfdr_div_r

    def sag(self, x: be.ndarray = 0.0, y: be.ndarray = 0.0) -> be.ndarray:
        """Calculates the surface sag at the given (x, y) coordinates."""
        r2 = x**2 + y**2
        r = be.sqrt(r2)
        conditions, choices = [], []
        
        conditions.append(r < self.semi_diameters[0])
        choices.append(
            self._zone_sag_formula(r2, self.radii[0], self.conics[0], self.coefficients[0])
            + self._offsets[0]
        )
        for i in range(1, self.num_zones):
            conditions.append(
                (r >= self.semi_diameters[i - 1]) & (r < self.semi_diameters[i])
            )
            choices.append(
                self._zone_sag_formula(r2, self.radii[i], self.conics[i], self.coefficients[i])
                + self._offsets[i]
            )
        
        default_choice = (
            self._zone_sag_formula(r2, self.radii[-1], self.conics[-1], self.coefficients[-1])
            + self._offsets[-1]
        )
        return be.select(conditions, choices, default=default_choice)

    def _surface_normal(self, x: be.ndarray, y: be.ndarray) -> tuple[be.ndarray, be.ndarray, be.ndarray]:
        """Calculates the surface normal vector at the given (x, y) coordinates."""
        r2 = x**2 + y**2
        r = be.where(r2 == 0, 1e-12, be.sqrt(r2))
        conditions, choices_deriv = [], []

        conditions.append(r < self.semi_diameters[0])
        choices_deriv.append(
            self._zone_derivative_r_div_r(r2, self.radii[0], self.conics[0], self.coefficients[0])
        )
        for i in range(1, self.num_zones):
            conditions.append(
                (r >= self.semi_diameters[i - 1]) & (r < self.semi_diameters[i])
            )
            choices_deriv.append(
                self._zone_derivative_r_div_r(r2, self.radii[i], self.conics[i], self.coefficients[i])
            )
        
        default_deriv = self._zone_derivative_r_div_r(r2, self.radii[-1], self.conics[-1], self.coefficients[-1])
        dfdr_div_r = be.select(conditions, choices_deriv, default=default_deriv)
        
        dfdx = x * dfdr_div_r
        dfdy = y * dfdr_div_r

        mag = be.sqrt(dfdx**2 + dfdy**2 + 1)
        nx, ny, nz = -dfdx / mag, -dfdy / mag, 1 / mag
        return nx, ny, nz

    def to_dict(self) -> dict:
        """Converts the geometry to a dictionary representation."""
        return {
            "type": "PiecewiseEvenAsphere",
            "coordinate_system": self.coordinate_system.to_dict(),
            "radii": self.radii,
            "semi_diameters": self.semi_diameters,
            "conics": self.conics,
            "coefficients": self.coefficients,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "PiecewiseEvenAsphere":
        """Creates a geometry instance from a dictionary representation."""
        if data.get("type") != "PiecewiseEvenAsphere":
            raise ValueError("Dictionary type does not match PiecewiseEvenAsphere.")

        cs = CoordinateSystem.from_dict(data["coordinate_system"])
        return cls(
            coordinate_system=cs,
            radii=data["radii"],
            semi_diameters=data["semi_diameters"],
            conics=data.get("conics"),
            coefficients=data.get("coefficients"),
        )

# --- Verification Example ---
if __name__ == '__main__':
    cs = CoordinateSystem()
    plt.style.use('seaborn-v0_8-whitegrid')

    # --- Example 1: Full definition with all parameters ---
    print("--- Running Example 1: Full Definition ---")
    try:
        p_asphere = PiecewiseEvenAsphere(
            coordinate_system=cs,
            radii=[50.0, 150.0, 1e9],
            semi_diameters=[10.0, 20.0, 30.0],
            conics=[0.0, -1.0, 0.0],
            coefficients=[[0.0, 1e-7], [], [0.0, 5e-8]]
        )
        print("Instance created successfully.")
        print(f"Offsets: {p_asphere._offsets}")

        max_diam = p_asphere.semi_diameters[-1]
        r_coords = np.linspace(-max_diam, max_diam, 500)
        sag_values = p_asphere.sag(x=be.array(r_coords), y=0.0)

        fig, ax = plt.subplots(figsize=(12, 7))
        ax.plot(r_coords, sag_values, label='Piecewise Asphere Sag', color='royalblue', linewidth=2)
        for i, diam in enumerate(p_asphere.semi_diameters):
            ax.axvline(x=diam, color='darkorange', linestyle='--', linewidth=1.2, label=f'Boundary {i+1}' if i==0 else None)
            ax.axvline(x=-diam, color='darkorange', linestyle='--', linewidth=1.2)
        
        ax.set_title('Example 1: Cross-Section of Fully Defined Asphere', fontsize=16)
        ax.set_xlabel('Radial Coordinate r (mm)', fontsize=12)
        ax.set_ylabel('Sag z (mm)', fontsize=12)
        ax.legend()
        ax.grid(True)
        ax.set_aspect('equal', adjustable='box')
        plt.show()

    except Exception as e:
        print(f"\nAn error occurred in Example 1: {e}")


    # --- Example 2: Simplified definition, omitting optional arguments ---
    print("\n--- Running Example 2: Omitting Optional Arguments ---")
    try:
        # Define a surface made of two simple spherical zones
        p_sphere = PiecewiseEvenAsphere(
            coordinate_system=cs,
            radii=[80.0, -120.0],
            semi_diameters=[15.0, 25.0]
            # conics and coefficients are omitted
        )
        print("Instance created successfully by omitting optional arguments.")
        print(f"Defaulted Conics: {p_sphere.conics}")
        print(f"Defaulted Coefficients: {p_sphere.coefficients}")
        print(f"Calculated Offsets: {p_sphere._offsets}")

        max_diam_sphere = p_sphere.semi_diameters[-1]
        r_coords_sphere = np.linspace(-max_diam_sphere, max_diam_sphere, 500)
        sag_values_sphere = p_sphere.sag(x=be.array(r_coords_sphere), y=0.0)

        fig, ax = plt.subplots(figsize=(12, 7))
        ax.plot(r_coords_sphere, sag_values_sphere, label='Piecewise Sphere Sag', color='seagreen', linewidth=2)
        for i, diam in enumerate(p_sphere.semi_diameters):
            ax.axvline(x=diam, color='darkorange', linestyle='--', linewidth=1.2, label=f'Boundary {i+1}' if i==0 else None)
            ax.axvline(x=-diam, color='darkorange', linestyle='--', linewidth=1.2)

        ax.set_title('Example 2: Cross-Section of Piecewise Sphere (Simplified Def.)', fontsize=16)
        ax.set_xlabel('Radial Coordinate r (mm)', fontsize=12)
        ax.set_ylabel('Sag z (mm)', fontsize=12)
        ax.legend()
        ax.grid(True)
        ax.set_aspect('equal', adjustable='box')
        plt.show()

    except Exception as e:
        print(f"\nAn error occurred in Example 2: {e}")

[HarrisonKramer]

Hi @goldengrape,

Thanks for sharing this! A PiecewiseEvenAsphere is a great idea and would be very useful.

This would be a welcome addition to the library if you're interested in opening a pull request. If you do, we might just need to tweak a couple of things like the flip() method and the serialization to fully align with the base classes. I am not necessarily opposed to using icontract, although I am generally hesitant to add new dependencies unless they're truly required.

I'd be happy to help out here if you want to collaborate on a PR. I could also take over eventually, but will have to add it to the backlog for now.

Thanks again,
Kramer