GRIN shell surface

[goldengrape]

I attempted to use a piecewise even asphere to assist in constructing a discretized gradient-index (GRIN) surface. This is only an initial trial, and I have not yet conducted any testing or validation.

# -*- coding: utf-8 -*-
"""
GRIN Medium and Discretized Lens Modeling System

This module provides a comprehensive system for defining a gradient-index (GRIN)
medium, calculating the precise geometry of its iso-index surfaces, and
constructing a discretized shell-based model of a GRIN lens.

The system is architected based on the following principles:

- Data-Oriented Programming:
  - `PiecewiseEvenAsphere`: Core is the list of geometric zone parameters.
  - `GrinProfile`: Purely encapsulates the coefficients defining the index profile.
  - The `GrinLensModel` orchestrates these data-centric objects.

- Functional Programming:
  - Core computational methods like `sag`, `normal`, and the quadratic solver for
    `IsoIndexSurface` are implemented as pure, stateless functions. Their output
    depends solely on their input parameters, ensuring predictability and testability.

- Axiomatic Design:
  - FR1: Define a physical surface geometry (`PiecewiseEvenAsphere`).
  - FR2: Define a GRIN refractive index profile (`GrinProfile`).
  - FR3: Represent a surface of constant index (`IsoIndexSurface`).
  - FR4: Assemble a discretized lens model (`GrinLensModel`).
  These functional requirements (FRs) are implemented in independent, decoupled
  modules (classes), ensuring maintainability and clarity.

- Design by Contract (DbC):
  - The `icontract` library is used to enforce preconditions and class invariants,
    such as ensuring that parameters for the quadratic solver are valid. This
    improves robustness and makes the code's assumptions explicit.

- Occam's Razor:
  - The `IsoIndexSurface` directly implements the analytical solution of the GRIN
    equation. This is the simplest and most direct correct approach, avoiding
    unnecessary complexity or approximation at the foundational level.
"""
from __future__ import annotations
from typing import List, Optional, Literal, Tuple


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

# ==============================================================================
# SECTION 1: USER-PROVIDED PIECEWISE EVEN ASPHERE CLASS
# ==============================================================================

class PiecewiseEvenAsphere(NewtonRaphsonGeometry):
    """
    Represents a piecewise even asphere geometry. (Implementation provided by user)
    """
    @icontract.require(lambda radii, semi_diameters: len(radii) == len(semi_diameters))
    @icontract.require(lambda radii, conics: conics is None or len(radii) == len(conics))
    @icontract.require(lambda radii, coefficients: coefficients is None or len(radii) == len(coefficients))
    @icontract.require(lambda semi_diameters: len(semi_diameters) > 0)
    @icontract.require(lambda semi_diameters: all(d > 0 for d in semi_diameters))
    @icontract.require(
        lambda semi_diameters: all(
            semi_diameters[i] < semi_diameters[i+1]
            for i in range(len(semi_diameters) - 1)
        )
    )
    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)
        processed_conics = [0.0] * num_zones if conics is None else conics
        processed_coefficients = [[] for _ in range(num_zones)] if coefficients is None else coefficients
        
        super().__init__(coordinate_system, radii[0], processed_conics[0], tol, max_iter)
        
        # --- FIX: Store coordinate_system as an instance attribute ---
        self.coordinate_system = coordinate_system
        # -----------------------------------------------------------

        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]:
        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:
        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:
        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:
        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]:
        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, dfdy = x * dfdr_div_r, y * dfdr_div_r
        mag = be.sqrt(dfdx**2 + dfdy**2 + 1)
        return -dfdx / mag, -dfdy / mag, 1 / mag
        
# ==============================================================================
# SECTION 2: GRIN PROFILE AND ISO-INDEX SURFACE DEFINITION
# ==============================================================================

