Grid Sag surface

[goldengrape]

Grid Sag surface.

untested

# -*- coding: utf-8 -*-
"""
Grid Sag Surface for Optiland Framework
=======================================

This module provides an implementation of the Zemax Grid Sag surface,
designed for integration into the `optiland` optical simulation framework.

Design Date: September 30, 2025
"""
from __future__ import annotations

import io
from dataclasses import dataclass
from typing import TextIO, Type

import icontract
import numpy as np
from scipy.interpolate import RectBivariateSpline
import optiland.backend as be
from optiland.coordinate_system import CoordinateSystem
# from optiland.geometries.newton_raphson import NewtonRaphsonGeometry
from optiland.geometries import EvenAsphere, NewtonRaphsonGeometry


@icontract.invariant(lambda self: self.nx > 0 and self.ny > 0)
@icontract.invariant(lambda self: self.delx > 0 and self.dely > 0)
@icontract.invariant(lambda self: self.z_grid.shape == (self.ny, self.nx))
@dataclass(frozen=True)
class GridSagData:
    """An immutable data class to store and validate data parsed from a Grid Sag file."""
    nx: int; ny: int; delx: float; dely: float; x_decenter: float
    y_decenter: float; z_grid: np.ndarray; raw_content: str
    x_coords: np.ndarray; y_coords: np.ndarray
    @staticmethod
    def from_stream(file_stream: TextIO) -> 'GridSagData':
        content = file_stream.read()
        lines = [line for line in content.splitlines() if line.strip() and not line.strip().startswith('!')]
        if not lines: raise ValueError("File is empty or contains only comments.")
        header_parts = lines[0].strip().split()
        if len(header_parts) != 7: raise ValueError(f"Header must contain 7 values, but found {len(header_parts)}.")
        nx, ny, delx, dely, unit_flag, xdec, ydec = header_parts
        nx, ny = int(nx), int(ny)
        delx, dely, xdec, ydec = float(delx), float(dely), float(xdec), float(ydec)
        scaler = {0: 1.0, 1: 10.0, 2: 25.4, 3: 1000.0}.get(int(unit_flag))
        if scaler is None: raise ValueError(f"Unknown unit flag: {unit_flag}.")
        delx, dely, xdec, ydec = delx*scaler, dely*scaler, xdec*scaler, ydec*scaler
        data_lines = lines[1:]
        if len(data_lines) != nx * ny: raise ValueError(f"Expected {nx*ny} data lines, but found {len(data_lines)}.")
        raw_data = np.loadtxt(io.StringIO("\n".join(data_lines)))
        if raw_data.ndim == 1: raw_data = raw_data.reshape(1, -1)
        if raw_data.shape[1] == 5 and np.any(raw_data[:, 4] == 1): raw_data[raw_data[:, 4] == 1, :4] = 0.0
        z_values = (raw_data[:, 0] * scaler).reshape((ny, nx))
        x_min = xdec - (nx - 1) / 2.0 * delx
        y_max = ydec + (ny - 1) / 2.0 * dely
        x_coords = np.linspace(x_min, x_min + (nx - 1) * delx, nx)
        y_coords = np.linspace(y_max, y_max - (ny - 1) * dely, ny)
        return GridSagData(nx=nx, ny=ny, delx=delx, dely=dely, x_decenter=xdec, y_decenter=ydec,
                           z_grid=z_values, raw_content=content, x_coords=x_coords, y_coords=y_coords)

