Surfaces API Reference

The src.geometric_surface module contains all geometric optical surface types. Each surface implements differentiable ray intersection and Snell's law refraction.

Base Class

Surface

Base class for all geometric surfaces. Implements Newton's method intersection (with one differentiable step) and vector Snell's law refraction.

src.geometric_surface.Surface

Surface(r, d, mat2, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: DeepObj

Base class for all geometric optical surfaces.

A surface sits at axial position d (mm) in the global coordinate system, has an aperture radius r (mm), and separates two optical media. Subclasses override :meth:_sag and :meth:_dfdxy to define their shape.

Ray–surface interaction is handled by three stages, implemented in :meth:ray_reaction:

Coordinate transform – ray is brought into the local surface frame.
Intersection – solved via Newton's method (:meth:newtons_method), using a non-differentiable iteration loop followed by a single differentiable Newton step to enable gradient flow.
Refraction / reflection – vector Snell's law (:meth:refract) or specular reflection (:meth:reflect).

Attributes:

Name	Type	Description
`d`	`Tensor`	Axial position of the surface vertex [mm].
`r`	`float`	Aperture radius [mm].
`mat2`	`Material`	Optical material on the transmission side.
`is_square`	`bool`	If `True` the aperture is square; otherwise circular.
`tolerancing`	`bool`	When `True`, manufacturing error offsets are applied during ray tracing.

Initialize a generic optical surface.

Parameters:

Name	Type	Description	Default
`r`	`float`	Aperture radius [mm].	required
`d`	`float`	Axial position of the surface vertex [mm].	required
`mat2`	`str or Material`	Material on the transmission side (e.g. `"N-BK7"`, `"air"`).	required
`pos_xy`	`list[float]`	Lateral offset `[x, y]` [mm]. Defaults to `[0.0, 0.0]`.	`[0.0, 0.0]`
`vec_local`	`list[float]`	Local normal direction. Defaults to `[0.0, 0.0, 1.0]` (on-axis).	`[0.0, 0.0, 1.0]`
`is_square`	`bool`	Use a square aperture. Defaults to `False`.	`False`
`device`	`str`	Compute device. Defaults to `"cpu"`.	`'cpu'`

Source code in src/geometric_surface/base.py

def __init__(
    self,
    r,
    d,
    mat2,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize a generic optical surface.

    Args:
        r (float): Aperture radius [mm].
        d (float): Axial position of the surface vertex [mm].
        mat2 (str or Material): Material on the transmission side
            (e.g. ``"N-BK7"``, ``"air"``).
        pos_xy (list[float], optional): Lateral offset ``[x, y]`` [mm].
            Defaults to ``[0.0, 0.0]``.
        vec_local (list[float], optional): Local normal direction.
            Defaults to ``[0.0, 0.0, 1.0]`` (on-axis).
        is_square (bool, optional): Use a square aperture.
            Defaults to ``False``.
        device (str, optional): Compute device. Defaults to ``"cpu"``.
    """
    super(Surface, self).__init__()

    # Global direction vector, always pointing to the positive z-axis
    self.vec_global = torch.tensor([0.0, 0.0, 1.0])

    # Surface position in global coordinate system
    self.d = torch.tensor(d)
    self.pos_x = torch.tensor(pos_xy[0])
    self.pos_y = torch.tensor(pos_xy[1])

    # Surface direction vector in global coordinate system
    self.vec_local = F.normalize(torch.tensor(vec_local), p=2, dim=-1)

    # Material after the surface
    self.mat2 = Material(mat2)

    # Surface aperture radius (non-differentiable)
    self.r = float(r)
    self.is_square = is_square
    if is_square:
        # r is the incircle radius
        self.h = 2 * self.r
        self.w = 2 * self.r

    # Newton method parameters (matched to HPC kernel settings)
    self.newton_maxiter = 8  # [int], maximum number of Newton iterations
    self.newton_convergence = 50.0 * 1e-6  # [mm], Newton method convergence threshold
    self.newton_step_bound = 5.0  # [mm], maximum step size in each iteration

    self.tolerancing = False
    self._R_tilt = None
    self._R_tilt_inv = None
    self.device = device if device is not None else torch.device("cpu")
    self.to(self.device)

    # Pre-compute rotation matrices (depends only on static vec_local/vec_global)
    self._cache_rotation_matrices()

init_from_dict `classmethod`

init_from_dict(surf_dict)

Initialize surface from a dict.

Source code in src/geometric_surface/base.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize surface from a dict."""
    raise NotImplementedError(
        f"init_from_dict() is not implemented for {cls.__name__}."
    )

ray_reaction

ray_reaction(ray, n1, n2, refraction=True)

Compute the output ray after intersection and refraction/reflection.

Transforms the ray to the local surface frame, solves the intersection via Newton's method, applies vector Snell's law (or specular reflection), then transforms back to global coordinates.

When tolerancing is active, mat2_n_error is added to n2 to simulate refractive-index manufacturing error.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	Incident ray bundle.	required
`n1`	`float`	Refractive index of the incident medium.	required
`n2`	`float`	Refractive index of the transmission medium.	required
`refraction`	`bool`	If `True` (default) refract the ray; if `False` reflect it.	`True`

Returns:

Name	Type	Description
`Ray`		Updated ray bundle after the surface interaction.

Source code in src/geometric_surface/base.py

def ray_reaction(self, ray, n1, n2, refraction=True):
    """Compute the output ray after intersection and refraction/reflection.

    Transforms the ray to the local surface frame, solves the intersection
    via Newton's method, applies vector Snell's law (or specular reflection),
    then transforms back to global coordinates.

    When tolerancing is active, ``mat2_n_error`` is added to ``n2`` to
    simulate refractive-index manufacturing error.

    Args:
        ray (Ray): Incident ray bundle.
        n1 (float): Refractive index of the incident medium.
        n2 (float): Refractive index of the transmission medium.
        refraction (bool, optional): If ``True`` (default) refract the ray;
            if ``False`` reflect it.

    Returns:
        Ray: Updated ray bundle after the surface interaction.
    """
    # Transform ray to local coordinate system
    ray = self.to_local_coord(ray)

    # Intersection
    ray = self.intersect(ray, n1)

    if refraction:
        # Apply refractive index tolerance error
        if self.tolerancing:
            n2 = n2 + self.mat2_n_error
        ray = self.refract(ray, n1 / n2)
    else:
        # Reflection
        ray = self.reflect(ray)

    # Transform ray to global coordinate system
    ray = self.to_global_coord(ray)

    return ray

intersect

intersect(ray, n=1.0)

Solve ray-surface intersection in local coordinate system.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	input ray.	required
`n`	`float`	refractive index. Defaults to 1.0.	`1.0`

Source code in src/geometric_surface/base.py

def intersect(self, ray, n=1.0):
    """Solve ray-surface intersection in local coordinate system.

    Args:
        ray (Ray): input ray.
        n (float, optional): refractive index. Defaults to 1.0.
    """
    # Solve ray-surface intersection time by Newton's method
    t, valid = self.newtons_method(ray)

    # Update ray
    new_o = ray.o + ray.d * t.unsqueeze(-1)
    ray.o = torch.where(valid.unsqueeze(-1), new_o, ray.o)
    ray.is_valid = ray.is_valid * valid

    if ray.coherent:
        if t.abs().max() > 100 and torch.get_default_dtype() == torch.float32:
            raise Exception(
                "Using float32 may cause precision problem for OPL calculation."
            )
        new_opl = ray.opl + n * t.unsqueeze(-1)
        ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)

    return ray

newtons_method

newtons_method(ray)

Solve intersection by Newton's method in local coordinate system.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	input ray.	required

Returns:

Name	Type	Description
`t`	`tensor`	intersection time.
`valid`	`tensor`	valid mask.

Source code in src/geometric_surface/base.py

def newtons_method(self, ray):
    """Solve intersection by Newton's method in local coordinate system.

    Args:
        ray (Ray): input ray.

    Returns:
        t (tensor): intersection time.
        valid (tensor): valid mask.
    """
    newton_maxiter = self.newton_maxiter
    newton_convergence = self.newton_convergence
    newton_step_bound = self.newton_step_bound

    # Ray direction components (reused across iterations)
    dxdt, dydt, dzdt = ray.d[..., 0], ray.d[..., 1], ray.d[..., 2]

    # Initial guess of t (can also use spherical surface for initial guess)
    t = -ray.o[..., 2] / dzdt

    # 1. Non-differentiable Newton's iterations to find the intersection
    #    Run (maxiter - 1) iterations; the differentiable step below acts as
    #    the final iteration while also enabling gradient flow.
    with torch.no_grad():
        for _ in range(newton_maxiter - 1):
            new_o = ray.o + ray.d * t.unsqueeze(-1)
            new_x, new_y = new_o[..., 0], new_o[..., 1]
            valid = self.is_within_data_range(new_x, new_y) & (ray.is_valid > 0)

            x, y = new_x * valid, new_y * valid
            ft = self._sag(x, y) - new_o[..., 2]
            dfdx, dfdy = self._dfdxy(x, y)
            dfdt = dfdx * dxdt + dfdy * dydt - dzdt
            t = t - torch.clamp(
                ft / (dfdt + EPSILON), -newton_step_bound, newton_step_bound
            )

    # 2. One differentiable Newton step (final iteration + gradient flow)
    new_o = ray.o + ray.d * t.unsqueeze(-1)
    new_x, new_y = new_o[..., 0], new_o[..., 1]
    valid = self.is_valid(new_x, new_y) & (ray.is_valid > 0)

    x, y = new_x * valid, new_y * valid
    ft = self._sag(x, y) - new_o[..., 2]
    dfdx, dfdy = self._dfdxy(x, y)
    dfdt = dfdx * dxdt + dfdy * dydt - dzdt
    t = t - torch.clamp(
        ft / (dfdt + EPSILON), -newton_step_bound, newton_step_bound
    )

    # 3. Determine valid solutions — reuse ft and valid from the diff step
    with torch.no_grad():
        valid = valid & (ft.abs() < newton_convergence)

    return t, valid

refract

refract(ray, eta)

Calculate refracted ray according to Snell's law in local coordinate system.

Normal vector points from the surface toward the side where the light is coming from. d is already normalized if both n and ray.d are normalized.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	incident ray.	required
`eta`	`float`	ratio of indices of refraction, eta = n_i / n_t	required

Returns:

Name	Type	Description
`ray`	`Ray`	refracted ray.

References

[1] https://registry.khronos.org/OpenGL-Refpages/gl4/html/refract.xhtml [2] https://en.wikipedia.org/wiki/Snell%27s_law, "Vector form" section.

Source code in src/geometric_surface/base.py

def refract(self, ray, eta):
    """Calculate refracted ray according to Snell's law in local coordinate system.

    Normal vector points from the surface toward the side where the light is coming from. d is already normalized if both n and ray.d are normalized.

    Args:
        ray (Ray): incident ray.
        eta (float): ratio of indices of refraction, eta = n_i / n_t

    Returns:
        ray (Ray): refracted ray.

    References:
        [1] https://registry.khronos.org/OpenGL-Refpages/gl4/html/refract.xhtml
        [2] https://en.wikipedia.org/wiki/Snell%27s_law, "Vector form" section.
    """
    # Compute normal vectors
    normal_vec = self.normal_vec(ray)

    # Compute refraction according to Snell's law, normal_vec * ray_d
    dot_product = (-normal_vec * ray.d).sum(-1).unsqueeze(-1)
    k = 1 - eta**2 * (1 - dot_product**2)

    # Total internal reflection
    valid = (k >= 0).squeeze(-1) & (ray.is_valid > 0)
    k = k * valid.unsqueeze(-1)

    # Update ray direction and obliquity
    new_d = eta * ray.d + (eta * dot_product - torch.sqrt(k + EPSILON)) * normal_vec
    # ==> Update obliq term to penalize steep rays in the later optimization.
    # obliq = cos(angle between old and new ray direction).
    # Threshold 0.9 (~25° refraction) activates for moderately steep
    # refractions, providing gradient signal before surfaces become extreme.
    obliq = torch.sum(new_d * ray.d, axis=-1).unsqueeze(-1)
    obliq_update_mask = valid.unsqueeze(-1) & (obliq < 0.9)
    ray.obliq = torch.where(obliq_update_mask, obliq * ray.obliq, ray.obliq)
    # ==>
    ray.d = torch.where(valid.unsqueeze(-1), new_d, ray.d)

    # Update ray valid mask
    ray.is_valid = ray.is_valid * valid

    return ray

reflect

reflect(ray)

Calculate reflected ray in local coordinate system.

Normal vector points from the surface toward the side where the light is coming from.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	incident ray.	required

Returns:

Name	Type	Description
`ray`	`Ray`	reflected ray.

References

[1] https://registry.khronos.org/OpenGL-Refpages/gl4/html/reflect.xhtml [2] https://en.wikipedia.org/wiki/Snell%27s_law, "Vector form" section.

Source code in src/geometric_surface/base.py

def reflect(self, ray):
    """Calculate reflected ray in local coordinate system.

    Normal vector points from the surface toward the side where the light is coming from.

    Args:
        ray (Ray): incident ray.

    Returns:
        ray (Ray): reflected ray.

    References:
        [1] https://registry.khronos.org/OpenGL-Refpages/gl4/html/reflect.xhtml
        [2] https://en.wikipedia.org/wiki/Snell%27s_law, "Vector form" section.
    """
    # Compute surface normal vectors
    normal_vec = self.normal_vec(ray)

    # Reflect
    dot_product = (normal_vec * ray.d).sum(-1).unsqueeze(-1)
    new_d = ray.d - 2 * dot_product * normal_vec
    new_d = F.normalize(new_d, p=2, dim=-1)

    # Update valid rays
    valid_mask = ray.is_valid > 0
    ray.d = torch.where(valid_mask.unsqueeze(-1), new_d, ray.d)

    return ray

normal_vec

normal_vec(ray)

Calculate surface normal vector at the intersection point in local coordinate system.

Normal vector points from the surface toward the side where the light is coming from.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	input ray.	required

Returns:

Name	Type	Description
`n_vec`	`tensor`	surface normal vector.

Source code in src/geometric_surface/base.py

def normal_vec(self, ray):
    """Calculate surface normal vector at the intersection point in local coordinate system.

    Normal vector points from the surface toward the side where the light is coming from.

    Args:
        ray (Ray): input ray.

    Returns:
        n_vec (tensor): surface normal vector.
    """
    x, y = ray.o[..., 0], ray.o[..., 1]
    nx, ny, nz = self.dfdxyz(x, y)
    n_vec = torch.stack((nx, ny, nz), axis=-1)
    n_vec = F.normalize(n_vec, p=2, dim=-1)

    is_forward = ray.d[..., 2].unsqueeze(-1) > 0
    n_vec = torch.where(is_forward, n_vec, -n_vec)
    return n_vec

to_local_coord

to_local_coord(ray)

Transform ray to local coordinate system.

When tolerancing is active, applies manufacturing error perturbations: d_error (axial shift), decenter_x/y_error (lateral shift), and tilt_error (rotation about the x-axis).

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	input ray in global coordinate system.	required

Returns:

Name	Type	Description
`ray`	`Ray`	transformed ray in local coordinate system.

Source code in src/geometric_surface/base.py

def to_local_coord(self, ray):
    """Transform ray to local coordinate system.

    When tolerancing is active, applies manufacturing error perturbations:
    d_error (axial shift), decenter_x/y_error (lateral shift), and
    tilt_error (rotation about the x-axis).

    Args:
        ray (Ray): input ray in global coordinate system.

    Returns:
        ray (Ray): transformed ray in local coordinate system.
    """
    # Shift ray origin to surface origin (with tolerance perturbations)
    if self.tolerancing:
        offset = torch.stack([
            self.pos_x + self.decenter_x_error,
            self.pos_y + self.decenter_y_error,
            self.d + self.d_error,
        ]).expand_as(ray.o)
    else:
        offset = torch.stack([self.pos_x, self.pos_y, self.d]).expand_as(ray.o)
    ray.o = ray.o - offset

    # Apply tilt rotation (tolerance-induced, using cached matrix)
    if self._R_tilt is not None:
        ray.o = self._apply_rotation(ray.o, self._R_tilt)
        ray.d = self._apply_rotation(ray.d, self._R_tilt)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    # Rotate ray origin and direction (using cached matrix for nominal orientation)
    if self._R_to_local is not None:
        ray.o = self._apply_rotation(ray.o, self._R_to_local)
        ray.d = self._apply_rotation(ray.d, self._R_to_local)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    return ray

to_global_coord

to_global_coord(ray)

Transform ray to global coordinate system.

When tolerancing is active, reverses the manufacturing error perturbations applied in :meth:to_local_coord.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	input ray in local coordinate system.	required

Returns:

Name	Type	Description
`ray`	`Ray`	transformed ray in global coordinate system.

Source code in src/geometric_surface/base.py

def to_global_coord(self, ray):
    """Transform ray to global coordinate system.

    When tolerancing is active, reverses the manufacturing error
    perturbations applied in :meth:`to_local_coord`.

    Args:
        ray (Ray): input ray in local coordinate system.

    Returns:
        ray (Ray): transformed ray in global coordinate system.
    """
    # Rotate ray origin and direction (using cached matrix for nominal orientation)
    if self._R_to_global is not None:
        ray.o = self._apply_rotation(ray.o, self._R_to_global)
        ray.d = self._apply_rotation(ray.d, self._R_to_global)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    # Reverse tilt rotation (tolerance-induced, using cached inverse matrix)
    if self._R_tilt_inv is not None:
        ray.o = self._apply_rotation(ray.o, self._R_tilt_inv)
        ray.d = self._apply_rotation(ray.d, self._R_tilt_inv)
        ray.d = F.normalize(ray.d, p=2, dim=-1)

    # Shift ray origin back to global coordinates (with tolerance perturbations)
    if self.tolerancing:
        offset = torch.stack([
            self.pos_x + self.decenter_x_error,
            self.pos_y + self.decenter_y_error,
            self.d + self.d_error,
        ]).expand_as(ray.o)
    else:
        offset = torch.stack([self.pos_x, self.pos_y, self.d]).expand_as(ray.o)
    ray.o = ray.o + offset

    return ray

sag

sag(x, y, valid=None)

Calculate sag (z) of the surface: z = f(x, y).

Valid term is used to avoid NaN when x, y exceed the data range, which happens in spherical and aspherical surfaces.

Calculating r = sqrt(x2, y2) may cause an NaN error during back-propagation. Because dr/dx = x / sqrt(x2 + y2), NaN will occur when x=y=0.

Source code in src/geometric_surface/base.py

def sag(self, x, y, valid=None):
    """Calculate sag (z) of the surface: z = f(x, y).

    Valid term is used to avoid NaN when x, y exceed the data range, which happens in spherical and aspherical surfaces.

    Calculating r = sqrt(x**2, y**2) may cause an NaN error during back-propagation. Because dr/dx = x / sqrt(x**2 + y**2), NaN will occur when x=y=0.
    """
    if valid is None:
        valid = self.is_valid(x, y)

    x, y = x * valid, y * valid
    return self._sag(x, y)

dfdxyz

dfdxyz(x, y, valid=None)

Compute derivatives of surface function. Surface function: f(x, y, z): sag(x, y) - z = 0. This function is used in Newton's method and normal vector calculation.

There are several methods to compute derivatives of surfaces

[1] Analytical derivatives: The current implementation is based on this method. But the implementation only works for surfaces which can be written as z = sag(x, y). For implicit surfaces, we need to compute derivatives (df/dx, df/dy, df/dz). [2] Numerical derivatives: Use finite difference method to compute derivatives. This can be used for those very complex surfaces, for example, NURBS. But it may suffer from numerical instability when the surface is very steep. [3] Automatic differentiation: Use torch.autograd to compute derivatives. This can work for almost all the surfaces and is accurate, but it requires an extra backward pass to compute the derivatives of the surface function.

Source code in src/geometric_surface/base.py

def dfdxyz(self, x, y, valid=None):
    """Compute derivatives of surface function. Surface function: f(x, y, z): sag(x, y) - z = 0. This function is used in Newton's method and normal vector calculation.

    There are several methods to compute derivatives of surfaces:
        [1] Analytical derivatives: The current implementation is based on this method. But the implementation only works for surfaces which can be written as z = sag(x, y). For implicit surfaces, we need to compute derivatives (df/dx, df/dy, df/dz).
        [2] Numerical derivatives: Use finite difference method to compute derivatives. This can be used for those very complex surfaces, for example, NURBS. But it may suffer from numerical instability when the surface is very steep.
        [3] Automatic differentiation: Use torch.autograd to compute derivatives. This can work for almost all the surfaces and is accurate, but it requires an extra backward pass to compute the derivatives of the surface function.
    """
    if valid is None:
        valid = self.is_valid(x, y)

    x, y = x * valid, y * valid
    dx, dy = self._dfdxy(x, y)
    return dx, dy, -torch.ones_like(x)

d2fdxyz2

d2fdxyz2(x, y, valid=None)

Compute second-order partial derivatives of the surface function f(x, y, z): sag(x, y) - z = 0. This function is currently only used for surfaces constraints.

Source code in src/geometric_surface/base.py

def d2fdxyz2(self, x, y, valid=None):
    """Compute second-order partial derivatives of the surface function f(x, y, z): sag(x, y) - z = 0. This function is currently only used for surfaces constraints."""
    if valid is None:
        valid = self.is_within_data_range(x, y)

    x, y = x * valid, y * valid

    # Compute second-order derivatives of sag(x, y)
    d2f_dx2, d2f_dxdy, d2f_dy2 = self._d2fdxy(x, y)

    # Mixed partial derivatives involving z are zero
    zeros = torch.zeros_like(x)
    d2f_dxdz = zeros  # ∂²f/∂x∂z = 0
    d2f_dydz = zeros  # ∂²f/∂y∂z = 0
    d2f_dz2 = zeros  # ∂²f/∂z² = 0

    return d2f_dx2, d2f_dxdy, d2f_dy2, d2f_dxdz, d2f_dydz, d2f_dz2

is_valid

is_valid(x, y)

Valid points within the data range and boundary of the surface.

Source code in src/geometric_surface/base.py

def is_valid(self, x, y):
    """Valid points within the data range and boundary of the surface."""
    return self.is_within_data_range(x, y) & self.is_within_boundary(x, y)

is_within_boundary

is_within_boundary(x, y)

Valid points within the boundary of the surface.

Source code in src/geometric_surface/base.py

def is_within_boundary(self, x, y):
    """Valid points within the boundary of the surface."""
    if self.is_square:
        valid = (torch.abs(x) <= (self.w / 2 + EPSILON)) & (
            torch.abs(y) <= (self.h / 2 + EPSILON)
        )
    else:
        if self.tolerancing:
            r = self.r + self.r_error
        else:
            r = self.r
        valid = (x**2 + y**2) <= (r**2 + EPSILON)

    return valid

is_within_data_range

is_within_data_range(x, y)

Valid points inside the data region of the sag function.

Source code in src/geometric_surface/base.py

def is_within_data_range(self, x, y):
    """Valid points inside the data region of the sag function."""
    return torch.ones_like(x, dtype=torch.bool)

max_height

max_height()

Maximum valid height.

Source code in src/geometric_surface/base.py

def max_height(self):
    """Maximum valid height."""
    return 10e3

surface_with_offset

surface_with_offset(x, y, valid_check=True)

Calculate z coordinate of the surface at (x, y).

This function is used in lens setup plotting and lens self-intersection detection.

Source code in src/geometric_surface/base.py

def surface_with_offset(self, x, y, valid_check=True):
    """Calculate z coordinate of the surface at (x, y).

    This function is used in lens setup plotting and lens self-intersection detection.
    """
    x = x if torch.is_tensor(x) else torch.tensor(x).to(self.device)
    y = y if torch.is_tensor(y) else torch.tensor(y).to(self.device)
    if valid_check:
        return self.sag(x, y) + self.d
    else:
        return self._sag(x, y) + self.d

surface_sag

surface_sag(x, y)

Calculate sag of the surface at (x, y).

This function is currently not used.

Source code in src/geometric_surface/base.py

def surface_sag(self, x, y):
    """Calculate sag of the surface at (x, y).

    This function is currently not used.
    """
    x = x if torch.is_tensor(x) else torch.tensor(x).to(self.device)
    y = y if torch.is_tensor(y) else torch.tensor(y).to(self.device)
    return self.sag(x, y).item()

get_optimizer_params

get_optimizer_params(lrs=[0.0001], optim_mat=False)

Get optimizer parameters for different parameters.

Parameters:

Name	Type	Description	Default
`lrs`	`list`	learning rates for different parameters.	`[0.0001]`
`optim_mat`	`bool`	whether to optimize material. Defaults to False.	`False`

Source code in src/geometric_surface/base.py

def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Get optimizer parameters for different parameters.

    Args:
        lrs (list): learning rates for different parameters.
        optim_mat (bool): whether to optimize material. Defaults to False.
    """
    raise NotImplementedError(
        "get_optimizer_params() is not implemented for {}".format(
            self.__class__.__name__
        )
    )

get_optimizer

get_optimizer(lrs=[0.0001], optim_mat=False)

Get optimizer for the surface.

Source code in src/geometric_surface/base.py

def get_optimizer(self, lrs=[1e-4], optim_mat=False):
    """Get optimizer for the surface."""
    params = self.get_optimizer_params(lrs, optim_mat=optim_mat)
    return torch.optim.Adam(params)

update_r

update_r(r)

Update surface radius.

Source code in src/geometric_surface/base.py

def update_r(self, r):
    """Update surface radius."""
    r_max = self.max_height()
    self.r = min(r, r_max)

init_tolerance

init_tolerance(tolerance_params=None)

Initialize tolerance parameters for the surface.

Parameters:

Name Type Description Default

tolerance_params

dict or None

Tolerance for surface parameters. Supported keys (all optional, default values shown):

.. code-block:: python

{
    "r_tole": 0.05,          # aperture radius [mm]
    "d_tole": 0.05,          # axial position [mm]
    "decenter_tole": 0.1,    # lateral decentre x & y [mm]
    "tilt_tole": 0.1,        # tilt [arcmin]
    "mat2_n_tole": 0.001,    # refractive index
}

None

References

[1] https://www.edmundoptics.com/knowledge-center/application-notes/optics/understanding-optical-specifications/ [2] https://wp.optics.arizona.edu/optomech/wp-content/uploads/sites/53/2016/08/8-Tolerancing-1.pdf [3] https://wp.optics.arizona.edu/jsasian/wp-content/uploads/sites/33/2016/03/L17_OPTI517_Lens-_Tolerancing.pdf

Source code in src/geometric_surface/base.py

def init_tolerance(self, tolerance_params=None):
    """Initialize tolerance parameters for the surface.

    Args:
        tolerance_params (dict or None): Tolerance for surface parameters.
            Supported keys (all optional, default values shown):

            .. code-block:: python

                {
                    "r_tole": 0.05,          # aperture radius [mm]
                    "d_tole": 0.05,          # axial position [mm]
                    "decenter_tole": 0.1,    # lateral decentre x & y [mm]
                    "tilt_tole": 0.1,        # tilt [arcmin]
                    "mat2_n_tole": 0.001,    # refractive index
                }

    References:
        [1] https://www.edmundoptics.com/knowledge-center/application-notes/optics/understanding-optical-specifications/
        [2] https://wp.optics.arizona.edu/optomech/wp-content/uploads/sites/53/2016/08/8-Tolerancing-1.pdf
        [3] https://wp.optics.arizona.edu/jsasian/wp-content/uploads/sites/33/2016/03/L17_OPTI517_Lens-_Tolerancing.pdf
    """
    if tolerance_params is None:
        tolerance_params = {}

    # Tolerance ranges
    self.r_tole = tolerance_params.get("r_tole", 0.05)
    self.d_tole = tolerance_params.get("d_tole", 0.05)
    self.decenter_tole = tolerance_params.get("decenter_tole", 0.1)
    self.tilt_tole = tolerance_params.get("tilt_tole", 0.1)
    self.mat2_n_tole = tolerance_params.get("mat2_n_tole", 0.001)

    # Initialize error values to zero (set to random values by sample_tolerance)
    self.r_error = 0.0
    self.d_error = 0.0
    self.decenter_x_error = 0.0
    self.decenter_y_error = 0.0
    self.tilt_error = 0.0
    self.mat2_n_error = 0.0
    # Cached tilt rotation matrices (populated by sample_tolerance)
    self._R_tilt = None
    self._R_tilt_inv = None

sample_tolerance

sample_tolerance()

Sample one set of random manufacturing errors for the surface.

Error distributions

r_error: Uniform[-r_tole, 0] (aperture only shrinks).
d_error: Normal(0, d_tole) axial position shift [mm].
decenter_x/y_error: Normal(0, decenter_tole) lateral shift [mm].
tilt_error: Normal(0, tilt_tole) tilt about x-axis [arcmin → rad].
mat2_n_error: Normal(0, mat2_n_tole) refractive index offset.

Source code in src/geometric_surface/base.py

@torch.no_grad()
def sample_tolerance(self):
    """Sample one set of random manufacturing errors for the surface.

    Error distributions:
        - r_error: Uniform[-r_tole, 0] (aperture only shrinks).
        - d_error: Normal(0, d_tole) axial position shift [mm].
        - decenter_x/y_error: Normal(0, decenter_tole) lateral shift [mm].
        - tilt_error: Normal(0, tilt_tole) tilt about x-axis [arcmin → rad].
        - mat2_n_error: Normal(0, mat2_n_tole) refractive index offset.
    """
    self.r_error = float(np.random.uniform(-self.r_tole, 0))  # [mm]
    self.d_error = float(np.random.randn() * self.d_tole)  # [mm]
    self.decenter_x_error = float(np.random.randn() * self.decenter_tole)  # [mm]
    self.decenter_y_error = float(np.random.randn() * self.decenter_tole)  # [mm]
    tilt_arcmin = float(np.random.randn() * self.tilt_tole)  # [arcmin]
    self.tilt_error = tilt_arcmin / 60.0 * np.pi / 180.0  # [rad]
    self.mat2_n_error = float(np.random.randn() * self.mat2_n_tole)

    # Cache tilt rotation matrices to avoid per-call tensor allocation
    if abs(self.tilt_error) > 1e-12:
        self._R_tilt = self._tilt_rotation_matrix(self.tilt_error, self.device)
        self._R_tilt_inv = self._tilt_rotation_matrix(-self.tilt_error, self.device)
    else:
        self._R_tilt = None
        self._R_tilt_inv = None

    self.tolerancing = True

zero_tolerance

zero_tolerance()

Reset all manufacturing errors to zero (nominal state).

Source code in src/geometric_surface/base.py

def zero_tolerance(self):
    """Reset all manufacturing errors to zero (nominal state)."""
    self.r_error = 0.0
    self.d_error = 0.0
    self.decenter_x_error = 0.0
    self.decenter_y_error = 0.0
    self.tilt_error = 0.0
    self.mat2_n_error = 0.0
    self._R_tilt = None
    self._R_tilt_inv = None
    self.tolerancing = False

sensitivity_score

sensitivity_score()

Compute first-order tolerance sensitivity scores via RSS formula.

For each parameter with a gradient, the score is: tolerance_range² × gradient², which approximates the variance of the loss contribution from that parameter's manufacturing error.

Returns:

Name	Type	Description
`dict`		Sensitivity gradients and RSS scores keyed by surface index.

Reference

[1] Page 10 from: https://wp.optics.arizona.edu/optomech/wp-content/uploads/sites/53/2016/08/8-Tolerancing-1.pdf

Source code in src/geometric_surface/base.py

def sensitivity_score(self):
    """Compute first-order tolerance sensitivity scores via RSS formula.

    For each parameter with a gradient, the score is:
    ``tolerance_range² × gradient²``, which approximates the variance of
    the loss contribution from that parameter's manufacturing error.

    Returns:
        dict: Sensitivity gradients and RSS scores keyed by surface index.

    Reference:
        [1] Page 10 from: https://wp.optics.arizona.edu/optomech/wp-content/uploads/sites/53/2016/08/8-Tolerancing-1.pdf
    """
    score_dict = {}
    idx = getattr(self, "surf_idx", id(self))

    if self.d.grad is not None:
        score_dict[f"surf{idx}_d_grad"] = round(self.d.grad.item(), 6)
        score_dict[f"surf{idx}_d_score"] = round(
            (self.d_tole**2 * self.d.grad**2).item(), 6
        )

    return score_dict

draw_r

draw_r()

Effective drawing radius, clamped to the valid data range.

Source code in src/geometric_surface/base.py

def draw_r(self):
    """Effective drawing radius, clamped to the valid data range."""
    return min(self.r, self.max_height())

draw_widget

draw_widget(ax, color='black', linestyle='solid')

Draw widget for the surface on the 2D plot.

Source code in src/geometric_surface/base.py

def draw_widget(self, ax, color="black", linestyle="solid"):
    """Draw widget for the surface on the 2D plot."""
    r_eff = self.draw_r()
    r = torch.linspace(-r_eff, r_eff, 128, device=self.device)
    z = self.surface_with_offset(
        r, torch.zeros(len(r), device=self.device), valid_check=False
    )
    ax.plot(
        z.cpu().detach().numpy(),
        r.cpu().detach().numpy(),
        color=color,
        linestyle=linestyle,
        linewidth=0.75,
    )

create_mesh

create_mesh(n_rings=32, n_arms=128, color=[0.06, 0.3, 0.6])

Create triangulated surface mesh.

Parameters:

Name	Type	Description	Default
`n_rings`	`int`	Number of concentric rings for sampling.	`32`
`n_arms`	`int`	Number of angular divisions.	`128`
`color`	`List[float]`	The color of the mesh.	`[0.06, 0.3, 0.6]`

Returns:

Name	Type	Description
`self`		The surface with mesh data.

Source code in src/geometric_surface/base.py

def create_mesh(self, n_rings=32, n_arms=128, color=[0.06, 0.3, 0.6]):
    """Create triangulated surface mesh.

    Args:
        n_rings (int): Number of concentric rings for sampling.
        n_arms (int): Number of angular divisions.
        color (List[float]): The color of the mesh.

    Returns:
        self: The surface with mesh data.
    """
    self.vertices = self._create_vertices(n_rings, n_arms)
    self.faces = self._create_faces(n_rings, n_arms)
    self.rim = self._create_rim(n_rings, n_arms)
    self.mesh_color = color
    return self

get_polydata

get_polydata()

Get PyVista PolyData object from previously generated vertices and faces.

PolyData object will be used to draw the surface and export as .obj file.

Source code in src/geometric_surface/base.py

def get_polydata(self):
    """Get PyVista PolyData object from previously generated vertices and faces.

    PolyData object will be used to draw the surface and export as .obj file.
    """
    from pyvista import PolyData

    face_vertex_n = 3  # vertices per triangle
    formatted_faces = np.hstack(
        [
            face_vertex_n * np.ones((self.faces.shape[0], 1), dtype=np.uint32),
            self.faces,
        ]
    )
    return PolyData(self.vertices, formatted_faces)

zmx_str

zmx_str(surf_idx, d_next)

Return Zemax surface string.

Source code in src/geometric_surface/base.py

def zmx_str(self, surf_idx, d_next):
    """Return Zemax surface string."""
    raise NotImplementedError(
        "zmx_str() is not implemented for {}".format(self.__class__.__name__)
    )

Surface(r, d, mat2, pos_xy=[0, 0], vec_local=[0, 0, 1], is_square=False, device="cpu")

Parameter	Type	Description
`r`	`float`	Aperture radius (mm)
`d`	`float`	Axial vertex position (mm), differentiable
`mat2`	`Material`	Material on the transmission side
`is_square`	`bool`	Square aperture if `True`, circular if `False`

Key methods:

Method	Description
`ray_reaction(ray, n1, n2)`	Main interface: intersect + refract/reflect
`intersect(ray, n)`	Ray-surface intersection (overridden by subclass)
`refract(ray, n1, n2)`	Vector Snell's law
`reflect(ray)`	Specular reflection
`sag(x, y)`	Surface height \(z(x, y)\)
`normal_vec(ray)`	Surface normal at intersection
`surf_dict()`	Serialize to dictionary
`get_optimizer_params(lrs)`	Return parameter groups for optimizer

Spheric

Spherical surface defined by curvature \(c = 1/R\).

src.geometric_surface.Spheric

Spheric(c, r, d, mat2, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: Surface

Spherical refractive surface parameterized by curvature.

The sag function is:

.. math::

z(x, y) = \frac{c \rho^2}{1 + \sqrt{1 - c^2 \rho^2}}, \quad
\rho^2 = x^2 + y^2

Attributes:

Name	Type	Description
`c`	`Tensor`	Surface curvature `1/R` [1/mm]. Differentiable with respect to gradient-based optimization.

Initialize a spherical surface.

Parameters:

Name	Type	Description	Default
`c`	`float`	Surface curvature `1/R` [1/mm]. Use `0` for a flat surface (equivalent to `Plane`).	required
`r`	`float`	Aperture radius [mm].	required
`d`	`float`	Axial vertex position [mm].	required
`mat2`	`str or Material`	Material on the transmission side.	required
`pos_xy`	`list[float]`	Lateral offset `[x, y]` [mm]. Defaults to `[0.0, 0.0]`.	`[0.0, 0.0]`
`vec_local`	`list[float]`	Local normal direction. Defaults to `[0.0, 0.0, 1.0]`.	`[0.0, 0.0, 1.0]`
`is_square`	`bool`	Square aperture flag. Defaults to `False`.	`False`
`device`	`str`	Compute device. Defaults to `"cpu"`.	`'cpu'`

Source code in src/geometric_surface/spheric.py

def __init__(
    self,
    c,
    r,
    d,
    mat2,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize a spherical surface.

    Args:
        c (float): Surface curvature ``1/R`` [1/mm].  Use ``0`` for a flat
            surface (equivalent to ``Plane``).
        r (float): Aperture radius [mm].
        d (float): Axial vertex position [mm].
        mat2 (str or Material): Material on the transmission side.
        pos_xy (list[float], optional): Lateral offset ``[x, y]`` [mm].
            Defaults to ``[0.0, 0.0]``.
        vec_local (list[float], optional): Local normal direction.
            Defaults to ``[0.0, 0.0, 1.0]``.
        is_square (bool, optional): Square aperture flag. Defaults to
            ``False``.
        device (str, optional): Compute device. Defaults to ``"cpu"``.
    """
    super(Spheric, self).__init__(
        r=r,
        d=d,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )
    self.c = torch.tensor(c)

    self.tolerancing = False
    self.to(device)

intersect

intersect(ray, n=1.0)

Solve ray-surface intersection in local coordinate system using analytical method.

Sphere equation: (x)^2 + (y)^2 + (z - R)^2 = R^2, where R = 1/c Ray equation: p(t) = o + t*d Solve quadratic equation for intersection parameter t.

Parameters:

Name	Type	Description	Default
`ray`	`Ray`	input ray.	required
`n`	`float`	refractive index. Defaults to 1.0.	`1.0`

Returns:

Name	Type	Description
`ray`	`Ray`	ray with updated position and opl.

Source code in src/geometric_surface/spheric.py

def intersect(self, ray, n=1.0):
    """Solve ray-surface intersection in local coordinate system using analytical method.

    Sphere equation: (x)^2 + (y)^2 + (z - R)^2 = R^2, where R = 1/c
    Ray equation: p(t) = o + t*d
    Solve quadratic equation for intersection parameter t.

    Args:
        ray (Ray): input ray.
        n (float, optional): refractive index. Defaults to 1.0.

    Returns:
        ray (Ray): ray with updated position and opl.
    """
    # Tolerance
    if self.tolerancing:
        c = self.c + self.c_error
    else:
        c = self.c

    if torch.abs(c) < EPSILON:
        # Handle flat surface as a plane
        t = (0.0 - ray.o[..., 2]) / ray.d[..., 2]
        new_o = ray.o + t.unsqueeze(-1) * ray.d
        valid = (new_o[..., 0] ** 2 + new_o[..., 1] ** 2 < self.r**2) & (
            ray.is_valid > 0
        )
    else:
        R = 1.0 / c

        # Vector from ray origin to sphere center at (0, 0, R)
        oc = ray.o.clone()
        oc[..., 2] = oc[..., 2] - R

        # Quadratic equation: a*t^2 + b*t + c = 0
        # a = d·d = 1 (since ray direction is normalized)
        # b = 2*(o-center)·d
        # c = (o-center)·(o-center) - R^2

        a = torch.sum(ray.d * ray.d, dim=-1)  # Should be 1 for normalized rays
        b = 2.0 * torch.sum(oc * ray.d, dim=-1)
        c_coeff = torch.sum(oc * oc, dim=-1) - R * R

        discriminant = b * b - 4 * a * c_coeff
        valid_intersect = discriminant >= 0

        sqrt_discriminant = torch.sqrt(torch.clamp(discriminant, min=EPSILON))
        t1 = (-b - sqrt_discriminant) / (2 * a + EPSILON)
        t2 = (-b + sqrt_discriminant) / (2 * a + EPSILON)

        # Choose intersection closest to z=0 (surface vertex)
        z1 = ray.o[..., 2] + t1 * ray.d[..., 2]
        z2 = ray.o[..., 2] + t2 * ray.d[..., 2]
        use_t1 = torch.abs(z1) < torch.abs(z2)
        t = torch.where(use_t1, t1, t2)

        new_o = ray.o + t.unsqueeze(-1) * ray.d

        # Check aperture
        r_squared = new_o[..., 0] ** 2 + new_o[..., 1] ** 2
        within_aperture = r_squared <= (self.r**2 + EPSILON)

        valid = valid_intersect & within_aperture & (ray.is_valid > 0)

    # Update ray position
    ray.o = torch.where(valid.unsqueeze(-1), new_o, ray.o)
    ray.is_valid = ray.is_valid * valid

    if ray.coherent:
        if t.abs().max() > 100 and torch.get_default_dtype() == torch.float32:
            raise Exception(
                "Using float32 may cause precision problem for OPL calculation."
            )
        new_opl = ray.opl + n * t.unsqueeze(-1)
        ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)

    return ray