@icontract.invariant(lambda self: self.nz2 != 0 or self.nz1 != 0, "Either nz1 or nz2 must be non-zero for z-dependence.")
class GrinProfile:
    """
    Encapsulates the data for a gradient-index (GRIN) medium profile.

    The refractive index n is defined by the formula:
    n(r, z) = n0 + nr2*r^2 + nr4*r^4 + nr6*r^6 + nz1*z + nz2*z^2
    where r^2 = x^2 + y^2.
    """
    def __init__(
        self, n0: float, nr2: float = 0.0, nr4: float = 0.0,
        nr6: float = 0.0, nz1: float = 0.0, nz2: float = 0.0
    ):
        self.n0, self.nr2, self.nr4 = n0, nr2, nr4
        self.nr6, self.nz1, self.nz2 = nr6, nz1, nz2

    def n(self, r: be.ndarray, z: be.ndarray) -> be.ndarray:
        """Calculates refractive index at specified (r, z) coordinates."""
        r2 = r**2
        return self.n0 + self.nr2*r2 + self.nr4*r2**2 + self.nr6*r2**3 + self.nz1*z + self.nz2*z**2

class IsoIndexSurface(NewtonRaphsonGeometry):
    """
    Represents a surface of constant refractive index within a GRIN medium.
    Its shape is the analytical solution to n(r, z) = n_iso.
    """
    @icontract.require(lambda profile: profile.nz2 != 0, "Quadratic solution requires nz2 to be non-zero.")
    def __init__(
        self,
        coordinate_system: CoordinateSystem,
        profile: GrinProfile,
        n_iso: float,
        branch: Literal['positive', 'negative'],
        tol: float = 1e-10,
        max_iter: int = 100
    ):
        super().__init__(coordinate_system, radius=1e9, conic=0.0, tol=tol, max_iter=max_iter)
        self.coordinate_system = coordinate_system
        self.profile = profile
        self.n_iso = n_iso
        self.branch = branch
        self.is_symmetric = False
        self.order = 6

    def _solve_quadratic_for_z(self, r2: be.ndarray) -> be.ndarray:
        """Solves the GRIN equation for z, returning the sag."""
        p = self.profile
        a, b = p.nz2, p.nz1
        c = (p.n0 - self.n_iso) + p.nr2 * r2 + p.nr4 * r2**2 + p.nr6 * r2**3
        discriminant = b**2 - 4 * a * c
        valid_discriminant = be.where(discriminant >= 0, discriminant, be.nan)
        sqrt_discriminant = be.sqrt(valid_discriminant)
        
        if self.branch == 'positive':
            return (-b + sqrt_discriminant) / (2 * a)
        return (-b - sqrt_discriminant) / (2 * a)

    def sag(self, x: be.ndarray = 0.0, y: be.ndarray = 0.0) -> be.ndarray:
        return self._solve_quadratic_for_z(x**2 + y**2)

    def _surface_normal(self, x: be.ndarray, y: be.ndarray) -> Tuple[be.ndarray, be.ndarray, be.ndarray]:
        r2, z = x**2 + y**2, self.sag(x, y)
        p = self.profile
        dn_dr_div_r = 2 * p.nr2 + 4 * p.nr4 * r2 + 6 * p.nr6 * r2**2
        dfdx, dfdy = x * dn_dr_div_r, y * dn_dr_div_r
        dfdz = p.nz1 + 2 * p.nz2 * z
        mag = be.sqrt(dfdx**2 + dfdy**2 + dfdz**2)
        
        is_invalid = be.isnan(z) | (mag == 0)
        safe_mag = be.where(is_invalid, 1.0, mag)
        nx, ny, nz = -dfdx / safe_mag, -dfdy / safe_mag, -dfdz / safe_mag
        
        return (be.where(is_invalid, be.nan, nx),
                be.where(is_invalid, be.nan, ny),
                be.where(is_invalid, be.nan, nz))

# ==============================================================================
# SECTION 3: DISCRETIZED GRIN LENS MODEL
# ==============================================================================

