Source code for geometries.standard

"""Standard Geometry

The Standard geometry represents a surface defined by a sphere or conic in two
dimensions. The surface is defined as:

z = r^2 / (R * (1 + sqrt(1 - (1 + k) * r^2 / R^2)))

where
- r^2 = x^2 + y^2
- R is the radius of curvature
- k is the conic constant

Kramer Harrison, 2024
"""

from __future__ import annotations

import warnings

import optiland.backend as be
from optiland.coordinate_system import CoordinateSystem
from optiland.geometries.base import BaseGeometry


def _is_radius_infinite(radius):
    """Checks if the given radius represents an infinite radius (a plane)."""
    is_inf_tensor = be.isinf(radius)
    if hasattr(is_inf_tensor, "ndim") and is_inf_tensor.ndim > 0:
        return bool(be.all(is_inf_tensor))
    return (
        bool(is_inf_tensor.item())
        if hasattr(is_inf_tensor, "item")
        else bool(is_inf_tensor)
    )


[docs] class StandardGeometry(BaseGeometry): """Represents a standard geometry with a given coordinate system, radius, and conic. Args: coordinate_system (CoordinateSystem): The coordinate system of the geometry. radius (float): The radius of curvature of the geometry. conic (float, optional): The conic constant of the geometry. Defaults to 0.0. Methods: sag(x=0, y=0): Calculates the surface sag of the geometry at the given coordinates. distance(rays): Finds the propagation distance to the geometry for the given rays. surface_normal(rays): Calculates the surface normal of the geometry at the given ray positions. """ def __init__(self, coordinate_system, radius, conic=0.0): super().__init__(coordinate_system) self.radius = be.array(radius) self.k = be.array(conic) self.is_symmetric = True def __str__(self): return "Standard"
[docs] def flip(self): """Flip the geometry. Changes the sign of the radius of curvature. The conic constant remains unchanged. """ self.radius = -self.radius
[docs] def scale(self, scale_factor: float): """Scale the geometry parameters. Args: scale_factor (float): The factor by which to scale the geometry. """ self.radius = self.radius * scale_factor
[docs] def sag(self, x=0, y=0): """Calculate the surface sag of the geometry at the given coordinates. Args: x (float or be.ndarray, optional): The x-coordinate(s). Defaults to 0. y (float or be.ndarray, optional): The y-coordinate(s). Defaults to 0. Returns: be.ndarray or float: The sag value(s) at the given coordinates. """ r2 = x**2 + y**2 return r2 / ( self.radius * (1 + be.sqrt(1 - (1 + self.k) * r2 / self.radius**2)) )
[docs] def distance(self, rays): """Find the propagation distance to the geometry for the given rays. Args: rays (RealRays): The rays for which to calculate the distance. Returns: be.ndarray: An array of distances from each ray's current position to its intersection point with the geometry. """ if _is_radius_infinite(self.radius): # intersection with the plane z=0 is z0 + t*Nz = 0 N_safe = be.where(be.abs(rays.N) > 1e-14, rays.N, 1e-14) return -rays.z / N_safe a = self.k * rays.N**2 + rays.L**2 + rays.M**2 + rays.N**2 b = ( 2 * self.k * rays.N * rays.z + 2 * rays.L * rays.x + 2 * rays.M * rays.y - 2 * rays.N * self.radius + 2 * rays.N * rays.z ) c = ( self.k * rays.z**2 - 2 * self.radius * rays.z + rays.x**2 + rays.y**2 + rays.z**2 ) # discriminant d = b**2 - 4 * a * c # two solutions for distance to conic with warnings.catch_warnings(): warnings.simplefilter("ignore") t1 = (-b + be.sqrt(d)) / (2 * a) t2 = (-b - be.sqrt(d)) / (2 * a) # find intersection points in z z1 = rays.z + t1 * rays.N z2 = rays.z + t2 * rays.N # take intersection closest to z = 0 (i.e., vertex of geometry) t = be.where(be.abs(z1) <= be.abs(z2), t1, t2) # handle case when a = 0 # Assumes b is not zero when a is zero, based on original logic. t = be.where(a == 0, -c / b, t) return t
[docs] def surface_normal(self, rays): """Calculate the surface normal of the geometry at the given points. Args: rays (RealRays): The rays, positioned at the surface, for which to calculate the surface normals. Returns: tuple[be.ndarray, be.ndarray, be.ndarray]: The x, y, and z components of the surface normal vectors. """ r2 = rays.x**2 + rays.y**2 denom = self.radius * be.sqrt(1 - (1 + self.k) * r2 / self.radius**2) dfdx = rays.x / denom dfdy = rays.y / denom dfdz = -1 mag = be.sqrt(dfdx**2 + dfdy**2 + dfdz**2) nx = dfdx / mag ny = dfdy / mag nz = dfdz / mag return nx, ny, nz
[docs] def to_dict(self): """Convert the geometry to a dictionary. Returns: dict: The dictionary representation of the geometry. """ geometry_dict = super().to_dict() geometry_dict.update({"radius": float(self.radius), "conic": float(self.k)}) return geometry_dict
[docs] @classmethod def from_dict(cls, data): """Create a geometry from a dictionary. Args: data (dict): The dictionary representation of the geometry. Returns: StandardGeometry: An instance of StandardGeometry. """ required_keys = {"cs", "radius"} if not required_keys.issubset(data): missing = required_keys - data.keys() raise ValueError(f"Missing required keys: {missing}") cs = CoordinateSystem.from_dict(data["cs"]) return cls(cs, data["radius"], data.get("conic", 0.0))