is_within_data_range

is_within_data_range(x, y)

Invalid when shape is non-defined.

Source code in src/geometric_surface/spheric.py

def is_within_data_range(self, x, y):
    """Invalid when shape is non-defined."""
    if self.tolerancing:
        c = self.c + self.c_error
    else:
        c = self.c

    valid = (x**2 + y**2) < 1 / c**2
    return valid

max_height

max_height()

Maximum valid height.

Source code in src/geometric_surface/spheric.py

def max_height(self):
    """Maximum valid height."""
    if self.tolerancing:
        c = self.c + self.c_error
    else:
        c = self.c

    max_height = torch.sqrt(1 / c**2).item() - 0.001
    return max_height

init_tolerance

init_tolerance(tolerance_params=None)

Initialize tolerance parameters for the surface.

Parameters:

Name	Type	Description	Default
`tolerance_params`	`dict`	Tolerance for surface parameters.	`None`

Source code in src/geometric_surface/spheric.py

def init_tolerance(self, tolerance_params=None):
    """Initialize tolerance parameters for the surface.

    Args:
        tolerance_params (dict): Tolerance for surface parameters.
    """
    super().init_tolerance(tolerance_params)
    if tolerance_params is None:
        tolerance_params = {}
    self.c_tole = tolerance_params.get("c_tole", 0.0001)
    self.c_error = 0.0