class GrinLensModel:
    """
    Constructs a discretized GRIN lens model from a physical front surface,
    a GRIN profile, and geometric constraints.
    
    The model consists of a series of iso-index surfaces that approximate the
    continuous gradient.
    """
    @icontract.require(lambda num_shells: num_shells > 0)
    def __init__(
        self,
        front_surface: PiecewiseEvenAsphere,
        grin_profile: GrinProfile,
        num_shells: int,
    ):
        self.front_surface = front_surface
        self.profile = grin_profile
        self.num_shells = num_shells
        self.iso_surfaces: List[Tuple[IsoIndexSurface, IsoIndexSurface]] = []
        self.shell_indices: List[float] = []
        
        self._build_model()

    def _build_model(self):
        """Generates the iso-index surfaces and shell index values."""
        max_radius = self.front_surface.semi_diameters[-1]
        
        z_at_vertex = self.front_surface.sag(0, 0)
        z_at_edge = self.front_surface.sag(max_radius, 0)
        n_max = self.profile.n(be.array(0.0), be.array(z_at_vertex))
        n_min = self.profile.n(be.array(max_radius), be.array(z_at_edge))

        n_boundaries = np.linspace(n_max, n_min, self.num_shells + 1)
        self.shell_indices = [(n_boundaries[i] + n_boundaries[i+1]) / 2.0 for i in range(self.num_shells)]
        
        print(f"Index range: [{n_max:.4f}, {n_min:.4f}]")
        print(f"Shell boundary indices: {[f'{n:.4f}' for n in n_boundaries]}")
        print(f"Average shell indices: {[f'{n:.4f}' for n in self.shell_indices]}")

        cs = self.front_surface.coordinate_system
        for n_iso in n_boundaries:
            pos_branch = IsoIndexSurface(cs, self.profile, n_iso, 'positive')
            neg_branch = IsoIndexSurface(cs, self.profile, n_iso, 'negative')
            self.iso_surfaces.append((pos_branch, neg_branch))

    def visualize_2d_cross_section(self):
        """Plots a 2D cross-section of the GRIN lens model."""
        max_radius = self.front_surface.semi_diameters[-1]
        r_coords = np.linspace(-max_radius, max_radius, 501)

        fig, ax = plt.subplots(figsize=(12, 8))

        front_sag = self.front_surface.sag(x=be.array(r_coords))
        ax.plot(r_coords, front_sag, 'k-', linewidth=2.5, label='Front Surface')

        colors = plt.cm.viridis(np.linspace(0, 1, len(self.iso_surfaces)))
        for i, (pos_surf, neg_surf) in enumerate(self.iso_surfaces):
            pos_sag = pos_surf.sag(x=be.array(r_coords))
            neg_sag = neg_surf.sag(x=be.array(r_coords))
            label = f'n = {pos_surf.n_iso:.4f}'
            ax.plot(r_coords, pos_sag, color=colors[i], linestyle='--', label=label)
            ax.plot(r_coords, neg_sag, color=colors[i], linestyle='--')
        
        ax.set_title('2D Cross-Section of Discretized GRIN Lens Model', fontsize=16)
        ax.set_xlabel('Radial Coordinate r (mm)', fontsize=12)
        ax.set_ylabel('Sag z (mm)', fontsize=12)
        ax.legend(title='Surface / Index Boundary')
        ax.grid(True)
        ax.set_aspect('equal', adjustable='box')
        plt.show()

# ==============================================================================
# SECTION 4: VERIFICATION EXAMPLE
# ==============================================================================

if __name__ == '__main__':
    cs = CoordinateSystem(x=0, y=0, z=0, rx=0, ry=0, rz=0, reference_cs=None)
    plt.style.use('seaborn-v0_8-whitegrid')

    # --- Step 1: Define the physical front surface of the lens ---
    front_surface = PiecewiseEvenAsphere(
        coordinate_system=cs,
        radii=[100.0, 150.0, 500.0],
        semi_diameters=[10.0, 20.0, 30.0],
        conics=[0.0, -0.5, -1.0],
        coefficients=[[], [1e-7], []]
    )

    # --- Step 2: Define the GRIN profile within the material ---
    grin_profile = GrinProfile(
        n0=1.6,
        nr2=-5.0e-4,
        nr4=-1.0e-8,
        nz1=1.0e-2,
        nz2=-2.0e-3
    )

    # --- Step 3: Create the discretized GRIN lens model ---
    print("--- Building GRIN Lens Model ---")
    grin_lens = GrinLensModel(
        front_surface=front_surface,
        grin_profile=grin_profile,
        num_shells=5
    )
    print("\n--- Model Construction Complete ---")

    # --- Step 4: Visualize the result ---
    grin_lens.visualize_2d_cross_section()

[goldengrape]

Attempt unsuccessful.
The shell-based GRIN lens model fails under sequential ray tracing due to aperture limitations. Specifically, rays are unable to traverse the outer regions of a small-aperture lens to reach the next interface. Discretizing the refractive index confines the GRIN region to a small ellipsoid. While a nominal surface can be constructed for this core, it blocks peripheral ray paths, severely restricting the system’s effective aperture.
Result: complete failure of the design concept.