class GridSag(NewtonRaphsonGeometry):
    """Represents a Grid Sag surface, integrated with the optiland framework."""
    def __init__(self, base_surface: NewtonRaphsonGeometry, grid_data: GridSagData):
        super().__init__(base_surface.cs, base_surface.radius, base_surface.k)
        self.base_surface = base_surface
        self.grid_data = grid_data
        kx = min(3, self.grid_data.nx - 1)
        ky = min(3, self.grid_data.ny - 1)
        if kx < 1 or ky < 1: raise ValueError(f"Grid size is too small for interpolation.")
        self._z_interpolator = RectBivariateSpline(
            self.grid_data.y_coords[::-1], self.grid_data.x_coords,
            self.grid_data.z_grid[::-1, :], kx=kx, ky=ky
        )
        self._x_min, self._x_max = self.grid_data.x_coords[0], self.grid_data.x_coords[-1]
        self._y_min, self._y_max = self.grid_data.y_coords[-1], self.grid_data.y_coords[0]

    def __str__(self): return f"Grid Sag on {self.base_surface}"
        
    def sag(self, x=0, y=0):
        was_scalar = np.isscalar(x)
        x_arr, y_arr = np.atleast_1d(x), np.atleast_1d(y)
        z_base = self.base_surface.sag(x_arr, y_arr)
        in_bounds = (x_arr >= self._x_min) & (x_arr <= self._x_max) & \
                    (y_arr >= self._y_min) & (y_arr <= self._y_max)
        z_grid = np.zeros_like(x_arr, dtype=float)
        if np.any(in_bounds):
            z_grid[in_bounds] = self._z_interpolator(y_arr[in_bounds], x_arr[in_bounds], grid=False)
        total_z = z_base + z_grid
        return total_z.item() if was_scalar else total_z

    def _surface_normal(self, x, y):
        was_scalar = np.isscalar(x)
        x_arr, y_arr = np.atleast_1d(x), np.atleast_1d(y)
        nx_b, ny_b, nz_b = self.base_surface._surface_normal(x_arr, y_arr)
        nz_b[nz_b == 0] = 1e-12
        dfdx_b, dfdy_b = -nx_b / nz_b, -ny_b / nz_b
        in_bounds = (x_arr >= self._x_min) & (x_arr <= self._x_max) & \
                    (y_arr >= self._y_min) & (y_arr <= self._y_max)
        dfdx_g = np.zeros_like(x_arr, dtype=float)
        dfdy_g = np.zeros_like(x_arr, dtype=float)
        if np.any(in_bounds):
            dfdx_g[in_bounds] = self._z_interpolator(y_arr[in_bounds], x_arr[in_bounds], dx=1, grid=False)
            dfdy_g[in_bounds] = self._z_interpolator(y_arr[in_bounds], x_arr[in_bounds], dy=1, grid=False)
        dfdx, dy = dfdx_b + dfdx_g, dfdy_b + dfdy_g
        mag = np.sqrt(dfdx**2 + dy**2 + 1)
        mag[mag == 0] = 1e-12
        nx, ny, nz = dfdx / mag, dy / mag, -1 / mag
        if was_scalar:
            return nx.item(), ny.item(), nz.item()
        return nx, ny, nz

    def to_dict(self):
        base_type_name = type(self.base_surface).__name__
        if base_type_name not in BASE_SURFACE_TYPES:
            raise TypeError(f"Base surface type '{base_type_name}' is not registered for serialization.")
        return {
            "type": "GridSag", "base_surface_type": base_type_name,
            "base_surface_data": self.base_surface.to_dict(),
            "grid_data_content": self.grid_data.raw_content
        }

    @classmethod
    def from_dict(cls, data: dict):
        if data.get("type") != "GridSag": raise ValueError("Dictionary type is not 'GridSag'.")
        base_type_name = data["base_surface_type"]
        base_surface_class = BASE_SURFACE_TYPES.get(base_type_name)
        if base_surface_class is None: raise TypeError(f"Unknown base surface type '{base_type_name}'.")
        base_surface = base_surface_class.from_dict(data["base_surface_data"])
        with io.StringIO(data["grid_data_content"]) as f:
            grid_data = GridSagData.from_stream(f)
        return cls(base_surface, grid_data)

BASE_SURFACE_TYPES: dict[str, Type[NewtonRaphsonGeometry]] = {
    "EvenAsphere": EvenAsphere
}

# --- Demo ---
if __name__ == '__main__':
    print("--- GridSag Surface (Scipy Interpolation & Framework Integration) Demo ---")

    # --- FIX: Use a correctly formatted multi-line string for the mock data ---
    mock_dat_content = """
! Example Grid Sag File
! 3x3 grid, 10mm spacing, 0.5mm bump at the center
3 3 10.0 10.0 0 0.0 0.0
0.0  0.0  0.0  0.0  0
0.0  0.0  0.0  0.0  0
0.0  0.0  0.0  0.0  0
0.0  0.0  0.0  0.0  0
0.5  0.0  0.0  0.0  0
0.0  0.0  0.0  0.0  0
0.0  0.0  0.0  0.0  0
0.0  0.0  0.0  0.0  0
0.0  0.0  0.0  0.0  0
"""
    
    cs = CoordinateSystem()
    base_sphere = EvenAsphere(coordinate_system=cs, radius=-200.0)
    
    try:
        # The parser now directly processes the correctly formatted string
        with io.StringIO(mock_dat_content) as f:
            grid_data_obj = GridSagData.from_stream(f)
        
        grid_surface = GridSag(base_sphere, grid_data_obj)
        print(f"Successfully created surface: {grid_surface}")

        print("\n--- Functionality Tests ---")
        x, y = 5.0, 5.0
        print(f"Total sag at ({x}, {y}) [scalar input]: {grid_surface.sag(x, y):.6f}")

        x_arr, y_arr = np.array([0., 5., 50.]), np.array([0., 5., 50.])
        sag_arr = grid_surface.sag(x_arr, y_arr)
        print(f"Total sag at multiple points [array input]: {[f'{z:.6f}' for z in sag_arr]}")

        print("\n--- Serialization Test ---")
        surface_dict = grid_surface.to_dict()
        print(f"Serialized dictionary: {{'type': '{surface_dict['type']}', 'base_surface_type': '{surface_dict['base_surface_type']}', ...}}")
        
        reconstructed_surface = GridSag.from_dict(surface_dict)
        print("Deserialization successful:", reconstructed_surface)
        
        assert np.isclose(grid_surface.sag(5.0, 5.0), reconstructed_surface.sag(5.0, 5.0))
        print("Serialization round-trip test passed.")

    except Exception as e:
        print(f"\nAn error occurred while creating or using the GridSag surface: {e}")

[HarrisonKramer]

Thanks for another solid proposal! I definitely think this is something we can add to the package. I will add this to the backlog.

Kramer