sample_tolerance

sample_tolerance()

Randomly perturb surface parameters to simulate manufacturing errors.

Source code in src/geometric_surface/spheric.py

def sample_tolerance(self):
    """Randomly perturb surface parameters to simulate manufacturing errors."""
    super().sample_tolerance()
    self.c_error = float(np.random.randn() * self.c_tole)

zero_tolerance

zero_tolerance()

Zero tolerance.

Source code in src/geometric_surface/spheric.py

def zero_tolerance(self):
    """Zero tolerance."""
    super().zero_tolerance()
    self.c_error = 0.0

sensitivity_score

sensitivity_score()

Tolerance squared sum.

Source code in src/geometric_surface/spheric.py

def sensitivity_score(self):
    """Tolerance squared sum."""
    score_dict = super().sensitivity_score()
    if self.c.grad is not None:
        idx = getattr(self, "surf_idx", id(self))
        score_dict[f"surf{idx}_c_grad"] = round(self.c.grad.item(), 6)
        score_dict[f"surf{idx}_c_score"] = round(
            (self.c_tole**2 * self.c.grad**2).item(), 6
        )
    return score_dict

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.0001], optim_mat=False)

Activate gradient computation for c and d and return optimizer parameters.

Source code in src/geometric_surface/spheric.py

def get_optimizer_params(self, lrs=[1e-4, 1e-4], optim_mat=False):
    """Activate gradient computation for c and d and return optimizer parameters."""
    self.c.requires_grad_(True)
    self.d.requires_grad_(True)

    params = []
    params.append({"params": [self.d], "lr": lrs[0]})
    params.append({"params": [self.c], "lr": lrs[1]})

    if optim_mat and self.mat2.get_name() != "air":
        params += self.mat2.get_optimizer_params()

    return params

surf_dict

surf_dict()

Return surface parameters.

Source code in src/geometric_surface/spheric.py

def surf_dict(self):
    """Return surface parameters."""
    roc = 1 / self.c.item() if self.c.item() != 0 else 0.0
    surf_dict = {
        "type": "Spheric",
        "r": round(self.r, 4),
        "(c)": round(self.c.item(), 4),
        "roc": round(roc, 4),
        "(d)": round(self.d.item(), 4),
        "mat2": self.mat2.get_name(),
    }

    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Return Zemax surface string.

Source code in src/geometric_surface/spheric.py

    def zmx_str(self, surf_idx, d_next):
        """Return Zemax surface string."""
        if self.mat2.get_name() == "air":
            zmx_str = f"""SURF {surf_idx} 
    TYPE STANDARD 
    CURV {self.c.item()} 
    DISZ {d_next.item()} 
    DIAM {self.r} 1 0 0 1 ""
"""
        else:
            zmx_str = f"""SURF {surf_idx} 
    TYPE STANDARD 
    CURV {self.c.item()} 
    DISZ {d_next.item()} 
    GLAS ___BLANK 1 0 {self.mat2.n} {self.mat2.V}
    DIAM {self.r} 1 0 0 1 ""
"""
        return zmx_str

\[z = \frac{c \rho^2}{1 + \sqrt{1 - c^2 \rho^2}}\]

where \(\rho^2 = x^2 + y^2\).

Spheric(c, r, d, mat2, device="cpu")

Parameter	Type	Description
`c`	`float`	Curvature \(= 1/R\) in mm\(^{-1}\) (differentiable)

Has an analytical intersect() solving the sphere-ray quadratic (faster than Newton's method).

Optimizer parameters: [lr_d, lr_c]

Aspheric

Even-asphere surface: spherical base with conic constant and polynomial corrections.

src.geometric_surface.Aspheric

Aspheric(r, d, c, k, ai, mat2, ai2=None, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: Surface

Even-order aspheric surface.

The sag function is:

.. math::

z(

ho) = rac{c\, ho^2}{1 + \sqrt{1-(1+k)c^2 ho^2}} + \sum_{i=2}^{n} a_{2i}\, ho^{2i}, \quad ho^2 = x^2 + y^2

The polynomial starts at the 4th-order term (a4) because the 2nd-order term competes with the base curvature c.

All coefficients c, k, and ai are differentiable torch tensors so they can be optimised with gradient descent.

Attributes:

Name	Type	Description
`c`	`Tensor`	Base curvature [1/mm].
`k`	`Tensor`	Conic constant.
`ai2`	`Tensor or None`	2nd-order aspheric coefficient (legacy).
`ai`	`Tensor`	Even-order aspheric coefficients `[a4, a6, a8, ...]`.

Initialize an aspheric surface.

Parameters:

Name	Type	Description	Default
`r`	`float`	Aperture radius [mm].	required
`d`	`float`	Axial vertex position [mm].	required
`c`	`float`	Base curvature `1/R` [1/mm].	required
`k`	`float`	Conic constant (`0` = sphere, `-1` = paraboloid).	required
`ai`	`list[float] or None`	Even-order aspheric coefficients starting from the 4th-order term: `[a4, a6, a8, ...]`. Pass `None` or an empty list for a pure conic.	required
`mat2`	`str or Material`	Material on the transmission side.	required
`ai2`	`float or None`	2nd-order aspheric coefficient from legacy data. Included in sag but not optimised. Defaults to `None`.	`None`
`pos_xy`	`list[float]`	Lateral offset `[x, y]` [mm]. Defaults to `[0.0, 0.0]`.	`[0.0, 0.0]`
`vec_local`	`list[float]`	Local normal direction. Defaults to `[0.0, 0.0, 1.0]`.	`[0.0, 0.0, 1.0]`
`is_square`	`bool`	Square aperture flag. Defaults to `False`.	`False`
`device`	`str`	Compute device. Defaults to `"cpu"`.	`'cpu'`

Source code in src/geometric_surface/aspheric.py

def __init__(
    self,
    r,
    d,
    c,
    k,
    ai,
    mat2,
    ai2=None,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize an aspheric surface.

    Args:
        r (float): Aperture radius [mm].
        d (float): Axial vertex position [mm].
        c (float): Base curvature ``1/R`` [1/mm].
        k (float): Conic constant (``0`` = sphere, ``-1`` = paraboloid).
        ai (list[float] or None): Even-order aspheric coefficients
            starting from the 4th-order term: ``[a4, a6, a8, ...]``.
            Pass ``None`` or an empty list for a pure conic.
        mat2 (str or Material): Material on the transmission side.
        ai2 (float or None, optional): 2nd-order aspheric coefficient
            from legacy data.  Included in sag but not optimised.
            Defaults to ``None``.
        pos_xy (list[float], optional): Lateral offset ``[x, y]`` [mm].
            Defaults to ``[0.0, 0.0]``.
        vec_local (list[float], optional): Local normal direction.
            Defaults to ``[0.0, 0.0, 1.0]``.
        is_square (bool, optional): Square aperture flag.
            Defaults to ``False``.
        device (str, optional): Compute device. Defaults to ``"cpu"``.
    """
    Surface.__init__(
        self,
        r=r,
        d=d,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    self._c = torch.tensor(c)
    self._k = torch.tensor(k)
    self._reparam = False

    # 2nd-order coefficient (legacy, not optimised)
    if ai2 is not None:
        self.ai2 = torch.tensor(float(ai2))
    else:
        self.ai2 = None

    self.ai_degree = 0
    self.ai = ai

    self.tolerancing = False
    self.to(device)

ai `property` `writable`

ai

Physical aspheric coefficients tensor [a4, a6, a8, ...].

c `property` `writable`

Physical curvature tensor.

k `property` `writable`

Physical conic constant tensor.

enable_reparam

enable_reparam()

Optimize normalized parameters while keeping physical coefficients readable.

Source code in src/geometric_surface/aspheric.py

def enable_reparam(self):
    """Optimize normalized parameters while keeping physical coefficients readable."""
    if self._reparam:
        return

    self._c_scale = max(abs(self._c.item()), 0.01)
    self._c_phi = self._scalar_tensor(self._c.item() / self._c_scale)
    self._c.requires_grad_(False)

    self._k_scale = max(abs(self._k.item()), 1.0)
    self._k_phi = self._scalar_tensor(self._k.item() / self._k_scale)
    self._k.requires_grad_(False)

    for i in range(self.ai_degree):
        order = 2 * (i + 2)
        p_name = f"ai{order}"
        raw_val = getattr(self, p_name)
        scale = self._get_ai_scale(order)
        phi = self._scalar_tensor(raw_val.item() / scale)
        setattr(self, f"_{p_name}_scale", scale)
        setattr(self, f"_{p_name}_phi", phi)
        raw_val.requires_grad_(False)

    self._reparam = True

disable_reparam

disable_reparam()

Write normalized parameters back to physical tensors and remove the transform.

Source code in src/geometric_surface/aspheric.py

@torch.no_grad()
def disable_reparam(self):
    """Write normalized parameters back to physical tensors and remove the transform."""
    if not self._reparam:
        return

    c_val = self.c.detach().clone()
    k_val = self.k.detach().clone()
    ai_vals = [self._get_ai_coeff(i).detach().clone() for i in range(self.ai_degree)]

    del self._c_scale, self._c_phi
    del self._k_scale, self._k_phi
    for i in range(self.ai_degree):
        order = 2 * (i + 2)
        delattr(self, f"_ai{order}_scale")
        delattr(self, f"_ai{order}_phi")

    self._reparam = False
    self._c = c_val
    self._k = k_val
    for i, coeff in enumerate(ai_vals):
        setattr(self, f"ai{2 * (i + 2)}", coeff)

update_r

update_r(r)

Update surface radius and preserve physical coefficients under reparam.

Source code in src/geometric_surface/aspheric.py

def update_r(self, r):
    """Update surface radius and preserve physical coefficients under reparam."""
    if not self._reparam:
        super().update_r(r)
        return

    old_r = self.r
    old_scales = {
        2 * (i + 2): self._get_ai_scale(2 * (i + 2), old_r)
        for i in range(self.ai_degree)
    }
    super().update_r(r)

    for order, old_scale in old_scales.items():
        new_scale = self._get_ai_scale(order)
        phi = getattr(self, f"_ai{order}_phi")
        with torch.no_grad():
            phi.mul_(old_scale / new_scale)
        setattr(self, f"_ai{order}_scale", new_scale)

is_within_data_range

is_within_data_range(x, y)

Invalid when shape is non-defined.

Source code in src/geometric_surface/aspheric.py

def is_within_data_range(self, x, y):
    """Invalid when shape is non-defined."""
    c, k = self._get_curvature_params()
    if k > -1:
        return (x**2 + y**2) < 1 / c**2 / (1 + k)
    return torch.ones_like(x, dtype=torch.bool)

max_height

max_height()

Maximum valid height.

Source code in src/geometric_surface/aspheric.py

def max_height(self):
    """Maximum valid height."""
    c, k = self._get_curvature_params()
    if k > -1:
        return torch.sqrt(1 / (k + 1) / (c**2)).item() - 0.001
    return 10e3

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.0001, 0.01, 0.0001], optim_mat=False, reparam=False)

Get optimizer parameters for the asphere.

When reparam=True the optimizer targets normalized hidden parameters so that c, k, and ai live on comparable scales — designed for single-lr optimizers like BFGS. With reparam=False (the default), per-order learning-rate scaling is applied to the raw physical parameters, which is correct for Adam and other per-parameter-lr optimizers.

Parameters:

Name	Type	Description	Default
`lrs`	`list`	learning rates for `[d, c, k, ai]`.	`[0.0001, 0.0001, 0.01, 0.0001]`
`optim_mat`	`bool`	whether to optimize material. Defaults to False.	`False`
`reparam`	`bool`	optimize normalized hidden parameters. Defaults to False.	`False`

Source code in src/geometric_surface/aspheric.py

def get_optimizer_params(
    self,
    lrs=[1e-4, 1e-4, 1e-2, 1e-4],
    optim_mat=False,
    reparam=False,
):
    """Get optimizer parameters for the asphere.

    When ``reparam=True`` the optimizer targets normalized hidden
    parameters so that ``c``, ``k``, and ``ai`` live on comparable
    scales — designed for single-lr optimizers like BFGS.  With
    ``reparam=False`` (the default), per-order learning-rate scaling
    is applied to the raw physical parameters, which is correct for
    Adam and other per-parameter-lr optimizers.

    Args:
        lrs (list, optional): learning rates for ``[d, c, k, ai]``.
        optim_mat (bool, optional): whether to optimize material.
            Defaults to False.
        reparam (bool, optional): optimize normalized hidden parameters.
            Defaults to False.
    """
    params = []

    if reparam:
        self.enable_reparam()
    elif self._reparam:
        self.disable_reparam()

    # Optimize distance
    self.d.requires_grad_(True)
    params.append({"params": [self.d], "lr": lrs[0]})

    if self._reparam:
        self._c_phi.requires_grad_(True)
        params.append({"params": [self._c_phi], "lr": lrs[1]})

        self._k_phi.requires_grad_(True)
        params.append({"params": [self._k_phi], "lr": lrs[2]})

        lr_base = lrs[3] if len(lrs) > 3 else 1e-4
        for i in range(self.ai_degree):
            order = 2 * (i + 2)
            phi = getattr(self, f"_ai{order}_phi")
            phi.requires_grad_(True)
            params.append({"params": [phi], "lr": lr_base})
    else:
        self.c.requires_grad_(True)
        params.append({"params": [self.c], "lr": lrs[1]})

        self.k.requires_grad_(True)
        params.append({"params": [self.k], "lr": lrs[2]})

        if self.ai is not None and self.ai_degree > 0:
            r_norm = self._get_reparam_radius()
            lr_base = lrs[3] if len(lrs) > 3 else 1e-4
            for i in range(self.ai_degree):
                p_name = f"ai{2 * (i + 2)}"
                p = getattr(self, p_name)
                p.requires_grad_(True)
                order = 2 * (i + 2)
                lr_ai = lr_base / r_norm**order
                params.append({"params": [p], "lr": lr_ai})

    # Optimize material parameters
    if optim_mat and self.mat2.get_name() != "air":
        params += self.mat2.get_optimizer_params()

    return params

init_tolerance

init_tolerance(tolerance_params=None)

Perturb the surface with some tolerance.

Parameters:

Name	Type	Description	Default
`tolerance_params`	`dict`	Tolerance for surface parameters.	`None`

References

[1] https://www.edmundoptics.com/capabilities/precision-optics/capabilities/aspheric-lenses/ [2] https://www.edmundoptics.com/knowledge-center/application-notes/optics/all-about-aspheric-lenses/?srsltid=AfmBOoon8AUXVALojol2s5K20gQk7W1qUisc6cE4WzZp3ATFY5T1pK8q

Source code in src/geometric_surface/aspheric.py

@torch.no_grad()
def init_tolerance(self, tolerance_params=None):
    """Perturb the surface with some tolerance.

    Args:
        tolerance_params (dict): Tolerance for surface parameters.

    References:
        [1] https://www.edmundoptics.com/capabilities/precision-optics/capabilities/aspheric-lenses/
        [2] https://www.edmundoptics.com/knowledge-center/application-notes/optics/all-about-aspheric-lenses/?srsltid=AfmBOoon8AUXVALojol2s5K20gQk7W1qUisc6cE4WzZp3ATFY5T1pK8q
    """
    super().init_tolerance(tolerance_params)
    if tolerance_params is None:
        tolerance_params = {}
    self.c_tole = tolerance_params.get("c_tole", 0.001)
    self.k_tole = tolerance_params.get("k_tole", 0.001)
    self.c_error = 0.0
    self.k_error = 0.0

sample_tolerance

sample_tolerance()

Randomly perturb surface parameters to simulate manufacturing errors.

Source code in src/geometric_surface/aspheric.py

def sample_tolerance(self):
    """Randomly perturb surface parameters to simulate manufacturing errors."""
    super().sample_tolerance()
    self.c_error = float(np.random.randn() * self.c_tole)
    self.k_error = float(np.random.randn() * self.k_tole)

zero_tolerance

zero_tolerance()

Zero tolerance.

Source code in src/geometric_surface/aspheric.py

def zero_tolerance(self):
    """Zero tolerance."""
    super().zero_tolerance()
    self.c_error = 0.0
    self.k_error = 0.0

sensitivity_score

sensitivity_score()

Tolerance squared sum.

Source code in src/geometric_surface/aspheric.py

def sensitivity_score(self):
    """Tolerance squared sum."""
    score_dict = super().sensitivity_score()
    idx = getattr(self, "surf_idx", id(self))

    c_grad = None
    if self._reparam and hasattr(self, "_c_phi") and self._c_phi.grad is not None:
        c_grad = self._c_phi.grad / self._c_scale
    elif not self._reparam and self.c.grad is not None:
        c_grad = self.c.grad

    k_grad = None
    if self._reparam and hasattr(self, "_k_phi") and self._k_phi.grad is not None:
        k_grad = self._k_phi.grad / self._k_scale
    elif not self._reparam and self.k.grad is not None:
        k_grad = self.k.grad

    if c_grad is not None:
        score_dict[f"surf{idx}_c_grad"] = round(c_grad.item(), 6)
        score_dict[f"surf{idx}_c_score"] = round(
            (self.c_tole**2 * c_grad**2).item(), 6
        )
    if k_grad is not None:
        score_dict[f"surf{idx}_k_grad"] = round(k_grad.item(), 6)
        score_dict[f"surf{idx}_k_score"] = round(
            (self.k_tole**2 * k_grad**2).item(), 6
        )
    return score_dict

surf_dict

surf_dict()

Return a dict of surface.

Source code in src/geometric_surface/aspheric.py

def surf_dict(self):
    """Return a dict of surface."""
    c, k = self._get_curvature_params()
    c_val = c.item()
    k_val = k.item()

    has_ai2 = self.ai2 is not None
    surf_dict = {
        "type": "Aspheric",
        "r": round(self.r, 4),
        "(c)": round(c_val, 4),
        "roc": round(1 / c_val, 4),
        "d": round(self.d.item(), 4),
        "k": round(k_val, 4),
        "ai": [],
        "use_ai2": has_ai2,
        "mat2": self.mat2.get_name(),
    }

    # Prepend a2 to ai list if present (ai2 key is informational;
    # deserialization reads ai[0] when use_ai2=True)
    if has_ai2:
        surf_dict["ai2"] = float(format(self.ai2.item(), ".6e"))
        surf_dict["ai"].append(float(format(self.ai2.item(), ".6e")))

    for i in range(self.ai_degree):
        order = i + 2
        coeff = self._get_ai_coeff(i)
        surf_dict[f"(ai{2 * order})"] = float(format(coeff.item(), ".6e"))
        surf_dict["ai"].append(float(format(coeff.item(), ".6e")))

    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Return Zemax surface string.

Source code in src/geometric_surface/aspheric.py

    def zmx_str(self, surf_idx, d_next):
        """Return Zemax surface string."""
        c, k = self._get_curvature_params()
        c_val = c.item()

        assert c_val != 0, (
            "Aperture surface is re-implemented in Aperture class."
        )
        assert self.ai is not None or k.item() != 0, (
            "Spheric surface is re-implemented in Spheric class."
        )

        # Collect absolute ai values, PARM 1 = a2, PARM 2+ = a4, a6, ...
        abs_ai = [self.ai2.item() if self.ai2 is not None else 0.0]
        for i in range(self.ai_degree):
            abs_ai.append(self._get_ai_coeff(i).item())

        # Pad with zeros for Zemax PARM format (needs 8 PARMs: a2..a16)
        while len(abs_ai) < 8:
            abs_ai.append(0.0)

        if self.mat2.get_name() == "air":
            zmx_str = f"""SURF {surf_idx}
    TYPE EVENASPH
    CURV {c_val}
    DISZ {d_next.item()}
    DIAM {self.r} 1 0 0 1 ""
    CONI {k}
    PARM 1 {abs_ai[0]}
    PARM 2 {abs_ai[1]}
    PARM 3 {abs_ai[2]}
    PARM 4 {abs_ai[3]}
    PARM 5 {abs_ai[4]}
    PARM 6 {abs_ai[5]}
    PARM 7 {abs_ai[6]}
    PARM 8 {abs_ai[7]}
"""
        else:
            zmx_str = f"""SURF {surf_idx}
    TYPE EVENASPH
    CURV {c_val}
    DISZ {d_next.item()}
    GLAS ___BLANK 1 0 {self.mat2.n} {self.mat2.V}
    DIAM {self.r} 1 0 0 1 ""
    CONI {k}
    PARM 1 {abs_ai[0]}
    PARM 2 {abs_ai[1]}
    PARM 3 {abs_ai[2]}
    PARM 4 {abs_ai[3]}
    PARM 5 {abs_ai[4]}
    PARM 6 {abs_ai[5]}
    PARM 7 {abs_ai[6]}
    PARM 8 {abs_ai[7]}
"""
        return zmx_str

\[z = \frac{c \rho^2}{1 + \sqrt{1 - (1+k) c^2 \rho^2}} + \sum_{i=2}^{N} a_{2i} \rho^{2i}\]

Aspheric(r, d, c, k, ai, mat2, device="cpu")

Parameter	Type	Description
`c`	`float`	Curvature (differentiable)
`k`	`float`	Conic constant (differentiable)
`ai`	`list[float]`	Polynomial coefficients \([a_4, a_6, a_8, \ldots]\) (differentiable)

Optimizer parameters: [lr_d, lr_c, lr_k, lr_ai]

Per-order learning rate scaling: \(\text{lr}_{a_{2n}} = \text{lr\_base} / r^{2n}\)

Reparametrization

For L-BFGS optimization, parameters are normalized to \(\sim\mathcal{O}(1)\):

# Enable: stores normalized phi tensors
surface.enable_reparam()

# Disable: restores physical values
surface.disable_reparam()

# Dynamic radius update (preserves physical coefficients)
surface.update_r(new_r)

Aperture

Aperture stop surface. Subclasses Plane with no refraction — only clips rays outside the clear aperture.

src.geometric_surface.Aperture

Aperture(r, d, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: Plane

Aperture surface.

Source code in src/geometric_surface/aperture.py

def __init__(
    self,
    r,
    d,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Aperture surface."""
    Plane.__init__(
        self,
        r=r,
        d=d,
        mat2="air",
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )
    self.tolerancing = False
    self.to(device)

ray_reaction

ray_reaction(ray, n1=1.0, n2=1.0, refraction=False)

Compute output ray after intersection and refraction.

Source code in src/geometric_surface/aperture.py

def ray_reaction(self, ray, n1=1.0, n2=1.0, refraction=False):
    """Compute output ray after intersection and refraction."""
    ray = self.to_local_coord(ray)
    ray = self.intersect(ray)
    ray = self.to_global_coord(ray)
    return ray

draw_widget

draw_widget(ax, color='orange', linestyle='solid')

Draw aperture wedge on the figure.

Source code in src/geometric_surface/aperture.py

def draw_widget(self, ax, color="orange", linestyle="solid"):
    """Draw aperture wedge on the figure."""
    d = self.d.item()
    aper_wedge_l = 0.05 * self.r  # [mm]
    aper_wedge_h = 0.15 * self.r  # [mm]

    # Parallel edges
    z = np.linspace(d - aper_wedge_l, d + aper_wedge_l, 3)
    x = -self.r * np.ones(3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)
    x = self.r * np.ones(3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)

    # Vertical edges
    z = d * np.ones(3)
    x = np.linspace(self.r, self.r + aper_wedge_h, 3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)
    x = np.linspace(-self.r - aper_wedge_h, -self.r, 3)
    ax.plot(z, x, color=color, linestyle=linestyle, linewidth=0.8)

draw_widget3D

draw_widget3D(ax, color='black')

Draw the aperture as a circle in a 3D plot.

Source code in src/geometric_surface/aperture.py

def draw_widget3D(self, ax, color="black"):
    """Draw the aperture as a circle in a 3D plot."""
    # Draw the edge circle
    theta = np.linspace(0, 2 * np.pi, 100)
    edge_x = self.r * np.cos(theta)
    edge_y = self.r * np.sin(theta)
    edge_z = np.full_like(edge_x, self.d.item())  # Constant z at aperture position

    # Plot the edge circle
    line = ax.plot(edge_z, edge_x, edge_y, color=color, linewidth=1.5)

    return line

create_mesh

create_mesh(n_rings=32, n_arms=128, color=[0.0, 0.0, 0.0])

Create triangulated surface mesh.

Parameters:

Name	Type	Description	Default
`n_rings`	`int`	Number of concentric rings for sampling.	`32`
`n_arms`	`int`	Number of angular divisions.	`128`
`color`	`List[float]`	The color of the mesh.	`[0.0, 0.0, 0.0]`

Returns:

Name	Type	Description
`self`		The surface with mesh data.

Source code in src/geometric_surface/aperture.py

def create_mesh(self, n_rings=32, n_arms=128, color=[0.0, 0.0, 0.0]):
    """Create triangulated surface mesh.

    Args:
        n_rings (int): Number of concentric rings for sampling.
        n_arms (int): Number of angular divisions.
        color (List[float]): The color of the mesh.

    Returns:
        self: The surface with mesh data.
    """
    self.vertices = self._create_vertices(n_rings, n_arms)
    self.faces = self._create_faces(n_rings, n_arms)
    self.rim = self._create_rim(n_rings, n_arms)
    self.mesh_color = color
    return self

get_optimizer_params

get_optimizer_params(lrs=[0.0001])

Activate gradient computation for d and return optimizer parameters.

Source code in src/geometric_surface/aperture.py

def get_optimizer_params(self, lrs=[1e-4]):
    """Activate gradient computation for d and return optimizer parameters."""
    self.d.requires_grad_(True)

    params = []
    params.append({"params": [self.d], "lr": lrs[0]})

    return params

surf_dict

surf_dict()

Dict of surface parameters.

Source code in src/geometric_surface/aperture.py

def surf_dict(self):
    """Dict of surface parameters."""
    surf_dict = {
        "type": "Aperture",
        "r": round(self.r, 4),
        "(d)": round(self.d.item(), 4),
        "mat2": "air",
        "is_square": self.is_square,
    }
    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Zemax surface string.

Source code in src/geometric_surface/aperture.py

    def zmx_str(self, surf_idx, d_next):
        """Zemax surface string."""
        zmx_str = f"""SURF {surf_idx}
    STOP
    TYPE STANDARD
    CURV 0.0
    DISZ {d_next.item()}
"""
        return zmx_str

Aperture(r, d, device="cpu")

Used as the aperture stop in lens systems. Has update_r() for curriculum aperture growth.

Plane

Flat surface (\(z = 0\) in local frame). Base class for Aperture, Mirror, and ThinLens.

src.geometric_surface.Plane

Plane(r, d, mat2, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: Surface

Plane surface.

Examples:

IR filter.
Lens cover glass.
DOE base.

The following surfaces inherit from Plane

Aperture.
Mirror.
ThinLens.

Source code in src/geometric_surface/plane.py

def __init__(
    self,
    r,
    d,
    mat2,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Plane surface.

    Examples:
        - IR filter.
        - Lens cover glass.
        - DOE base.

    The following surfaces inherit from Plane:
        - Aperture.
        - Mirror.
        - ThinLens.
    """
    Surface.__init__(
        self,
        r=r,
        d=d,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

intersect

intersect(ray, n=1.0)

Solve ray-surface intersection in local coordinate system and update ray data.

Source code in src/geometric_surface/plane.py

def intersect(self, ray, n=1.0):
    """Solve ray-surface intersection in local coordinate system and update ray data."""
    # Solve intersection
    t = (0.0 - ray.o[..., 2]) / ray.d[..., 2]
    new_o = ray.o + t.unsqueeze(-1) * ray.d

    # Aperture mask
    if self.is_square:
        valid = (
            (torch.abs(new_o[..., 0]) < self.w / 2)
            & (torch.abs(new_o[..., 1]) < self.h / 2)
            & (ray.is_valid > 0)
        )
    else:
        valid = (new_o[..., 0] ** 2 + new_o[..., 1] ** 2 < self.r**2) & (
            ray.is_valid > 0
        )

    # Update rays
    new_o = ray.o + ray.d * t.unsqueeze(-1)
    ray.o = torch.where(valid.unsqueeze(-1), new_o, ray.o)
    ray.is_valid = ray.is_valid * valid

    if ray.coherent:
        ray.opl = torch.where(
            valid.unsqueeze(-1), ray.opl + n * t.unsqueeze(-1), ray.opl
        )

    return ray

normal_vec

normal_vec(ray)

Calculate surface normal vector at intersection points in local coordinate system.

Normal vector points from the surface toward the side where the light is coming from.

Source code in src/geometric_surface/plane.py

def normal_vec(self, ray):
    """Calculate surface normal vector at intersection points in local coordinate system.

    Normal vector points from the surface toward the side where the light is coming from.
    """
    normal_vec = torch.zeros_like(ray.d)
    normal_vec[..., 2] = -1

    is_forward = ray.d[..., 2].unsqueeze(-1) > 0
    normal_vec = torch.where(is_forward, normal_vec, -normal_vec)
    return normal_vec

get_optimizer_params

get_optimizer_params(lrs=[0.0001], optim_mat=False)

Activate gradient computation for d and return optimizer parameters.

Source code in src/geometric_surface/plane.py

def get_optimizer_params(self, lrs=[1e-4], optim_mat=False):
    """Activate gradient computation for d and return optimizer parameters."""
    params = []

    # Optimize d
    self.d.requires_grad_(True)
    params.append({"params": [self.d], "lr": lrs[0]})

    # Optimize material parameters
    if optim_mat and self.mat2.get_name() != "air":
        params += self.mat2.get_optimizer_params()

    return params

Mirror

Flat mirror surface. Subclasses Plane with reflect() instead of refract().

src.geometric_surface.Mirror

Mirror(r, d, mat2='air', pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=True, device='cpu')

Bases: Plane

Mirror surface.

Source code in src/geometric_surface/mirror.py

def __init__(
    self,
    r,
    d,
    mat2="air",
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=True,
    device="cpu",
):
    """Mirror surface."""
    Surface.__init__(
        self,
        r=r,
        d=d,
        mat2=mat2,
        is_square=is_square,
        pos_xy=pos_xy,
        vec_local=vec_local,
        device=device,
    )

ray_reaction

ray_reaction(ray, n1=None, n2=None)

Compute output ray after intersection and reflection with the mirror surface.

Source code in src/geometric_surface/mirror.py

def ray_reaction(self, ray, n1=None, n2=None):
    """Compute output ray after intersection and reflection with the mirror surface."""
    ray = self.to_local_coord(ray)
    ray = self.intersect(ray)
    ray = self.reflect(ray)
    ray = self.to_global_coord(ray)
    return ray

surf_dict

surf_dict()

Return surface parameters.

Source code in src/geometric_surface/mirror.py

def surf_dict(self):
    """Return surface parameters."""
    surf_dict = {
        "type": self.__class__.__name__,
        "r": self.r,
        "d": self.d,
        "mat2": self.mat2.get_name(),
    }
    return surf_dict

ThinLens

Ideal thin lens with focal length \(f\). Implements paraxial phase approximation.

src.geometric_surface.ThinLens

ThinLens(r, d, f=100.0, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: Plane

Thin lens surface.

Source code in src/geometric_surface/thinlens.py

def __init__(
    self,
    r,
    d,
    f=100.0,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Thin lens surface."""
    Plane.__init__(
        self,
        r=r,
        d=d,
        mat2="air",
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )
    self.f = torch.tensor(f)

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.0001], optim_mat=False)

Activate gradient computation for f and d and return optimizer parameters.

Source code in src/geometric_surface/thinlens.py

def get_optimizer_params(self, lrs=[1e-4, 1e-4], optim_mat=False):
    """Activate gradient computation for f and d and return optimizer parameters."""
    params = []

    self.d.requires_grad_(True)
    params.append({"params": [self.d], "lr": lrs[0]})

    self.f.requires_grad_(True)
    params.append({"params": [self.f], "lr": lrs[1]})

    return params

refract

refract(ray, eta=1.0)

For a thin lens, all rays will converge to z = f plane. Therefore we trace the chief-ray (parallel-shift to surface center) to find the final convergence point for each ray.

For coherent ray tracing, we can think it as a Fresnel lens with infinite refractive index. (1) Lens maker's equation (2) Spherical lens function

Source code in src/geometric_surface/thinlens.py

def refract(self, ray, eta=1.0):
    """For a thin lens, all rays will converge to z = f plane. Therefore we trace the chief-ray (parallel-shift to surface center) to find the final convergence point for each ray.

    For coherent ray tracing, we can think it as a Fresnel lens with infinite refractive index.
    (1) Lens maker's equation
    (2) Spherical lens function
    """
    forward = (ray.d * ray.is_valid.unsqueeze(-1))[..., 2].sum() > 0

    # Calculate convergence point
    if forward:
        t0 = self.f / ray.d[..., 2]
        xy_final = ray.d[..., :2] * t0.unsqueeze(-1)
        z_final = (
            (self.d + self.f).view(1).expand_as(xy_final[..., 0].unsqueeze(-1))
        )
        o_final = torch.cat([xy_final, z_final], dim=-1)
    else:
        t0 = -self.f / ray.d[..., 2]
        xy_final = ray.d[..., :2] * t0.unsqueeze(-1)
        z_final = (
            (self.d - self.f).view(1).expand_as(xy_final[..., 0].unsqueeze(-1))
        )
        o_final = torch.cat([xy_final, z_final], dim=-1)

    # New ray direction
    if self.f > 0:
        new_d = o_final - ray.o
    else:
        new_d = ray.o - o_final
    new_d = F.normalize(new_d, p=2, dim=-1)
    ray.d = new_d

    # Optical path length change
    if ray.coherent:
        valid = ray.is_valid > 0
        if forward:
            new_opl = (
                ray.opl
                - (ray.o[..., 0] ** 2 + ray.o[..., 1] ** 2)
                / self.f
                / 2
                / ray.d[..., 2]
            )
            ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)
        else:
            new_opl = (
                ray.opl
                + (ray.o[..., 0] ** 2 + ray.o[..., 1] ** 2)
                / self.f
                / 2
                / ray.d[..., 2]
            )
            ray.opl = torch.where(valid.unsqueeze(-1), new_opl, ray.opl)

    return ray

ThinLens(r, d, f=100.0, device="cpu")

Parameter	Type	Description
`f`	`float`	Focal length in mm (differentiable)

QTypeFreeform

Forbes Q-polynomial freeform surface (Qbfs). Orthogonal polynomial basis for reduced sensitivity to coefficient errors.

src.geometric_surface.QTypeFreeform

QTypeFreeform(r, d, c, k, qm, mat2, r_norm=None, pos_xy=[0.0, 0.0], vec_local=[0.0, 0.0, 1.0], is_square=False, device='cpu')

Bases: Surface

Q-type (Forbes Qbfs polynomial) freeform surface.

This surface type uses Forbes Q-polynomials to represent rotationally symmetric aspheric departures from a base conic surface. The representation is well-suited for optimization due to the orthogonality of the basis functions.

Attributes:

Name	Type	Description
`c`	`tensor`	Curvature of the base surface (1/radius of curvature)
`k`	`tensor`	Conic constant
`r_norm`	`float`	Normalization radius for Q polynomials (typically equals r)
`qm`	`tensor`	Q polynomial coefficients [a_0, a_1, ..., a_{n-1}]

Initialize Q-type freeform surface.

Parameters:

Name	Type	Description	Default
`r`	`float`	Aperture radius of the surface	required
`d`	`float`	Distance from origin to surface vertex	required
`c`	`float`	Curvature of the base surface (1/radius of curvature)	required
`k`	`float`	Conic constant (k=0 for sphere, k=-1 for paraboloid)	required
`qm`	`list`	Q polynomial coefficients [a_0, a_1, ..., a_{n-1}]	required
`mat2`	`str or Material`	Material after the surface	required
`r_norm`	`float`	Normalization radius. Defaults to r.	`None`
`pos_xy`	`list`	Surface center position [x, y]. Defaults to [0, 0].	`[0.0, 0.0]`
`vec_local`	`list`	Local surface normal. Defaults to [0, 0, 1].	`[0.0, 0.0, 1.0]`
`is_square`	`bool`	Whether aperture is square. Defaults to False.	`False`
`device`	`str`	Torch device. Defaults to "cpu".	`'cpu'`

Source code in src/geometric_surface/qtype.py

def __init__(
    self,
    r,
    d,
    c,
    k,
    qm,
    mat2,
    r_norm=None,
    pos_xy=[0.0, 0.0],
    vec_local=[0.0, 0.0, 1.0],
    is_square=False,
    device="cpu",
):
    """Initialize Q-type freeform surface.

    Args:
        r (float): Aperture radius of the surface
        d (float): Distance from origin to surface vertex
        c (float): Curvature of the base surface (1/radius of curvature)
        k (float): Conic constant (k=0 for sphere, k=-1 for paraboloid)
        qm (list): Q polynomial coefficients [a_0, a_1, ..., a_{n-1}]
        mat2 (str or Material): Material after the surface
        r_norm (float, optional): Normalization radius. Defaults to r.
        pos_xy (list): Surface center position [x, y]. Defaults to [0, 0].
        vec_local (list): Local surface normal. Defaults to [0, 0, 1].
        is_square (bool): Whether aperture is square. Defaults to False.
        device (str): Torch device. Defaults to "cpu".
    """
    Surface.__init__(
        self,
        r=r,
        d=d,
        mat2=mat2,
        pos_xy=pos_xy,
        vec_local=vec_local,
        is_square=is_square,
        device=device,
    )

    self.c = torch.tensor(c)
    self.k = torch.tensor(k)
    self.r_norm = r_norm if r_norm is not None else r

    # Store Q polynomial coefficients
    if qm is not None and len(qm) > 0:
        self.qm = torch.tensor(qm, dtype=torch.float64)
        self.n_qterms = len(qm)
        # Also store individual coefficients for optimization
        for i, coef in enumerate(qm):
            setattr(self, f"q{i}", torch.tensor(coef))
    else:
        self.qm = None
        self.n_qterms = 0

    self.tolerancing = False
    self.to(device)

init_from_dict `classmethod`

init_from_dict(surf_dict)

Initialize surface from a dictionary specification.

Source code in src/geometric_surface/qtype.py

@classmethod
def init_from_dict(cls, surf_dict):
    """Initialize surface from a dictionary specification."""
    if "roc" in surf_dict:
        c = 1 / surf_dict["roc"]
    else:
        c = surf_dict["c"]

    return cls(
        r=surf_dict["r"],
        d=surf_dict["d"],
        c=c,
        k=surf_dict.get("k", 0.0),
        qm=surf_dict.get("qm", []),
        mat2=surf_dict["mat2"],
        r_norm=surf_dict.get("r_norm", None),
    )

is_within_data_range

is_within_data_range(x, y)

Check if points are within valid surface data range.

For conic surfaces with k > -1, there's a maximum valid radius.

Parameters:

Name	Type	Description	Default
`x`	`tensor`	x coordinates	required
`y`	`tensor`	y coordinates	required

Returns:

Name	Type	Description
`tensor`		Boolean mask of valid points

Source code in src/geometric_surface/qtype.py

def is_within_data_range(self, x, y):
    """Check if points are within valid surface data range.

    For conic surfaces with k > -1, there's a maximum valid radius.

    Args:
        x (tensor): x coordinates
        y (tensor): y coordinates

    Returns:
        tensor: Boolean mask of valid points
    """
    if self.tolerancing:
        c = self.c + self.c_error
        k = self.k + self.k_error
    else:
        c = self.c
        k = self.k

    r2 = x**2 + y**2

    # Check conic validity
    if k > -1 and abs(c) > EPSILON:
        valid_conic = r2 < 1 / (c**2 * (1 + k)) - EPSILON
    else:
        valid_conic = torch.ones_like(x, dtype=torch.bool)

    # Check normalized radius (should be <= 1 for Q polynomials)
    u2 = r2 / (self.r_norm**2)
    valid_qpoly = u2 <= 1 + EPSILON

    return valid_conic & valid_qpoly

max_height

max_height()

Maximum valid height (radial distance).

Source code in src/geometric_surface/qtype.py

def max_height(self):
    """Maximum valid height (radial distance)."""
    if self.tolerancing:
        c = self.c + self.c_error
        k = self.k + self.k_error
    else:
        c = self.c
        k = self.k

    # Conic limit
    if k > -1 and abs(c) > EPSILON:
        max_conic = np.sqrt(1 / ((k + 1) * c**2)) - 0.001
    else:
        max_conic = 10e3

    # Q polynomial limit (normalization radius)
    max_q = self.r_norm

    return min(max_conic, max_q)

get_optimizer_params

get_optimizer_params(lrs=[0.0001, 0.0001, 0.01, 1e-06], decay=0.1, optim_mat=False)

Get optimizer parameters for different surface parameters.

Parameters:

Name	Type	Description	Default
`lrs`	`list`	Learning rates for [d, c, k, q_coefficients].	`[0.0001, 0.0001, 0.01, 1e-06]`
`decay`	`float`	Decay factor for higher-order Q coefficients.	`0.1`
`optim_mat`	`bool`	Whether to optimize material parameters.	`False`

Returns:

Name	Type	Description
`list`		Parameter groups for optimizer

Source code in src/geometric_surface/qtype.py

def get_optimizer_params(
    self, lrs=[1e-4, 1e-4, 1e-2, 1e-6], decay=0.1, optim_mat=False
):
    """Get optimizer parameters for different surface parameters.

    Args:
        lrs (list): Learning rates for [d, c, k, q_coefficients].
        decay (float): Decay factor for higher-order Q coefficients.
        optim_mat (bool): Whether to optimize material parameters.

    Returns:
        list: Parameter groups for optimizer
    """
    params = []

    # Distance
    self.d.requires_grad_(True)
    params.append({"params": [self.d], "lr": lrs[0]})

    # Curvature
    self.c.requires_grad_(True)
    params.append({"params": [self.c], "lr": lrs[1]})

    # Conic constant
    self.k.requires_grad_(True)
    params.append({"params": [self.k], "lr": lrs[2]})

    # Q polynomial coefficients
    if self.n_qterms > 0:
        base_lr = lrs[3] if len(lrs) > 3 else 1e-6
        for m in range(self.n_qterms):
            qm = getattr(self, f"q{m}")
            qm.requires_grad_(True)
            # Decay learning rate for higher order terms
            lr = base_lr * (decay**m)
            params.append({"params": [qm], "lr": lr})

    # Material parameters
    if optim_mat and self.mat2.get_name() != "air":
        params += self.mat2.get_optimizer_params()

    return params

init_tolerance

init_tolerance(tolerance_params=None)

Initialize tolerance parameters for manufacturing error simulation.

Source code in src/geometric_surface/qtype.py

@torch.no_grad()
def init_tolerance(self, tolerance_params=None):
    """Initialize tolerance parameters for manufacturing error simulation."""
    super().init_tolerance(tolerance_params)
    if tolerance_params is None:
        tolerance_params = {}
    self.c_tole = tolerance_params.get("c_tole", 0.001)
    self.k_tole = tolerance_params.get("k_tole", 0.001)
    self.c_error = 0.0
    self.k_error = 0.0

sample_tolerance

sample_tolerance()

Sample random manufacturing errors.

Source code in src/geometric_surface/qtype.py

def sample_tolerance(self):
    """Sample random manufacturing errors."""
    super().sample_tolerance()
    self.c_error = float(np.random.randn() * self.c_tole)
    self.k_error = float(np.random.randn() * self.k_tole)

zero_tolerance

zero_tolerance()

Reset all tolerances to zero.

Source code in src/geometric_surface/qtype.py

def zero_tolerance(self):
    """Reset all tolerances to zero."""
    super().zero_tolerance()
    self.c_error = 0.0
    self.k_error = 0.0

surf_dict

surf_dict()

Return dictionary representation of surface.

Source code in src/geometric_surface/qtype.py

def surf_dict(self):
    """Return dictionary representation of surface."""
    surf_dict = {
        "type": "QTypeFreeform",
        "r": round(self.r, 4),
        "d": round(self.d.item(), 4),
        "(c)": round(self.c.item(), 6),
        "roc": round(1 / self.c.item(), 4)
        if abs(self.c.item()) > EPSILON
        else float("inf"),
        "k": round(self.k.item(), 6),
        "r_norm": round(self.r_norm, 4),
        "qm": [],
        "mat2": self.mat2.get_name(),
    }

    for m in range(self.n_qterms):
        qm = getattr(self, f"q{m}")
        surf_dict["qm"].append(float(format(qm.item(), ".6e")))
        surf_dict[f"(q{m})"] = float(format(qm.item(), ".6e"))

    return surf_dict

zmx_str

zmx_str(surf_idx, d_next)

Return Zemax surface string.

Note: Zemax uses a different Q-type representation (QTYPE surface). This export is approximate and may need adjustment for specific Zemax versions.

Source code in src/geometric_surface/qtype.py

    def zmx_str(self, surf_idx, d_next):
        """Return Zemax surface string.

        Note: Zemax uses a different Q-type representation (QTYPE surface).
        This export is approximate and may need adjustment for specific Zemax versions.
        """
        if self.mat2.get_name() == "air":
            zmx_str = f"""SURF {surf_idx}
    TYPE QTYPE
    CURV {self.c.item()}
    DISZ {d_next.item()}
    DIAM {self.r} 1 0 0 1 ""
    CONI {self.k.item()}
    PARM 1 {self.r_norm}
"""
        else:
            zmx_str = f"""SURF {surf_idx}
    TYPE QTYPE
    CURV {self.c.item()}
    DISZ {d_next.item()}
    GLAS ___BLANK 1 0 {self.mat2.n} {self.mat2.V}
    DIAM {self.r} 1 0 0 1 ""
    CONI {self.k.item()}
    PARM 1 {self.r_norm}
"""

        # Add Q coefficients
        for m in range(self.n_qterms):
            qm = getattr(self, f"q{m}")
            zmx_str += f"    PARM {m + 2} {qm.item()}\n"

        return zmx_str

\[z = z_{\text{conic}} + u^4 \sum_{m=0}^{M} a_m \, Q_m^{\text{bfs}}(u^2)\]

where \(u = \rho / r_{\text{norm}}\).

Diffractive Surfaces

Phase surfaces for diffractive optical elements (DOEs) and metalenses are in src.phase_surface:

Surface	Description
`Binary2Phase`	Zemax Binary 2 even-order polynomial phase
`FresnelPhase`	Fresnel lens phase profile
`ZernikePhase`	Zernike polynomial phase
`PolyPhase`	General polynomial phase
`GratingPhase`	Diffraction grating phase
`CubicPhase`	Cubic phase mask
`NURBSPhase`	NURBS-based freeform phase

Surfaces API Reference

Base Class

Surface

src.geometric_surface.Surface

init_from_dict classmethod

ray_reaction

intersect

newtons_method

refract

reflect

normal_vec

to_local_coord

to_global_coord

sag

dfdxyz

d2fdxyz2

is_valid

is_within_boundary

is_within_data_range

max_height

surface_with_offset

surface_sag

get_optimizer_params

get_optimizer

update_r

init_tolerance

sample_tolerance

zero_tolerance

sensitivity_score

draw_r

draw_widget

create_mesh

get_polydata

zmx_str

Spheric

src.geometric_surface.Spheric

intersect

is_within_data_range

max_height

init_tolerance

sample_tolerance

zero_tolerance

sensitivity_score

get_optimizer_params

surf_dict

zmx_str

Aspheric

src.geometric_surface.Aspheric

ai property writable

c property writable

k property writable

enable_reparam

disable_reparam

update_r

is_within_data_range

max_height

get_optimizer_params

init_tolerance

sample_tolerance

zero_tolerance

sensitivity_score

surf_dict

zmx_str

Reparametrization

Aperture

src.geometric_surface.Aperture

ray_reaction

draw_widget

draw_widget3D

create_mesh

get_optimizer_params

surf_dict

zmx_str

Plane

src.geometric_surface.Plane

intersect

normal_vec

get_optimizer_params

Mirror

src.geometric_surface.Mirror

init_from_dict `classmethod`

ai `property` `writable`

c `property` `writable`

k `property` `writable`

init_from_dict `classmethod`