GeoLens

Differentiable multi-element refractive lens via geometric ray tracing. GeoLens is the primary lens model in DeepLens: it ray-traces through a stack of optical surfaces to compute PSFs, render images, and optimize lens geometry end-to-end.

deeplens.GeoLens

GeoLens(
    filename=None,
    device=None,
    dtype=torch.float32,
    primary_wvln=DEFAULT_WAVE,
    wvln_rgb=WAVE_RGB,
    obj_depth=DEPTH,
)

Bases: GeoLensPSF, GeoLensRender, GeoLensEval, GeoLensOptim, GeoLensOps, GeoLensVis, GeoLensIO, GeoLensVis3D, Lens

Differentiable geometric lens using vectorised ray tracing.

The primary lens model in DeepLens. Supports multi-element refractive (and partially reflective) systems loaded from JSON, Zemax .zmx, or Code V .seq files. Accuracy is aligned with Zemax OpticStudio.

Uses a mixin architecture: eight specialised mixin classes are composed at class-definition time to keep each concern isolated: GeoLensPSF (PSF computation), GeoLensRender (image simulation: render dispatch and reverse ray tracing), GeoLensEval (spot/MTF/distortion/vignetting evaluation), GeoLensOptim (losses and gradient-based optimisation), GeoLensOps (in-place lens operations), GeoLensVis (2-D layout/ray visualisation), GeoLensIO (JSON/Zemax read-write), and GeoLensVis3D (3-D mesh visualisation).

Attributes:

Name	Type	Description
surfaces	list[Surface]	Ordered list of optical surfaces.
d_sensor	Tensor	Derived distance from the first surface to the sensor plane [mm].
foclen	float	Effective focal length [mm].
fnum	float	F-number.
rfov	float	Real half-diagonal field of view [radians].
sensor_size	tuple	Physical sensor size (W, H) [mm].
sensor_res	tuple	Sensor resolution (W, H) [pixels].
pixel_size	float	Pixel pitch [mm].

Reference

Xinge Yang et al., "Curriculum learning for ab initio deep learned refractive optics," Nature Communications 2024.

Initialize a refractive lens.

There are two ways to initialize a GeoLens

1. Read a lens from .json/.zmx/.seq file
2. Initialize a lens with no lens file, then manually add surfaces

Parameters:

Name	Type	Description	Default
filename	str	Path to lens file (.json, .zmx, or .seq). Defaults to None.	None
device	device	Device for tensor computations. Defaults to None.	None
dtype	dtype	Data type for computations. Defaults to torch.float32.	float32
primary_wvln	float	Primary design wavelength [µm]. Used as fallback when a method is called without an explicit wvln. Defaults to DEFAULT_WAVE.	DEFAULT_WAVE
wvln_rgb	sequence of float	Three wavelengths used for RGB computations, ordered [R, G, B] in µm. Defaults to WAVE_RGB.	WAVE_RGB
obj_depth	float	Default object depth [mm], used when a method is called without an explicit depth. Defaults to DEPTH.	DEPTH

Source code in deeplens-src/deeplens/geolens.py

def __init__(
    self,
    filename=None,
    device=None,
    dtype=torch.float32,
    primary_wvln=DEFAULT_WAVE,
    wvln_rgb=WAVE_RGB,
    obj_depth=DEPTH,
):
    """Initialize a refractive lens.

    There are two ways to initialize a GeoLens:
        1. Read a lens from .json/.zmx/.seq file
        2. Initialize a lens with no lens file, then manually add surfaces

    Args:
        filename (str, optional): Path to lens file (.json, .zmx, or .seq). Defaults to None.
        device (torch.device, optional): Device for tensor computations. Defaults to None.
        dtype (torch.dtype, optional): Data type for computations. Defaults to torch.float32.
        primary_wvln (float, optional): Primary design wavelength [µm].
            Used as fallback when a method is called without an explicit
            ``wvln``.  Defaults to ``DEFAULT_WAVE``.
        wvln_rgb (sequence of float, optional): Three wavelengths used
            for RGB computations, ordered ``[R, G, B]`` in µm.  Defaults
            to ``WAVE_RGB``.
        obj_depth (float, optional): Default object depth [mm], used
            when a method is called without an explicit ``depth``.
            Defaults to ``DEPTH``.
    """
    super().__init__(
        device=device,
        dtype=dtype,
        primary_wvln=primary_wvln,
        wvln_rgb=wvln_rgb,
        obj_depth=obj_depth,
    )
    # Load lens file
    if filename is not None:
        self.read_lens(filename)
    else:
        self.surfaces = []
        # Placeholder until the caller adds surfaces and calls
        # post_computation(), which sets aper_idx via calc_pupil().
        self.aper_idx = None
        # Set default sensor size and resolution
        self.sensor_size = (8.0, 8.0)
        self.sensor_res = (2000, 2000)
        self.to(self.device)

d_sensor property writable

d_sensor

Global axial position of the sensor plane [mm].

Surface 0 is the origin. The sensor position is derived by summing each surface's d_next, so the image plane remains part of the same differentiable sequential thickness chain.

read_lens

read_lens(filename)

Read a GeoLens from a file.

Supported file formats

• .json: DeepLens native JSON format
• .zmx: Zemax lens file format
• .seq: CODE V sequence file format

Parameters:

Name	Type	Description	Default
filename	str	Path to the lens file.	required

Note

Sensor size and resolution will usually be overwritten by values from the file.

Source code in deeplens-src/deeplens/geolens.py

def read_lens(self, filename):
    """Read a GeoLens from a file.

    Supported file formats:
        - .json: DeepLens native JSON format
        - .zmx: Zemax lens file format
        - .seq: CODE V sequence file format

    Args:
        filename (str): Path to the lens file.

    Note:
        Sensor size and resolution will usually be overwritten by values from the file.
    """
    # Load lens file
    if filename[-5:] == ".json":
        self.read_lens_json(filename)
    elif filename[-4:] == ".zmx":
        self.read_lens_zmx(filename)
    elif filename[-4:] == ".seq":
        self.read_lens_seq(filename)
    else:
        raise ValueError(f"File format {filename[-4:]} not supported.")

    # After loading lens, compute foclen, fov and fnum
    self.create_dummy_sensor()
    self.to(self.device)
    self.astype(self.dtype)
    self.post_computation()

post_computation

post_computation()

Compute derived optical properties after loading or modifying lens.

Calculates and caches

• Effective focal length (EFL)
• Entrance and exit pupil positions and radii
• Field of view (FoV) in horizontal, vertical, and diagonal directions
• F-number
• Lens design constraints (edge/center thickness bounds, etc.)

Note

This method should be called after any changes to the lens geometry.

Source code in deeplens-src/deeplens/geolens.py

def post_computation(self):
    """Compute derived optical properties after loading or modifying lens.

    Calculates and caches:
        - Effective focal length (EFL)
        - Entrance and exit pupil positions and radii
        - Field of view (FoV) in horizontal, vertical, and diagonal directions
        - F-number
        - Lens design constraints (edge/center thickness bounds, etc.)

    Note:
        This method should be called after any changes to the lens geometry.
    """
    self.calc_foclen()
    self.calc_pupil()
    self.calc_fov()
    self.init_constraints()

__call__

__call__(ray)

Trace rays through the lens system (callable shorthand for trace).

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object to trace.	required

Returns:

Name	Type	Description
ray_out	Ray	Ray after propagation through the surfaces.
ray_o_record	list or None	Recorded ray positions, or None.

Source code in deeplens-src/deeplens/geolens.py

def __call__(self, ray):
    """Trace rays through the lens system (callable shorthand for `trace`).

    Args:
        ray (Ray): Ray object to trace.

    Returns:
        ray_out (Ray): Ray after propagation through the surfaces.
        ray_o_record (list or None): Recorded ray positions, or None.
    """
    return self.trace(ray)

sample_grid_rays

sample_grid_rays(
    depth=float("inf"),
    num_grid=(11, 11),
    num_rays=SPP_PSF,
    wvln=None,
    uniform_fov=True,
    sample_more_off_axis=False,
    scale_pupil=1.0,
)

Sample a grid of rays spanning the field of view from object space.

If depth is infinite, samples collimated rays at evenly-spaced field angles; if depth is finite, samples diverging point-source rays from a grid of object points. Used for PSF maps, RMS error maps, and spot diagrams.

Parameters:

Name	Type	Description	Default
depth	float	Object distance in mm. Use float("inf") for collimated light. Defaults to float("inf").	float('inf')
num_grid	int or tuple	Number of grid points as (num_x, num_y), or a single int for both. Defaults to (11, 11).	(11, 11)
num_rays	int	Number of rays per grid point. Defaults to SPP_PSF.	SPP_PSF
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
uniform_fov	bool	If True, sample uniform FoV angles; otherwise sample a uniform object grid. Defaults to True.	True
sample_more_off_axis	bool	If True, concentrate grid samples toward off-axis fields. Defaults to False.	False
scale_pupil	float	Scale factor for pupil radius. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
rays	Ray	Sampled rays with shape [num_grid[1], num_grid[0], num_rays, 3].

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def sample_grid_rays(
    self,
    depth=float("inf"),
    num_grid=(11, 11),
    num_rays=SPP_PSF,
    wvln=None,
    uniform_fov=True,
    sample_more_off_axis=False,
    scale_pupil=1.0,
):
    """Sample a grid of rays spanning the field of view from object space.

    If `depth` is infinite, samples collimated rays at evenly-spaced field
    angles; if `depth` is finite, samples diverging point-source rays from a
    grid of object points. Used for PSF maps, RMS error maps, and spot
    diagrams.

    Args:
        depth (float, optional): Object distance in mm. Use `float("inf")`
            for collimated light. Defaults to `float("inf")`.
        num_grid (int or tuple, optional): Number of grid points as
            (num_x, num_y), or a single int for both. Defaults to (11, 11).
        num_rays (int, optional): Number of rays per grid point. Defaults to SPP_PSF.
        wvln (float, optional): Wavelength in µm. When None (default),
            falls back to `self.primary_wvln`.
        uniform_fov (bool, optional): If True, sample uniform FoV angles;
            otherwise sample a uniform object grid. Defaults to True.
        sample_more_off_axis (bool, optional): If True, concentrate grid
            samples toward off-axis fields. Defaults to False.
        scale_pupil (float, optional): Scale factor for pupil radius. Defaults to 1.0.

    Returns:
        rays (Ray): Sampled rays with shape [num_grid[1], num_grid[0], num_rays, 3].
    """
    wvln = self.primary_wvln if wvln is None else wvln

    # Normalize num_grid to a tuple if it's an int
    if isinstance(num_grid, int):
        num_grid = (num_grid, num_grid)

    # Calculate field angles for grid source. Top-left field has positive fov_x and negative fov_y
    x_list = [x for x in np.linspace(1, -1, num_grid[0])]
    y_list = [y for y in np.linspace(-1, 1, num_grid[1])]
    if sample_more_off_axis:
        x_list = [np.sign(x) * np.abs(x) ** 0.5 for x in x_list]
        y_list = [np.sign(y) * np.abs(y) ** 0.5 for y in y_list]

    # Calculate FoV_x and FoV_y
    if uniform_fov:
        # Sample uniform FoV angles
        fov_x_list = [x * self.hfov / 2 for x in x_list]
        fov_y_list = [y * self.vfov / 2 for y in y_list]
        fov_x_list = [float(np.rad2deg(fov_x)) for fov_x in fov_x_list]
        fov_y_list = [float(np.rad2deg(fov_y)) for fov_y in fov_y_list]
    else:
        # Sample uniform object grid
        fov_x_list = [np.arctan(x * np.tan(self.hfov / 2)) for x in x_list]
        fov_y_list = [np.arctan(y * np.tan(self.vfov / 2)) for y in y_list]
        fov_x_list = [float(np.rad2deg(fov_x)) for fov_x in fov_x_list]
        fov_y_list = [float(np.rad2deg(fov_y)) for fov_y in fov_y_list]

    # Sample rays (collimated or point source via unified API)
    rays = self.sample_from_fov(
        fov_x=fov_x_list,
        fov_y=fov_y_list,
        depth=depth,
        num_rays=num_rays,
        wvln=wvln,
        scale_pupil=scale_pupil,
    )
    return rays

sample_radial_rays

sample_radial_rays(
    num_field=5,
    depth=float("inf"),
    num_rays=SPP_PSF,
    wvln=None,
    direction="y",
    fov_max=None,
)

Sample radial rays at evenly-spaced field angles along a chosen direction.

The sampled angles are radial field angles: for "diagonal" the per-axis components are atan(tan(fov) / sqrt(2)) so that a field listed as fov really lands at radial angle fov, matching the "x" and "y" directions.

Parameters:

Name	Type	Description	Default
num_field	int	Number of field angles from on-axis to full-field. Defaults to 5.	5
depth	float	Object distance in mm. Use float('inf') for collimated light. Defaults to float('inf').	float('inf')
num_rays	int	Rays per field position. Defaults to SPP_PSF.	SPP_PSF
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
direction	str	Sampling direction — "y" (meridional, default), "x" (sagittal), "diagonal" (45° azimuth).	'y'
fov_max	float	Full-field radial angle [radians]. When None (default), falls back to self.rfov_eff (paraxial pinhole FoV). Pass self.rfov for the real ray-traced FoV, which reaches the sensor corner on a distorting lens.	None

Returns:

Name	Type	Description
ray	Ray	Ray object with shape [num_field, num_rays, 3].

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def sample_radial_rays(
    self,
    num_field=5,
    depth=float("inf"),
    num_rays=SPP_PSF,
    wvln=None,
    direction="y",
    fov_max=None,
):
    """Sample radial rays at evenly-spaced field angles along a chosen direction.

    The sampled angles are *radial* field angles: for ``"diagonal"`` the
    per-axis components are ``atan(tan(fov) / sqrt(2))`` so that a field
    listed as ``fov`` really lands at radial angle ``fov``, matching the
    ``"x"`` and ``"y"`` directions.

    Args:
        num_field (int): Number of field angles from on-axis to full-field.
            Defaults to 5.
        depth (float): Object distance in mm. Use ``float('inf')`` for
            collimated light. Defaults to ``float('inf')``.
        num_rays (int): Rays per field position. Defaults to ``SPP_PSF``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.
        direction (str): Sampling direction —
            ``"y"`` (meridional, default),
            ``"x"`` (sagittal),
            ``"diagonal"`` (45° azimuth).
        fov_max (float): Full-field radial angle [radians]. When ``None``
            (default), falls back to ``self.rfov_eff`` (paraxial pinhole
            FoV). Pass ``self.rfov`` for the real ray-traced FoV, which
            reaches the sensor corner on a distorting lens.

    Returns:
        ray (Ray): Ray object with shape ``[num_field, num_rays, 3]``.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    device = self.device
    fov_max = self.rfov_eff if fov_max is None else fov_max
    fov_deg = fov_max * 180 / torch.pi
    fov_list = torch.linspace(
        0, fov_deg, num_field, device=device, dtype=self.dtype
    )

    if direction == "y":
        ray = self.sample_from_fov(
            fov_x=0.0, fov_y=fov_list, depth=depth, num_rays=num_rays, wvln=wvln
        )
    elif direction == "x":
        ray = self.sample_from_fov(
            fov_x=fov_list, fov_y=0.0, depth=depth, num_rays=num_rays, wvln=wvln
        )
    elif direction == "diagonal":
        # Split the radial angle across both axes so the field lands at
        # radial angle `f`, not at atan(sqrt(2) * tan(f)).
        fov_axis = torch.rad2deg(
            torch.atan(torch.tan(torch.deg2rad(fov_list)) / math.sqrt(2))
        )
        # sample_from_fov creates a meshgrid; for pairwise diagonal, loop
        rays = [
            self.sample_from_fov(
                fov_x=f.item(),
                fov_y=f.item(),
                depth=depth,
                num_rays=num_rays,
                wvln=wvln,
            )
            for f in fov_axis
        ]
        ray_o = torch.stack([r.o for r in rays], dim=0)
        ray_d = torch.stack([r.d for r in rays], dim=0)
        ray = Ray(ray_o, ray_d, wvln, device=device)
    else:
        raise ValueError(
            f"Invalid direction: {direction!r}. Use 'x', 'y', or 'diagonal'."
        )
    return ray

sample_from_points

sample_from_points(
    points=[[0.0, 0.0, -10000.0]],
    num_rays=SPP_PSF,
    wvln=None,
    entrance_pupil=True,
    scale_pupil=1.0,
)

Sample rays from point sources in object space (absolute physical coordinates).

Rays originate at the given object points and fan out toward the entrance pupil. Used for PSF and chief-ray calculation.

Parameters:

Name	Type	Description	Default
points	list or Tensor	Object-space ray origins [mm] with shape [3], [N, 3], or [Nx, Ny, 3]. Defaults to [[0.0, 0.0, -10000.0]].	[[0.0, 0.0, -10000.0]]
num_rays	int	Number of rays per point. Defaults to SPP_PSF.	SPP_PSF
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
entrance_pupil	bool	If True (default), aim rays at the entrance pupil; otherwise at surface 0.	True
scale_pupil	float	Scale factor for pupil radius. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
rays	Ray	Sampled rays with shape [*points.shape[:-1], num_rays, 3].

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def sample_from_points(
    self,
    points=[[0.0, 0.0, -10000.0]],
    num_rays=SPP_PSF,
    wvln=None,
    entrance_pupil=True,
    scale_pupil=1.0,
):
    """Sample rays from point sources in object space (absolute physical coordinates).

    Rays originate at the given object points and fan out toward the
    entrance pupil. Used for PSF and chief-ray calculation.

    Args:
        points (list or torch.Tensor): Object-space ray origins [mm] with
            shape [3], [N, 3], or [Nx, Ny, 3]. Defaults to [[0.0, 0.0, -10000.0]].
        num_rays (int): Number of rays per point. Defaults to SPP_PSF.
        wvln (float): Wavelength in µm. When None (default), falls back to
            `self.primary_wvln`.
        entrance_pupil (bool): If True (default), aim rays at the entrance
            pupil; otherwise at surface 0.
        scale_pupil (float): Scale factor for pupil radius. Defaults to 1.0.

    Returns:
        rays (Ray): Sampled rays with shape [*points.shape[:-1], num_rays, 3].
    """
    wvln = self.primary_wvln if wvln is None else wvln

    # Ray origin is given
    if not torch.is_tensor(points):
        ray_o = torch.as_tensor(points, device=self.device, dtype=self.dtype)
    else:
        ray_o = points.to(device=self.device, dtype=self.dtype)

    # Sample points on the pupil
    if entrance_pupil:
        pupilz, pupilr = self.get_entrance_pupil()
    else:
        pupilz, pupilr = self.surf_d(0).item(), self.surfaces[0].r
    pupilr *= scale_pupil
    ray_o2 = self.sample_circle(
        r=pupilr, z=pupilz, shape=(*ray_o.shape[:-1], num_rays)
    )

    # Compute ray directions
    if len(ray_o.shape) == 1:
        # Input point shape is [3]
        ray_o = ray_o.unsqueeze(0).repeat(num_rays, 1)  # shape [num_rays, 3]
        ray_d = ray_o2 - ray_o

    elif len(ray_o.shape) == 2:
        # Input point shape is [N, 3]
        ray_o = ray_o.unsqueeze(1).repeat(1, num_rays, 1)  # shape [N, num_rays, 3]
        ray_d = ray_o2 - ray_o

    elif len(ray_o.shape) == 3:
        # Input point shape is [Nx, Ny, 3]
        ray_o = ray_o.unsqueeze(2).repeat(
            1, 1, num_rays, 1
        )  # shape [Nx, Ny, num_rays, 3]
        ray_d = ray_o2 - ray_o

    else:
        raise Exception("The shape of input object positions is not supported.")

    # The physical-stop distance (`ray.stop_dist`) is stamped only when
    # tracing crosses the aperture; entrance-pupil sampling is not the
    # chief-ray definition.
    rays = Ray(ray_o, ray_d, wvln, device=self.device)
    return rays

sample_from_fov

sample_from_fov(
    fov_x=[0.0],
    fov_y=[0.0],
    depth=float("inf"),
    num_rays=SPP_CALC,
    wvln=None,
    entrance_pupil=True,
    scale_pupil=1.0,
)

Sample rays from object space at given field angles.

For infinite depth, generates collimated parallel rays: origins are distributed on the entrance pupil and all rays in a field share the same direction determined by the FOV angle.

For finite depth, generates diverging point-source rays: the point source position is determined by FOV angle and depth, and rays fan out toward the entrance pupil.

Parameters:

Name	Type	Description	Default
fov_x	float or list	Field angle(s) in the xz plane (degrees).	[0.0]
fov_y	float or list	Field angle(s) in the yz plane (degrees).	[0.0]
depth	float	Object distance in mm. float('inf') for collimated rays, finite for point-source rays.	float('inf')
num_rays	int	Number of rays per field point.	SPP_CALC
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
entrance_pupil	bool	If True, sample on entrance pupil; otherwise on surface 0. Default: True.	True
scale_pupil	float	Scale factor for pupil radius.	1.0

Returns:

Name	Type	Description
rays	Ray	Rays with shape [..., num_rays, 3], where leading dims are squeezed when the corresponding fov input is scalar.

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def sample_from_fov(
    self,
    fov_x=[0.0],
    fov_y=[0.0],
    depth=float("inf"),
    num_rays=SPP_CALC,
    wvln=None,
    entrance_pupil=True,
    scale_pupil=1.0,
):
    """Sample rays from object space at given field angles.

    For infinite depth, generates collimated parallel rays: origins are
    distributed on the entrance pupil and all rays in a field share the
    same direction determined by the FOV angle.

    For finite depth, generates diverging point-source rays: the point
    source position is determined by FOV angle and depth, and rays fan
    out toward the entrance pupil.

    Args:
        fov_x (float or list): Field angle(s) in the xz plane (degrees).
        fov_y (float or list): Field angle(s) in the yz plane (degrees).
        depth (float): Object distance in mm. ``float('inf')`` for
            collimated rays, finite for point-source rays.
        num_rays (int): Number of rays per field point.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.
        entrance_pupil (bool): If True, sample on entrance pupil;
            otherwise on surface 0. Default: True.
        scale_pupil (float): Scale factor for pupil radius.

    Returns:
        rays (Ray): Rays with shape ``[..., num_rays, 3]``, where leading dims
            are squeezed when the corresponding fov input is scalar.
    """
    wvln = self.primary_wvln if wvln is None else wvln

    # Track which inputs were scalar for output shape
    x_scalar = isinstance(fov_x, (float, int))
    y_scalar = isinstance(fov_y, (float, int))
    if x_scalar:
        fov_x = [float(fov_x)]
    if y_scalar:
        fov_y = [float(fov_y)]

    fov_x_rad = torch.as_tensor(fov_x, device=self.device, dtype=self.dtype) * (
        math.pi / 180.0
    )
    fov_y_rad = torch.as_tensor(fov_y, device=self.device, dtype=self.dtype) * (
        math.pi / 180.0
    )
    fov_x_grid, fov_y_grid = torch.meshgrid(fov_x_rad, fov_y_rad, indexing="xy")

    if depth == float("inf"):
        # Collimated rays: origins on pupil, uniform direction per field
        if entrance_pupil:
            pupilz, pupilr = self.get_entrance_pupil()
        else:
            pupilz, pupilr = self.surf_d(0).item(), self.surfaces[0].r
        pupilr *= scale_pupil
        ray_o = self.sample_circle(
            r=pupilr, z=pupilz, shape=[len(fov_y), len(fov_x), num_rays]
        )
        dx = torch.tan(fov_x_grid).unsqueeze(-1).expand_as(ray_o[..., 0])
        dy = torch.tan(fov_y_grid).unsqueeze(-1).expand_as(ray_o[..., 1])
        dz = torch.ones_like(ray_o[..., 2])
        ray_d = torch.stack((dx, dy, dz), dim=-1)

        if x_scalar:
            ray_o = ray_o.squeeze(1)
            ray_d = ray_d.squeeze(1)
        if y_scalar:
            ray_o = ray_o.squeeze(0)
            ray_d = ray_d.squeeze(0)

        rays = Ray(ray_o, ray_d, wvln, device=self.device)
        rays.prop_to(-1.0)

    else:
        # Point-source rays: origin at object point, fan toward pupil
        x = torch.tan(fov_x_grid) * depth
        y = torch.tan(fov_y_grid) * depth
        z = torch.full_like(x, depth)
        points = torch.stack((x, y, z), dim=-1)

        if x_scalar:
            points = points.squeeze(-2)
        if y_scalar:
            points = points.squeeze(0)

        rays = self.sample_from_points(
            points=points,
            num_rays=num_rays,
            wvln=wvln,
            entrance_pupil=entrance_pupil,
            scale_pupil=scale_pupil,
        )

    return rays

sample_circle

sample_circle(r, z, shape=[16, 16, 512])

Sample points uniformly inside a circle on a constant-z plane.

Parameters:

Name	Type	Description	Default
r	float	Radius of the circle [mm].	required
z	float	Z-coordinate shared by all sampled points [mm].	required
shape	list	Shape of the point grid (excluding the trailing coordinate dimension). Defaults to [16, 16, 512].	[16, 16, 512]

Returns:

Name	Type	Description
points	Tensor	Sampled points with shape [*shape, 3].

Source code in deeplens-src/deeplens/geolens.py

def sample_circle(self, r, z, shape=[16, 16, 512]):
    """Sample points uniformly inside a circle on a constant-z plane.

    Args:
        r (float): Radius of the circle [mm].
        z (float): Z-coordinate shared by all sampled points [mm].
        shape (list): Shape of the point grid (excluding the trailing
            coordinate dimension). Defaults to [16, 16, 512].

    Returns:
        points (torch.Tensor): Sampled points with shape [*shape, 3].
    """
    device = self.device

    # Generate random angles and radii
    theta = torch.rand(*shape, device=device, dtype=self.dtype) * 2 * torch.pi
    r2 = torch.rand(*shape, device=device, dtype=self.dtype) * r**2
    radius = torch.sqrt(r2)

    # Stack to form 3D points
    x = radius * torch.cos(theta)
    y = radius * torch.sin(theta)
    z_tensor = torch.full_like(x, z)
    points = torch.stack((x, y, z_tensor), dim=-1)

    # Manually sample chief ray
    # points[..., 0, :2] = 0.0

    return points

trace

trace(ray, surf_range=None, record=False)

Trace rays through the lens.

Forward or backward tracing is selected automatically from the sign of the ray z-direction.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object to trace.	required
surf_range	range	Range of surface indices to trace through. When None (default), traces through all surfaces.	None
record	bool	If True, record ray positions at each surface. Defaults to False.	False

Returns:

Name	Type	Description
ray_out	Ray	Ray after propagation through the surfaces.
ray_o_record	list or None	Recorded ray positions at each surface, or None when record is False.

Source code in deeplens-src/deeplens/geolens.py

def trace(self, ray, surf_range=None, record=False):
    """Trace rays through the lens.

    Forward or backward tracing is selected automatically from the sign of
    the ray z-direction.

    Args:
        ray (Ray): Ray object to trace.
        surf_range (range, optional): Range of surface indices to trace
            through. When None (default), traces through all surfaces.
        record (bool): If True, record ray positions at each surface. Defaults to False.

    Returns:
        ray_out (Ray): Ray after propagation through the surfaces.
        ray_o_record (list or None): Recorded ray positions at each surface,
            or None when record is False.
    """
    if surf_range is None:
        surf_range = range(0, len(self.surfaces))

    forward = ray.d[..., 2] > 0
    backward = ray.d[..., 2] < 0
    if bool(forward.all().item()):
        ray_out, ray_o_rec = self.forward_tracing(ray, surf_range, record=record)
    elif bool(backward.all().item()):
        ray_out, ray_o_rec = self.backward_tracing(ray, surf_range, record=record)
    else:
        raise ValueError(
            "A ray bundle must have one tracing direction; split mixed "
            "forward/backward or zero-z-direction rays before tracing."
        )

    return ray_out, ray_o_rec

trace2obj

trace2obj(ray)

Trace rays through the lens toward object space.

Convenience wrapper around trace that discards the position record. Typically called with sensor-side (backward-propagating) rays.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object to trace.	required

Returns:

Name	Type	Description
ray	Ray	Ray after propagation through the lens.

Source code in deeplens-src/deeplens/geolens.py

def trace2obj(self, ray):
    """Trace rays through the lens toward object space.

    Convenience wrapper around `trace` that discards the position record.
    Typically called with sensor-side (backward-propagating) rays.

    Args:
        ray (Ray): Ray object to trace.

    Returns:
        ray (Ray): Ray after propagation through the lens.
    """
    ray, _ = self.trace(ray)
    return ray

trace2sensor

trace2sensor(ray, record=False)

Forward trace rays through the lens and propagate them to the sensor plane.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object to trace.	required
record	bool	If True, record ray positions at each surface. Defaults to False.	False

Returns:

Name	Type	Description
ray	Ray	Ray propagated to the sensor plane. When record is True, returns a tuple (ray, ray_o_record) where ray_o_record is the list of recorded ray positions at each surface (invalid points set to NaN).

Source code in deeplens-src/deeplens/geolens.py

def trace2sensor(self, ray, record=False):
    """Forward trace rays through the lens and propagate them to the sensor plane.

    Args:
        ray (Ray): Ray object to trace.
        record (bool): If True, record ray positions at each surface. Defaults to False.

    Returns:
        ray (Ray): Ray propagated to the sensor plane. When record is True,
            returns a tuple (ray, ray_o_record) where ray_o_record is the list
            of recorded ray positions at each surface (invalid points set to NaN).
    """
    # Trace rays
    ray, ray_o_record = self.trace(ray, record=record)
    ray = ray.prop_to(self.d_sensor)

    if record:
        ray_o = ray.o.clone().detach()
        # Set to NaN to be skipped in 2d layout visualization
        ray_o[ray.is_valid == 0] = float("nan")
        ray_o_record.append(ray_o)
        return ray, ray_o_record
    else:
        return ray

trace2exit_pupil

trace2exit_pupil(ray)

Forward trace rays through the lens to exit pupil plane.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object to trace.	required

Returns:

Name	Type	Description
ray	Ray	Ray object propagated to the exit pupil plane.

Source code in deeplens-src/deeplens/geolens.py

def trace2exit_pupil(self, ray):
    """Forward trace rays through the lens to exit pupil plane.

    Args:
        ray (Ray): Ray object to trace.

    Returns:
        ray (Ray): Ray object propagated to the exit pupil plane.
    """
    ray = self.trace2sensor(ray)
    pupil_z, _ = self.get_exit_pupil()
    ray = ray.prop_to(pupil_z)
    return ray

forward_tracing

forward_tracing(ray, surf_range, record)

Trace forward using sequential per-surface reference frames.

Rays enter and leave in global coordinates. Interactions happen with each vertex at local z=0; after a surface, the ray origin is shifted by -d_next to express it in the next surface's frame.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object to trace.	required
surf_range	range	Range of surface indices to trace through.	required
record	bool	If True, record ray positions at each surface.	required

Returns:

Name	Type	Description
ray_out	Ray	Ray after propagation through all surfaces.
ray_o_record	list or None	Ray positions at each surface, or None if record is False.

Source code in deeplens-src/deeplens/geolens.py

def forward_tracing(self, ray, surf_range, record):
    """Trace forward using sequential per-surface reference frames.

    Rays enter and leave in global coordinates. Interactions happen with
    each vertex at local z=0; after a surface, the ray origin is shifted by
    `-d_next` to express it in the next surface's frame.

    Args:
        ray (Ray): Ray object to trace.
        surf_range (range): Range of surface indices to trace through.
        record (bool): If True, record ray positions at each surface.

    Returns:
        ray_out (Ray): Ray after propagation through all surfaces.
        ray_o_record (list or None): Ray positions at each surface, or None
            if record is False.
    """
    if record:
        ray_o_record = []
        ray_o_record.append(ray.o.clone().detach())
    else:
        ray_o_record = None

    surf_indices = self._validate_surf_range(surf_range)
    if not surf_indices:
        return ray, ray_o_record
    first, last = min(surf_indices), max(surf_indices)
    surf_set = set(surf_indices)

    z_frame = self.surf_d(first)
    ray.o[..., 2] -= z_frame

    mat1 = Material("air") if first == 0 else self.surfaces[first - 1].mat2

    # Re-anchor a far object-space bundle once per trace, after entering the
    # first selected surface's local frame: a float32 origin at large |z|
    # loses the low-order sag when the origin and the intersection distance
    # cancel. Anchoring at -10 mm (not -1 mm) keeps rays clear of a concave
    # first surface, whose edge can reach several mm of negative sag.
    if bool((ray.o[..., 2] < -10.0).any()):
        ray.prop_to(-10.0, n=mat1.ior(ray.wvln))

    for i in range(first, last + 1):
        surf = self.surfaces[i]
        if i in surf_set:
            n1 = mat1.ior(ray.wvln)
            n2 = surf.mat2.ior(ray.wvln)
            ray = surf.ray_reaction(ray, n1, n2)
            mat1 = surf.mat2

            if i == self.aper_idx:
                self._assign_stop_distance(ray)

            if record:
                ray_out_o = ray.o.clone().detach()
                ray_out_o[..., 2] += z_frame.detach()
                ray_out_o[ray.is_valid == 0] = float("nan")
                ray_o_record.append(ray_out_o)

        dz = surf._get_effective_d_next()
        ray.o[..., 2] -= dz
        z_frame = z_frame + dz

    ray.o[..., 2] += z_frame

    return ray, ray_o_record

backward_tracing

backward_tracing(ray, surf_range, record)

Trace backward through the inverse sequential frame steps.

Parameters:

Name	Type	Description	Default
ray	Ray	Ray object to trace.	required
surf_range	range	Range of surface indices to trace through.	required
record	bool	If True, record ray positions at each surface.	required

Returns:

Name	Type	Description
ray_out	Ray	Ray after backward propagation through all surfaces.
ray_o_record	list or None	Ray positions at each surface, or None if record is False.

Source code in deeplens-src/deeplens/geolens.py

def backward_tracing(self, ray, surf_range, record):
    """Trace backward through the inverse sequential frame steps.

    Args:
        ray (Ray): Ray object to trace.
        surf_range (range): Range of surface indices to trace through.
        record (bool): If True, record ray positions at each surface.

    Returns:
        ray_out (Ray): Ray after backward propagation through all surfaces.
        ray_o_record (list or None): Ray positions at each surface, or None
            if record is False.
    """
    if record:
        ray_o_record = []
        ray_o_record.append(ray.o.clone().detach())
    else:
        ray_o_record = None

    surf_indices = self._validate_surf_range(surf_range)
    if not surf_indices:
        return ray, ray_o_record
    first, last = min(surf_indices), max(surf_indices)
    surf_set = set(surf_indices)

    z_frame = self.surf_d(last + 1)
    ray.o[..., 2] -= z_frame

    # The medium on the image side of the last traced surface is always
    # that surface's `mat2`, including when it is the final lens surface.
    # This preserves the pre-refactor material convention for cover glass
    # and other non-air image-space media.
    mat1 = self.surfaces[last].mat2

    # Mirror of the re-anchor in `forward_tracing`: a far float32 origin
    # loses the low-order sag when it cancels against the intersection
    # distance. Local z=0 here is the plane *after* surface `last`, so +10
    # is on the incoming side of every surface in the range. `sample_sensor`
    # enters at local z=0, so the library's own reverse-rendering path never
    # trips this; only a caller-supplied bundle from far +z does.
    if bool((ray.o[..., 2] > 10.0).any()):
        ray.prop_to(10.0, n=mat1.ior(ray.wvln))

    for i in range(last, first - 1, -1):
        surf = self.surfaces[i]
        dz = surf._get_effective_d_next()
        ray.o[..., 2] += dz
        z_frame = z_frame - dz

        if i in surf_set:
            n1 = mat1.ior(ray.wvln)
            mat2 = Material("air") if i == 0 else self.surfaces[i - 1].mat2
            n2 = mat2.ior(ray.wvln)
            ray = surf.ray_reaction(ray, n1, n2)
            mat1 = mat2

            if i == self.aper_idx:
                self._assign_stop_distance(ray)

            if record:
                ray_out_o = ray.o.clone().detach()
                ray_out_o[..., 2] += z_frame.detach()
                ray_out_o[ray.is_valid == 0] = float("nan")
                ray_o_record.append(ray_out_o)

    ray.o[..., 2] += z_frame

    return ray, ray_o_record

calc_foclen

calc_foclen()

Compute effective focal length (EFL) by paraxial ray tracing.

Traces the paraxial marginal ray of an object at infinity through the surfaces with the y-nu recursion, using the reduced slope \(\omega = n u\):

\[\omega' = \omega - y \phi, \qquad y_{next} = y + \omega' t / n'\]

where the surface power \(\phi\) comes from each surface's paraxial_power. This is the first-order convention used by Zemax and CODE V: only vertex geometry contributes, so conic constants, aspheric coefficients and freeform departures do not affect the result, and the sensor position is irrelevant. Launching \(y = 1\), \(\omega = 0\) gives \(EFL = -n'_k / \omega'_k\) and \(BFL = -y_k n'_k / \omega'_k\).

Returns:

Name	Type	Description
eff_foclen	float	Effective focal length [mm].

Raises:

Type	Description
ValueError	If the system is afocal, so the focal length is undefined. Failing here avoids writing an infinite foclen, which would silently poison calc_fov/calc_scale/set_fnum.

Note

Also caches self.efl (effective focal length [mm]), self.foclen (alias of self.efl), and self.bfl (paraxial back focal length, the distance from the last surface to the rear focal point [mm]).

Reference

[1] W. Smith, "Modern Optical Engineering", the y-nu paraxial raytrace. [2] https://optics.ansys.com/hc/en-us/articles/42661756008083-Understanding-paraxial-ray-tracing

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_foclen(self):
    """Compute effective focal length (EFL) by paraxial ray tracing.

    Traces the paraxial marginal ray of an object at infinity through the
    surfaces with the y-nu recursion, using the reduced slope
    $\\omega = n u$:

    $$\\omega' = \\omega - y \\phi, \\qquad y_{next} = y + \\omega' t / n'$$

    where the surface power $\\phi$ comes from each surface's
    `paraxial_power`. This is the first-order convention used by Zemax and
    CODE V: only vertex geometry contributes, so conic constants, aspheric
    coefficients and freeform departures do not affect the result, and the
    sensor position is irrelevant. Launching $y = 1$, $\\omega = 0$ gives
    $EFL = -n'_k / \\omega'_k$ and $BFL = -y_k n'_k / \\omega'_k$.

    Returns:
        eff_foclen (float): Effective focal length [mm].

    Raises:
        ValueError: If the system is afocal, so the focal length is
            undefined. Failing here avoids writing an infinite `foclen`,
            which would silently poison `calc_fov`/`calc_scale`/`set_fnum`.

    Note:
        Also caches `self.efl` (effective focal length [mm]), `self.foclen`
        (alias of `self.efl`), and `self.bfl` (paraxial back focal length,
        the distance from the last surface to the rear focal point [mm]).

    Reference:
        [1] W. Smith, "Modern Optical Engineering", the y-nu paraxial
            raytrace.
        [2] https://optics.ansys.com/hc/en-us/articles/42661756008083-Understanding-paraxial-ray-tracing
    """
    wvln = self.primary_wvln
    n = Material("air").ior(wvln).double()
    y = torch.ones((), dtype=torch.float64, device=self.device)
    omega = torch.zeros((), dtype=torch.float64, device=self.device)

    for surf in self.surfaces:
        n2 = surf.mat2.ior(wvln).double()
        omega = omega - y * surf.paraxial_power(n, n2).double()
        y_last = y
        y = y + omega * (surf._get_effective_d_next().double() / n2)
        n = n2

    if omega.abs() < EPSILON:
        raise ValueError(
            "calc_foclen: the system is afocal (zero paraxial power); the "
            "effective focal length is undefined."
        )

    self.efl = float(-n / omega)
    self.foclen = self.efl
    self.bfl = float(-y_last * n / omega)
    return self.efl

calc_numerical_aperture

calc_numerical_aperture(n=1.0)

Compute numerical aperture (NA).

Parameters:

Name	Type	Description	Default
n	float	Refractive index. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
NA	float	Numerical aperture.

Reference

[1] https://en.wikipedia.org/wiki/Numerical_aperture

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_numerical_aperture(self, n=1.0):
    """Compute numerical aperture (NA).

    Args:
        n (float, optional): Refractive index. Defaults to 1.0.

    Returns:
        NA (float): Numerical aperture.

    Reference:
        [1] https://en.wikipedia.org/wiki/Numerical_aperture
    """
    return n * math.sin(math.atan(1 / 2 / self.fnum))

calc_focal_plane

calc_focal_plane(wvln=None)

Compute the focus distance in the object space. Ray starts from sensor center and traces to the object space.

Parameters:

Name	Type	Description	Default
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None

Returns:

Name	Type	Description
focal_plane	float	Object-space focus distance [mm] (negative z, in front of the lens).

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_focal_plane(self, wvln=None):
    """Compute the focus distance in the object space. Ray starts from sensor center and traces to the object space.

    Args:
        wvln (float, optional): Wavelength in µm. When ``None`` (default),
            falls back to ``self.primary_wvln``.

    Returns:
        focal_plane (float): Object-space focus distance [mm] (negative z, in front of the lens).
    """
    wvln = self.primary_wvln if wvln is None else wvln
    device = self.device

    # Sample point source rays from sensor center
    o1 = torch.zeros(SPP_CALC, 3, device=device, dtype=self.dtype)
    o1[:, 2] = self.d_sensor

    # Sample the first surface as pupil
    # o2 = self.sample_circle(self.surfaces[0].r, z=0.0, shape=[SPP_CALC])
    # o2 *= 0.5  # Shrink sample region to improve accuracy
    pupilz, pupilr = self.get_exit_pupil()
    o2 = self.sample_circle(pupilr, pupilz, shape=[SPP_CALC])
    d = o2 - o1
    ray = Ray(o1, d, wvln, device=device)

    # Trace rays to object space
    ray = self.trace2obj(ray)

    # Optical axis intersection
    t = (ray.d[..., 0] * ray.o[..., 0] + ray.d[..., 1] * ray.o[..., 1]) / (
        ray.d[..., 0] ** 2 + ray.d[..., 1] ** 2
    )
    focus_z = (ray.o[..., 2] - ray.d[..., 2] * t)[ray.is_valid > 0].cpu().numpy()
    focus_z = focus_z[~np.isnan(focus_z) & (focus_z < 0)]

    if len(focus_z) > 0:
        focal_plane = float(np.mean(focus_z))
    else:
        raise ValueError(
            "No valid rays found, focal plane in the image space cannot be computed."
        )

    return focal_plane

calc_sensor_plane

calc_sensor_plane(depth=float('inf'))

Calculate in-focus sensor plane.

Parameters:

Name	Type	Description	Default
depth	float	Depth of the object plane. Defaults to float("inf").	float('inf')

Returns:

Name	Type	Description
d_sensor	Tensor	In-focus sensor z-position [mm] in image space (scalar tensor).

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_sensor_plane(self, depth=float("inf")):
    """Calculate in-focus sensor plane.

    Args:
        depth (float, optional): Depth of the object plane. Defaults to float("inf").

    Returns:
        d_sensor (torch.Tensor): In-focus sensor z-position [mm] in image space (scalar tensor).
    """
    # Sample and trace rays, shape [SPP_CALC, 3]
    ray = self.sample_from_fov(fov_x=0.0, fov_y=0.0, depth=depth, num_rays=SPP_CALC)
    ray = self.trace2sensor(ray)

    # Calculate in-focus sensor position
    t = (ray.d[:, 0] * ray.o[:, 0] + ray.d[:, 1] * ray.o[:, 1]) / (
        ray.d[:, 0] ** 2 + ray.d[:, 1] ** 2
    )
    focus_z = ray.o[:, 2] - ray.d[:, 2] * t
    focus_z = focus_z[ray.is_valid > 0]
    focus_z = focus_z[~torch.isnan(focus_z) & (focus_z > 0)]
    d_sensor = torch.mean(focus_z)
    return d_sensor

calc_fov

calc_fov()

Compute field of view (FoV) of the lens in radians.

Calculates FoV using two methods

1. Perspective projection — from focal length and sensor size (effective FoV, ignoring distortion).
2. Forward ray tracing — sweeps FOV angles from object side, traces to sensor, and finds the angle whose centroid image height matches the sensor half-diagonal. This avoids the failure of the old backward-tracing approach on wide-angle lenses where pupil aberration at full field leaves zero valid rays.

Note

Caches the following attributes (all FoV values in radians): self.vfov (vertical FoV), self.hfov (horizontal FoV), self.dfov (diagonal FoV), self.rfov_eff (effective paraxial half-diagonal FoV, ignoring distortion), self.rfov (real half-diagonal FoV from ray tracing, accounts for distortion), self.real_dfov (real diagonal FoV from ray tracing), and self.eqfl (35 mm equivalent focal length [mm]).

Reference

[1] https://en.wikipedia.org/wiki/Angle_of_view_(photography)

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_fov(self):
    """Compute field of view (FoV) of the lens in radians.

    Calculates FoV using two methods:
        1. **Perspective projection** — from focal length and sensor size
           (effective FoV, ignoring distortion).
        2. **Forward ray tracing** — sweeps FOV angles from object side,
           traces to sensor, and finds the angle whose centroid image height
           matches the sensor half-diagonal. This avoids the failure of the
           old backward-tracing approach on wide-angle lenses where pupil
           aberration at full field leaves zero valid rays.

    Note:
        Caches the following attributes (all FoV values in radians):
        `self.vfov` (vertical FoV), `self.hfov` (horizontal FoV),
        `self.dfov` (diagonal FoV), `self.rfov_eff` (effective paraxial
        half-diagonal FoV, ignoring distortion), `self.rfov` (real
        half-diagonal FoV from ray tracing, accounts for distortion),
        `self.real_dfov` (real diagonal FoV from ray tracing), and
        `self.eqfl` (35 mm equivalent focal length [mm]).

    Reference:
        [1] https://en.wikipedia.org/wiki/Angle_of_view_(photography)
    """
    if not hasattr(self, "foclen"):
        return

    # 1. Perspective projection (effective FoV)
    self.hfov = 2 * math.atan(self.sensor_size[0] / 2 / self.foclen)
    self.vfov = 2 * math.atan(self.sensor_size[1] / 2 / self.foclen)
    self.dfov = 2 * math.atan(self.r_sensor / self.foclen)
    self.rfov_eff = self.dfov / 2  # effective (paraxial) half-diagonal FoV

    # 2. Forward ray tracing to calculate real FoV (distortion-affected)
    # Sweep FOV angles from object side, trace to sensor, and find which
    # angle produces an image height matching r_sensor.
    num_fov = 64
    fov_lo = float(np.rad2deg(self.rfov_eff)) * 0.5
    fov_hi = min(float(np.rad2deg(self.rfov_eff)) * 1.8, 89.0)
    fov_samples = torch.linspace(
        fov_lo, fov_hi, num_fov, device=self.device, dtype=self.dtype
    )

    ray = self.sample_from_fov(fov_x=0.0, fov_y=fov_samples.tolist(), num_rays=256)
    ray = self.trace2sensor(ray)

    # Centroid image height per FOV angle, shape [num_fov]
    valid = ray.is_valid > 0  # [num_fov, num_rays]
    masked_y = ray.o[..., 1] * valid
    n_valid = valid.sum(dim=-1).clamp(min=1)
    imgh = (masked_y.sum(dim=-1) / n_valid).abs()

    # Find the FOV angle whose image height is closest to r_sensor
    has_valid = valid.sum(dim=-1) > 10
    if has_valid.any():
        imgh[~has_valid] = float("inf")
        diff = (imgh - self.r_sensor).abs()
        best_idx = diff.argmin().item()
        rfov = fov_samples[best_idx].item() * math.pi / 180.0
        self.rfov = rfov
        self.real_dfov = 2 * rfov
    else:
        self.rfov = self.rfov_eff
        self.real_dfov = self.dfov

    # 3. Compute 35mm equivalent focal length. 35mm sensor: 36mm * 24mm
    self.eqfl = 21.63 / math.tan(self.rfov_eff)

calc_scale

calc_scale(depth)

Calculate the scale factor (object height / image height).

Uses the pinhole camera model to compute magnification.

Parameters:

Name	Type	Description	Default
depth	float	Object distance from the lens (negative z direction).	required

Returns:

Name	Type	Description
scale	float	Scale factor relating object height to image height.

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_scale(self, depth):
    """Calculate the scale factor (object height / image height).

    Uses the pinhole camera model to compute magnification.

    Args:
        depth (float): Object distance from the lens (negative z direction).

    Returns:
        scale (float): Scale factor relating object height to image height.
    """
    return -depth / self.foclen

calc_pupil

calc_pupil()

Compute entrance and exit pupil positions and radii.

The entrance and exit pupils must be recalculated whenever

• First-order parameters change (e.g., field of view, object height, image height),
• Lens geometry or materials change (e.g., surface curvatures, refractive indices, thicknesses),
• Or generally, any time the lens configuration is modified.

Note

Caches self.aper_idx (aperture surface index), self.exit_pupilz/self.exit_pupilr (real exit pupil position and radius [mm]), self.entr_pupilz/self.entr_pupilr (real entrance pupil position and radius [mm]), self.exit_pupilz_parax/self.exit_pupilr_parax and self.entr_pupilz_parax/self.entr_pupilr_parax (paraxial pupils), and self.fnum (F-number from focal length and entrance pupil).

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_pupil(self):
    """Compute entrance and exit pupil positions and radii.

    The entrance and exit pupils must be recalculated whenever:
        - First-order parameters change (e.g., field of view, object height, image height),
        - Lens geometry or materials change (e.g., surface curvatures, refractive indices, thicknesses),
        - Or generally, any time the lens configuration is modified.

    Note:
        Caches `self.aper_idx` (aperture surface index),
        `self.exit_pupilz`/`self.exit_pupilr` (real exit pupil position and
        radius [mm]), `self.entr_pupilz`/`self.entr_pupilr` (real entrance
        pupil position and radius [mm]),
        `self.exit_pupilz_parax`/`self.exit_pupilr_parax` and
        `self.entr_pupilz_parax`/`self.entr_pupilr_parax` (paraxial pupils),
        and `self.fnum` (F-number from focal length and entrance pupil).
    """
    # Find aperture
    self.aper_idx = None
    for i in range(len(self.surfaces)):
        if getattr(self.surfaces[i], "is_aperture", False):
            self.aper_idx = i
            break

    if self.aper_idx is None:
        for i in range(len(self.surfaces)):
            if isinstance(self.surfaces[i], Aperture):
                self.aper_idx = i
                break

    if self.aper_idx is None:
        self.aper_idx = np.argmin([s.r for s in self.surfaces])
        print("No aperture found, use the smallest surface as aperture.")

    # Compute entrance and exit pupil
    self.exit_pupilz, self.exit_pupilr = self.calc_exit_pupil_rayaiming()
    self.entr_pupilz, self.entr_pupilr = self.calc_entrance_pupil_rayaiming()
    self.exit_pupilz_parax, self.exit_pupilr_parax = self.calc_pupil_paraxial(
        reverse=False
    )
    self.entr_pupilz_parax, self.entr_pupilr_parax = self.calc_pupil_paraxial(
        reverse=True
    )

    for name, radius in (
        ("entrance", self.entr_pupilr),
        ("exit", self.exit_pupilr),
        ("paraxial entrance", self.entr_pupilr_parax),
        ("paraxial exit", self.exit_pupilr_parax),
    ):
        try:
            radius_value = float(radius)
        except (TypeError, ValueError) as exc:
            raise ValueError(
                f"Invalid {name} pupil radius {radius!r}; check lens surfaces "
                "and aperture data."
            ) from exc
        if not math.isfinite(radius_value) or radius_value <= 0:
            raise ValueError(
                f"Invalid {name} pupil radius {radius_value!r}; check lens "
                "surfaces and aperture data."
            )

    # Compute F-number
    self.fnum = self.foclen / (2 * self.entr_pupilr)

get_entrance_pupil

get_entrance_pupil(paraxial=False)

Get entrance pupil location and radius.

Parameters:

Name	Type	Description	Default
paraxial	bool	If True, return paraxial approximation values. If False, return real ray-traced values. Defaults to False.	False

Returns:

Name	Type	Description
pupilz	float	Entrance pupil z-position [mm].
pupilr	float	Entrance pupil radius [mm].

Source code in deeplens-src/deeplens/geolens.py

def get_entrance_pupil(self, paraxial=False):
    """Get entrance pupil location and radius.

    Args:
        paraxial (bool, optional): If True, return paraxial approximation values.
            If False, return real ray-traced values. Defaults to False.

    Returns:
        pupilz (float): Entrance pupil z-position [mm].
        pupilr (float): Entrance pupil radius [mm].
    """
    if paraxial:
        return self.entr_pupilz_parax, self.entr_pupilr_parax
    else:
        return self.entr_pupilz, self.entr_pupilr

get_exit_pupil

get_exit_pupil(paraxial=False)

Get exit pupil location and radius.

Parameters:

Name	Type	Description	Default
paraxial	bool	If True, return paraxial approximation values. If False, return real ray-traced values. Defaults to False.	False

Returns:

Name	Type	Description
pupilz	float	Exit pupil z-position [mm].
pupilr	float	Exit pupil radius [mm].

Source code in deeplens-src/deeplens/geolens.py

def get_exit_pupil(self, paraxial=False):
    """Get exit pupil location and radius.

    Args:
        paraxial (bool, optional): If True, return paraxial approximation values.
            If False, return real ray-traced values. Defaults to False.

    Returns:
        pupilz (float): Exit pupil z-position [mm].
        pupilr (float): Exit pupil radius [mm].
    """
    if paraxial:
        return self.exit_pupilz_parax, self.exit_pupilr_parax
    else:
        return self.exit_pupilz, self.exit_pupilr

calc_pupil_paraxial

calc_pupil_paraxial(reverse=False)

Image the aperture stop through the surfaces on one side of it.

The pupils are the first-order images of the stop, so they follow from the same y-nu recursion as calc_foclen rather than from tracing real rays. A ray is launched from the axial point of the stop (\(y = 0\), \(\omega = n u = 1\)) and propagated through the surfaces after the stop (exit pupil) or backwards through those before it (entrance pupil); where it re-crosses the axis is the pupil plane. Stop and pupil are conjugate, so the transverse magnification is given by the Smith-Helmholtz relation

\[m = \frac{n_0 u_0}{n_k u_k} = \frac{\omega_0}{\omega_k}.\]

Parameters:

Name	Type	Description	Default
reverse	bool	Trace backwards through the surfaces before the stop (entrance pupil) instead of forwards through those after it (exit pupil). Defaults to False.	False

Returns:

Name	Type	Description
pupilz	float	Pupil z-position [mm].
pupilr	float	Pupil radius [mm].

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_pupil_paraxial(self, reverse=False):
    """Image the aperture stop through the surfaces on one side of it.

    The pupils are the first-order images of the stop, so they follow from
    the same y-nu recursion as `calc_foclen` rather than from tracing real
    rays. A ray is launched from the axial point of the stop ($y = 0$,
    $\\omega = n u = 1$) and propagated through the surfaces after the stop
    (exit pupil) or backwards through those before it (entrance pupil);
    where it re-crosses the axis is the pupil plane. Stop and pupil are
    conjugate, so the transverse magnification is given by the
    Smith-Helmholtz relation

    $$m = \\frac{n_0 u_0}{n_k u_k} = \\frac{\\omega_0}{\\omega_k}.$$

    Args:
        reverse (bool, optional): Trace backwards through the surfaces
            before the stop (entrance pupil) instead of forwards through
            those after it (exit pupil). Defaults to False.

    Returns:
        pupilz (float): Pupil z-position [mm].
        pupilr (float): Pupil radius [mm].
    """
    wvln = self.primary_wvln
    air = Material("air").ior(wvln).double()
    surfaces = self.surfaces
    aper_idx = self.aper_idx
    aper_surf = surfaces[aper_idx]
    aper_r = aper_surf.r

    z = self.surf_d(aper_idx).double()
    y = torch.zeros((), dtype=torch.float64, device=self.device)
    omega = torch.ones((), dtype=torch.float64, device=self.device)

    if reverse:
        n = surfaces[aper_idx - 1].mat2.ior(wvln).double() if aper_idx > 0 else air
        for i in range(aper_idx - 1, -1, -1):
            d = surfaces[i]._get_effective_d_next().double()
            y = y - omega * (d / n)
            z = z - d
            n1 = surfaces[i - 1].mat2.ior(wvln).double() if i > 0 else air
            omega = omega + y * surfaces[i].paraxial_power(n1, n).double()
            n = n1
    else:
        n = aper_surf.mat2.ior(wvln).double()
        for i in range(aper_idx, len(surfaces) - 1):
            d = surfaces[i]._get_effective_d_next().double()
            y = y + omega * (d / n)
            z = z + d
            surf = surfaces[i + 1]
            n2 = surf.mat2.ior(wvln).double()
            omega = omega - y * surf.paraxial_power(n, n2).double()
            n = n2

    # A stop conjugate to infinity (telecentric side) has no finite pupil.
    if omega.abs() < EPSILON:
        idx = 0 if reverse else -1
        print("Stop images to infinity, use the end surface as pupil.")
        return self.surf_d(idx).item(), surfaces[idx].r

    return float(z - y * n / omega), float(aper_r / omega.abs())

calc_exit_pupil_rayaiming

calc_exit_pupil_rayaiming()

Calculate exit pupil location and radius from real rays.

Rays are emitted from the edge of the aperture stop in large quantities and traced to the last surface; the exit pupil position and radius come from the intersection points of those rays. Slower than calc_pupil_paraxial and affected by aperture-related aberrations.

Returns:

Name	Type	Description
avg_pupilz	float	z coordinate of exit pupil.
avg_pupilr	float	radius of exit pupil.

Reference

[1] Exit pupil: how many rays can come from sensor to object space. [2] https://en.wikipedia.org/wiki/Exit_pupil

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_exit_pupil_rayaiming(self):
    """Calculate exit pupil location and radius from real rays.

    Rays are emitted from the edge of the aperture stop in large quantities
    and traced to the last surface; the exit pupil position and radius come
    from the intersection points of those rays. Slower than
    `calc_pupil_paraxial` and affected by aperture-related aberrations.

    Returns:
        avg_pupilz (float): z coordinate of exit pupil.
        avg_pupilr (float): radius of exit pupil.

    Reference:
        [1] Exit pupil: how many rays can come from sensor to object space.
        [2] https://en.wikipedia.org/wiki/Exit_pupil
    """
    if self.aper_idx is None:
        print("No aperture, use the last surface as exit pupil.")
        return self.surf_d(-1).item(), self.surfaces[-1].r

    # Sample rays from the aperture edge
    aper_z = self.surf_d(self.aper_idx).item()
    aper_r = self.surfaces[self.aper_idx].r
    ray_o = torch.tensor(
        [[aper_r, 0, aper_z]], device=self.device, dtype=self.dtype
    ).repeat(128, 1)
    rfov = float(np.arctan(self.r_sensor / self.foclen))
    phi_rad = torch.linspace(
        -rfov / 2, rfov / 2, 128, device=self.device, dtype=self.dtype
    )

    d = torch.stack(
        (torch.sin(phi_rad), torch.zeros_like(phi_rad), torch.cos(phi_rad)), axis=-1
    )
    ray = Ray(ray_o, d, wvln=self.primary_wvln, device=self.device)

    # Ray tracing from aperture edge to last surface
    surf_range = range(self.aper_idx + 1, len(self.surfaces))
    ray, _ = self.trace(ray, surf_range=surf_range)

    # Compute intersection points, solving the equation: o1+d1*t1 = o2+d2*t2
    ray_o = torch.stack(
        [ray.o[ray.is_valid != 0][:, 0], ray.o[ray.is_valid != 0][:, 2]], dim=-1
    )
    ray_d = torch.stack(
        [ray.d[ray.is_valid != 0][:, 0], ray.d[ray.is_valid != 0][:, 2]], dim=-1
    )
    intersection_points = self.compute_intersection_points_2d(ray_o, ray_d)

    # Handle the case where no intersection points are found or small pupil
    if len(intersection_points) == 0:
        print("No intersection points found, use the last surface as exit pupil.")
        avg_pupilr = self.surfaces[-1].r
        avg_pupilz = self.surf_d(-1).item()
    else:
        avg_pupilr = torch.mean(intersection_points[:, 0]).item()
        avg_pupilz = torch.mean(intersection_points[:, 1]).item()

        if avg_pupilr < EPSILON:
            print(
                "Zero or negative exit pupil is detected, use the last surface as pupil."
            )
            avg_pupilr = self.surfaces[-1].r
            avg_pupilz = self.surf_d(-1).item()

    return avg_pupilz, avg_pupilr

calc_entrance_pupil_rayaiming

calc_entrance_pupil_rayaiming()

Calculate entrance pupil of the lens from real rays.

The entrance pupil is the optical image of the physical aperture stop, as seen through the optical elements in front of the stop. We sample backward rays from the aperture stop edge and trace them to the first surface, then find the intersection points of the reverse extension of the rays. The average of the intersection points defines the entrance pupil position and radius. Slower than calc_pupil_paraxial and affected by aperture-related aberrations.

Returns:

Name	Type	Description
avg_pupilz	float	Entrance pupil z-position [mm].
avg_pupilr	float	Entrance pupil radius [mm].

Note

[1] Use calc_pupil_paraxial unless precise ray aiming is required. [2] This function only works for object at a far distance. For microscopes, this function usually returns a negative entrance pupil.

Reference

[1] Entrance pupil: how many rays can come from object space to sensor. [2] https://en.wikipedia.org/wiki/Entrance_pupil: "In an optical system, the entrance pupil is the optical image of the physical aperture stop, as 'seen' through the optical elements in front of the stop." [3] Zemax LLC, OpticStudio User Manual, Version 19.4, Document No. 2311, 2019.

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def calc_entrance_pupil_rayaiming(self):
    """Calculate entrance pupil of the lens from real rays.

    The entrance pupil is the optical image of the physical aperture stop, as seen through the optical elements in front of the stop. We sample backward rays from the aperture stop edge and trace them to the first surface, then find the intersection points of the reverse extension of the rays. The average of the intersection points defines the entrance pupil position and radius. Slower than `calc_pupil_paraxial` and affected by aperture-related aberrations.

    Returns:
        avg_pupilz (float): Entrance pupil z-position [mm].
        avg_pupilr (float): Entrance pupil radius [mm].

    Note:
        [1] Use `calc_pupil_paraxial` unless precise ray aiming is required.
        [2] This function only works for object at a far distance. For microscopes, this function usually returns a negative entrance pupil.

    Reference:
        [1] Entrance pupil: how many rays can come from object space to sensor.
        [2] https://en.wikipedia.org/wiki/Entrance_pupil: "In an optical system, the entrance pupil is the optical image of the physical aperture stop, as 'seen' through the optical elements in front of the stop."
        [3] Zemax LLC, *OpticStudio User Manual*, Version 19.4, Document No. 2311, 2019.
    """
    if self.aper_idx is None:
        print("No aperture stop, use the first surface as entrance pupil.")
        return self.surf_d(0).item(), self.surfaces[0].r

    # Sample rays from edge of aperture stop
    aper_z = self.surf_d(self.aper_idx).item()
    aper_r = self.surfaces[self.aper_idx].r

    ray_o = torch.tensor(
        [[aper_r, 0, aper_z]], device=self.device, dtype=self.dtype
    ).repeat(128, 1)
    rfov = float(np.arctan(self.r_sensor / self.foclen))
    phi = torch.linspace(
        -rfov / 2, rfov / 2, 128, device=self.device, dtype=self.dtype
    )

    d = torch.stack(
        (torch.sin(phi), torch.zeros_like(phi), -torch.cos(phi)), axis=-1
    )
    ray = Ray(ray_o, d, wvln=self.primary_wvln, device=self.device)

    # Ray tracing from aperture edge to first surface
    surf_range = range(0, self.aper_idx)
    ray, _ = self.trace(ray, surf_range=surf_range)

    # Compute intersection points, solving the equation: o1+d1*t1 = o2+d2*t2
    ray_o = torch.stack(
        [ray.o[ray.is_valid > 0][:, 0], ray.o[ray.is_valid > 0][:, 2]], dim=-1
    )
    ray_d = torch.stack(
        [ray.d[ray.is_valid > 0][:, 0], ray.d[ray.is_valid > 0][:, 2]], dim=-1
    )
    intersection_points = self.compute_intersection_points_2d(ray_o, ray_d)

    # Handle the case where no intersection points are found or small entrance pupil
    if len(intersection_points) == 0:
        print(
            "No intersection points found, use the first surface as entrance pupil."
        )
        avg_pupilr = self.surfaces[0].r
        avg_pupilz = self.surf_d(0).item()
    else:
        avg_pupilr = torch.mean(intersection_points[:, 0]).item()
        avg_pupilz = torch.mean(intersection_points[:, 1]).item()

        if avg_pupilr < EPSILON:
            print(
                "Zero or negative entrance pupil is detected, use the first surface as entrance pupil."
            )
            avg_pupilr = self.surfaces[0].r
            avg_pupilz = self.surf_d(0).item()

    return avg_pupilz, avg_pupilr

compute_intersection_points_2d staticmethod

compute_intersection_points_2d(origins, directions)

Compute the intersection points of 2D lines.

Parameters:

Name	Type	Description	Default
origins	Tensor	Origins of the lines. Shape: [N, 2]	required
directions	Tensor	Directions of the lines. Shape: [N, 2]	required

Returns:

Name	Type	Description
points	Tensor	Intersection points. Shape: [N*(N-1)/2, 2]

Source code in deeplens-src/deeplens/geolens.py

@staticmethod
def compute_intersection_points_2d(origins, directions):
    """Compute the intersection points of 2D lines.

    Args:
        origins (torch.Tensor): Origins of the lines. Shape: [N, 2]
        directions (torch.Tensor): Directions of the lines. Shape: [N, 2]

    Returns:
        points (torch.Tensor): Intersection points. Shape: [N*(N-1)/2, 2]
    """
    N = origins.shape[0]

    # Create pairwise combinations of indices
    idx = torch.arange(N)
    idx_i, idx_j = torch.combinations(idx, r=2).unbind(1)

    Oi = origins[idx_i]  # Shape: [N*(N-1)/2, 2]
    Oj = origins[idx_j]  # Shape: [N*(N-1)/2, 2]
    Di = directions[idx_i]  # Shape: [N*(N-1)/2, 2]
    Dj = directions[idx_j]  # Shape: [N*(N-1)/2, 2]

    # A pair of non-parallel 2-D lines has a closed-form intersection.
    # Filter parallel/near-parallel pairs before division: batched
    # ``torch.linalg.lstsq`` rejects an entire CUDA batch when even one
    # pair is rank deficient, which made valid real prescriptions fail
    # during pupil estimation.
    b = Oj - Oi  # Shape: [N*(N-1)/2, 2]
    cross_d = Di[:, 0] * Dj[:, 1] - Di[:, 1] * Dj[:, 0]
    direction_scale = torch.linalg.vector_norm(
        Di, dim=-1
    ) * torch.linalg.vector_norm(Dj, dim=-1)
    tolerance = 100 * torch.finfo(directions.dtype).eps * direction_scale
    valid = cross_d.abs() > tolerance
    if not valid.any():
        return origins.new_empty((0, 2))

    Oi = Oi[valid]
    Oj = Oj[valid]
    Di = Di[valid]
    Dj = Dj[valid]
    b = b[valid]
    cross_d = cross_d[valid]
    s = (b[:, 0] * Dj[:, 1] - b[:, 1] * Dj[:, 0]) / cross_d
    t = (b[:, 0] * Di[:, 1] - b[:, 1] * Di[:, 0]) / cross_d

    # Calculate the intersection points using either rays
    P_i = Oi + s.unsqueeze(-1) * Di  # Shape: [N*(N-1)/2, 2]
    P_j = Oj + t.unsqueeze(-1) * Dj  # Shape: [N*(N-1)/2, 2]

    # Take the average to mitigate numerical precision issues
    P = (P_i + P_j) / 2

    return P

surf_d

surf_d(idx)

Return the derived global vertex position of surface idx [mm].

surf_d(0) is zero and surf_d(i) is the differentiable prefix sum of d_next for surfaces before i. Passing len(surfaces) returns the sensor position. Negative indices follow Python surface indexing.

Source code in deeplens-src/deeplens/geolens.py

def surf_d(self, idx):
    """Return the derived global vertex position of surface `idx` [mm].

    `surf_d(0)` is zero and `surf_d(i)` is the differentiable prefix sum
    of `d_next` for surfaces before `i`. Passing `len(surfaces)` returns the
    sensor position. Negative indices follow Python surface indexing.
    """
    num_surfs = len(self.surfaces)
    if idx < 0:
        idx += num_surfs
    if idx < 0 or idx > num_surfs:
        raise IndexError(f"Surface index {idx} out of range [0, {num_surfs}].")
    if idx == 0 or num_surfs == 0:
        return torch.zeros((), device=self.device, dtype=self.dtype)
    return torch.stack(
        [surface._get_effective_d_next() for surface in self.surfaces[:idx]]
    ).sum()

refocus

refocus(foc_dist=float('inf'))

Refocus the lens to a depth distance by changing sensor position.

Parameters:

Name	Type	Description	Default
foc_dist	float	Object focus distance [mm]. Use float('inf') for infinity focus. Defaults to float('inf').	float('inf')

Note

In DSLR, phase detection autofocus (PDAF) is a popular and efficient method. But here we simplify the problem by calculating the in-focus position of green light.

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def refocus(self, foc_dist=float("inf")):
    """Refocus the lens to a depth distance by changing sensor position.

    Args:
        foc_dist (float, optional): Object focus distance [mm].
            Use ``float('inf')`` for infinity focus. Defaults to ``float('inf')``.

    Note:
        In DSLR, phase detection autofocus (PDAF) is a popular and efficient method. But here we simplify the problem by calculating the in-focus position of green light.
    """
    # Calculate in-focus sensor position
    d_sensor_new = self.calc_sensor_plane(depth=foc_dist)

    # Update sensor position
    assert d_sensor_new > 0, "Obtained negative sensor position."
    self.d_sensor = d_sensor_new

    # FoV will be slightly changed
    self.post_computation()

set_fnum

set_fnum(fnum)

Set F-number and aperture radius using binary search.

Parameters:

Name	Type	Description	Default
fnum	float	target F-number.	required

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def set_fnum(self, fnum):
    """Set F-number and aperture radius using binary search.

    Args:
        fnum (float): target F-number.
    """
    target_pupil_r = self.foclen / fnum / 2
    aper_r = self.surfaces[self.aper_idx].r
    lo, hi = 0.1 * aper_r, 5.0 * aper_r

    pupilr = None
    for _ in range(40):
        mid = 0.5 * (lo + hi)
        self.surfaces[self.aper_idx].update_r(float(mid))
        _, pupilr = self.calc_entrance_pupil_rayaiming()
        if abs(pupilr - target_pupil_r) / target_pupil_r < 1e-3:
            break
        if pupilr > target_pupil_r:
            hi = mid
        else:
            lo = mid
    else:
        logging.warning(
            f"set_fnum: did not converge, pupil_r={pupilr:.4f}, target={target_pupil_r:.4f}"
        )

    self.calc_pupil()

set_target_fov_fnum

set_target_fov_fnum(rfov, fnum)

Set FoV, image height, and F-number as design targets.

Only use this method to assign design targets (it overwrites the cached first-order quantities directly rather than measuring them).

Parameters:

Name	Type	Description	Default
rfov	float	Half-diagonal FoV. Interpreted as radians; if the value is greater than \(\pi\) it is treated as degrees and converted to radians.	required
fnum	float	Target F-number.	required

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def set_target_fov_fnum(self, rfov, fnum):
    """Set FoV, image height, and F-number as design targets.

    Only use this method to assign design targets (it overwrites the
    cached first-order quantities directly rather than measuring them).

    Args:
        rfov (float): Half-diagonal FoV. Interpreted as radians; if the
            value is greater than $\\pi$ it is treated as degrees and
            converted to radians.
        fnum (float): Target F-number.
    """
    if rfov > math.pi:
        self.rfov_eff = rfov / 180.0 * math.pi
    else:
        self.rfov_eff = rfov

    self.rfov = self.rfov_eff
    self.real_dfov = 2 * self.rfov
    self.foclen = self.r_sensor / math.tan(self.rfov_eff)
    self.eqfl = 21.63 / math.tan(self.rfov_eff)
    self.fnum = fnum
    aper_r = self.foclen / fnum / 2
    self.surfaces[self.aper_idx].update_r(float(aper_r))

    # Update pupil after setting aperture radius
    self.calc_pupil()

set_fov

set_fov(rfov)

Set half-diagonal field of view as a design target.

Unlike calc_fov() which derives FoV from focal length and sensor size, this method directly assigns the target FoV for lens optimisation.

Parameters:

Name	Type	Description	Default
rfov	float	Half-diagonal FoV in radians.	required

Source code in deeplens-src/deeplens/geolens.py

@torch.no_grad()
def set_fov(self, rfov):
    """Set half-diagonal field of view as a design target.

    Unlike ``calc_fov()`` which derives FoV from focal length and sensor
    size, this method directly assigns the target FoV for lens optimisation.

    Args:
        rfov (float): Half-diagonal FoV in radians.
    """
    self.rfov_eff = rfov
    self.rfov = rfov
    self.real_dfov = 2 * self.rfov
    self.eqfl = 21.63 / math.tan(self.rfov_eff)

Components

GeoLens uses a mixin architecture — its functionality is split across the focused classes below. You normally interact only with GeoLens itself; these are documented for reference.

deeplens.geolens_pkg.psf_compute.GeoLensPSF

Mixin providing PSF computation for GeoLens.

Exposes three PSF models through a single psf dispatcher: incoherent geometric ray tracing, coherent exit-pupil diffraction (ASM propagation), and Huygens-Fresnel integration. The geometric and coherent models are differentiable; Huygens is not. This class is not instantiated directly; it is mixed into GeoLens.

psf

psf(points, wvln=None, ks=PSF_KS, **kwargs)

Compute the Point Spread Function (PSF) for given point sources.

Dispatches to one of three PSF models

• geometric: incoherent intensity ray tracing (fast, differentiable).
• coherent: coherent tracing to exit pupil + ASM propagation (accurate, differentiable, single point).
• huygens: Huygens-Fresnel integration (accurate, not differentiable, single point).

Parameters:

Name	Type	Description	Default
points	Tensor	Normalized point source positions. Shape [N, 3] with x, y in [-1, 1] and z in [-Inf, 0]. The coherent and huygens models accept only a single point ([3] or [1, 3]).	required
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
ks	int	Output kernel size in pixels. Defaults to PSF_KS.	PSF_KS
**kwargs		Model-specific options: spp (int): Rays sampled per source. If None, uses the model-specific default (SPP_PSF / SPP_COHERENT). recenter (bool): If True (default), center the PSF on the chief ray; otherwise on the pinhole projection. model (str): One of 'geometric' (default), 'coherent', 'huygens'. return_field (bool): For the coherent and Huygens models, return the energy-normalized complex sensor-plane field instead of intensity. Defaults to False.	{}

Returns:

Type	Description
	torch.Tensor: Intensity PSF normalized to sum to 1, or an energy-normalized complex field when return_field=True. Shape [ks, ks] for a single point, or [N, ks, ks] for the geometric model with N points.

Raises:

Type	Description
ValueError	If model is not one of the supported names.

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

def psf(self, points, wvln=None, ks=PSF_KS, **kwargs):
    """Compute the Point Spread Function (PSF) for given point sources.

    Dispatches to one of three PSF models:
        - geometric: incoherent intensity ray tracing (fast, differentiable).
        - coherent: coherent tracing to exit pupil + ASM propagation (accurate, differentiable, single point).
        - huygens: Huygens-Fresnel integration (accurate, not differentiable, single point).

    Args:
        points (torch.Tensor): Normalized point source positions. Shape [N, 3]
            with x, y in [-1, 1] and z in [-Inf, 0]. The coherent and huygens
            models accept only a single point ([3] or [1, 3]).
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`.
        ks (int, optional): Output kernel size in pixels. Defaults to PSF_KS.
        **kwargs: Model-specific options:
            spp (int): Rays sampled per source. If None, uses the
            model-specific default (SPP_PSF / SPP_COHERENT).
            recenter (bool): If True (default), center the PSF on the chief
            ray; otherwise on the pinhole projection.
            model (str): One of 'geometric' (default), 'coherent', 'huygens'.
            return_field (bool): For the coherent and Huygens models, return
                the energy-normalized complex sensor-plane field instead of
                intensity. Defaults to False.

    Returns:
        torch.Tensor: Intensity PSF normalized to sum to 1, or an
            energy-normalized complex field when ``return_field=True``.
            Shape [ks, ks] for a single point, or [N, ks, ks] for the
            geometric model with N points.

    Raises:
        ValueError: If `model` is not one of the supported names.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    spp = kwargs.get("spp", None)
    recenter = kwargs.get("recenter", True)
    return_field = kwargs.get("return_field", False)
    model = kwargs.get("model", "geometric")
    if model == "geometric":
        spp = SPP_PSF if spp is None else spp
        return self.psf_geometric(points, ks, wvln, spp, recenter)
    elif model == "coherent":
        spp = SPP_COHERENT if spp is None else spp
        return self.psf_coherent(points, ks, wvln, spp, recenter, return_field)
    elif model == "huygens":
        spp = SPP_COHERENT if spp is None else spp
        return self.psf_huygens(points, ks, wvln, spp, recenter, return_field)
    else:
        raise ValueError(f"Unknown PSF model: {model}")

psf_geometric

psf_geometric(points, ks=PSF_KS, wvln=None, spp=SPP_PSF, recenter=True)

Compute the single-wavelength geometric PSF by incoherent ray binning.

Samples rays from each object point, traces them incoherently to the sensor, and bins the hit positions into a ks × ks intensity kernel. This model is fast and differentiable.

Parameters:

Name	Type	Description	Default
points	Tensor	Normalized point source positions. Shape [N, 3] with x, y in [-1, 1] and z in [-Inf, 0].	required
ks	int	Output kernel size in pixels. Defaults to PSF_KS.	PSF_KS
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
spp	int	Rays sampled per source. Defaults to SPP_PSF.	SPP_PSF
recenter	bool	If True (default), center on the chief ray; otherwise on the pinhole projection.	True

Returns:

Name	Type	Description
psf	Tensor	PSF normalized to sum to 1. Shape [ks, ks] for a single point, or [N, ks, ks] for N points.

Reference

[1] https://optics.ansys.com/hc/en-us/articles/42661723066515-What-is-a-Point-Spread-Function

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

def psf_geometric(self, points, ks=PSF_KS, wvln=None, spp=SPP_PSF, recenter=True):
    """Compute the single-wavelength geometric PSF by incoherent ray binning.

    Samples rays from each object point, traces them incoherently to the
    sensor, and bins the hit positions into a `ks × ks` intensity kernel.
    This model is fast and differentiable.

    Args:
        points (torch.Tensor): Normalized point source positions. Shape [N, 3]
            with x, y in [-1, 1] and z in [-Inf, 0].
        ks (int, optional): Output kernel size in pixels. Defaults to PSF_KS.
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`.
        spp (int, optional): Rays sampled per source. Defaults to SPP_PSF.
        recenter (bool, optional): If True (default), center on the chief ray;
            otherwise on the pinhole projection.

    Returns:
        psf (torch.Tensor): PSF normalized to sum to 1. Shape [ks, ks] for a
            single point, or [N, ks, ks] for N points.

    Reference:
        [1] https://optics.ansys.com/hc/en-us/articles/42661723066515-What-is-a-Point-Spread-Function
    """
    wvln = self.primary_wvln if wvln is None else wvln
    sensor_w, sensor_h = self.sensor_size
    pixel_size = self.pixel_size
    device = self.device

    # Points shape of [N, 3]
    if not torch.is_tensor(points):
        points = torch.as_tensor(points, device=device, dtype=self.dtype)
    else:
        points = points.to(device=device, dtype=self.dtype)

    if len(points.shape) == 1:
        single_point = True
        points = points.unsqueeze(0)
    else:
        single_point = False

    # Sample rays. Ray position in the object space by perspective projection
    depth = points[:, 2]
    scale = self.calc_scale(depth)
    point_obj_x = points[..., 0] * scale * sensor_w / 2
    point_obj_y = points[..., 1] * scale * sensor_h / 2
    point_obj = torch.stack([point_obj_x, point_obj_y, points[..., 2]], dim=-1)
    ray = self.sample_from_points(points=point_obj, num_rays=spp, wvln=wvln)

    # Trace rays to sensor plane (incoherent)
    ray.is_coherent = False
    ray = self.trace2sensor(ray)

    # Calculate PSF center, shape [N, 2]. At the primary wavelength the
    # bundle just traced is reused for the chief ray instead of tracing a
    # second one. Other wavelengths still trace at the primary wavelength so
    # every channel shares one reference center and lateral color stays in
    # the kernels. Sparse bundles are not reused: chief-ray accuracy scales
    # as 1/sqrt(spp).
    if recenter:
        reuse_ray = ray if wvln == self.primary_wvln and spp >= SPP_CALC else None
        pointc = self.psf_center(point_obj, method="chief_ray", ray=reuse_ray)
    else:
        pointc = self.psf_center(point_obj, method="pinhole")

    # Monte Carlo integration
    psf = forward_integral(ray.flip_xy(), ps=pixel_size, ks=ks, pointc=pointc)

    # Intensity normalization
    psf = psf / (torch.sum(psf, dim=(-2, -1), keepdim=True) + EPSILON)

    if single_point:
        psf = psf.squeeze(0)

    return diff_float(psf)

psf_coherent

psf_coherent(
    points, ks=PSF_KS, wvln=None, spp=SPP_COHERENT, recenter=True, return_field=False
)

Compute the coherent exit-pupil PSF (alias for psf_pupil_prop).

Traces coherent rays to the exit pupil and propagates the wavefront to the sensor with the Angular Spectrum Method (ASM). See psf_pupil_prop for full argument and return documentation.

Parameters:

Name	Type	Description	Default
points	Tensor	Single normalized point source [3] or [1, 3] with x, y in [-1, 1] and z in [-Inf, 0].	required
ks	int	Output kernel size in pixels. Defaults to PSF_KS.	PSF_KS
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
spp	int	Rays sampled. Defaults to SPP_COHERENT.	SPP_COHERENT
recenter	bool	If True (default), center on the chief ray.	True
return_field	bool	Return the energy-normalized complex sensor-plane field instead of intensity. Defaults to False.	False

Returns:

Type	Description
	torch.Tensor: Intensity PSF normalized to sum to 1, or an energy-normalized complex field when return_field=True. Shape [ks, ks].

Note

The complex field has an arbitrary global phase because optical path length is measured relative to a reference; relative phase is meaningful.

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

def psf_coherent(
    self,
    points,
    ks=PSF_KS,
    wvln=None,
    spp=SPP_COHERENT,
    recenter=True,
    return_field=False,
):
    """Compute the coherent exit-pupil PSF (alias for `psf_pupil_prop`).

    Traces coherent rays to the exit pupil and propagates the wavefront to
    the sensor with the Angular Spectrum Method (ASM). See `psf_pupil_prop`
    for full argument and return documentation.

    Args:
        points (torch.Tensor): Single normalized point source [3] or [1, 3]
            with x, y in [-1, 1] and z in [-Inf, 0].
        ks (int, optional): Output kernel size in pixels. Defaults to PSF_KS.
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`.
        spp (int, optional): Rays sampled. Defaults to SPP_COHERENT.
        recenter (bool, optional): If True (default), center on the chief ray.
        return_field (bool, optional): Return the energy-normalized complex
            sensor-plane field instead of intensity. Defaults to False.

    Returns:
        torch.Tensor: Intensity PSF normalized to sum to 1, or an
            energy-normalized complex field when ``return_field=True``.
            Shape [ks, ks].

    Note:
        The complex field has an arbitrary global phase because optical path
        length is measured relative to a reference; relative phase is meaningful.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    return self.psf_pupil_prop(
        points,
        ks=ks,
        wvln=wvln,
        spp=spp,
        recenter=recenter,
        return_field=return_field,
    )

psf_pupil_prop

psf_pupil_prop(
    points, ks=PSF_KS, wvln=None, spp=SPP_COHERENT, recenter=True, return_field=False
)

Compute the single-point monochromatic PSF via the exit-pupil diffraction model.

Steps

1. Compute the complex wavefront at the exit-pupil plane by coherent ray tracing.
2. Propagate to the sensor plane with the Angular Spectrum Method (ASM) and return either its field or intensity. This function is differentiable.

Parameters:

Name	Type	Description	Default
points	Tensor or list	Single normalized point source [3] or [1, 3] with x, y in [-1, 1] and z in [-Inf, 0].	required
ks	int	Size of the output PSF patch in pixels. If None, the full propagated intensity field is returned uncropped. Defaults to PSF_KS.	PSF_KS
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
spp	int	Number of rays to sample. Defaults to SPP_COHERENT.	SPP_COHERENT
recenter	bool	If True (default), center on the chief ray; otherwise on the pinhole projection.	True
return_field	bool	Return the energy-normalized complex sensor-plane field instead of intensity. Defaults to False.	False

Returns:

Type	Description
	torch.Tensor: Intensity PSF normalized to sum to 1, or an energy-normalized complex field when return_field=True. Shape [ks, ks] when ks is given, or [1, 1, 2H, 2H] when ks is None (twice the exit-pupil grid after zero-padding).

Reference

[1] "End-to-End Hybrid Refractive-Diffractive Lens Design with Differentiable Ray-Wave Model", SIGGRAPH Asia 2024.

Note

Similar to the ZEMAX FFT PSF, but free-space propagation uses the Angular Spectrum Method (ASM) instead of a single FFT. ASM is more accurate because the FFT approach assumes a far-field condition (e.g., chief ray perpendicular to the image plane). The complex field has an arbitrary global phase because optical path length is measured relative to a reference; relative phase is meaningful.

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

def psf_pupil_prop(
    self,
    points,
    ks=PSF_KS,
    wvln=None,
    spp=SPP_COHERENT,
    recenter=True,
    return_field=False,
):
    """Compute the single-point monochromatic PSF via the exit-pupil diffraction model.

    Steps:
        1. Compute the complex wavefront at the exit-pupil plane by coherent ray tracing.
        2. Propagate to the sensor plane with the Angular Spectrum Method (ASM)
           and return either its field or intensity. This function is differentiable.

    Args:
        points (torch.Tensor or list): Single normalized point source [3] or
            [1, 3] with x, y in [-1, 1] and z in [-Inf, 0].
        ks (int, optional): Size of the output PSF patch in pixels. If None,
            the full propagated intensity field is returned uncropped.
            Defaults to PSF_KS.
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`.
        spp (int, optional): Number of rays to sample. Defaults to SPP_COHERENT.
        recenter (bool, optional): If True (default), center on the chief ray;
            otherwise on the pinhole projection.
        return_field (bool, optional): Return the energy-normalized complex
            sensor-plane field instead of intensity. Defaults to False.

    Returns:
        torch.Tensor: Intensity PSF normalized to sum to 1, or an
            energy-normalized complex field when ``return_field=True``.
            Shape [ks, ks] when `ks` is given, or [1, 1, 2H, 2H] when
            `ks` is None (twice the exit-pupil grid after zero-padding).

    Reference:
        [1] "End-to-End Hybrid Refractive-Diffractive Lens Design with Differentiable Ray-Wave Model", SIGGRAPH Asia 2024.

    Note:
        Similar to the ZEMAX FFT PSF, but free-space propagation uses the
        Angular Spectrum Method (ASM) instead of a single FFT. ASM is more
        accurate because the FFT approach assumes a far-field condition
        (e.g., chief ray perpendicular to the image plane).
        The complex field has an arbitrary global phase because optical path
        length is measured relative to a reference; relative phase is meaningful.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    # Pupil field by coherent ray tracing
    wavefront, psfc = self.pupil_field(
        points=points, wvln=wvln, spp=spp, recenter=recenter
    )

    # Propagate to sensor plane and get intensity
    pupilz, pupilr = self.get_exit_pupil()
    h, w = wavefront.shape
    # Manually pad wave field
    wavefront = F.pad(
        wavefront.unsqueeze(0).unsqueeze(0),
        [h // 2, h // 2, w // 2, w // 2],
        mode="constant",
        value=0,
    )
    # Free-space propagation using Angular Spectrum Method (ASM)
    sensor_field = AngularSpectrumMethod(
        wavefront,
        z=self.d_sensor - pupilz,
        wvln=wvln,
        ps=self.pixel_size,
        padding=False,
    )
    # Get intensity
    psf_inten = sensor_field.abs() ** 2

    # Calculate PSF center
    h, w = psf_inten.shape[-2:]
    # consider both interplation and padding
    psfc_idx_i = ((2 - psfc[1]) * h / 4).round().long()
    psfc_idx_j = ((2 + psfc[0]) * w / 4).round().long()

    if return_field:
        field = sensor_field
        if ks is not None:
            field = (
                F.pad(
                    field,
                    [ks // 2, ks // 2, ks // 2, ks // 2],
                    mode="constant",
                    value=0,
                )
                .squeeze(0)
                .squeeze(0)
            )
            field = field[
                psfc_idx_i : psfc_idx_i + ks,
                psfc_idx_j : psfc_idx_j + ks,
            ]
        energy = field.abs().square().sum(dim=(-2, -1), keepdim=True)
        return field / torch.sqrt(energy + EPSILON)

    # Crop valid PSF region and normalize
    if ks is not None:
        psf_inten_pad = (
            F.pad(
                psf_inten,
                [ks // 2, ks // 2, ks // 2, ks // 2],
                mode="constant",
                value=0,
            )
            .squeeze(0)
            .squeeze(0)
        )
        psf = psf_inten_pad[
            psfc_idx_i : psfc_idx_i + ks, psfc_idx_j : psfc_idx_j + ks
        ]
    else:
        psf = psf_inten

    # Intensity normalization, shape of [ks, ks] or [h, w]
    psf = psf / (torch.sum(psf, dim=(-2, -1), keepdim=True) + EPSILON)

    return diff_float(psf)

pupil_field

pupil_field(points, wvln=None, spp=SPP_COHERENT, recenter=True)

Compute the complex wavefront at the exit-pupil plane by coherent ray tracing.

The wavefront is xy-flipped for subsequent PSF calculation and binned at the sensor pixel size onto a square [H, H] grid (H = sensor height in pixels). This function is differentiable.

Parameters:

Name	Type	Description	Default
points	Tensor or list	Single normalized point source [3] or [1, 3] with x, y in [-1, 1] and z in [-Inf, 0].	required
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
spp	int	Number of rays to sample. Must be at least 1,000,000 for accurate coherent simulation. Defaults to SPP_COHERENT.	SPP_COHERENT
recenter	bool	If True (default), center on the chief ray; otherwise on the pinhole projection.	True

Returns:

Name	Type	Description
wavefront	Tensor	Complex wavefront at the exit pupil, binned at pixel size. Shape [H, H].
psf_center	list	Normalized PSF center [x, y] on the sensor in [-1, 1].

Note

Default dtype must be torch.float64 for accurate phase calculation.

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

def pupil_field(self, points, wvln=None, spp=SPP_COHERENT, recenter=True):
    """Compute the complex wavefront at the exit-pupil plane by coherent ray tracing.

    The wavefront is xy-flipped for subsequent PSF calculation and binned at
    the sensor pixel size onto a square `[H, H]` grid (H = sensor height in
    pixels). This function is differentiable.

    Args:
        points (torch.Tensor or list): Single normalized point source [3] or
            [1, 3] with x, y in [-1, 1] and z in [-Inf, 0].
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`.
        spp (int, optional): Number of rays to sample. Must be at least
            1,000,000 for accurate coherent simulation. Defaults to SPP_COHERENT.
        recenter (bool, optional): If True (default), center on the chief ray;
            otherwise on the pinhole projection.

    Returns:
        wavefront (torch.Tensor): Complex wavefront at the exit pupil, binned
            at pixel size. Shape [H, H].
        psf_center (list): Normalized PSF center [x, y] on the sensor in [-1, 1].

    Note:
        Default dtype must be torch.float64 for accurate phase calculation.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    assert spp >= 1_000_000, (
        f"Ray sampling {spp} is too small for coherent ray tracing, which may lead to inaccurate simulation."
    )
    if self.dtype != torch.float64:
        raise ValueError(
            "Coherent pupil propagation requires a float64 lens; call lens.astype(torch.float64)."
        )

    sensor_w, sensor_h = self.sensor_size
    device = self.device

    if isinstance(points, list):
        points = torch.as_tensor(points, device=device, dtype=self.dtype).unsqueeze(
            0
        )
    elif torch.is_tensor(points) and len(points.shape) == 1:
        points = points.unsqueeze(0).to(device=device, dtype=self.dtype)
    elif torch.is_tensor(points) and len(points.shape) == 2:
        points = points.to(device=device, dtype=self.dtype)
        assert points.shape[0] == 1, (
            f"pupil_field only supports single point input, got shape {points.shape}"
        )
    else:
        raise ValueError(f"Unsupported point type {points.type()}.")

    assert points.shape[0] == 1, (
        "Only one point is supported for pupil field calculation."
    )

    # Ray origin in the object space
    scale = self.calc_scale(points[:, 2].item())
    point_obj_x = points[:, 0] * scale * sensor_w / 2
    point_obj_y = points[:, 1] * scale * sensor_h / 2
    points_obj = torch.stack([point_obj_x, point_obj_y, points[:, 2]], dim=-1)

    # Ray center determined by chief ray
    # Shape of [N, 2], un-normalized physical coordinates
    if recenter:
        pointc = self.psf_center(points_obj, method="chief_ray")
    else:
        pointc = self.psf_center(points_obj, method="pinhole")

    # Ray-tracing to exit_pupil
    ray = self.sample_from_points(points=points_obj, num_rays=spp, wvln=wvln)
    ray.is_coherent = True
    ray = self.trace2exit_pupil(ray)

    # Calculate complex field (same physical size and resolution as the sensor)
    # Complex field is flipped here for further PSF calculation
    pointc_ref = torch.zeros_like(points[:, :2])  # [N, 2]
    wavefront = forward_integral(
        ray.flip_xy(),
        ps=self.pixel_size,
        ks=self.sensor_res[1],
        pointc=pointc_ref,
    )
    wavefront = wavefront.squeeze(0)  # [H, H]

    # PSF center (on the sensor plane).
    pointc = pointc[0, :]
    psf_center = [
        pointc[0] / sensor_w * 2,
        pointc[1] / sensor_h * 2,
    ]

    return wavefront, psf_center

psf_huygens

psf_huygens(
    points, ks=PSF_KS, wvln=None, spp=SPP_COHERENT, recenter=True, return_field=False
)

Compute the single-wavelength Huygens PSF by spherical-wave integration.

Not differentiable, due to the heavy computational cost.

Steps

1. Trace coherent rays to the exit-pupil plane.
2. Treat every ray as a secondary point source emitting a spherical wave, and coherently sum these waves over the PSF pixel grid. Each contribution uses the Huygens-Fresnel obliquity factor \(0.5 (1 + \cos\theta)\) and \(1/r\) spherical-wave amplitude decay.

Parameters:

Name	Type	Description	Default
points	Tensor	Single normalized point source [3] or [1, 3] with x, y in [-1, 1] and z in [-Inf, 0].	required
ks	int	Output kernel size in pixels. Defaults to PSF_KS.	PSF_KS
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
spp	int	Rays sampled. Defaults to SPP_COHERENT.	SPP_COHERENT
recenter	bool	If True (default), center on the chief ray; otherwise on the pinhole projection.	True
return_field	bool	Return the energy-normalized complex sensor-plane field instead of intensity. Defaults to False.	False

Returns:

Type	Description
	torch.Tensor: Intensity PSF normalized to sum to 1, or an energy-normalized complex field when return_field=True. Shape [ks, ks].

Reference

[1] "Optical Aberrations Correction in Postprocessing Using Imaging Simulation", TOG 2021.

Note

Different from the ZEMAX Huygens PSF, which traces rays to the image plane and performs plane-wave integration. The complex field has an arbitrary global phase because optical path length is measured relative to a reference; relative phase is meaningful.

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

def psf_huygens(
    self,
    points,
    ks=PSF_KS,
    wvln=None,
    spp=SPP_COHERENT,
    recenter=True,
    return_field=False,
):
    """Compute the single-wavelength Huygens PSF by spherical-wave integration.

    Not differentiable, due to the heavy computational cost.

    Steps:
        1. Trace coherent rays to the exit-pupil plane.
        2. Treat every ray as a secondary point source emitting a spherical
           wave, and coherently sum these waves over the PSF pixel grid. Each
           contribution uses the Huygens-Fresnel obliquity factor
           $0.5 (1 + \\cos\\theta)$ and $1/r$ spherical-wave amplitude decay.

    Args:
        points (torch.Tensor): Single normalized point source [3] or [1, 3]
            with x, y in [-1, 1] and z in [-Inf, 0].
        ks (int, optional): Output kernel size in pixels. Defaults to PSF_KS.
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`.
        spp (int, optional): Rays sampled. Defaults to SPP_COHERENT.
        recenter (bool, optional): If True (default), center on the chief ray;
            otherwise on the pinhole projection.
        return_field (bool, optional): Return the energy-normalized complex
            sensor-plane field instead of intensity. Defaults to False.

    Returns:
        torch.Tensor: Intensity PSF normalized to sum to 1, or an
            energy-normalized complex field when ``return_field=True``.
            Shape [ks, ks].

    Reference:
        [1] "Optical Aberrations Correction in Postprocessing Using Imaging Simulation", TOG 2021.

    Note:
        Different from the ZEMAX Huygens PSF, which traces rays to the image
        plane and performs plane-wave integration.
        The complex field has an arbitrary global phase because optical path
        length is measured relative to a reference; relative phase is meaningful.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    if self.dtype != torch.float64:
        raise ValueError(
            "Huygens propagation requires a float64 lens; call lens.astype(torch.float64)."
        )

    sensor_w, sensor_h = self.sensor_size
    pixel_size = self.pixel_size
    device = self.device
    wvln_mm = wvln * 1e-3  # Convert wavelength to mm

    # Points shape of [N, 3]
    if not torch.is_tensor(points):
        points = torch.as_tensor(points, device=device, dtype=self.dtype)
    else:
        points = points.to(device=device, dtype=self.dtype)

    if len(points.shape) == 1:
        single_point = True
        points = points.unsqueeze(0)
    elif len(points.shape) == 2 and points.shape[0] == 1:
        single_point = True
    else:
        raise ValueError(
            f"Points must be of shape [3] or [1, 3], got {points.shape}."
        )

    # Sample rays from object point
    depth = points[:, 2]
    scale = self.calc_scale(depth)
    point_obj_x = points[..., 0] * scale * sensor_w / 2
    point_obj_y = points[..., 1] * scale * sensor_h / 2
    point_obj = torch.stack([point_obj_x, point_obj_y, points[..., 2]], dim=-1)
    ray = self.sample_from_points(points=point_obj, num_rays=spp, wvln=wvln)

    # Trace rays coherently through the lens to exit pupil
    ray.is_coherent = True
    ray = self.trace2exit_pupil(ray)

    # Calculate PSF center (not flipped here)
    if recenter:
        pointc = -self.psf_center(point_obj, method="chief_ray")
    else:
        pointc = -self.psf_center(point_obj, method="pinhole")

    # Build PSF pixel coordinates (sensor plane at z = d_sensor)
    sensor_z = self.d_sensor.item()
    psf_half_size = (ks / 2) * pixel_size  # Physical half-size of PSF region
    x_coords = torch.linspace(
        -psf_half_size + pixel_size / 2,
        psf_half_size - pixel_size / 2,
        ks,
        device=device,
        dtype=self.dtype,
    )
    y_coords = torch.linspace(
        psf_half_size - pixel_size / 2,
        -psf_half_size + pixel_size / 2,
        ks,
        device=device,
        dtype=self.dtype,
    )
    psf_x, psf_y = torch.meshgrid(
        pointc[0, 0] + x_coords, pointc[0, 1] + y_coords, indexing="xy"
    )  # [ks, ks] each

    # Get valid rays only
    valid_mask = ray.is_valid > 0
    valid_pos = ray.o[valid_mask]  # [num_valid, 3]
    valid_dir = ray.d[valid_mask]  # [num_valid, 3]
    valid_opl = ray.opl[valid_mask]  # [num_valid]
    num_valid = valid_pos.shape[0]

    # Huygens integration: sum spherical waves from each secondary source
    psf_complex = torch.zeros(ks, ks, dtype=torch.complex128, device=device)
    opl_min = valid_opl.min()

    # Compute distance from each secondary source to each pixel
    batch_size = min(num_valid, 10_000)  # Process rays in batches
    for batch_start in range(0, num_valid, batch_size):
        batch_end = min(batch_start + batch_size, num_valid)

        # Batch ray data
        batch_pos = valid_pos[batch_start:batch_end]  # [batch, 3]
        batch_dir = valid_dir[batch_start:batch_end]  # [batch, 3]
        batch_opl = valid_opl[batch_start:batch_end].squeeze(-1)  # [batch]

        # Distance from each secondary source to each pixel
        # batch_pos: [batch, 3], psf_x: [ks, ks]
        dx = psf_x.unsqueeze(-1) - batch_pos[:, 0]  # [ks, ks, batch]
        dy = psf_y.unsqueeze(-1) - batch_pos[:, 1]  # [ks, ks, batch]
        dz = sensor_z - batch_pos[:, 2]  # [batch]

        # Distance r from secondary source to pixel
        r = torch.sqrt(dx**2 + dy**2 + dz**2)  # [ks, ks, batch]

        # Obliquity factor: cos(theta) where theta is angle from normal
        # Using ray direction at exit pupil (dz component)
        obliq = torch.abs(batch_dir[:, 2])  # [batch]
        amp = 0.5 * (1.0 + obliq)  # Huygens–Fresnel obliquity factor

        # Total optical path = OPL through lens + distance to pixel
        total_opl = batch_opl + r  # [ks, ks, batch]

        # Phase relative to reference
        phase = torch.fmod((total_opl - opl_min) / wvln_mm, 1.0) * (
            2 * torch.pi
        )  # [ks, ks, batch]

        # Complex amplitude: A * exp(i * phase) / r (spherical wave decay)
        # We use 1/r for spherical wave amplitude decay
        complex_amp = (amp / r) * torch.exp(1j * phase)  # [ks, ks, batch]

        # Sum contributions from this batch
        psf_complex += complex_amp.sum(dim=-1)  # [ks, ks]

    if return_field:
        field = torch.flip(psf_complex, [-2, -1])
        energy = field.abs().square().sum(dim=(-2, -1), keepdim=True)
        return field / torch.sqrt(energy + EPSILON)

    # Convert complex field to intensity
    psf = psf_complex.abs() ** 2

    # Intensity normalization
    psf = psf / (torch.sum(psf, dim=(-2, -1), keepdim=True) + EPSILON)

    # Flip PSF
    psf = torch.flip(psf, [-2, -1])

    if single_point:
        psf = psf.squeeze(0)

    return diff_float(psf)

psf_map

psf_map(depth=None, grid=(7, 7), ks=PSF_KS, spp=SPP_PSF, wvln=None, recenter=True)

Compute the geometric PSF map across the field of view at a given depth.

Overrides the base Lens method to improve efficiency by tracing all field points in parallel.

Parameters:

Name	Type	Description	Default
depth	float	Object plane depth [mm]. When None (default), falls back to self.obj_depth.	None
grid	int or tuple	Grid size (grid_w, grid_h); an int is broadcast to a square grid. Defaults to (7, 7).	(7, 7)
ks	int	Output kernel size in pixels. Defaults to PSF_KS.	PSF_KS
spp	int	Rays sampled per source. Defaults to SPP_PSF.	SPP_PSF
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
recenter	bool	If True (default), center on the chief ray.	True

Returns:

Name	Type	Description
psf_map	Tensor	PSF map. Shape [grid_h, grid_w, 1, ks, ks].

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

def psf_map(
    self,
    depth=None,
    grid=(7, 7),
    ks=PSF_KS,
    spp=SPP_PSF,
    wvln=None,
    recenter=True,
):
    """Compute the geometric PSF map across the field of view at a given depth.

    Overrides the base `Lens` method to improve efficiency by tracing all
    field points in parallel.

    Args:
        depth (float, optional): Object plane depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        grid (int or tuple, optional): Grid size (grid_w, grid_h); an int is
            broadcast to a square grid. Defaults to (7, 7).
        ks (int, optional): Output kernel size in pixels. Defaults to PSF_KS.
        spp (int, optional): Rays sampled per source. Defaults to SPP_PSF.
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`.
        recenter (bool, optional): If True (default), center on the chief ray.

    Returns:
        psf_map (torch.Tensor): PSF map. Shape [grid_h, grid_w, 1, ks, ks].
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth
    if isinstance(grid, int):
        grid = (grid, grid)
    points = self.point_source_grid(depth=depth, grid=grid)
    points = points.reshape(-1, 3)
    psfs = self.psf(
        points=points, ks=ks, recenter=recenter, spp=spp, wvln=wvln
    ).unsqueeze(1)  # [grid_h * grid_w, 1, ks, ks]

    psf_map = psfs.reshape(grid[1], grid[0], 1, ks, ks)
    return psf_map

psf_center

psf_center(points_obj, method='chief_ray', ray=None)

Compute the reference PSF center on the sensor for a given point source.

With method "chief_ray" it returns the negated sensor intercept of the sampled real ray closest to the physical aperture-stop centre. Invalid chief rays fall back to the pinhole model independently per field, as does the whole computation when the lens has no aperture stop. With "pinhole" it uses an ideal perspective projection (no distortion).

Parameters:

Name	Type	Description	Default
points_obj	Tensor	Un-normalized object-plane point(s), shape [..., 3][mm], spanning [-Inf, Inf] x [-Inf, Inf] x [-Inf, 0].	required
method	str	"chief_ray" or "pinhole". Defaults to "chief_ray".	'chief_ray'
ray	Ray	Bundle already traced to the sensor from points_obj at the primary wavelength, shape [..., num_rays, 3]. When given, the chief ray is extracted from it instead of tracing a new bundle. Ignored for "pinhole".	None

Returns:

Name	Type	Description
psf_center	Tensor	Un-normalized PSF center on the sensor plane [mm], shape [..., 2].

Raises:

Type	Description
ValueError	If method is neither "chief_ray" nor "pinhole".

Source code in deeplens-src/deeplens/geolens_pkg/psf_compute.py

@torch.no_grad()
def psf_center(self, points_obj, method="chief_ray", ray=None):
    """Compute the reference PSF center on the sensor for a given point source.

    With method "chief_ray" it returns the negated sensor intercept of the
    sampled real ray closest to the physical aperture-stop centre. Invalid
    chief rays fall back to the pinhole model independently per field, as
    does the whole computation when the lens has no aperture stop. With
    "pinhole" it uses an ideal perspective projection (no distortion).

    Args:
        points_obj (torch.Tensor): Un-normalized object-plane point(s), shape
            [..., 3] [mm], spanning [-Inf, Inf] x [-Inf, Inf] x [-Inf, 0].
        method (str, optional): "chief_ray" or "pinhole". Defaults to "chief_ray".
        ray (Ray, optional): Bundle already traced to the sensor from
            `points_obj` at the primary wavelength, shape [..., num_rays, 3].
            When given, the chief ray is extracted from it instead of
            tracing a new bundle. Ignored for "pinhole".

    Returns:
        psf_center (torch.Tensor): Un-normalized PSF center on the sensor
            plane [mm], shape [..., 2].

    Raises:
        ValueError: If `method` is neither "chief_ray" nor "pinhole".
    """
    if method == "chief_ray":
        if self.aper_idx is None:
            logger.warning(
                "Lens has no aperture stop; using the pinhole model for "
                "PSF centers."
            )
            return self.psf_center(points_obj, method="pinhole")

        if ray is None:
            chief_ray = self.calc_chief_ray(points_obj, num_rays=SPP_PSF)
        else:
            chief_ray = self.calc_chief_ray(ray=ray)
        psf_center = -chief_ray.o[..., 0, :2]

        # Info level: a vignetted chief ray is routine mid-optimization, and
        # psf()/psf_map() run every iteration, so a warning would flood logs.
        valid = chief_ray.is_valid[..., 0].bool()
        if not valid.all():
            failed = int((~valid).sum().item())
            logger.info(
                "%d of %d chief rays are invalid; using the pinhole model "
                "for those PSF centers.",
                failed,
                valid.numel(),
            )
            pinhole_center = self.psf_center(points_obj, method="pinhole")
            psf_center = torch.where(
                valid.unsqueeze(-1), psf_center, pinhole_center
            )

    elif method == "pinhole":
        # Pinhole camera perspective projection, distortion not considered
        if points_obj[..., 2].min().abs() < 100:
            print(
                "Point source is too close, pinhole model may be inaccurate for PSF center calculation."
            )
        tan_point_fov_x = -points_obj[..., 0] / points_obj[..., 2]
        tan_point_fov_y = -points_obj[..., 1] / points_obj[..., 2]
        psf_center_x = self.foclen * tan_point_fov_x
        psf_center_y = self.foclen * tan_point_fov_y
        psf_center = torch.stack([psf_center_x, psf_center_y], dim=-1).to(
            self.device
        )

    else:
        raise ValueError(
            f"Unsupported method for PSF center calculation: {method}."
        )

    return psf_center

For coherent and Huygens models, pass return_field=True to psf() to get the energy-normalized complex sensor-plane field instead of intensity. Its global phase is arbitrary; relative phase is meaningful.

field = lens.psf(points, model="coherent", return_field=True)

deeplens.geolens_pkg.eval.GeoLensEval

Mixin that adds classical optical evaluation methods to GeoLens.

This class is never instantiated on its own. It is mixed into GeoLens via multiple inheritance, so every method can access lens geometry (self.d_sensor, self.rfov, …) and ray-tracing routines (self.trace(), self.trace2sensor(), …) directly through self.

All evaluation functions follow the same pattern

1. Sample rays from object space (parallel / grid / radial).
2. Trace rays through the lens (self.trace or self.trace2sensor).
3. Analyze ray positions / directions at the sensor plane.
4. Optionally produce a matplotlib figure saved to disk.

Results are accuracy-aligned with Zemax OpticStudio for the same lens prescriptions and ray-sampling densities.

Attributes consumed from GeoLens (via self): d_sensor (float): Axial position of the sensor plane (mm). sensor_size (tuple[float, float]): Sensor (width, height) in mm. pixel_size (float): Pixel pitch in mm. sensor_res (tuple[int, int]): Sensor resolution (W, H) in pixels. rfov (float): Half field-of-view in radians. foclen (float): Equivalent focal length in mm. fnum (float): F-number. aper_idx (int): Index of the aperture stop surface. device (torch.device): Compute device (CPU / CUDA).

calc_chief_ray

calc_chief_ray(
    points_obj=None,
    *,
    ray=None,
    fov=None,
    plane="meridional",
    num_rays=SPP_CALC,
    wvln=None,
    scale_pupil=1.0,
    record=False,
)

Calculate one real, physical-stop-centred chief ray per field.

DeepLens uses the classical physical definition: the chief ray is the real ray through the centre of the aperture stop. Because the tracer samples a finite bundle, this method returns the sampled ray closest to the stop centre rather than a synthetic bundle centroid. The selection residual scales as ~1/sqrt(num_rays) of the stop radius; raise num_rays when tighter stop centring is needed (e.g. relative distortion at small field angles, where the residual is amplified by the small ideal image height).

Exactly one input mode must be supplied:

1. points_obj samples and traces rays from object-space points to the sensor before extracting the chief ray (finite conjugates).
2. fov traces a collimated bundle at the given field angles (infinite conjugates) in the plane direction.
3. ray extracts from an already traced bundle that crossed the aperture stop, avoiding duplicate tracing.

Selection happens before final validity is checked. If the closest stop-centred ray is clipped after the stop, it remains the selected ray and is returned as invalid rather than being replaced by a farther ray. The result's stop_dist is its distance from the stop centre in stop radii. If no sample has a finite recorded distance, including an untraced bundle, the result is invalid with stop_dist=inf. Multiply a recorded distance by the stop radius used for tracing to obtain the distance in millimetres.

Parameters:

Name	Type	Description	Default
points_obj	Tensor | list | None	Physical object points in millimetres, shape [..., 3].	None
ray	Ray | None	Already traced ray bundle that crossed the aperture stop, shape [..., num_rays, 3].	None
fov	float | list | Tensor | None	Field angle(s) in degrees for infinite-conjugate chief rays, shape [N] (a scalar becomes [1]).	None
plane	str	'meridional' (y-axis) or 'sagittal' (x-axis); used by fov mode only. Defaults to 'meridional'.	'meridional'
num_rays	int	Samples per field in points_obj / fov mode. Defaults to SPP_CALC.	SPP_CALC
wvln	float | None	Wavelength in micrometres for the sampling modes. Defaults to the lens primary wavelength.	None
scale_pupil	float	Entrance-pupil sampling-radius multiplier in the sampling modes. Defaults to 1.0.	1.0
record	bool	Return (chief_ray, path) with a singleton ray dimension in every path entry. Not valid in ray mode.	False

Returns:

Type	Description
	Ray | tuple[Ray, list[torch.Tensor]]: One selected chief ray per
	field, with shape [..., 1, 3], and optionally its traced path.

Raises:

Type	Description
ValueError	If the input modes are invalid, sampling parameters are invalid, the supplied bundle is empty, or path recording is requested for a supplied bundle.
TypeError	If ray is not a Ray.

Note

This method is discrete and decorated with torch.no_grad(). It is intended for evaluation and PSF centring, not gradient losses.

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def calc_chief_ray(
    self,
    points_obj=None,
    *,
    ray=None,
    fov=None,
    plane="meridional",
    num_rays=SPP_CALC,
    wvln=None,
    scale_pupil=1.0,
    record=False,
):
    """Calculate one real, physical-stop-centred chief ray per field.

    DeepLens uses the classical physical definition: the chief ray is the
    real ray through the centre of the aperture stop. Because the tracer
    samples a finite bundle, this method returns the sampled ray closest to
    the stop centre rather than a synthetic bundle centroid. The selection
    residual scales as ``~1/sqrt(num_rays)`` of the stop radius; raise
    ``num_rays`` when tighter stop centring is needed (e.g. relative
    distortion at small field angles, where the residual is amplified by
    the small ideal image height).

    Exactly one input mode must be supplied:

    1. ``points_obj`` samples and traces rays from object-space points to
       the sensor before extracting the chief ray (finite conjugates).
    2. ``fov`` traces a collimated bundle at the given field angles
       (infinite conjugates) in the ``plane`` direction.
    3. ``ray`` extracts from an already traced bundle that crossed the
       aperture stop, avoiding duplicate tracing.

    Selection happens before final validity is checked. If the closest
    stop-centred ray is clipped after the stop, it remains the selected ray
    and is returned as invalid rather than being replaced by a farther ray.
    The result's ``stop_dist`` is its distance from the stop centre in
    stop radii. If no sample has a finite recorded distance, including an
    untraced bundle, the result is invalid with ``stop_dist=inf``. Multiply
    a recorded distance by the stop radius used for tracing to obtain
    the distance in millimetres.

    Args:
        points_obj (torch.Tensor | list | None): Physical object points in
            millimetres, shape ``[..., 3]``.
        ray (Ray | None): Already traced ray bundle that crossed the
            aperture stop, shape ``[..., num_rays, 3]``.
        fov (float | list | torch.Tensor | None): Field angle(s) in
            degrees for infinite-conjugate chief rays, shape ``[N]`` (a
            scalar becomes ``[1]``).
        plane (str): ``'meridional'`` (y-axis) or ``'sagittal'`` (x-axis);
            used by ``fov`` mode only. Defaults to ``'meridional'``.
        num_rays (int): Samples per field in ``points_obj`` / ``fov``
            mode. Defaults to ``SPP_CALC``.
        wvln (float | None): Wavelength in micrometres for the sampling
            modes. Defaults to the lens primary wavelength.
        scale_pupil (float): Entrance-pupil sampling-radius multiplier in
            the sampling modes. Defaults to 1.0.
        record (bool): Return ``(chief_ray, path)`` with a singleton ray
            dimension in every path entry. Not valid in ``ray`` mode.

    Returns:
        Ray | tuple[Ray, list[torch.Tensor]]: One selected chief ray per
        field, with shape ``[..., 1, 3]``, and optionally its traced path.

    Raises:
        ValueError: If the input modes are invalid, sampling parameters are
            invalid, the supplied bundle is empty, or path
            recording is requested for a supplied bundle.
        TypeError: If ``ray`` is not a :class:`Ray`.

    Note:
        This method is discrete and decorated with ``torch.no_grad()``. It
        is intended for evaluation and PSF centring, not gradient losses.
    """
    has_points = points_obj is not None
    has_ray = ray is not None
    has_fov = fov is not None
    if int(has_points) + int(has_ray) + int(has_fov) != 1:
        raise ValueError("Provide exactly one of points_obj, fov, or ray.")
    if not isinstance(num_rays, int) or num_rays < 1:
        raise ValueError("num_rays must be a positive integer.")
    if not math.isfinite(float(scale_pupil)) or float(scale_pupil) <= 0:
        raise ValueError("scale_pupil must be a finite positive number.")
    if plane not in ("meridional", "sagittal"):
        raise ValueError(f"Invalid plane: {plane}")
    wvln = self.primary_wvln if wvln is None else wvln

    ray_path = None
    if has_points:
        sampled_ray = self.sample_from_points(
            points=points_obj,
            num_rays=num_rays,
            wvln=wvln,
            scale_pupil=scale_pupil,
        )
        if record:
            ray, ray_path = self.trace2sensor(sampled_ray, record=True)
        else:
            ray = self.trace2sensor(sampled_ray)
    elif has_fov:
        fov_t = torch.as_tensor(fov, device=self.device, dtype=self.dtype)
        if fov_t.ndim == 0:
            fov_t = fov_t.unsqueeze(0)
        if fov_t.ndim != 1:
            raise ValueError("fov must be a scalar or a 1-D sequence.")
        if plane == "sagittal":
            fov_x, fov_y = fov_t.tolist(), 0.0
        else:
            fov_x, fov_y = 0.0, fov_t.tolist()
        sampled_ray = self.sample_from_fov(
            fov_x=fov_x,
            fov_y=fov_y,
            depth=float("inf"),
            num_rays=num_rays,
            wvln=wvln,
            scale_pupil=scale_pupil,
        )
        if record:
            ray, ray_path = self.trace2sensor(sampled_ray, record=True)
        else:
            ray = self.trace2sensor(sampled_ray)
    else:
        if not isinstance(ray, Ray):
            raise TypeError("ray must be a deeplens.light.Ray instance.")
        if record:
            raise ValueError(
                "record=True requires a sampling mode; a supplied traced "
                "bundle does not contain its earlier path."
            )

    if ray.o.ndim < 2 or ray.o.shape[-2] < 1:
        raise ValueError("ray must contain at least one sample ray.")

    # Rays invalid at the stop carry inf, so the minimum is the closest
    # valid sample and a finite minimum means the field reached the stop.
    stop_dist = ray.stop_dist.masked_fill(ray.stop_dist.isnan(), float("inf"))
    residual_normalized, sample_index = stop_dist.min(dim=-1, keepdim=True)
    stop_reached = torch.isfinite(residual_normalized)

    def gather_scalar(value):
        return torch.gather(value, dim=-1, index=sample_index)

    def gather_vector(value):
        index = sample_index.unsqueeze(-1).expand(
            *sample_index.shape, value.shape[-1]
        )
        return torch.gather(value, dim=-2, index=index)

    chief_ray = ray.clone()
    chief_ray.o = gather_vector(ray.o)
    chief_ray.d = gather_vector(ray.d)
    chief_ray.en = gather_vector(ray.en)
    chief_ray.bend_penalty = gather_vector(ray.bend_penalty)
    chief_ray.opl = gather_vector(ray.opl)
    chief_ray.is_valid = gather_scalar(ray.is_valid) * stop_reached.to(
        dtype=ray.is_valid.dtype
    )
    chief_ray.stop_dist = residual_normalized
    chief_ray.shape = chief_ray.o.shape[:-1]

    if record:
        chief_path = [gather_vector(path_entry) for path_entry in ray_path]
        return chief_ray, chief_path
    return chief_ray

spot_points

spot_points(points, num_rays=SPP_PSF, wvln=None)

Trace rays from object points to sensor and return the traced Ray.

Samples rays from each physical object point toward the entrance pupil, traces through all lens surfaces (refraction + clipping), and returns the resulting Ray object on the sensor plane.

This is the shared computational core for spot diagrams (draw_spot_radial, draw_spot_map) and RMS error maps (rms_map, rms_map_rgb).

Algorithm

1. self.sample_from_points(points, num_rays, wvln) generates a fan of num_rays rays per object point, aimed at the entrance pupil.
2. self.trace2sensor() propagates through all surfaces and clips vignetted rays.

Parameters:

Name	Type	Description	Default
points	Tensor	Physical 3D object-space coordinates with shape [..., 3] (mm). Supported layouts: - [3] — single point. - [N, 3] — N points (e.g. radial field positions). - [H, W, 3] — 2-D field grid. Generated by self.point_source_grid(normalized=False) for grid sampling, or self.point_source_radial(normalized=False) for radial sampling.	required
num_rays	int	Number of rays sampled per object point. Defaults to SPP_PSF.	SPP_PSF
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None

Returns:

Name	Type	Description
ray	Ray	Traced ray on the sensor plane, with shape [..., num_rays, 3] for positions and [..., num_rays] for validity mask. Use ray.o[..., :2] for transverse positions and ray.is_valid for the validity mask. ray.centroid() gives the centroid over valid rays.

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def spot_points(self, points, num_rays=SPP_PSF, wvln=None):
    """Trace rays from object points to sensor and return the traced Ray.

    Samples rays from each physical object point toward the entrance pupil,
    traces through all lens surfaces (refraction + clipping), and returns
    the resulting Ray object on the sensor plane.

    This is the shared computational core for spot diagrams
    (``draw_spot_radial``, ``draw_spot_map``) and RMS error maps
    (``rms_map``, ``rms_map_rgb``).

    Algorithm:
        1. ``self.sample_from_points(points, num_rays, wvln)`` generates a
           fan of ``num_rays`` rays per object point, aimed at the entrance
           pupil.
        2. ``self.trace2sensor()`` propagates through all surfaces and
           clips vignetted rays.

    Args:
        points (torch.Tensor): Physical 3D object-space coordinates with
            shape ``[..., 3]`` (mm).  Supported layouts:
            - ``[3]`` — single point.
            - ``[N, 3]`` — N points (e.g. radial field positions).
            - ``[H, W, 3]`` — 2-D field grid.
            Generated by ``self.point_source_grid(normalized=False)`` for
            grid sampling, or ``self.point_source_radial(normalized=False)``
            for radial sampling.
        num_rays (int): Number of rays sampled per object point.
            Defaults to ``SPP_PSF``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.

    Returns:
        ray (Ray): Traced ray on the sensor plane, with shape
            ``[..., num_rays, 3]`` for positions and ``[..., num_rays]``
            for validity mask. Use ``ray.o[..., :2]`` for transverse
            positions and ``ray.is_valid`` for the validity mask.
            ``ray.centroid()`` gives the centroid over valid rays.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    ray = self.sample_from_points(points=points, num_rays=num_rays, wvln=wvln)
    return self.trace2sensor(ray)

draw_spot_radial

draw_spot_radial(
    save_name="./lens_spot_radial.png",
    num_fov=5,
    depth=None,
    num_rays=SPP_PSF,
    wvln_list=None,
    direction="y",
    show=False,
)

Draw spot diagrams at evenly-spaced field angles along a chosen direction.

A spot diagram visualizes the transverse ray-intercept distribution on the sensor plane for a point source at a given field angle and depth. It reveals the combined effect of all aberrations (spherical, coma, astigmatism, field curvature, chromatic, …).

Field positions are sampled by field angle uniformly from on-axis (0) to full-field (self.rfov), so the FoV 1.0 subplot reaches the full image height — consistent with analysis_spot().

Algorithm

For each wavelength in wvln_list: 1. self.sample_radial_rays(direction) samples rays at num_fov field angles in [0, self.rfov] along the chosen direction. 2. self.trace2sensor() traces them to the sensor. 3. Valid ray (x, y) positions are scatter-plotted per subplot. All wavelengths are overlaid in a single figure with RGB coloring.

Parameters:

Name	Type	Description	Default
save_name	str	File path for the output PNG. Defaults to './lens_spot_radial.png'.	'./lens_spot_radial.png'
num_fov	int	Number of field positions sampled uniformly from on-axis (0) to full-field. Defaults to 5.	5
depth	float	Object distance in mm (negative = real object). When None or float('inf') (default None), falls back to self.obj_depth.	None
num_rays	int	Rays per field position per wavelength. Defaults to SPP_PSF.	SPP_PSF
wvln_list	list[float]	Wavelengths in µm. When None (default), falls back to self.wvln_rgb.	None
direction	str	Sampling direction — "y" (meridional, default), "x" (sagittal), "diagonal" (45°).	'y'
show	bool	If True, display the figure interactively instead of saving to disk. Defaults to False.	False

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def draw_spot_radial(
    self,
    save_name="./lens_spot_radial.png",
    num_fov=5,
    depth=None,
    num_rays=SPP_PSF,
    wvln_list=None,
    direction="y",
    show=False,
):
    """Draw spot diagrams at evenly-spaced field angles along a chosen direction.

    A *spot diagram* visualizes the transverse ray-intercept distribution on
    the sensor plane for a point source at a given field angle and depth.
    It reveals the combined effect of all aberrations (spherical, coma,
    astigmatism, field curvature, chromatic, …).

    Field positions are sampled by **field angle** uniformly from on-axis
    (0) to full-field (``self.rfov``), so the ``FoV 1.0`` subplot reaches
    the full image height — consistent with ``analysis_spot()``.

    Algorithm:
        For each wavelength in ``wvln_list``:
            1. ``self.sample_radial_rays(direction)`` samples rays at
               ``num_fov`` field angles in ``[0, self.rfov]`` along the
               chosen direction.
            2. ``self.trace2sensor()`` traces them to the sensor.
            3. Valid ray (x, y) positions are scatter-plotted per subplot.
        All wavelengths are overlaid in a single figure with RGB coloring.

    Args:
        save_name (str): File path for the output PNG.
            Defaults to ``'./lens_spot_radial.png'``.
        num_fov (int): Number of field positions sampled uniformly from
            on-axis (0) to full-field. Defaults to 5.
        depth (float): Object distance in mm (negative = real object).
            When ``None`` or ``float('inf')`` (default ``None``), falls
            back to ``self.obj_depth``.
        num_rays (int): Rays per field position per wavelength.
            Defaults to ``SPP_PSF``.
        wvln_list (list[float]): Wavelengths in µm.  When ``None``
            (default), falls back to ``self.wvln_rgb``.
        direction (str): Sampling direction —
            ``"y"`` (meridional, default), ``"x"`` (sagittal),
            ``"diagonal"`` (45°).
        show (bool): If ``True``, display the figure interactively instead
            of saving to disk. Defaults to ``False``.
    """
    wvln_list = self.wvln_rgb if wvln_list is None else wvln_list
    assert isinstance(wvln_list, list), "wvln_list must be a list"
    if depth is None or depth == float("inf"):
        depth = self.obj_depth

    # Field fractions for subplot titles (0 -> on-axis, 1 -> full field)
    fov_fracs = torch.linspace(0, 1, num_fov)

    # Prepare figure
    fig, axs = plt.subplots(1, num_fov, figsize=(num_fov * 3.5, 3))
    axs = np.atleast_1d(axs)

    # Trace and draw each wavelength separately, overlaying results
    for wvln_idx, wvln in enumerate(wvln_list):
        # Sample rays by field angle (0 .. self.rfov) so FoV 1.0 reaches the
        # full image height, matching analysis_spot().
        ray = self.sample_radial_rays(
            num_field=num_fov,
            depth=depth,
            num_rays=num_rays,
            wvln=wvln,
            direction=direction,
            fov_max=self.rfov,
        )
        ray = self.trace2sensor(ray)
        ray_o = ray.o[..., :2].cpu().numpy()
        ray_valid_np = ray.is_valid.cpu().numpy()

        color = RGB_COLORS[wvln_idx % len(RGB_COLORS)]

        # Plot multiple spot diagrams in one figure
        for i in range(num_fov):
            valid = ray_valid_np[i, :]
            xi, yi = ray_o[i, :, 0], ray_o[i, :, 1]

            # Filter valid rays
            mask = valid > 0
            x_valid, y_valid = xi[mask], yi[mask]

            # Plot points and center of mass for this wavelength
            axs[i].scatter(x_valid, y_valid, 2, color=color, alpha=0.5)
            axs[i].set_aspect("equal", adjustable="datalim")
            axs[i].tick_params(axis="both", which="major", labelsize=6)
            if wvln_idx == 0:
                axs[i].set_title(f"FoV {fov_fracs[i].item():.2f}", fontsize=8)

    if show:
        plt.show()
    else:
        assert save_name.endswith(".png"), "save_name must end with .png"
        plt.savefig(save_name, bbox_inches="tight", format="png", dpi=300)
    plt.close(fig)

draw_spot_map

draw_spot_map(
    save_name="./lens_spot_map.png",
    num_grid=5,
    depth=None,
    num_rays=SPP_PSF,
    wvln_list=None,
    show=False,
)

Draw a 2-D grid of spot diagrams across the full field of view.

Unlike draw_spot_radial (which samples only a radial slice), this method samples a num_grid × num_grid grid of field positions covering both the x (sagittal) and y (meridional) axes, revealing off-axis aberrations that are invisible in a 1-D radial scan.

Grid field positions are sampled by field angle, spanning the full field on both axes so the corner cells reach the full image height — consistent with draw_spot_radial / analysis_spot.

Algorithm

For each wavelength in wvln_list: 1. self.sample_grid_rays() samples a grid_h × grid_w field-angle grid spanning horizontal [-hfov/2, hfov/2] and vertical [-vfov/2, vfov/2] angles. 2. self.trace2sensor() traces them to the sensor. 3. Valid (x, y) positions are scatter-plotted in the corresponding subplot of the num_grid × num_grid figure. All wavelengths are overlaid with RGB coloring.

Parameters:

Name	Type	Description	Default
save_name	str	File path for the output PNG. Defaults to './lens_spot_map.png'.	'./lens_spot_map.png'
num_grid	int | tuple[int, int]	Number of grid points along each axis. Total subplots = grid_w * grid_h. Defaults to 5.	5
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
num_rays	int	Rays per grid cell per wavelength. Defaults to SPP_PSF.	SPP_PSF
wvln_list	list[float]	Wavelengths in µm. When None (default), falls back to self.wvln_rgb.	None
show	bool	If True, display interactively. Defaults to False.	False

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def draw_spot_map(
    self,
    save_name="./lens_spot_map.png",
    num_grid=5,
    depth=None,
    num_rays=SPP_PSF,
    wvln_list=None,
    show=False,
):
    """Draw a 2-D grid of spot diagrams across the full field of view.

    Unlike ``draw_spot_radial`` (which samples only a radial slice),
    this method samples a ``num_grid × num_grid`` grid of field positions
    covering both the x (sagittal) and y (meridional) axes, revealing
    off-axis aberrations that are invisible in a 1-D radial scan.

    Grid field positions are sampled by **field angle**, spanning the full
    field on both axes so the corner cells reach the full image height —
    consistent with ``draw_spot_radial`` / ``analysis_spot``.

    Algorithm:
        For each wavelength in ``wvln_list``:
            1. ``self.sample_grid_rays()`` samples a ``grid_h × grid_w``
               field-angle grid spanning horizontal ``[-hfov/2, hfov/2]``
               and vertical ``[-vfov/2, vfov/2]`` angles.
            2. ``self.trace2sensor()`` traces them to the sensor.
            3. Valid (x, y) positions are scatter-plotted in the
               corresponding subplot of the ``num_grid × num_grid`` figure.
        All wavelengths are overlaid with RGB coloring.

    Args:
        save_name (str): File path for the output PNG.
            Defaults to ``'./lens_spot_map.png'``.
        num_grid (int | tuple[int, int]): Number of grid points along each
            axis. Total subplots = ``grid_w * grid_h``. Defaults to 5.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        num_rays (int): Rays per grid cell per wavelength.
            Defaults to ``SPP_PSF``.
        wvln_list (list[float]): Wavelengths in µm.  When ``None``
            (default), falls back to ``self.wvln_rgb``.
        show (bool): If ``True``, display interactively. Defaults to ``False``.
    """
    wvln_list = self.wvln_rgb if wvln_list is None else wvln_list
    depth = self.obj_depth if depth is None else depth
    assert isinstance(wvln_list, list), "wvln_list must be a list"
    if isinstance(num_grid, int):
        num_grid = (num_grid, num_grid)

    grid_w, grid_h = num_grid
    fig, axs = plt.subplots(grid_h, grid_w, figsize=(grid_w * 3, grid_h * 3))
    axs = np.atleast_2d(axs)

    # Loop wavelengths and overlay scatters
    for wvln_idx, wvln in enumerate(wvln_list):
        # Sample a field-angle grid spanning the full field so the corner
        # cells reach full image height. Shape [grid_h, grid_w, num_rays, 3].
        ray = self.sample_grid_rays(
            depth=depth, num_grid=num_grid, num_rays=num_rays, wvln=wvln
        )
        ray = self.trace2sensor(ray)

        # Convert to numpy, shape [grid_h, grid_w, num_rays, 2]
        ray_o = ray.o[..., :2].cpu().numpy()
        ray_valid_np = ray.is_valid.cpu().numpy()

        color = RGB_COLORS[wvln_idx % len(RGB_COLORS)]

        # Draw per grid cell
        for i in range(grid_h):
            for j in range(grid_w):
                valid = ray_valid_np[i, j, :]
                xi, yi = ray_o[i, j, :, 0], ray_o[i, j, :, 1]

                # Filter valid rays
                mask = valid > 0
                x_valid, y_valid = xi[mask], yi[mask]

                # Plot points for this wavelength
                axs[i, j].scatter(x_valid, y_valid, 2, color=color, alpha=0.5)
                axs[i, j].set_aspect("equal", adjustable="datalim")
                axs[i, j].tick_params(axis="both", which="major", labelsize=6)

    if show:
        plt.show()
    else:
        assert save_name.endswith(".png"), "save_name must end with .png"
        plt.savefig(save_name, bbox_inches="tight", format="png", dpi=300)
    plt.close(fig)

rms_map

rms_map(num_grid=32, depth=None, wvln=None, center=None)

Compute per-field-position RMS spot radius for a single wavelength.

Traces SPP_PSF rays per grid cell and computes the root-mean-square distance of valid ray hits from a reference centroid. When center is None, each cell uses its own centroid (monochromatic blur). When an external center is provided (e.g. the green-channel centroid), the RMS includes the chromatic shift from that reference.

Algorithm

1. self.point_source_grid(normalized=False) generates physical object points on a [num_grid, num_grid] field grid.
2. self.spot_points() samples SPP_PSF rays per point and traces to sensor.
3. If center is None, compute per-cell centroid c = mean(valid ray_xy); otherwise use the provided center.
4. RMS = sqrt( mean( ||ray_xy - c||^2 ) ).

Parameters:

Name	Type	Description	Default
num_grid	int | tuple[int, int]	Spatial resolution of the field sampling grid. Defaults to 32.	32
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
center	Tensor | None	External reference centroid with shape [grid_h, grid_w, 2]. If None, each cell's own centroid is used. Defaults to None.	None

Returns:

Name	Type	Description
rms	Tensor	RMS spot error map, shape [grid_h, grid_w], in mm.
centroid	Tensor	Per-cell centroid used as reference, shape [grid_h, grid_w, 2]. Useful for passing as center to subsequent calls (e.g. in rms_map_rgb).

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def rms_map(self, num_grid=32, depth=None, wvln=None, center=None):
    """Compute per-field-position RMS spot radius for a single wavelength.

    Traces ``SPP_PSF`` rays per grid cell and computes the root-mean-square
    distance of valid ray hits from a reference centroid.  When ``center``
    is ``None``, each cell uses its own centroid (monochromatic blur).
    When an external ``center`` is provided (e.g. the green-channel
    centroid), the RMS includes the chromatic shift from that reference.

    Algorithm:
        1. ``self.point_source_grid(normalized=False)`` generates physical
           object points on a ``[num_grid, num_grid]`` field grid.
        2. ``self.spot_points()`` samples ``SPP_PSF`` rays per point and
           traces to sensor.
        3. If ``center`` is ``None``, compute per-cell centroid
           ``c = mean(valid ray_xy)``; otherwise use the provided ``center``.
        4. ``RMS = sqrt( mean( ||ray_xy - c||^2 ) )``.

    Args:
        num_grid (int | tuple[int, int]): Spatial resolution of the field
            sampling grid. Defaults to 32.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.
        center (torch.Tensor | None): External reference centroid with shape
            ``[grid_h, grid_w, 2]``.  If ``None``, each cell's own
            centroid is used. Defaults to ``None``.

    Returns:
        rms (torch.Tensor): RMS spot error map, shape ``[grid_h, grid_w]``,
            in mm.
        centroid (torch.Tensor): Per-cell centroid used as reference, shape
            ``[grid_h, grid_w, 2]``.  Useful for passing as
            ``center`` to subsequent calls (e.g. in ``rms_map_rgb``).
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth
    if isinstance(num_grid, int):
        num_grid = (num_grid, num_grid)

    # Generate physical grid points and trace rays to sensor
    points = self.point_source_grid(depth=depth, grid=num_grid, normalized=False)
    ray = self.spot_points(points, num_rays=SPP_PSF, wvln=wvln)

    # Reuse Ray.centroid() — shape [grid_h, grid_w, 3], slice to [grid_h, grid_w, 2]
    centroid = ray.centroid()[..., :2]

    # Use external center if provided, otherwise own centroid
    ref = center if center is not None else centroid

    # RMS relative to reference, shape [grid_h, grid_w]
    ray_xy = ray.o[..., :2]
    ray_valid = ray.is_valid
    rms = torch.sqrt(
        (((ray_xy - ref.unsqueeze(-2)) ** 2).sum(-1) * ray_valid).sum(-1)
        / (ray_valid.sum(-1) + EPSILON)
    )

    return rms, centroid

rms_map_rgb

rms_map_rgb(num_grid=32, depth=None)

Compute per-field-position RMS spot radius for R, G, B wavelengths.

The RMS spot radius is a standard measure of geometrical image quality. For each field position in a num_grid × num_grid grid, this method traces SPP_PSF rays per wavelength and computes the root-mean-square distance of valid ray hits from a common reference centroid.

The reference centroid is the green-channel centroid. Using a common reference means the returned RMS values include lateral chromatic aberration (the shift between R/G/B centroids), making the map useful as a polychromatic image-quality metric.

Algorithm

1. Call rms_map(wvln=green) to get the green RMS map and the green centroid.
2. Call rms_map(wvln=red, center=green_centroid) and rms_map(wvln=blue, center=green_centroid) to measure R/B blur relative to the green reference.
3. Stack as [R, G, B].

Parameters:

Name	Type	Description	Default
num_grid	int or tuple[int, int]	Spatial resolution of the field sampling grid. Defaults to 32.	32
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None

Returns:

Name	Type	Description
rms_rgb	Tensor	RMS spot error map with shape [3, grid_h, grid_w] (channels ordered R, G, B). Units are mm (same as sensor coordinates).

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def rms_map_rgb(self, num_grid=32, depth=None):
    """Compute per-field-position RMS spot radius for R, G, B wavelengths.

    The RMS spot radius is a standard measure of geometrical image quality.
    For each field position in a ``num_grid × num_grid`` grid, this method
    traces ``SPP_PSF`` rays per wavelength and computes the root-mean-square
    distance of valid ray hits from a **common** reference centroid.

    The reference centroid is the green-channel centroid.  Using a common
    reference means the returned RMS values include *lateral chromatic
    aberration* (the shift between R/G/B centroids), making the map useful
    as a polychromatic image-quality metric.

    Algorithm:
        1. Call ``rms_map(wvln=green)`` to get the green RMS map **and**
           the green centroid.
        2. Call ``rms_map(wvln=red, center=green_centroid)`` and
           ``rms_map(wvln=blue, center=green_centroid)`` to measure R/B
           blur relative to the green reference.
        3. Stack as ``[R, G, B]``.

    Args:
        num_grid (int or tuple[int, int]): Spatial resolution of the field
            sampling grid. Defaults to 32.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.

    Returns:
        rms_rgb (torch.Tensor): RMS spot error map with shape
            ``[3, grid_h, grid_w]`` (channels ordered R, G, B). Units
            are mm (same as sensor coordinates).
    """
    depth = self.obj_depth if depth is None else depth
    # Green first to obtain the shared reference centroid
    rms_g, green_centroid = self.rms_map(
        num_grid=num_grid, depth=depth, wvln=self.wvln_rgb[1]
    )

    # Red and blue relative to the green centroid
    rms_r, _ = self.rms_map(
        num_grid=num_grid, depth=depth, wvln=self.wvln_rgb[0], center=green_centroid
    )
    rms_b, _ = self.rms_map(
        num_grid=num_grid, depth=depth, wvln=self.wvln_rgb[2], center=green_centroid
    )

    return torch.stack([rms_r, rms_g, rms_b], dim=0)

calc_distortion_radial

calc_distortion_radial(
    num_points=GEO_GRID, wvln=None, plane="meridional", num_rays=SPP_CALC
)

Compute fractional distortion at evenly-spaced field angles along the meridional direction.

Distortion is defined as (h_actual - h_ideal) / h_ideal, where h_ideal = f * tan(theta) (rectilinear projection) and h_actual is the chief-ray image height on the sensor. A positive value means pincushion distortion; negative means barrel distortion.

This is the computational counterpart to draw_spot_radial: it samples num_points field angles uniformly from 0 to self.rfov and returns both the sampled angles and the corresponding distortion values, making it easy to pair with other radial evaluation functions.

Algorithm

1. Derive rfov_deg from self.rfov (radians → degrees).
2. Sample num_points field angles uniformly in [0, rfov_deg]. The on-axis sample (0°) is replaced by a tiny positive angle to avoid 0/0.
3. Compute h_ideal = foclen * tan(angle) for each sample.
4. Compute the physical-stop chief ray at the sensor via calc_chief_ray(fov=...).
5. Extract h_actual from the appropriate transverse coordinate (x for sagittal, y for meridional).
6. Return (h_actual - h_ideal) / h_ideal.

Parameters:

Name	Type	Description	Default
num_points	int	Number of evenly-spaced field-angle samples from on-axis (0°) to full-field (self.rfov). Defaults to GEO_GRID.	GEO_GRID
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
plane	str	'meridional' (y-axis) or 'sagittal' (x-axis). Defaults to 'meridional'.	'meridional'
num_rays	int	Pupil samples per field angle for chief-ray selection. The chief-ray residual (and hence the distortion noise floor, amplified at small field angles by the small ideal image height) scales as ~1/sqrt(num_rays). Defaults to SPP_CALC.	SPP_CALC

Returns:

Name	Type	Description
rfov_samples	ndarray	Field angles in degrees, shape [num_points].
distortions	ndarray	Fractional distortion at each angle, shape [num_points]. Dimensionless (multiply by 100 for percent).

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def calc_distortion_radial(
    self,
    num_points=GEO_GRID,
    wvln=None,
    plane="meridional",
    num_rays=SPP_CALC,
):
    """Compute fractional distortion at evenly-spaced field angles along the meridional direction.

    Distortion is defined as ``(h_actual - h_ideal) / h_ideal``, where
    ``h_ideal = f * tan(theta)`` (rectilinear projection) and ``h_actual``
    is the chief-ray image height on the sensor.  A positive value means
    pincushion distortion; negative means barrel distortion.

    This is the computational counterpart to ``draw_spot_radial``: it
    samples ``num_points`` field angles uniformly from 0 to ``self.rfov``
    and returns both the sampled angles and the corresponding distortion
    values, making it easy to pair with other radial evaluation functions.

    Algorithm:
        1. Derive ``rfov_deg`` from ``self.rfov`` (radians → degrees).
        2. Sample ``num_points`` field angles uniformly in
           ``[0, rfov_deg]``.  The on-axis sample (0°) is replaced by a
           tiny positive angle to avoid 0/0.
        3. Compute ``h_ideal = foclen * tan(angle)`` for each sample.
        4. Compute the physical-stop chief ray at the sensor via
           ``calc_chief_ray(fov=...)``.
        5. Extract ``h_actual`` from the appropriate transverse coordinate
           (x for sagittal, y for meridional).
        6. Return ``(h_actual - h_ideal) / h_ideal``.

    Args:
        num_points (int): Number of evenly-spaced field-angle samples from
            on-axis (0°) to full-field (``self.rfov``).
            Defaults to ``GEO_GRID``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.
        plane (str): ``'meridional'`` (y-axis) or ``'sagittal'`` (x-axis).
            Defaults to ``'meridional'``.
        num_rays (int): Pupil samples per field angle for chief-ray
            selection. The chief-ray residual (and hence the distortion
            noise floor, amplified at small field angles by the small
            ideal image height) scales as ``~1/sqrt(num_rays)``.
            Defaults to ``SPP_CALC``.

    Returns:
        rfov_samples (np.ndarray): Field angles in degrees, shape
            ``[num_points]``.
        distortions (np.ndarray): Fractional distortion at each angle, shape
            ``[num_points]``.  Dimensionless (multiply by 100 for
            percent).
    """
    wvln = self.primary_wvln if wvln is None else wvln
    rfov_deg = self.rfov * 180 / torch.pi

    # Sample field angles uniformly from 0 to rfov_deg.
    # For the on-axis point (FOV=0), distortion is 0/0.  We compute it at a
    # tiny positive angle to obtain the correct limit, which may be non-zero
    # when the sensor is not at the paraxial focus.
    rfov_samples = torch.linspace(0, rfov_deg, num_points)
    rfov_compute = rfov_samples.clone()
    if rfov_compute[0] == 0:
        # Guard rfov_samples[1] for the single-sample case (num_points == 1).
        tiny = (
            rfov_samples[1].item() * 0.01
            if len(rfov_samples) > 1
            else min(0.01, float(rfov_deg) * 0.01)
        )
        rfov_compute[0] = min(0.01, tiny)

    # Ideal image height: h_ideal = f * tan(theta)
    eff_foclen = float(self.foclen)
    ideal_imgh = eff_foclen * np.tan(rfov_compute.numpy() * np.pi / 180)

    # Actual image height from the physical-stop chief ray at the sensor.
    if plane == "sagittal":
        axis = 0
    elif plane == "meridional":
        axis = 1
    else:
        raise ValueError(f"Invalid plane: {plane}")
    chief_ray = self.calc_chief_ray(
        fov=rfov_compute, plane=plane, wvln=wvln, num_rays=num_rays
    )
    actual_imgh = chief_ray.o[..., 0, axis].abs().cpu().numpy()

    # Fractional distortion, with safe handling of the on-axis singularity
    ideal_imgh = np.asarray(ideal_imgh)
    mask = np.abs(ideal_imgh) < EPSILON
    distortions = np.where(
        mask, 0.0, (actual_imgh - ideal_imgh) / np.where(mask, 1.0, ideal_imgh)
    )

    return rfov_samples.numpy(), distortions

draw_distortion_radial

draw_distortion_radial(
    save_name=None, num_points=GEO_GRID, wvln=None, plane="meridional", show=False
)

Draw distortion-vs-field-angle curve in Zemax style.

Produces a plot with field angle on the y-axis and percent distortion on the x-axis, matching the layout convention used in Zemax OpticStudio. Useful for quick visual assessment of barrel / pincushion distortion.

Algorithm

1. Call calc_distortion_radial to obtain field angles and fractional distortion values.
2. Convert distortion to percent and plot.

Parameters:

Name	Type	Description	Default
save_name	str | None	File path for the output PNG. If None, auto-generates './{plane}_distortion_inf.png'.	None
num_points	int	Number of field-angle samples. Defaults to GEO_GRID.	GEO_GRID
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
plane	str	'meridional' or 'sagittal'. Defaults to 'meridional'.	'meridional'
show	bool	If True, display interactively. Defaults to False.	False

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def draw_distortion_radial(
    self,
    save_name=None,
    num_points=GEO_GRID,
    wvln=None,
    plane="meridional",
    show=False,
):
    """Draw distortion-vs-field-angle curve in Zemax style.

    Produces a plot with field angle on the y-axis and percent distortion
    on the x-axis, matching the layout convention used in Zemax OpticStudio.
    Useful for quick visual assessment of barrel / pincushion distortion.

    Algorithm:
        1. Call ``calc_distortion_radial`` to obtain field angles and
           fractional distortion values.
        2. Convert distortion to percent and plot.

    Args:
        save_name (str | None): File path for the output PNG.  If ``None``,
            auto-generates ``'./{plane}_distortion_inf.png'``.
        num_points (int): Number of field-angle samples.
            Defaults to ``GEO_GRID``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.
        plane (str): ``'meridional'`` or ``'sagittal'``.
            Defaults to ``'meridional'``.
        show (bool): If ``True``, display interactively. Defaults to ``False``.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    rfov_deg = self.rfov * 180 / torch.pi

    # Calculate distortion at evenly-spaced field angles
    rfov_samples, distortions = self.calc_distortion_radial(
        num_points=num_points, wvln=wvln, plane=plane
    )

    # Convert to percentage and handle NaN
    values = np.nan_to_num(distortions * 100, nan=0.0).tolist()

    # Create figure
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_title(f"{plane} Surface Distortion")

    # Draw distortion curve
    ax.plot(values, rfov_samples, linestyle="-", color="g", linewidth=1.5)

    # Draw reference line (vertical line)
    ax.axvline(x=0, color="k", linestyle="-", linewidth=0.8)

    # Set grid
    ax.grid(True, color="gray", linestyle="-", linewidth=0.5, alpha=1)

    # Dynamically adjust x-axis range
    value = max(abs(v) for v in values)
    margin = value * 0.2  # 20% margin
    x_min, x_max = -max(0.2, value + margin), max(0.2, value + margin)

    # Set ticks
    x_ticks = np.linspace(-value, value, 3)
    y_ticks = np.linspace(0, rfov_deg, 3)

    ax.set_xticks(x_ticks)
    ax.set_yticks(y_ticks)

    # Format tick labels
    x_labels = [f"{x:.1f}%" for x in x_ticks]
    y_labels = [f"{y:.1f}" for y in y_ticks]

    ax.set_xticklabels(x_labels)
    ax.set_yticklabels(y_labels)

    # Set axis labels
    ax.set_xlabel("Distortion (%)")
    ax.set_ylabel("Field of View (degrees)")

    # Set axis range
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(0, rfov_deg)

    if show:
        plt.show()
    else:
        if save_name is None:
            save_name = f"./{plane}_distortion_inf.png"
        plt.savefig(save_name, bbox_inches="tight", format="png", dpi=300)
    plt.close(fig)

calc_distortion_map

calc_distortion_map(num_grid=16, depth=None, wvln=None)

Compute a 2-D distortion grid mapping ideal to actual image positions.

For each cell in a num_grid × num_grid field grid, rays are traced to the sensor and the physical-stop chief-ray position is extracted (fields without a valid chief ray fall back to the bundle centroid). The position is then normalized to [-1, 1] sensor coordinates, producing a map that shows how each ideal image point is displaced by lens distortion.

This map can be used with torch.nn.functional.grid_sample to warp or unwarp rendered images.

Parameters:

Name	Type	Description	Default
num_grid	int	Grid resolution along each axis. Defaults to 16.	16
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None

Returns:

Name	Type	Description
distortion_grid	Tensor	Distortion grid with shape [grid_h, grid_w, 2]. Each entry (x, y) is in normalized sensor coordinates [-1, 1], representing the actual chief-ray image position for the corresponding ideal grid position.

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def calc_distortion_map(self, num_grid=16, depth=None, wvln=None):
    """Compute a 2-D distortion grid mapping ideal to actual image positions.

    For each cell in a ``num_grid × num_grid`` field grid, rays are traced
    to the sensor and the physical-stop chief-ray position is extracted
    (fields without a valid chief ray fall back to the bundle centroid).
    The position is then normalized to ``[-1, 1]`` sensor coordinates,
    producing a map that shows how each ideal image point is displaced by
    lens distortion.

    This map can be used with ``torch.nn.functional.grid_sample`` to warp
    or unwarp rendered images.

    Args:
        num_grid (int): Grid resolution along each axis. Defaults to 16.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.

    Returns:
        distortion_grid (torch.Tensor): Distortion grid with shape
            ``[grid_h, grid_w, 2]``. Each entry ``(x, y)`` is in
            normalized sensor coordinates ``[-1, 1]``, representing the
            actual chief-ray image position for the corresponding ideal
            grid position.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth
    # Sample and trace rays, shape (grid_size, grid_size, num_rays, 3)
    ray = self.sample_grid_rays(
        depth=depth, num_grid=num_grid, wvln=wvln, uniform_fov=False
    )
    ray = self.trace2sensor(ray)

    # Reference position per cell, shape (grid_size, grid_size, 2): the
    # physical-stop chief ray (centroid fallback for invalid fields).
    # Normalize each axis by its own half-extent so non-square sensors
    # map correctly to [-1, 1]: x by sensor_size[0] (width, W),
    # y by sensor_size[1] (height, H).  Sign is flipped on both axes to
    # undo image inversion, matching ``distortion_center``.
    sensor_w, sensor_h = self.sensor_size
    ray_xy = -self._chief_or_centroid_xy(ray)
    x_dist = ray_xy[..., 0] / (sensor_w / 2)
    y_dist = ray_xy[..., 1] / (sensor_h / 2)
    distortion_grid = torch.stack((x_dist, y_dist), dim=-1)
    return distortion_grid

calc_inv_distortion_map

calc_inv_distortion_map(num_grid=16, depth=None, wvln=None)

Compute a grid for applying lens distortion with grid_sample.

For each point on the distorted sensor grid, backward rays are traced through the lens to the target object-depth plane. The traced object intersections are converted to normalized ideal image coordinates. Passing this grid to torch.nn.functional.grid_sample samples an undistorted image and produces a distorted image.

Parameters:

Name	Type	Description	Default
num_grid	int or tuple	Grid resolution. If a tuple is supplied, it is interpreted as (grid_w, grid_h).	16
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None

Returns:

Name	Type	Description
inv_distortion_grid	Tensor	Inverse distortion grid with shape [grid_h, grid_w, 2] in grid_sample coordinates.

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def calc_inv_distortion_map(self, num_grid=16, depth=None, wvln=None):
    """Compute a grid for applying lens distortion with ``grid_sample``.

    For each point on the distorted sensor grid, backward rays are traced
    through the lens to the target object-depth plane. The traced object
    intersections are converted to normalized ideal image coordinates.
    Passing this grid to ``torch.nn.functional.grid_sample`` samples an
    undistorted image and produces a distorted image.

    Args:
        num_grid (int or tuple): Grid resolution. If a tuple is supplied,
            it is interpreted as ``(grid_w, grid_h)``.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.

    Returns:
        inv_distortion_grid (torch.Tensor): Inverse distortion grid with
            shape ``[grid_h, grid_w, 2]`` in ``grid_sample`` coordinates.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth
    if isinstance(num_grid, int):
        num_grid = (num_grid, num_grid)

    grid_w, grid_h = num_grid
    sensor_w, sensor_h = self.sensor_size
    device = self.device

    # Convert grid_sample output coordinates to physical sensor positions.
    # Distortion maps use negated sensor positions as image coordinates.
    x, y = torch.meshgrid(
        torch.linspace(sensor_w / 2, -sensor_w / 2, grid_w, device=device),
        torch.linspace(sensor_h / 2, -sensor_h / 2, grid_h, device=device),
        indexing="xy",
    )
    z = torch.full_like(x, self.d_sensor.item())

    pupilz, pupilr = self.get_exit_pupil()
    ray_o2 = self.sample_circle(
        r=pupilr, z=pupilz, shape=(grid_h, grid_w, SPP_CALC)
    )
    ray_o = torch.stack((x, y, z), dim=-1).unsqueeze(2).repeat(1, 1, SPP_CALC, 1)
    ray = Ray(ray_o, ray_o2 - ray_o, wvln, device=device)

    ray = self.trace2obj(ray)
    ray = ray.prop_to(depth)
    # Backward tracing stamps stop weights at the aperture, so the chief
    # backward ray (through the stop centre) gives the object point whose
    # forward chief ray lands on this sensor cell — the exact inverse of
    # the forward distortion map. Centroid fallback for invalid fields.
    point_obj = self._chief_or_centroid_xy(ray)

    scale = self.calc_scale(depth)
    x_ideal = point_obj[..., 0] / (scale * sensor_w / 2)
    y_ideal = point_obj[..., 1] / (scale * sensor_h / 2)
    inv_distortion_grid = torch.stack((x_ideal, y_ideal), dim=-1)
    return inv_distortion_grid

distortion_center

distortion_center(points)

Compute the distorted image position for arbitrary normalized object points.

Given object points in normalized coordinates, this method converts them to physical object-space positions, traces rays from each point through the lens, and returns the physical-stop chief-ray position on the sensor (bundle-centroid fallback for fields without a valid chief ray) in normalized [-1, 1] coordinates. This is the inverse mapping needed for distortion correction (unwarping).

Algorithm

1. Convert normalized (x, y) ∈ [-1, 1] to physical object-space positions using self.calc_scale(depth) and self.sensor_size.
2. self.sample_from_points() generates rays from each point.
3. self.trace2sensor() propagates rays.
4. Extract the chief-ray position and normalize back to [-1, 1].

Parameters:

Name	Type	Description	Default
points	Tensor	Normalized point source positions with shape [N, 3] or [..., 3]. x, y ∈ [-1, 1] encode the field position; z ∈ (-∞, 0] is the object depth in mm.	required

Returns:

Name	Type	Description
distortion_center	Tensor	Normalized distortion centroid positions with shape [N, 2] or [..., 2]. x, y ∈ [-1, 1].

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

def distortion_center(self, points):
    """Compute the distorted image position for arbitrary normalized object points.

    Given object points in normalized coordinates, this method converts them
    to physical object-space positions, traces rays from each point through
    the lens, and returns the physical-stop chief-ray position on the sensor
    (bundle-centroid fallback for fields without a valid chief ray) in
    normalized ``[-1, 1]`` coordinates.  This is the inverse mapping needed
    for distortion correction (unwarping).

    Algorithm:
        1. Convert normalized ``(x, y)`` ∈ [-1, 1] to physical object-space
           positions using ``self.calc_scale(depth)`` and ``self.sensor_size``.
        2. ``self.sample_from_points()`` generates rays from each point.
        3. ``self.trace2sensor()`` propagates rays.
        4. Extract the chief-ray position and normalize back to ``[-1, 1]``.

    Args:
        points (torch.Tensor): Normalized point source positions with shape
            ``[N, 3]`` or ``[..., 3]``.  ``x, y`` ∈ [-1, 1] encode the
            field position; ``z`` ∈ (-∞, 0] is the object depth in mm.

    Returns:
        distortion_center (torch.Tensor): Normalized distortion centroid
            positions with shape ``[N, 2]`` or ``[..., 2]``.
            ``x, y`` ∈ [-1, 1].
    """
    sensor_w, sensor_h = self.sensor_size

    # Convert normalized points to object space coordinates
    depth = points[..., 2]
    scale = self.calc_scale(depth)
    points_obj_x = points[..., 0] * scale * sensor_w / 2
    points_obj_y = points[..., 1] * scale * sensor_h / 2
    points_obj = torch.stack([points_obj_x, points_obj_y, depth], dim=-1)

    # Sample rays and trace to sensor
    ray = self.sample_from_points(points=points_obj)
    ray = self.trace2sensor(ray)

    # Chief-ray position (centroid fallback), normalized to [-1, 1]
    ray_center = -self._chief_or_centroid_xy(ray)  # shape [..., 2]
    distortion_center_x = ray_center[..., 0] / (sensor_w / 2)
    distortion_center_y = ray_center[..., 1] / (sensor_h / 2)
    distortion_center = torch.stack(
        (distortion_center_x, distortion_center_y), dim=-1
    )
    return distortion_center

draw_distortion_map

draw_distortion_map(save_name=None, num_grid=16, depth=None, wvln=None, show=False)

Draw a scatter plot of the distortion grid.

Visualizes the output of calc_distortion_map() as a scatter plot on [-1, 1] normalized sensor coordinates. An undistorted lens would show a perfect rectilinear grid; deviations reveal barrel or pincushion distortion.

Parameters:

Name	Type	Description	Default
save_name	str | None	File path for the output PNG. If None, auto-generates './distortion_{depth}.png'.	None
num_grid	int	Grid resolution per axis. Defaults to 16.	16
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None
show	bool	If True, display interactively. Defaults to False.	False

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def draw_distortion_map(
    self, save_name=None, num_grid=16, depth=None, wvln=None, show=False
):
    """Draw a scatter plot of the distortion grid.

    Visualizes the output of ``calc_distortion_map()`` as a scatter plot on
    ``[-1, 1]`` normalized sensor coordinates.  An undistorted lens would
    show a perfect rectilinear grid; deviations reveal barrel or pincushion
    distortion.

    Args:
        save_name (str | None): File path for the output PNG.  If ``None``,
            auto-generates ``'./distortion_{depth}.png'``.
        num_grid (int): Grid resolution per axis. Defaults to 16.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.
        show (bool): If ``True``, display interactively. Defaults to ``False``.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth
    # Ray tracing to calculate distortion map
    distortion_grid = self.calc_distortion_map(
        num_grid=num_grid, depth=depth, wvln=wvln
    )
    # Scale axes so the plot preserves the physical sensor aspect ratio:
    # longer side → ±1, shorter side → ±(shorter/longer).
    sensor_w, sensor_h = self.sensor_size
    max_half = max(sensor_w, sensor_h) / 2
    aspect_x = (sensor_w / 2) / max_half
    aspect_y = (sensor_h / 2) / max_half
    x1 = distortion_grid[..., 0].cpu().numpy() * aspect_x
    y1 = distortion_grid[..., 1].cpu().numpy() * aspect_y

    # Draw image
    fig, ax = plt.subplots()
    ax.set_axisbelow(True)
    ax.grid(True)
    ax.scatter(x1, y1, s=20, zorder=3)
    ax.axis("scaled")

    # Grid lines based on grid_size, scaled per axis so the overlay
    # matches the data extent (±aspect_x × ±aspect_y).
    ax.set_xticks(np.linspace(-aspect_x, aspect_x, num_grid))
    ax.set_yticks(np.linspace(-aspect_y, aspect_y, num_grid))
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.tick_params(length=0)
    for spine in ax.spines.values():
        spine.set_visible(False)

    if show:
        plt.show()
    else:
        depth_str = "inf" if depth == float("inf") else f"{-depth}mm"
        if save_name is None:
            save_name = f"./distortion_{depth_str}.png"
        plt.savefig(save_name, bbox_inches="tight", format="png", dpi=300)
    plt.close(fig)

mtf

mtf(fov, wvln=None)

Compute the geometric MTF at a single field position.

The Modulation Transfer Function describes how well the lens preserves contrast as a function of spatial frequency. MTF = 1 at low frequencies (perfect contrast) and falls toward 0 near the diffraction limit or the Nyquist frequency of the sensor.

This implementation uses the geometric (ray-based) approach: 1. Compute the PSF at the given field position via self.psf(). 2. Convert PSF → MTF via psf2mtf() (project onto tangential and sagittal axes, then take the magnitude of the 1-D FFT).

Tangential MTF captures resolution in the meridional (radial) direction; sagittal MTF captures resolution perpendicular to it. The difference between the two indicates astigmatism.

Parameters:

Name	Type	Description	Default
fov	float	Field angle in radians. Internally mapped to a normalized point [0, -fov/rfov, self.obj_depth].	required
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln.	None

Returns:

Name	Type	Description
freq	ndarray	Spatial frequency axis in cycles/mm (positive frequencies only, excluding DC).
mtf_tan	ndarray	Tangential (meridional) MTF values, normalized so that MTF → 1 at low frequency.
mtf_sag	ndarray	Sagittal MTF values, same normalization.

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

def mtf(self, fov, wvln=None):
    """Compute the geometric MTF at a single field position.

    The *Modulation Transfer Function* describes how well the lens preserves
    contrast as a function of spatial frequency.  MTF = 1 at low frequencies
    (perfect contrast) and falls toward 0 near the diffraction limit or the
    Nyquist frequency of the sensor.

    This implementation uses the *geometric* (ray-based) approach:
        1. Compute the PSF at the given field position via ``self.psf()``.
        2. Convert PSF → MTF via ``psf2mtf()`` (project onto tangential and
           sagittal axes, then take the magnitude of the 1-D FFT).

    Tangential MTF captures resolution in the meridional (radial) direction;
    sagittal MTF captures resolution perpendicular to it.  The difference
    between the two indicates astigmatism.

    Args:
        fov (float): Field angle in radians.  Internally mapped to a
            normalized point ``[0, -fov/rfov, self.obj_depth]``.
        wvln (float): Wavelength in µm. When ``None`` (default), falls
            back to ``self.primary_wvln``.

    Returns:
        freq (np.ndarray): Spatial frequency axis in cycles/mm (positive
            frequencies only, excluding DC).
        mtf_tan (np.ndarray): Tangential (meridional) MTF values, normalized
            so that MTF → 1 at low frequency.
        mtf_sag (np.ndarray): Sagittal MTF values, same normalization.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    point = [0, -fov / self.rfov, self.obj_depth]
    psf = self.psf(points=point, recenter=True, wvln=wvln)
    freq, mtf_tan, mtf_sag = self.psf2mtf(psf, pixel_size=self.pixel_size)
    return freq, mtf_tan, mtf_sag

psf2mtf staticmethod

psf2mtf(psf, pixel_size)

Convert a 2-D point-spread function to tangential and sagittal MTF curves.

The MTF is the magnitude of the optical transfer function (OTF), which is the Fourier transform of the PSF. For separable 1-D analysis: 1. Integrate the PSF along the x-axis → tangential line-spread function (LSF_tan). 2. Integrate the PSF along the y-axis → sagittal LSF_sag. 3. Take |FFT(LSF)| and normalize by the DC component so that MTF(0) = 1.

Only positive frequencies (excluding DC) are returned, following the convention used in Zemax MTF plots.

Parameters:

Name	Type	Description	Default
psf	Tensor | ndarray	2-D PSF with shape [H, W]. The array's y-axis (rows) corresponds to the tangential (meridional) direction; x-axis (columns) to the sagittal direction.	required
pixel_size	float	Pixel pitch in mm. Determines the frequency axis scaling: Nyquist = 0.5 / pixel_size cycles/mm.	required

Returns:

Name	Type	Description
freq	ndarray	Spatial frequency in cycles/mm (positive, excluding DC). Length is roughly W // 2.
mtf_tan	ndarray	Tangential MTF, normalized to 1 at DC.
mtf_sag	ndarray	Sagittal MTF, normalized to 1 at DC.

References

• https://en.wikipedia.org/wiki/Optical_transfer_function
• Edmund Optics: Introduction to Modulation Transfer Function.

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@staticmethod
def psf2mtf(psf, pixel_size):
    """Convert a 2-D point-spread function to tangential and sagittal MTF curves.

    The MTF is the magnitude of the optical transfer function (OTF), which
    is the Fourier transform of the PSF.  For separable 1-D analysis:
        1. Integrate the PSF along the x-axis → *tangential* line-spread
           function (LSF_tan).
        2. Integrate the PSF along the y-axis → *sagittal* LSF_sag.
        3. Take ``|FFT(LSF)|`` and normalize by the DC component so that
           MTF(0) = 1.

    Only positive frequencies (excluding DC) are returned, following the
    convention used in Zemax MTF plots.

    Args:
        psf (torch.Tensor | np.ndarray): 2-D PSF with shape ``[H, W]``.
            The array's y-axis (rows) corresponds to the **tangential**
            (meridional) direction; x-axis (columns) to the **sagittal**
            direction.
        pixel_size (float): Pixel pitch in mm.  Determines the frequency
            axis scaling: ``Nyquist = 0.5 / pixel_size`` cycles/mm.

    Returns:
        freq (np.ndarray): Spatial frequency in cycles/mm (positive,
            excluding DC).  Length is roughly ``W // 2``.
        mtf_tan (np.ndarray): Tangential MTF, normalized to 1 at DC.
        mtf_sag (np.ndarray): Sagittal MTF, normalized to 1 at DC.

    References:
        - https://en.wikipedia.org/wiki/Optical_transfer_function
        - Edmund Optics: Introduction to Modulation Transfer Function.
    """
    # Convert to numpy (supports torch tensors and numpy arrays)
    try:
        psf_np = psf.detach().cpu().numpy()
    except AttributeError:
        try:
            psf_np = psf.cpu().numpy()
        except AttributeError:
            psf_np = np.asarray(psf)

    # Compute line spread functions (integrate PSF over orthogonal axes)
    # y-axis corresponds to tangential; x-axis corresponds to sagittal
    lsf_sagittal = psf_np.sum(axis=0)  # function of x
    lsf_tangential = psf_np.sum(axis=1)  # function of y

    # One-sided spectra (for real inputs)
    mtf_sag = np.abs(np.fft.rfft(lsf_sagittal))
    mtf_tan = np.abs(np.fft.rfft(lsf_tangential))

    # Normalize by DC to ensure MTF(0) == 1
    dc_sag = mtf_sag[0] if mtf_sag.size > 0 else 1.0
    dc_tan = mtf_tan[0] if mtf_tan.size > 0 else 1.0
    if dc_sag != 0:
        mtf_sag = mtf_sag / dc_sag
    if dc_tan != 0:
        mtf_tan = mtf_tan / dc_tan

    # Frequency axis in cycles/mm (one-sided)
    fx = np.fft.rfftfreq(lsf_sagittal.size, d=pixel_size)
    freq = fx
    positive_freq_idx = freq > 0

    return (
        freq[positive_freq_idx],
        mtf_tan[positive_freq_idx],
        mtf_sag[positive_freq_idx],
    )

draw_mtf

draw_mtf(
    save_name="./lens_mtf.png",
    relative_fov_list=[0.0, 0.7, 1.0],
    depth_list=None,
    psf_ks=128,
    show=False,
)

Draw a grid of MTF curves for multiple depths and field positions.

Produces a len(depth_list) × len(relative_fov_list) subplot grid. Each subplot shows the tangential (T, solid) and sagittal (S, dashed) MTF for R, G, B wavelengths plus a vertical line at the sensor Nyquist frequency (0.5 / pixel_size cycles/mm).

Algorithm per subplot

1. Compute the RGB PSF via self.psf_rgb() at the specified (depth, relative_fov) with kernel size psf_ks.
2. For each wavelength channel, call psf2mtf() to obtain the tangential and sagittal MTF curves.
3. Plot frequency vs MTF with RGB coloring (T solid, S dashed).

Parameters:

Name	Type	Description	Default
save_name	str	File path for the output PNG. Defaults to './lens_mtf.png'.	'./lens_mtf.png'
relative_fov_list	list[float]	Relative field positions in [0, 1], where 0 = on-axis and 1 = full field. Defaults to [0.0, 0.7, 1.0].	[0.0, 0.7, 1.0]
depth_list	list[float]	Object distances in mm. float('inf') is automatically replaced by self.obj_depth. When None (default), uses [self.obj_depth].	None
psf_ks	int	PSF kernel size in pixels (controls frequency resolution of the resulting MTF). Defaults to 128.	128
show	bool	If True, display interactively. Defaults to False.	False

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def draw_mtf(
    self,
    save_name="./lens_mtf.png",
    relative_fov_list=[0.0, 0.7, 1.0],
    depth_list=None,
    psf_ks=128,
    show=False,
):
    """Draw a grid of MTF curves for multiple depths and field positions.

    Produces a ``len(depth_list) × len(relative_fov_list)`` subplot grid.
    Each subplot shows the tangential (T, solid) and sagittal (S, dashed)
    MTF for R, G, B wavelengths plus a vertical line at the sensor Nyquist
    frequency (``0.5 / pixel_size`` cycles/mm).

    Algorithm per subplot:
        1. Compute the RGB PSF via ``self.psf_rgb()`` at the specified
           ``(depth, relative_fov)`` with kernel size ``psf_ks``.
        2. For each wavelength channel, call ``psf2mtf()`` to obtain the
           tangential and sagittal MTF curves.
        3. Plot frequency vs MTF with RGB coloring (T solid, S dashed).

    Args:
        save_name (str): File path for the output PNG.
            Defaults to ``'./lens_mtf.png'``.
        relative_fov_list (list[float]): Relative field positions in
            ``[0, 1]``, where 0 = on-axis and 1 = full field.
            Defaults to ``[0.0, 0.7, 1.0]``.
        depth_list (list[float]): Object distances in mm.
            ``float('inf')`` is automatically replaced by
            ``self.obj_depth``.  When ``None`` (default), uses
            ``[self.obj_depth]``.
        psf_ks (int): PSF kernel size in pixels (controls frequency
            resolution of the resulting MTF). Defaults to 128.
        show (bool): If ``True``, display interactively. Defaults to ``False``.
    """
    if depth_list is None:
        depth_list = [self.obj_depth]
    pixel_size = self.pixel_size
    nyquist_freq = 0.5 / pixel_size
    num_fovs = len(relative_fov_list)
    if float("inf") in depth_list:
        depth_list = [
            self.obj_depth if x == float("inf") else x for x in depth_list
        ]
    num_depths = len(depth_list)

    # Create figure and subplots (num_depths * num_fovs subplots)
    fig, axs = plt.subplots(
        num_depths, num_fovs, figsize=(num_fovs * 3, num_depths * 3), squeeze=False
    )

    # Iterate over depth and field of view
    for depth_idx, depth in enumerate(depth_list):
        for fov_idx, fov_relative in enumerate(relative_fov_list):
            # Calculate rgb PSF
            point = [0, -fov_relative, depth]
            psf_rgb = self.psf_rgb(points=point, ks=psf_ks, recenter=True)

            # Calculate MTF curves for rgb wavelengths
            for wvln_idx, wvln in enumerate(self.wvln_rgb):
                # Calculate tangential + sagittal MTF curves from PSF
                psf = psf_rgb[wvln_idx]
                freq, mtf_tan, mtf_sag = self.psf2mtf(psf, pixel_size)

                # Plot MTF curves (tangential solid, sagittal dashed)
                ax = axs[depth_idx, fov_idx]
                color = RGB_COLORS[wvln_idx % len(RGB_COLORS)]
                wvln_label = RGB_LABELS[wvln_idx % len(RGB_LABELS)]
                wvln_nm = int(wvln * 1000)
                ax.plot(
                    freq,
                    mtf_tan,
                    color=color,
                    linestyle="-",
                    label=f"{wvln_label}({wvln_nm}nm)-T",
                )
                ax.plot(
                    freq,
                    mtf_sag,
                    color=color,
                    linestyle="--",
                    label=f"{wvln_label}({wvln_nm}nm)-S",
                )

            # Draw Nyquist frequency
            ax.axvline(
                x=nyquist_freq,
                color="k",
                linestyle=":",
                linewidth=1.2,
                label="Nyquist",
            )

            # Set title and label for subplot
            fov_deg = round(fov_relative * self.rfov * 180 / np.pi, 1)
            depth_str = "inf" if depth == float("inf") else f"{depth}"
            ax.set_title(f"FOV: {fov_deg}deg, Depth: {depth_str}mm", fontsize=8)
            ax.set_xlabel("Spatial Frequency [cycles/mm]", fontsize=8)
            ax.set_ylabel("MTF", fontsize=8)
            ax.legend(fontsize=6)
            ax.tick_params(axis="both", which="major", labelsize=7)
            ax.grid(True)
            ax.set_ylim(0, 1.05)

    plt.tight_layout()
    if show:
        plt.show()
    else:
        assert save_name.endswith(".png"), "save_name must end with .png"
        plt.savefig(save_name, bbox_inches="tight", format="png", dpi=300)
    plt.close(fig)

vignetting

vignetting(depth=None, num_grid=32, num_rays=512)

Compute the relative-illumination (vignetting) map across the field.

Vignetting measures how much light is lost at each field position due to rays being clipped by lens apertures or barrel edges. It is computed as the fraction of traced rays that remain valid (not vignetted) at each grid cell, normalized by the total number of launched rays.

A value of 1.0 means all rays reach the sensor (no vignetting); 0.0 means complete light blockage. Real lenses typically show 1.0 on-axis and fall off toward the field edges due to mechanical vignetting and the cos⁴ illumination law.

Algorithm

1. self.sample_grid_rays() with uniform_fov=False (uniform image-space sampling) to ensure correct sensor-plane mapping.
2. self.trace2sensor() propagates rays and marks clipped ones as invalid.
3. Per-cell throughput = count(valid) / num_rays.

Parameters:

Name	Type	Description	Default
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
num_grid	int	Grid resolution per axis. Defaults to 32.	32
num_rays	int	Rays launched per grid cell. Higher values reduce Monte-Carlo noise. Defaults to 512.	512

Returns:

Name	Type	Description
vignetting	Tensor	Vignetting map with shape [num_grid, num_grid], values in [0, 1].

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def vignetting(self, depth=None, num_grid=32, num_rays=512):
    """Compute the relative-illumination (vignetting) map across the field.

    Vignetting measures how much light is lost at each field position due to
    rays being clipped by lens apertures or barrel edges.  It is computed as
    the fraction of traced rays that remain valid (not vignetted) at each
    grid cell, normalized by the total number of launched rays.

    A value of 1.0 means all rays reach the sensor (no vignetting); 0.0
    means complete light blockage.  Real lenses typically show 1.0 on-axis
    and fall off toward the field edges due to mechanical vignetting and the
    cos⁴ illumination law.

    Algorithm:
        1. ``self.sample_grid_rays()`` with ``uniform_fov=False`` (uniform
           image-space sampling) to ensure correct sensor-plane mapping.
        2. ``self.trace2sensor()`` propagates rays and marks clipped ones as
           invalid.
        3. Per-cell throughput = ``count(valid) / num_rays``.

    Args:
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        num_grid (int): Grid resolution per axis. Defaults to 32.
        num_rays (int): Rays launched per grid cell.  Higher values reduce
            Monte-Carlo noise. Defaults to 512.

    Returns:
        vignetting (torch.Tensor): Vignetting map with shape
            ``[num_grid, num_grid]``, values in ``[0, 1]``.
    """
    depth = self.obj_depth if depth is None else depth
    # Sample rays in uniform image space (not FOV angles) for correct sensor mapping
    # shape [num_grid, num_grid, num_rays, 3]
    ray = self.sample_grid_rays(
        depth=depth, num_grid=num_grid, num_rays=num_rays, uniform_fov=False
    )

    # Trace rays to sensor
    ray = self.trace2sensor(ray)

    # Calculate vignetting map
    vignetting = ray.is_valid.sum(-1) / (ray.is_valid.shape[-1])
    return vignetting

draw_vignetting

draw_vignetting(filename=None, depth=None, resolution=512, show=False)

Draw the vignetting map as a grayscale image with a colorbar.

Computes the vignetting map via self.vignetting(), bilinearly upsamples it to resolution × resolution, and displays it as a grayscale image where white = no vignetting and black = fully vignetted.

Parameters:

Name	Type	Description	Default
filename	str | None	File path for the output PNG. If None, auto-generates './vignetting_{depth}.png'.	None
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
resolution	int	Output image size in pixels (square). Defaults to 512.	512
show	bool	If True, display interactively. Defaults to False.	False

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def draw_vignetting(self, filename=None, depth=None, resolution=512, show=False):
    """Draw the vignetting map as a grayscale image with a colorbar.

    Computes the vignetting map via ``self.vignetting()``, bilinearly
    upsamples it to ``resolution × resolution``, and displays it as a
    grayscale image where white = no vignetting and black = fully vignetted.

    Args:
        filename (str | None): File path for the output PNG.  If ``None``,
            auto-generates ``'./vignetting_{depth}.png'``.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        resolution (int): Output image size in pixels (square).
            Defaults to 512.
        show (bool): If ``True``, display interactively. Defaults to ``False``.
    """
    depth = self.obj_depth if depth is None else depth
    # Calculate vignetting map
    vignetting = self.vignetting(depth=depth)

    # Interpolate vignetting map to desired resolution
    vignetting = F.interpolate(
        vignetting.unsqueeze(0).unsqueeze(0),
        size=(resolution, resolution),
        mode="bilinear",
        align_corners=False,
    ).squeeze()

    fig, ax = plt.subplots()
    ax.set_title("Relative Illumination (Vignetting)")
    im = ax.imshow(vignetting.cpu().numpy(), cmap="gray", vmin=0.0, vmax=1.0)
    fig.colorbar(im, ax=ax, ticks=[0.0, 0.25, 0.5, 0.75, 1.0])

    if show:
        plt.show()
    else:
        if filename is None:
            filename = f"./vignetting_{depth}.png"
        plt.savefig(filename, bbox_inches="tight", format="png", dpi=300)
    plt.close(fig)

analysis_rendering

analysis_rendering(
    img_org,
    save_name=None,
    depth=None,
    spp=SPP_RENDER,
    unwarp=False,
    method="ray_tracing",
    show=False,
)

Render a test image through the lens and report PSNR / SSIM.

Simulates what the sensor would capture if the given image were placed at the specified object distance. The rendering accounts for all geometric aberrations (blur, distortion, vignetting, chromatic effects). Optionally applies an inverse distortion warp (unwarp) and reports quality metrics for both the raw and unwarped renderings.

Algorithm

1. Convert img_org to a [1, 3, H, W] float tensor and temporarily set the sensor resolution to match.
2. Call self.render() with the chosen method (ray tracing or PSF-map / PSF-patch convolution).
3. Compute PSNR and SSIM between the original and rendered images.
4. If unwarp=True, apply self.unwarp() to correct geometric distortion and report metrics again.
5. Restore the original sensor resolution.

Parameters:

Name	Type	Description	Default
img_org	ndarray | Tensor	Source image with shape [H, W, 3], either uint8 [0, 255] or float [0, 1].	required
save_name	str | None	Path prefix for saved PNGs. If not None, saves '{save_name}.png' and (if unwarped) '{save_name}_unwarped.png'. Defaults to None.	None
depth	float	Object distance in mm. When None (default), falls back to self.obj_depth.	None
spp	int	Samples (rays) per pixel for rendering. Defaults to SPP_RENDER.	SPP_RENDER
unwarp	bool	If True, apply distortion correction after rendering. Defaults to False.	False
method	str	Rendering backend — 'ray_tracing', 'psf_map', or 'psf_patch'. Defaults to 'ray_tracing'.	'ray_tracing'
show	bool	If True, display the result with matplotlib. Defaults to False.	False

Returns:

Name	Type	Description
img_render	Tensor	Rendered (and optionally unwarped) image with shape [1, 3, H, W], float values in [0, 1].

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def analysis_rendering(
    self,
    img_org,
    save_name=None,
    depth=None,
    spp=SPP_RENDER,
    unwarp=False,
    method="ray_tracing",
    show=False,
):
    """Render a test image through the lens and report PSNR / SSIM.

    Simulates what the sensor would capture if the given image were placed
    at the specified object distance.  The rendering accounts for all
    geometric aberrations (blur, distortion, vignetting, chromatic effects).
    Optionally applies an inverse distortion warp (``unwarp``) and reports
    quality metrics for both the raw and unwarped renderings.

    Algorithm:
        1. Convert ``img_org`` to a ``[1, 3, H, W]`` float tensor and
           temporarily set the sensor resolution to match.
        2. Call ``self.render()`` with the chosen method (ray tracing or
           PSF-map / PSF-patch convolution).
        3. Compute PSNR and SSIM between the original and rendered images.
        4. If ``unwarp=True``, apply ``self.unwarp()`` to correct geometric
           distortion and report metrics again.
        5. Restore the original sensor resolution.

    Args:
        img_org (np.ndarray | torch.Tensor): Source image with shape
            ``[H, W, 3]``, either uint8 ``[0, 255]`` or float ``[0, 1]``.
        save_name (str | None): Path prefix for saved PNGs.  If not
            ``None``, saves ``'{save_name}.png'`` and (if unwarped)
            ``'{save_name}_unwarped.png'``. Defaults to ``None``.
        depth (float): Object distance in mm. When ``None`` (default),
            falls back to ``self.obj_depth``.
        spp (int): Samples (rays) per pixel for rendering.
            Defaults to ``SPP_RENDER``.
        unwarp (bool): If ``True``, apply distortion correction after
            rendering. Defaults to ``False``.
        method (str): Rendering backend — ``'ray_tracing'``, ``'psf_map'``,
            or ``'psf_patch'``. Defaults to ``'ray_tracing'``.
        show (bool): If ``True``, display the result with matplotlib.
            Defaults to ``False``.

    Returns:
        img_render (torch.Tensor): Rendered (and optionally unwarped) image
            with shape ``[1, 3, H, W]``, float values in ``[0, 1]``.
    """
    from skimage.metrics import peak_signal_noise_ratio, structural_similarity
    from torchvision.utils import save_image

    depth = self.obj_depth if depth is None else depth
    # Change sensor resolution to match the image
    sensor_res_original = self.sensor_res
    if isinstance(img_org, np.ndarray):
        img = (
            torch.from_numpy(img_org).permute(2, 0, 1).unsqueeze(0).float() / 255.0
        )
    elif torch.is_tensor(img_org):
        img = img_org.permute(2, 0, 1).unsqueeze(0).float()
        if img.max() > 1.0:
            img = img / 255.0
    img = img.to(self.device)
    self.set_sensor_res(sensor_res=img.shape[-2:])

    # Image rendering
    img_render = self.render(img, depth=depth, method=method, spp=spp)

    # Compute PSNR and SSIM
    img_np = img.squeeze(0).permute(1, 2, 0).cpu().numpy()
    render_np = (
        img_render.squeeze(0).permute(1, 2, 0).clamp(0, 1).cpu().detach().numpy()
    )
    render_psnr = round(
        peak_signal_noise_ratio(img_np, render_np, data_range=1.0), 3
    )
    render_ssim = round(
        structural_similarity(img_np, render_np, channel_axis=2, data_range=1.0), 4
    )
    print(f"Rendered image: PSNR={render_psnr:.3f}, SSIM={render_ssim:.4f}")

    # Save image
    if save_name is not None:
        save_image(img_render, f"{save_name}.png")

    # Unwarp to correct geometry distortion
    if unwarp:
        img_render = self.unwarp(img_render, depth)

        # Compute PSNR and SSIM
        render_np = (
            img_render.squeeze(0)
            .permute(1, 2, 0)
            .clamp(0, 1)
            .cpu()
            .detach()
            .numpy()
        )
        render_psnr = round(
            peak_signal_noise_ratio(img_np, render_np, data_range=1.0), 3
        )
        render_ssim = round(
            structural_similarity(
                img_np, render_np, channel_axis=2, data_range=1.0
            ),
            4,
        )
        print(
            f"Rendered image (unwarped): PSNR={render_psnr:.3f}, SSIM={render_ssim:.4f}"
        )

        if save_name is not None:
            save_image(img_render, f"{save_name}_unwarped.png")

    # Change the sensor resolution back
    self.set_sensor_res(sensor_res=sensor_res_original)

    # Show image
    if show:
        plt.imshow(img_render.cpu().squeeze(0).permute(1, 2, 0).numpy())
        plt.title("Rendered image")
        plt.axis("off")
        plt.show()
        plt.close()

    return img_render

analysis_spot

analysis_spot(num_field=3, depth=float('inf'))

Compute RMS and geometric spot radii at multiple field positions for RGB.

Traces rays at num_field evenly-spaced field positions along the meridional direction for three wavelengths (R, G, B), and computes polychromatic RMS and geometric spot radii referenced to the combined centroid across all wavelengths (matching Zemax's default "RMS Spot Radius w.r.t. Centroid").

This provides a quick polychromatic spot-size summary used for design comparisons and printed to stdout during analysis().

Algorithm (per field point): 1. Trace R, G, B rays through the lens to the sensor. 2. Pool all valid ray intercepts (across all three wavelengths) and compute one combined centroid c. 3. RMS = sqrt(mean(||xy - c||²)) over all pooled rays — a single polychromatic RMS that includes lateral chromatic aberration. 4. radius = max(||xy - c||) over all pooled rays. 5. Convert from mm to μm (× 1000).

Parameters:

Name	Type	Description	Default
num_field	int	Number of field positions sampled from on-axis to full-field. Defaults to 3.	3
depth	float	Object distance in mm. Use float('inf') for collimated light. Defaults to float('inf').	float('inf')

Returns:

Name	Type	Description
rms_results	dict[str, dict[str, float]]	Spot analysis results keyed by field position string (e.g., 'fov0.0', 'fov0.5', 'fov1.0'). Each value is a dict with: - 'rms': Polychromatic RMS spot radius in μm. - 'radius': Polychromatic geometric spot radius in μm.

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def analysis_spot(self, num_field=3, depth=float("inf")):
    """Compute RMS and geometric spot radii at multiple field positions for RGB.

    Traces rays at ``num_field`` evenly-spaced field positions along the
    meridional direction for three wavelengths (R, G, B), and computes
    polychromatic RMS and geometric spot radii referenced to the
    **combined centroid across all wavelengths** (matching Zemax's
    default "RMS Spot Radius w.r.t. Centroid").

    This provides a quick polychromatic spot-size summary used for design
    comparisons and printed to stdout during ``analysis()``.

    Algorithm (per field point):
        1. Trace R, G, B rays through the lens to the sensor.
        2. Pool all valid ray intercepts (across all three wavelengths)
           and compute one combined centroid ``c``.
        3. RMS = sqrt(mean(||xy - c||²)) over all pooled rays — a single
           polychromatic RMS that includes lateral chromatic aberration.
        4. radius = max(||xy - c||) over all pooled rays.
        5. Convert from mm to μm (× 1000).

    Args:
        num_field (int): Number of field positions sampled from on-axis
            to full-field. Defaults to 3.
        depth (float): Object distance in mm.  Use ``float('inf')`` for
            collimated light. Defaults to ``float('inf')``.

    Returns:
        rms_results (dict[str, dict[str, float]]): Spot analysis results
            keyed by field position string (e.g., ``'fov0.0'``,
            ``'fov0.5'``, ``'fov1.0'``). Each value is a dict with:
                - ``'rms'``: Polychromatic RMS spot radius in μm.
                - ``'radius'``: Polychromatic geometric spot radius in μm.
    """
    # Trace each wavelength and pool rays across wavelengths per field
    xy_list = []
    valid_list = []
    for wvln in self.wvln_rgb:
        ray = self.sample_radial_rays(
            num_field=num_field,
            depth=depth,
            num_rays=SPP_PSF,
            wvln=wvln,
            fov_max=self.rfov,
        )
        ray = self.trace2sensor(ray)
        xy_list.append(ray.o[..., :2])
        valid_list.append(ray.is_valid)

    # Pool over wavelengths, shape [num_field, 3*num_rays, 2] and [num_field, 3*num_rays]
    xy_all = torch.cat(xy_list, dim=-2)
    valid_all = torch.cat(valid_list, dim=-1)

    # Combined polychromatic centroid per field, shape [num_field, 1, 2]
    valid_mask = valid_all.unsqueeze(-1)
    center = (xy_all * valid_mask).sum(-2) / (
        valid_all.sum(-1, keepdim=True) + EPSILON
    )
    center = center.unsqueeze(-2)

    # Squared distance to combined centroid, shape [num_field, 3*num_rays]
    dist_sq = ((xy_all - center) ** 2).sum(-1)

    # Polychromatic RMS spot radius per field, shape [num_field]
    spot_rms = (
        (dist_sq * valid_all).sum(-1) / (valid_all.sum(-1) + EPSILON)
    ).sqrt()
    # Geometric spot radius (max distance among valid rays)
    dist_masked = torch.where(
        valid_all > 0, dist_sq, torch.full_like(dist_sq, -1.0)
    )
    spot_radius = dist_masked.max(dim=-1).values.clamp(min=0.0).sqrt()

    # Convert mm → μm
    avg_rms_radius_um = spot_rms * 1000.0
    avg_geo_radius_um = spot_radius * 1000.0

    # Print results
    print(f"Ray spot analysis results for depth {depth}:")
    print(
        f"RMS radius: FoV (0.0) {avg_rms_radius_um[0]:.3f} um, FoV (0.5) {avg_rms_radius_um[num_field // 2]:.3f} um, FoV (1.0) {avg_rms_radius_um[-1]:.3f} um"
    )
    print(
        f"Geo radius: FoV (0.0) {avg_geo_radius_um[0]:.3f} um, FoV (0.5) {avg_geo_radius_um[num_field // 2]:.3f} um, FoV (1.0) {avg_geo_radius_um[-1]:.3f} um"
    )

    # Save to dict
    rms_results = {}
    fov_ls = torch.linspace(0, 1, num_field)
    for i in range(num_field):
        fov = round(fov_ls[i].item(), 2)
        rms_results[f"fov{fov}"] = {
            "rms": round(avg_rms_radius_um[i].item(), 4),
            "radius": round(avg_geo_radius_um[i].item(), 4),
        }

    return rms_results

analysis

analysis(
    save_name="./lens",
    depth=float("inf"),
    full_eval=False,
    render=False,
    render_unwarp=False,
    lens_title=None,
    show=False,
)

Run a comprehensive optical analysis pipeline for the lens.

This is the main entry point for evaluating a lens design. It chains multiple evaluation steps in order, saving all plots with a common save_name prefix.

Execution flow

1. Always: draw the lens layout (draw_layout) and compute polychromatic spot RMS/radius (analysis_spot).
2. If full_eval=True: additionally generate:
3. Spot diagram (draw_spot_radial).
4. MTF grid (draw_mtf).
5. Distortion curve (draw_distortion_radial).
6. Vignetting map (draw_vignetting).
7. If render=True: render a test chart image through the lens and report PSNR/SSIM (analysis_rendering).

Parameters:

Name	Type	Description	Default
save_name	str	Path prefix for all output files. Each plot appends a suffix (e.g., '_spot.png', '_mtf.png'). Defaults to './lens'.	'./lens'
depth	float	Object distance in mm. float('inf') is replaced by self.obj_depth for rendering and vignetting. Defaults to float('inf').	float('inf')
full_eval	bool	If True, run all evaluation plots. If False, only layout + spot RMS. Defaults to False.	False
render	bool	If True, render a test image through the lens. Defaults to False.	False
render_unwarp	bool	If True (and render=True), also produce an unwarped rendering. Defaults to False.	False
lens_title	str | None	Title string for the layout plot. Defaults to None.	None
show	bool	If True, display all plots interactively. Defaults to False.	False

Source code in deeplens-src/deeplens/geolens_pkg/eval.py

@torch.no_grad()
def analysis(
    self,
    save_name="./lens",
    depth=float("inf"),
    full_eval=False,
    render=False,
    render_unwarp=False,
    lens_title=None,
    show=False,
):
    """Run a comprehensive optical analysis pipeline for the lens.

    This is the main entry point for evaluating a lens design.  It chains
    multiple evaluation steps in order, saving all plots with a common
    ``save_name`` prefix.

    Execution flow:
        1. **Always**: draw the lens layout (``draw_layout``) and compute
           polychromatic spot RMS/radius (``analysis_spot``).
        2. **If** ``full_eval=True``: additionally generate:
           - Spot diagram (``draw_spot_radial``).
           - MTF grid (``draw_mtf``).
           - Distortion curve (``draw_distortion_radial``).
           - Vignetting map (``draw_vignetting``).
        3. **If** ``render=True``: render a test chart image through the
           lens and report PSNR/SSIM (``analysis_rendering``).

    Args:
        save_name (str): Path prefix for all output files.  Each plot
            appends a suffix (e.g., ``'_spot.png'``, ``'_mtf.png'``).
            Defaults to ``'./lens'``.
        depth (float): Object distance in mm.  ``float('inf')`` is replaced
            by ``self.obj_depth`` for rendering and vignetting.
            Defaults to ``float('inf')``.
        full_eval (bool): If ``True``, run all evaluation plots.  If
            ``False``, only layout + spot RMS. Defaults to ``False``.
        render (bool): If ``True``, render a test image through the lens.
            Defaults to ``False``.
        render_unwarp (bool): If ``True`` (and ``render=True``), also
            produce an unwarped rendering. Defaults to ``False``.
        lens_title (str | None): Title string for the layout plot.
            Defaults to ``None``.
        show (bool): If ``True``, display all plots interactively.
            Defaults to ``False``.
    """
    # Draw lens layout and ray path
    self.draw_layout(
        filename=f"{save_name}.png",
        lens_title=lens_title,
        depth=depth,
        show=show,
    )

    # Calculate RMS error
    self.analysis_spot(depth=depth)

    # Comprehensive optical evaluation
    if full_eval:
        # Draw spot diagram
        self.draw_spot_radial(
            save_name=f"{save_name}_spot.png",
            depth=depth,
            show=show,
        )

        # Draw MTF
        if depth == float("inf"):
            self.draw_mtf(
                depth_list=[self.obj_depth],
                save_name=f"{save_name}_mtf.png",
                show=show,
            )
        else:
            self.draw_mtf(
                depth_list=[depth],
                save_name=f"{save_name}_mtf.png",
                show=show,
            )

        # Draw distortion
        self.draw_distortion_radial(
            save_name=f"{save_name}_distortion.png",
            show=show,
        )

        # Draw vignetting
        eval_depth = self.obj_depth if depth == float("inf") else depth
        self.draw_vignetting(
            filename=f"{save_name}_vignetting.png",
            depth=eval_depth,
            show=show,
        )

    # Render an image, compute PSNR and SSIM
    if render:
        depth = self.obj_depth if depth == float("inf") else depth
        img_org = Image.open("./datasets/charts/NBS_1963_1k.png").convert("RGB")
        img_org = np.array(img_org)
        self.analysis_rendering(
            img_org,
            depth=depth,
            spp=SPP_RENDER,
            unwarp=render_unwarp,
            save_name=f"{save_name}_render",
            show=show,
        )

deeplens.geolens_pkg.optim.GeoLensOptim

Mixin providing differentiable optimisation for GeoLens.

Implements gradient-based lens design using PyTorch autograd:

• Loss functions – RMS spot error, focus, surface regularity, gap constraints, material validity.
• Constraint initialisation – edge-thickness and self-intersection guards.
• Optimizer helpers – parameter groups with per-type learning rates and cosine annealing schedules.
• High-level optimize() – curriculum-learning training loop.

This class is not instantiated directly; it is mixed into GeoLens.

References

Xinge Yang et al., "Curriculum learning for ab initio deep learned refractive optics," Nature Communications 2024.

init_constraints

init_constraints(constraint_params=None)

Initialize geometry, ray-angle, and distortion constraints for the lens.

Selects a cellphone or camera constraint preset based on whether the sensor radius is below 12 mm, sets the air-gap, thickness, BFL, TTL, surface-shape, CRA, bend-angle, and distortion limits, and propagates the bend-angle limit onto every surface.

Parameters:

Name	Type	Description	Default
constraint_params	dict	Constraint parameters. Currently unused (reserved for future overrides). Defaults to None.	None

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def init_constraints(self, constraint_params=None):
    """Initialize geometry, ray-angle, and distortion constraints for the lens.

    Selects a cellphone or camera constraint preset based on whether the
    sensor radius is below 12 mm, sets the air-gap, thickness, BFL, TTL,
    surface-shape, CRA, bend-angle, and distortion limits, and propagates
    the bend-angle limit onto every surface.

    Args:
        constraint_params (dict, optional): Constraint parameters. Currently
            unused (reserved for future overrides). Defaults to None.
    """
    # In the future, we want to use constraint_params to set the constraints.
    if constraint_params is None:
        constraint_params = {}

    if self.r_sensor < 12.0:
        self.is_cellphone = True

        self.air_edge_min = 0.05
        self.air_edge_max = 5.0
        self.air_center_min = 0.05
        self.air_center_max = 5.0

        self.thick_edge_min = 0.25
        self.thick_edge_max = 5.0
        self.thick_center_min = 0.25
        self.thick_center_max = 5.0

        self.bfl_min = 0.8
        self.bfl_max = 5.0

        self.ttl_min = 0.0
        self.ttl_max = 50.0

        # Surface shape constraints
        self.sag2diam_max = 0.5
        self.diam2thick_max = 15.0
        self.tmax2tmin_max = 5.0
        self.surf_angle_max = 45.0  # deg

        # Ray angle constraints
        self.chief_ray_angle_max = 45.0  # deg
        self.bend_angle_max = 30.0  # deg

        # Distortion constraint
        self.distortion_max = 0.10  # 10 % relative distortion

    else:
        self.is_cellphone = False

        self.air_edge_min = 0.1
        self.air_edge_max = 100.0  # float("inf")
        self.air_center_min = 0.1
        self.air_center_max = 100.0  # float("inf")

        self.thick_edge_min = 1.0
        self.thick_edge_max = 20.0
        self.thick_center_min = 2.0
        self.thick_center_max = 20.0

        self.bfl_min = 5.0
        self.bfl_max = 100.0  # float("inf")

        self.ttl_min = 0.0  # disabled by default
        self.ttl_max = 300.0  # float("inf")

        # Surface shape constraints
        self.sag2diam_max = 0.5
        self.diam2thick_max = 20.0
        self.tmax2tmin_max = 10.0
        self.surf_angle_max = 45.0  # deg

        # Ray angle constraints
        self.chief_ray_angle_max = 45.0  # deg
        self.bend_angle_max = 30.0  # deg

        # Distortion constraint
        self.distortion_max = 0.02  # 2 % relative distortion

    # Propagate bend angle limit onto every surface so refract() reads it.
    for s in self.surfaces:
        s.bend_angle_max = self.bend_angle_max

loss_reg

loss_reg(
    w_focus=1.0,
    w_cra=1.0,
    w_ray_bend=1.0,
    w_clearance=1.0,
    w_envelope=1.0,
    w_profile=1.0,
)

Compute combined regularization loss for lens design.

Aggregates multiple constraint losses to keep the lens physically valid during gradient-based optimisation.

Parameters:

Name	Type	Description	Default
w_focus	float	Weight for focus loss. Defaults to 1.0.	1.0
w_cra	float	Weight for chief ray angle loss. Defaults to 1.0.	1.0
w_ray_bend	float	Weight for per-surface bend penalty. Defaults to 1.0.	1.0
w_clearance	float	Weight for the clearance penalty (min air gap, min thickness, min BFL, min TTL). Defaults to 1.0.	1.0
w_envelope	float	Weight for the envelope penalty (max air gap, max thickness, max BFL, max TTL). Defaults to 1.0.	1.0
w_profile	float	Weight for per-surface profile feasibility (sag, slope). Defaults to 1.0.	1.0

Returns:

Name	Type	Description
loss_reg	Tensor	Scalar combined regularization loss.
loss_dict	dict	Per-component loss values for logging.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_reg(
    self,
    w_focus=1.0,
    w_cra=1.0,
    w_ray_bend=1.0,
    w_clearance=1.0,
    w_envelope=1.0,
    w_profile=1.0,
):
    """Compute combined regularization loss for lens design.

    Aggregates multiple constraint losses to keep the lens physically valid
    during gradient-based optimisation.

    Args:
        w_focus (float, optional): Weight for focus loss. Defaults to 1.0.
        w_cra (float, optional): Weight for chief ray angle loss. Defaults to 1.0.
        w_ray_bend (float, optional): Weight for per-surface bend penalty. Defaults to 1.0.
        w_clearance (float, optional): Weight for the clearance penalty
            (min air gap, min thickness, min BFL, min TTL). Defaults to 1.0.
        w_envelope (float, optional): Weight for the envelope penalty
            (max air gap, max thickness, max BFL, max TTL). Defaults to 1.0.
        w_profile (float, optional): Weight for per-surface profile
            feasibility (sag, slope). Defaults to 1.0.

    Returns:
        loss_reg (torch.Tensor): Scalar combined regularization loss.
        loss_dict (dict): Per-component loss values for logging.
    """
    # Loss functions for regularization
    # loss_focus = self.loss_infocus()
    loss_cra = self.loss_cra()
    loss_ray_bend = self.loss_ray_bend()
    loss_clearance, loss_envelope = self.loss_bound()
    loss_profile = self.loss_profile()
    # loss_mat = self.loss_mat()
    loss_reg = (
        # w_focus * loss_focus
        +w_clearance * loss_clearance
        + w_envelope * loss_envelope
        + w_profile * loss_profile
        + w_cra * loss_cra
        + w_ray_bend * loss_ray_bend
        # w_mat * loss_mat
    )

    # Return loss and loss dictionary
    loss_dict = {
        # "loss_focus": loss_focus.item(),
        "loss_clearance": loss_clearance.item(),
        "loss_envelope": loss_envelope.item(),
        "loss_profile": loss_profile.item(),
        "loss_cra": loss_cra.item(),
        "loss_ray_bend": loss_ray_bend.item(),
        # 'loss_mat': loss_mat.item(),
    }
    return loss_reg, loss_dict

loss_infocus

loss_infocus(target=0.005, wvln=None)

Sample on-axis parallel rays and penalize the sensor-plane spot RMS.

Traces a zero-field ray bundle to the sensor and applies a one-sided penalty \(\mathrm{relu}(\text{rms} - \text{target})\) that activates only when the RMS spot radius exceeds the target.

Parameters:

Name	Type	Description	Default
target	float	Target on-axis RMS spot radius in mm. Defaults to 0.005.	0.005
wvln	float	Wavelength in µm. When None (default), falls back to the green channel of self.wvln_rgb. Defaults to None.	None

Returns:

Name	Type	Description
loss	Tensor	Scalar focus penalty (at least 0).

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_infocus(self, target=0.005, wvln=None):
    """Sample on-axis parallel rays and penalize the sensor-plane spot RMS.

    Traces a zero-field ray bundle to the sensor and applies a one-sided
    penalty $\\mathrm{relu}(\\text{rms} - \\text{target})$ that activates
    only when the RMS spot radius exceeds the target.

    Args:
        target (float, optional): Target on-axis RMS spot radius in mm.
            Defaults to 0.005.
        wvln (float, optional): Wavelength in µm. When None (default),
            falls back to the green channel of `self.wvln_rgb`. Defaults to None.

    Returns:
        loss (torch.Tensor): Scalar focus penalty (at least 0).
    """
    if wvln is None:
        wvln = self.wvln_rgb[1]
    loss = torch.tensor(0.0, device=self.device)

    # Ray tracing and calculate RMS error
    ray = self.sample_from_fov(fov_x=0.0, fov_y=0.0, wvln=wvln, num_rays=SPP_CALC)
    ray = self.trace2sensor(ray)
    rms_error = ray.rms_error()

    # Smooth penalty: activates when rms_error exceeds target
    loss += relu(rms_error - target)

    return loss

loss_profile

loss_profile()

Penalize infeasible per-surface profile shapes.

The "profile" is the z(r) curve of a single surface. This loss makes sure each surface is physically manufacturable by checking: 1. Sag-to-diameter ratio exceeding sag2diam_max. 2. Maximum surface slope angle exceeding surf_angle_max (deg).

Returns:

Name	Type	Description
loss	Tensor	Scalar profile feasibility penalty.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_profile(self):
    """Penalize infeasible per-surface profile shapes.

    The "profile" is the z(r) curve of a single surface. This loss makes
    sure each surface is physically manufacturable by checking:
        1. Sag-to-diameter ratio exceeding ``sag2diam_max``.
        2. Maximum surface slope angle exceeding ``surf_angle_max`` (deg).

    Returns:
        loss (torch.Tensor): Scalar profile feasibility penalty.
    """
    sag2diam_max = self.sag2diam_max
    grad_max = math.tan(math.radians(self.surf_angle_max))

    loss_grad = torch.tensor(0.0, device=self.device)
    loss_sag2diam = torch.tensor(0.0, device=self.device)
    for i in self.find_diff_surf():
        # Sample points on the surface
        x_ls = torch.linspace(0.0, 1.0, 32, device=self.device) * self.surfaces[i].r
        y_ls = torch.zeros_like(x_ls)

        # Sag
        sag_ls = self.surfaces[i].sag(x_ls, y_ls)
        sag2diam = sag_ls.abs().max() / self.surfaces[i].r / 2
        loss_sag2diam += relu((sag2diam - sag2diam_max) / sag2diam_max)

        # 1st-order derivative
        grad_ls = self.surfaces[i].dfdxyz(x_ls, y_ls)[0]
        grad = grad_ls.abs().max()
        loss_grad += relu((grad - grad_max) / grad_max)

        # # Diameter to thickness ratio, thick_max to thick_min ratio
        # if not self.surfaces[i].mat2.name == "air":
        #     surf2 = self.surfaces[i + 1]
        #     surf1 = self.surfaces[i]

        #     # Penalize diameter to thickness ratio
        #     diam2thick = 2 * max(surf2.r, surf1.r) / surf1.d_next
        #     loss_diam2thick += torch.nn.functional.relu(diam2thick - diam2thick_max)

        #     # Penalize thick_max to thick_min ratio.
        #     # Use torch.maximum/minimum for differentiable max/min.
        #     r_edge = min(surf2.r, surf1.r)
        #     thick_center = surf1.d_next
        #     thick_edge = surf2.surface_with_offset(
        #         r_edge, 0.0, d=self.surf_d(i + 1)
        #     ) - surf1.surface_with_offset(r_edge, 0.0, d=self.surf_d(i))
        #     thick_max = torch.maximum(thick_center, thick_edge)
        #     thick_min = torch.minimum(thick_center, thick_edge).clamp(min=0.01)
        #     tmax2tmin = thick_max / thick_min

        #     loss_tmax2tmin += torch.nn.functional.relu(tmax2tmin - tmax2tmin_max)

    return loss_sag2diam + loss_grad

loss_bound

loss_bound()

Penalize geometry-bound violations in a single surface-sampling pass.

Each surface pair is sampled once and its distances feed both the clearance (min) and envelope (max) relu penalties for air gaps, glass thickness, BFL, and TTL.

Returns:

Name	Type	Description
loss_clearance	Tensor	Scalar clearance penalty for parts that are too close / too thin.
loss_envelope	Tensor	Scalar envelope penalty for the overall assembly growing beyond its spatial budget. Returned separately so callers can weight the two independently.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_bound(self):
    """Penalize geometry-bound violations in a single surface-sampling pass.

    Each surface pair is sampled once and its distances feed both the
    clearance (min) and envelope (max) relu penalties for air gaps,
    glass thickness, BFL, and TTL.

    Returns:
        loss_clearance (torch.Tensor): Scalar clearance penalty for parts
            that are too close / too thin.
        loss_envelope (torch.Tensor): Scalar envelope penalty for the
            overall assembly growing beyond its spatial budget.  Returned
            separately so callers can weight the two independently.
    """
    # Min bounds (clearance)
    air_center_min = self.air_center_min
    air_edge_min = self.air_edge_min
    thick_center_min = self.thick_center_min
    thick_edge_min = self.thick_edge_min
    bfl_min = self.bfl_min
    ttl_min = self.ttl_min

    # Max bounds (envelope)
    air_center_max = self.air_center_max
    air_edge_max = self.air_edge_max
    thick_center_max = self.thick_center_max
    thick_edge_max = self.thick_edge_max
    bfl_max = self.bfl_max
    ttl_max = self.ttl_max

    loss_clearance = torch.tensor(0.0, device=self.device)
    loss_envelope = torch.tensor(0.0, device=self.device)
    air_c_range = air_center_max - air_center_min
    air_e_range = air_edge_max - air_edge_min
    thick_c_range = thick_center_max - thick_center_min
    thick_e_range = thick_edge_max - thick_edge_min
    bfl_range = bfl_max - bfl_min
    ttl_range = ttl_max - ttl_min

    for i in range(len(self.surfaces) - 1):
        current_surf = self.surfaces[i]
        next_surf = self.surfaces[i + 1]

        # Sample surfaces once and reuse for both clearance and envelope
        r_center = torch.tensor(0.0, device=self.device) * current_surf.r
        z_prev_center = current_surf.surface_with_offset(
            r_center, 0.0, valid_check=False, d=self.surf_d(i)
        )
        z_next_center = next_surf.surface_with_offset(
            r_center, 0.0, valid_check=False, d=self.surf_d(i + 1)
        )

        r_edge = torch.linspace(0.5, 1.0, 16, device=self.device) * current_surf.r
        z_prev_edge = current_surf.surface_with_offset(
            r_edge, 0.0, valid_check=False, d=self.surf_d(i)
        )
        z_next_edge = next_surf.surface_with_offset(
            r_edge, 0.0, valid_check=False, d=self.surf_d(i + 1)
        )

        dist_center = z_next_center - z_prev_center
        dist_edges = z_next_edge - z_prev_edge
        dist_edge_lo = torch.min(dist_edges)
        dist_edge_hi = torch.max(dist_edges)

        if current_surf.mat2.name == "air":
            loss_clearance += relu((air_center_min - dist_center) / air_c_range)
            loss_clearance += relu((air_edge_min - dist_edge_lo) / air_e_range)
            loss_envelope += relu((dist_center - air_center_max) / air_c_range)
            loss_envelope += relu((dist_edge_hi - air_edge_max) / air_e_range)
        else:
            loss_clearance += relu((thick_center_min - dist_center) / thick_c_range)
            loss_clearance += relu((thick_edge_min - dist_edge_lo) / thick_e_range)
            loss_envelope += relu((dist_center - thick_center_max) / thick_c_range)
            loss_envelope += relu((dist_edge_hi - thick_edge_max) / thick_e_range)

    # Back focal length
    last_surf = self.surfaces[-1]
    r = torch.linspace(0.0, 1.0, 32, device=self.device) * last_surf.r
    z_last_surf = self.d_sensor - last_surf.surface_with_offset(
        r, 0.0, d=self.surf_d(-1)
    )
    bfl_lo = torch.min(z_last_surf)
    bfl_hi = torch.max(z_last_surf)
    loss_clearance += relu((bfl_min - bfl_lo) / bfl_range)
    loss_envelope += relu((bfl_hi - bfl_max) / bfl_range)

    # Total track length
    ttl = self.d_sensor
    loss_clearance += relu((ttl_min - ttl) / ttl_range)
    loss_envelope += relu((ttl - ttl_max) / ttl_range)

    return loss_clearance, loss_envelope

loss_cra

loss_cra()

Penalize chief ray angle at sensor exceeding chief_ray_angle_max.

Uses a near-paraxial pupil sample (scale_pupil=0.2) over the full FoV. The penalty is \(\mathrm{relu}(\cos\theta_\text{ref} - \cos\theta_\text{CRA})\), valid-ray averaged, where $\cos\theta = $ ray.d[..., 2].

Returns:

Name	Type	Description
loss	Tensor	Scalar CRA penalty (at least 0).

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_cra(self):
    """Penalize chief ray angle at sensor exceeding chief_ray_angle_max.

    Uses a near-paraxial pupil sample (scale_pupil=0.2) over the full FoV.
    The penalty is $\\mathrm{relu}(\\cos\\theta_\\text{ref} - \\cos\\theta_\\text{CRA})$,
    valid-ray averaged, where $\\cos\\theta = $ `ray.d[..., 2]`.

    Returns:
        loss (torch.Tensor): Scalar CRA penalty (at least 0).
    """
    cos_cra_ref = float(np.cos(np.deg2rad(self.chief_ray_angle_max)))

    ray = self.sample_ring_arm_rays(
        num_ring=8, num_arm=2, spp=SPP_CALC, scale_pupil=0.2
    )
    ray = self.trace2sensor(ray)
    cos_cra = ray.d[..., 2]
    valid = ray.is_valid > 0
    penalty_cra = relu(cos_cra_ref - cos_cra)
    return (penalty_cra * valid).sum() / (valid.sum() + EPSILON)

loss_ray_bend

loss_ray_bend()

Penalize accumulated per-surface bend angles exceeding bend_angle_max.

Reads ray.bend_penalty, an additive sum of per-surface relu contributions collected during trace2sensor. Each surface contributes independently, so large bends at one surface are not hidden by small bends at another. Uses a full-pupil sample (scale_pupil=1.0).

Returns:

Name	Type	Description
loss	Tensor	Scalar bend penalty (at least 0).

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_ray_bend(self):
    """Penalize accumulated per-surface bend angles exceeding bend_angle_max.

    Reads ``ray.bend_penalty``, an additive sum of per-surface relu
    contributions collected during ``trace2sensor``.  Each surface
    contributes independently, so large bends at one surface are not hidden
    by small bends at another.  Uses a full-pupil sample (scale_pupil=1.0).

    Returns:
        loss (torch.Tensor): Scalar bend penalty (at least 0).
    """
    ray = self.sample_ring_arm_rays(
        num_ring=8, num_arm=2, spp=SPP_CALC, scale_pupil=1.0
    )
    ray = self.trace2sensor(ray)
    bend_penalty = ray.bend_penalty.squeeze(-1)
    valid = ray.is_valid > 0
    return (bend_penalty * valid).sum() / (valid.sum() + EPSILON)

loss_mat

loss_mat()

Penalize material parameters outside manufacturable ranges.

Constrains refractive index n to [1.5, 1.9] and Abbe number V to [30, 70] for each non-air surface material.

Returns:

Name	Type	Description
loss_mat	Tensor	Scalar material penalty loss.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_mat(self):
    """Penalize material parameters outside manufacturable ranges.

    Constrains refractive index *n* to [1.5, 1.9] and Abbe number *V* to
    [30, 70] for each non-air surface material.

    Returns:
        loss_mat (torch.Tensor): Scalar material penalty loss.
    """
    n_max = 1.9
    n_min = 1.5
    V_max = 70
    V_min = 30
    loss_mat = torch.tensor(0.0, device=self.device)
    for i in range(len(self.surfaces)):
        if self.surfaces[i].mat2.name != "air":
            if self.surfaces[i].mat2.n > n_max:
                loss_mat += (self.surfaces[i].mat2.n - n_max) / (n_max - n_min)
            if self.surfaces[i].mat2.n < n_min:
                loss_mat += (n_min - self.surfaces[i].mat2.n) / (n_max - n_min)
            if self.surfaces[i].mat2.V > V_max:
                loss_mat += (self.surfaces[i].mat2.V - V_max) / (V_max - V_min)
            if self.surfaces[i].mat2.V < V_min:
                loss_mat += (V_min - self.surfaces[i].mat2.V) / (V_max - V_min)

    return loss_mat

loss_rms

loss_rms(num_grid=GEO_GRID, depth=None, num_rays=SPP_PSF, sample_more_off_axis=False)

Compute the RGB spot-size RMS loss over a grid of field points.

Traces R, G, B ray bundles (green first) to the sensor and measures the spot radius against the green pinhole center. The green spot error sets a detached per-field weight mask that emphasises harder fields.

Parameters:

Name	Type	Description	Default
num_grid	int	Number of field-grid points per axis. Defaults to GEO_GRID.	GEO_GRID
depth	float	Object-plane depth in mm. When None (default), falls back to self.obj_depth. Defaults to None.	None
num_rays	int	Number of rays per field point. Defaults to SPP_PSF.	SPP_PSF
sample_more_off_axis	bool	If True, concentrate field samples toward the field edge. Defaults to False.	False

Returns:

Name	Type	Description
avg_rms_error	Tensor	Scalar RMS spot error in mm, averaged over the R, G, B wavelengths.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def loss_rms(
    self,
    num_grid=GEO_GRID,
    depth=None,
    num_rays=SPP_PSF,
    sample_more_off_axis=False,
):
    """Compute the RGB spot-size RMS loss over a grid of field points.

    Traces R, G, B ray bundles (green first) to the sensor and measures the
    spot radius against the green pinhole center. The green spot error sets
    a detached per-field weight mask that emphasises harder fields.

    Args:
        num_grid (int, optional): Number of field-grid points per axis.
            Defaults to GEO_GRID.
        depth (float, optional): Object-plane depth in mm. When None
            (default), falls back to `self.obj_depth`. Defaults to None.
        num_rays (int, optional): Number of rays per field point.
            Defaults to SPP_PSF.
        sample_more_off_axis (bool, optional): If True, concentrate field
            samples toward the field edge. Defaults to False.

    Returns:
        avg_rms_error (torch.Tensor): Scalar RMS spot error in mm, averaged
            over the R, G, B wavelengths.
    """
    depth = self.obj_depth if depth is None else depth
    # Iterate green first so the error-adaptive weight mask is anchored
    # on the reference (green) wavelength.
    loss_rms_ls = []
    w_mask = None
    for i, wvln in enumerate(
        [self.wvln_rgb[1], self.wvln_rgb[0], self.wvln_rgb[2]]
    ):
        ray = self.sample_grid_rays(
            depth=depth,
            num_grid=num_grid,
            num_rays=num_rays,
            wvln=wvln,
            sample_more_off_axis=sample_more_off_axis,
        )

        # Reference center from green chief-ray (pinhole), broadcast to rays.
        if i == 0:
            with torch.no_grad():
                center_ref = -self.psf_center(
                    points_obj=ray.o[:, :, 0, :], method="pinhole"
                )
            center_ref = center_ref.unsqueeze(-2)

        ray = self.trace2sensor(ray)

        # Per-FOV MSE → RMS, zeroing invalid rays before squaring to
        # avoid Inf*0 = NaN.
        ray_xy = ray.o[..., :2]
        ray_valid = ray.is_valid
        ray_err = ray_xy - center_ref
        ray_err = torch.where(
            ray_valid.bool().unsqueeze(-1), ray_err, torch.zeros_like(ray_err)
        )
        mse = (ray_err**2).sum(-1).sum(-1) / (ray_valid.sum(-1) + EPSILON)
        l_rms = (mse + EPSILON).sqrt()

        # First wavelength (green) defines the detached weight mask.
        if w_mask is None:
            w_mask = mse.detach()
            w_mask = w_mask / (w_mask.mean() + EPSILON)

        l_rms_weighted = (l_rms * w_mask).sum() / (w_mask.sum() + EPSILON)
        loss_rms_ls.append(l_rms_weighted)

    avg_rms_error = torch.stack(loss_rms_ls).mean(dim=0)
    return avg_rms_error

sample_ring_arm_rays

sample_ring_arm_rays(
    num_ring=8,
    num_arm=2,
    spp=2048,
    depth=None,
    wvln=None,
    scale_pupil=1.0,
    sample_more_off_axis=True,
)

Sample rays from object space using a ring-arm pattern.

This method distributes sampling points (origins of ray bundles) on a polar grid in the object plane, defined by field of view. This is useful for capturing lens performance across the full field. The points include the center and num_ring rings with num_arm points on each.

Uses self.rfov (ray-traced real FoV, accounts for distortion) rather than self.rfov_eff (paraxial pinhole FoV) so the full distorted field is covered.

Parameters:

Name	Type	Description	Default
num_ring	int	Number of rings to sample in the field of view. Defaults to 8.	8
num_arm	int	Number of arms (spokes) sampled per ring. Defaults to 2.	2
spp	int	Number of rays sampled per field point. Defaults to 2048.	2048
depth	float	Depth of the object plane in mm. When None (default), falls back to self.obj_depth. Defaults to None.	None
wvln	float	Wavelength in µm. When None (default), falls back to self.primary_wvln. Defaults to None.	None
scale_pupil	float	Scale factor for the entrance pupil radius. Defaults to 1.0.	1.0
sample_more_off_axis	bool	If True, warp ring field angles by a square-root profile to concentrate samples toward the field edge. Defaults to True.	True

Returns:

Name	Type	Description
rays	Ray	Ray bundle with field points laid out as [num_ring, num_arm] and spp rays each.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def sample_ring_arm_rays(
    self,
    num_ring=8,
    num_arm=2,
    spp=2048,
    depth=None,
    wvln=None,
    scale_pupil=1.0,
    sample_more_off_axis=True,
):
    """Sample rays from object space using a ring-arm pattern.

    This method distributes sampling points (origins of ray bundles) on a polar grid in the object plane,
    defined by field of view. This is useful for capturing lens performance across the full field.
    The points include the center and `num_ring` rings with `num_arm` points on each.

    Uses ``self.rfov`` (ray-traced real FoV, accounts for distortion) rather than
    ``self.rfov_eff`` (paraxial pinhole FoV) so the full distorted field is covered.

    Args:
        num_ring (int, optional): Number of rings to sample in the field
            of view. Defaults to 8.
        num_arm (int, optional): Number of arms (spokes) sampled per ring.
            Defaults to 2.
        spp (int, optional): Number of rays sampled per field point.
            Defaults to 2048.
        depth (float, optional): Depth of the object plane in mm. When None
            (default), falls back to `self.obj_depth`. Defaults to None.
        wvln (float, optional): Wavelength in µm. When None (default), falls
            back to `self.primary_wvln`. Defaults to None.
        scale_pupil (float, optional): Scale factor for the entrance pupil
            radius. Defaults to 1.0.
        sample_more_off_axis (bool, optional): If True, warp ring field
            angles by a square-root profile to concentrate samples toward
            the field edge. Defaults to True.

    Returns:
        rays (Ray): Ray bundle with field points laid out as
            [num_ring, num_arm] and `spp` rays each.
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth
    # Create points on rings and arms
    max_fov_rad = self.rfov
    if sample_more_off_axis:
        beta_values = torch.linspace(0.0, 1.0, num_ring, device=self.device)
        beta_transformed = beta_values**0.5
        ring_fovs = max_fov_rad * beta_transformed
    else:
        ring_fovs = max_fov_rad * torch.linspace(
            0.0, 1.0, num_ring, device=self.device
        )

    arm_angles = torch.linspace(0.0, 2 * torch.pi, num_arm + 1, device=self.device)[
        :-1
    ]
    ring_grid, arm_grid = torch.meshgrid(ring_fovs, arm_angles, indexing="ij")
    x = depth * torch.tan(ring_grid) * torch.cos(arm_grid)
    y = depth * torch.tan(ring_grid) * torch.sin(arm_grid)
    z = torch.full_like(x, depth)
    points = torch.stack([x, y, z], dim=-1)  # shape: [num_ring, num_arm, 3]

    # Sample rays
    rays = self.sample_from_points(
        points=points, num_rays=spp, wvln=wvln, scale_pupil=scale_pupil
    )
    return rays

optimize

optimize(
    lrs=[0.001, 0.0001, 0.1, 0.0001],
    iterations=5000,
    test_per_iter=100,
    optim_mat=False,
    shape_control=True,
    sample_more_off_axis=False,
    result_dir=None,
)

Optimise the lens by minimising RGB RMS spot errors.

Runs a curriculum-learning training loop with Adam optimiser and cosine annealing. Periodically evaluates the lens, saves intermediate results, and optionally corrects surface shapes.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates for [d_next, c, k, a] parameter groups. Defaults to [1e-3, 1e-4, 1e-1, 1e-4].	[0.001, 0.0001, 0.1, 0.0001]
iterations	int	Total training iterations. Defaults to 5000.	5000
test_per_iter	int	Evaluate and save every N iterations. Defaults to 100.	100
optim_mat	bool	If True, include material parameters (n, V) in optimisation. Defaults to False.	False
shape_control	bool	If True, call correct_shape() at each evaluation step. Defaults to True.	True
sample_more_off_axis	bool	If True, concentrate ray samples toward the edge of the field to improve off-axis correction. Passed directly to sample_ring_arm_rays. Defaults to False.	False
result_dir	str	Directory to save results. If None, auto-generates a timestamped directory. Defaults to None.	None

Note

Debug hints: 1. Slowly optimise with small learning rate. 2. FoV and thickness should match well. 3. Keep parameter ranges reasonable. 4. Higher aspheric order is better but more sensitive. 5. More iterations with larger ray sampling improves convergence.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def optimize(
    self,
    lrs=[1e-3, 1e-4, 1e-1, 1e-4],
    iterations=5000,
    test_per_iter=100,
    optim_mat=False,
    shape_control=True,
    sample_more_off_axis=False,
    result_dir=None,
):
    """Optimise the lens by minimising RGB RMS spot errors.

    Runs a curriculum-learning training loop with Adam optimiser and cosine
    annealing. Periodically evaluates the lens, saves intermediate results,
    and optionally corrects surface shapes.

    Args:
        lrs (list, optional): Learning rates for [d_next, c, k, a] parameter groups.
            Defaults to [1e-3, 1e-4, 1e-1, 1e-4].
        iterations (int, optional): Total training iterations. Defaults to 5000.
        test_per_iter (int, optional): Evaluate and save every N iterations.
            Defaults to 100.
        optim_mat (bool, optional): If True, include material parameters (n, V)
            in optimisation. Defaults to False.
        shape_control (bool, optional): If True, call ``correct_shape()`` at each
            evaluation step. Defaults to True.
        sample_more_off_axis (bool, optional): If True, concentrate ray samples
            toward the edge of the field to improve off-axis correction.
            Passed directly to ``sample_ring_arm_rays``. Defaults to False.
        result_dir (str, optional): Directory to save results. If None,
            auto-generates a timestamped directory. Defaults to None.

    Note:
        Debug hints:
            1. Slowly optimise with small learning rate.
            2. FoV and thickness should match well.
            3. Keep parameter ranges reasonable.
            4. Higher aspheric order is better but more sensitive.
            5. More iterations with larger ray sampling improves convergence.
    """
    # Experiment settings
    depth = self.obj_depth
    num_ring = 32
    num_arm = 8
    spp = 2048

    # Result directory and logger
    if result_dir is None:
        result_dir = (
            f"./results/{datetime.now().strftime('%m%d-%H%M%S')}-DesignLens"
        )

    os.makedirs(result_dir, exist_ok=True)
    if not logging.getLogger().hasHandlers():
        logger = logging.getLogger()
        logger.setLevel("DEBUG")
        fmt = logging.Formatter(
            "%(asctime)s:%(levelname)s:%(message)s", "%Y-%m-%d %H:%M:%S"
        )
        sh = logging.StreamHandler()
        sh.setFormatter(fmt)
        sh.setLevel("INFO")
        fh = logging.FileHandler(f"{result_dir}/output.log")
        fh.setFormatter(fmt)
        fh.setLevel("INFO")
        logger.addHandler(sh)
        logger.addHandler(fh)
    logging.info(
        f"lr:{lrs}, iterations:{iterations}, num_ring:{num_ring}, num_arm:{num_arm}, rays_per_fov:{spp}."
    )
    logging.info(
        "If Out-of-Memory, try to reduce num_ring, num_arm, and rays_per_fov."
    )

    # Optimizer and scheduler
    optimizer = self.get_optimizer(lrs, optim_mat=optim_mat)
    scheduler = get_cosine_schedule_with_warmup(
        optimizer, num_warmup_steps=100, num_training_steps=iterations
    )

    # Training loop
    pbar = tqdm(
        total=iterations + 1,
        desc="Progress",
        postfix={"loss_rms": 0},
    )
    for i in range(iterations + 1):
        # ===> Evaluate the lens
        if i % test_per_iter == 0:
            with torch.no_grad():
                if shape_control and i > 0:
                    self.correct_shape()

                self.write_lens_json(f"{result_dir}/iter{i}.json")
                self.analysis(f"{result_dir}/iter{i}")

                # Sample rays
                self.calc_pupil()
                rays_backup = []
                for wv in self.wvln_rgb:
                    ray = self.sample_ring_arm_rays(
                        num_ring=num_ring,
                        num_arm=num_arm,
                        spp=spp,
                        depth=depth,
                        wvln=wv,
                        scale_pupil=1.05,
                        sample_more_off_axis=sample_more_off_axis,
                    )
                    rays_backup.append(ray)

                # Pinhole ideal for distortion reference (distortion-free).
                pinhole_ref = -self.psf_center(
                    points_obj=ray.o[:, :, 0, :], method="pinhole"
                )

        # Iteration 0 is the baseline state. Stop after evaluating iteration
        # ``iterations`` so the public count equals the number of optimizer
        # updates rather than performing one extra step.
        if i == iterations:
            pbar.update(1)
            break

        # ===> Optimize lens by minimizing RMS
        # Green is traced first: its centroid sets center_ref and drives
        # the distortion penalty; red and blue reuse the same center_ref.
        loss_rms_ls = []
        loss_distortion = torch.tensor(0.0, device=self.device)
        w_mask = None
        center_ref = None
        wvln_order = [1, 0, 2]  # green, red, blue
        for wv_idx in wvln_order:
            # Ray tracing to sensor, [num_ring, num_arm, num_rays, 3]
            ray = rays_backup[wv_idx].clone()
            ray = self.trace2sensor(ray)

            if center_ref is None:
                # Green centroid at sensor, shape [num_ring, num_arm, 2].
                centroid_xy = ray.centroid()[..., :2]

                # Distortion: relative displacement of green centroid from
                # pinhole ideal, averaged equally over all off-axis fields.
                ideal_height = pinhole_ref.norm(dim=-1)
                field_mask = ideal_height > EPSILON
                distortion = (centroid_xy - pinhole_ref).norm(dim=-1)
                distortion = distortion / ideal_height.clamp_min(EPSILON)
                violation = distortion - self.distortion_max
                penalty = relu(violation / self.distortion_max)
                n_fields = field_mask.sum().clamp_min(1)
                loss_distortion = (penalty * field_mask.float()).sum() / n_fields

                # Detach so RMS gradient moves spot shape, not its
                # position; distortion loss handles placement.
                center_ref = centroid_xy.detach().unsqueeze(-2)

            # Ray error to center and valid mask
            ray_valid = ray.is_valid
            ray_err = ray.o[..., :2] - center_ref
            ray_err = torch.where(
                ray_valid.bool().unsqueeze(-1), ray_err, torch.zeros_like(ray_err)
            )

            # MSE per field point, shape [num_ring, num_arm]
            mse = (ray_err**2).sum(-1).sum(-1) / (ray_valid.sum(-1) + EPSILON)

            # Weight mask
            if w_mask is None:
                w_mask = mse.detach().sqrt().clone()
                w_mask = w_mask / (w_mask.mean() + EPSILON)
                w_mask[0, :] = 1.0

            # RMS and weighted loss
            l_rms = torch.clamp(mse, min=EPSILON).sqrt()
            l_rms_weighted = (l_rms * w_mask).sum() / (w_mask.sum() + EPSILON)
            loss_rms_ls.append(l_rms_weighted)

        # RMS loss for all wavelengths
        loss_rms = sum(loss_rms_ls) / len(loss_rms_ls)

        # Total loss
        w_reg = 0.1
        loss_reg, loss_dict = self.loss_reg()
        L_total = loss_rms + w_reg * (loss_reg + loss_distortion)

        # Back-propagation
        optimizer.zero_grad()
        L_total.backward()
        optimizer.step()
        scheduler.step()

        pbar.set_postfix(
            loss_rms=loss_rms.item(),
            loss_dist=loss_distortion.item(),
            **loss_dict,
        )
        pbar.update(1)

    pbar.close()

find_diff_surf

find_diff_surf()

Get differentiable/optimizable surface indices.

Returns a list of surface indices that can be optimized during lens design. Excludes the aperture surface from optimization.

Returns:

Name	Type	Description
diff_surf_range	list or range	Surface indices excluding the aperture.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def find_diff_surf(self):
    """Get differentiable/optimizable surface indices.

    Returns a list of surface indices that can be optimized during lens design.
    Excludes the aperture surface from optimization.

    Returns:
        diff_surf_range (list or range): Surface indices excluding the
            aperture.
    """
    if self.aper_idx is None:
        diff_surf_range = range(len(self.surfaces))
    else:
        diff_surf_range = list(range(0, self.aper_idx)) + list(
            range(self.aper_idx + 1, len(self.surfaces))
        )
    return diff_surf_range

get_optimizer_params

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

Build per-surface Adam parameter groups with per-type learning rates.

Collects trainable parameters for every surface (dispatching on surface type), plus the sensor distance, into a list of optimizer param groups.

Recommendation

For cellphone lens: [d_next, c, k, a], [1e-4, 1e-4, 1e-1, 1e-4]. For camera lens: [d_next, c, 0, 0], [1e-3, 1e-4, 0, 0].

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates for the [d_next, c, k, a] parameter groups. Defaults to [1e-4, 1e-4, 1e-2, 1e-4].	[0.0001, 0.0001, 0.01, 0.0001]
optim_mat	bool	Whether to optimize material parameters. Defaults to False.	False
optim_surf_range	list or None	Surface indices to optimize. When None, all surfaces are used. Defaults to None.	None

Returns:

Name	Type	Description
params	list	List of optimizer parameter-group dicts.

Raises:

Type	Description
Exception	If a surface type is not supported for optimization.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def get_optimizer_params(
    self,
    lrs=[1e-4, 1e-4, 1e-2, 1e-4],
    optim_mat=False,
    optim_surf_range=None,
):
    """Build per-surface Adam parameter groups with per-type learning rates.

    Collects trainable parameters for every surface (dispatching on surface
    type), plus the sensor distance, into a list of optimizer param groups.

    Recommendation:
        For cellphone lens: [d_next, c, k, a], [1e-4, 1e-4, 1e-1, 1e-4].
        For camera lens: [d_next, c, 0, 0], [1e-3, 1e-4, 0, 0].

    Args:
        lrs (list, optional): Learning rates for the [d_next, c, k, a] parameter
            groups. Defaults to [1e-4, 1e-4, 1e-2, 1e-4].
        optim_mat (bool, optional): Whether to optimize material parameters.
            Defaults to False.
        optim_surf_range (list or None, optional): Surface indices to
            optimize. When None, all surfaces are used. Defaults to None.

    Returns:
        params (list): List of optimizer parameter-group dicts.

    Raises:
        Exception: If a surface type is not supported for optimization.
    """
    # Find surfaces to be optimized
    if optim_surf_range is None:
        # optim_surf_range = self.find_diff_surf()
        optim_surf_range = range(len(self.surfaces))

    # Optimize lens surface parameters
    params = []
    for surf_idx in optim_surf_range:
        surf = self.surfaces[surf_idx]

        if isinstance(surf, Aperture):
            params += surf.get_optimizer_params(lrs=[lrs[0]])

        elif isinstance(surf, Aspheric):
            params += surf.get_optimizer_params(lrs=lrs[:4], optim_mat=optim_mat)

        elif isinstance(surf, Phase):
            # Phase surfaces take [d_lr, coeff_lr]. Use a dedicated 5th lr
            # when provided, otherwise fall back to the last lr so the
            # standard 4-element lrs convention does not IndexError.
            coeff_lr = lrs[4] if len(lrs) > 4 else lrs[-1]
            params += surf.get_optimizer_params(lrs=[lrs[0], coeff_lr])

        # elif isinstance(surf, GaussianRBF):
        #     params += surf.get_optimizer_params(lrs=lr, optim_mat=optim_mat)

        # elif isinstance(surf, NURBS):
        #     params += surf.get_optimizer_params(lrs=lr, optim_mat=optim_mat)

        elif isinstance(surf, Plane):
            params += surf.get_optimizer_params(lrs=[lrs[0]], optim_mat=optim_mat)

        # elif isinstance(surf, PolyEven):
        #     params += surf.get_optimizer_params(lrs=lr, optim_mat=optim_mat)

        elif isinstance(surf, Spheric):
            params += surf.get_optimizer_params(
                lrs=[lrs[0], lrs[1]], optim_mat=optim_mat
            )

        elif isinstance(surf, ThinLens):
            params += surf.get_optimizer_params(
                lrs=[lrs[0], lrs[1]], optim_mat=optim_mat
            )

        else:
            raise Exception(
                f"Surface type {surf.__class__.__name__} is not supported for optimization yet."
            )

    return params

get_optimizer

get_optimizer(
    lrs=[0.0001, 0.0001, 0.1, 0.0001], optim_surf_range=None, optim_mat=False
)

Build an Adam optimizer over all trainable lens parameters.

Parameters:

Name	Type	Description	Default
lrs	list	Learning rates for the [d_next, c, k, ai] parameter groups. Defaults to [1e-4, 1e-4, 1e-1, 1e-4].	[0.0001, 0.0001, 0.1, 0.0001]
optim_surf_range	list or None	Surface indices to optimize. When None, all surfaces are included. Defaults to None.	None
optim_mat	bool	Whether to include material parameters (n, V). Defaults to False.	False

Returns:

Name	Type	Description
optimizer	Adam	Configured Adam optimizer.

Source code in deeplens-src/deeplens/geolens_pkg/optim.py

def get_optimizer(
    self,
    lrs=[1e-4, 1e-4, 1e-1, 1e-4],
    optim_surf_range=None,
    optim_mat=False,
):
    """Build an Adam optimizer over all trainable lens parameters.

    Args:
        lrs (list, optional): Learning rates for the [d_next, c, k, ai] parameter
            groups. Defaults to [1e-4, 1e-4, 1e-1, 1e-4].
        optim_surf_range (list or None, optional): Surface indices to
            optimize. When None, all surfaces are included. Defaults to None.
        optim_mat (bool, optional): Whether to include material parameters
            (n, V). Defaults to False.

    Returns:
        optimizer (torch.optim.Adam): Configured Adam optimizer.
    """
    # Get optimizer
    params = self.get_optimizer_params(
        lrs=lrs, optim_surf_range=optim_surf_range, optim_mat=optim_mat
    )
    optimizer = torch.optim.Adam(params)
    # optimizer = torch.optim.SGD(params)
    return optimizer

deeplens.geolens_pkg.ops.GeoLensOps

Mixin providing in-place lens operations for GeoLens.

Bundles methods that modify a lens during design optimization: sizing clear apertures by ray tracing (pruning) and correcting lens geometry. Intended to be mixed into the GeoLens class, so all methods access lens state (self.surfaces, self.d_sensor, self.rfov, etc.) on the host.

Key methods

prune_surf: Size clear apertures by ray tracing. correct_shape: Fix lens geometry during optimization.

prune_surf

prune_surf(mounting_margin=None)

Prune surface radii so all valid rays pass through, then enforce manufacturability.

Traces 16 meridional fields from 0 to the full FoV to find the maximum ray height [mm] on each surface, expands it by a mounting margin, then caps the proposed radii to satisfy an edge-sag limit and edge-clearance (air-gap / edge-thickness) constraints with neighbouring surfaces. Aperture surfaces are not resized. The capped radii are committed via each surface's update_r.

Parameters:

Name	Type	Description	Default
mounting_margin	float or None	Absolute mounting margin [mm] added to the ray-traced clear-aperture radius. If None, the margin is auto-selected per surface: 5% of the ray-traced radius when that radius is below 5 mm, otherwise 1 mm. Defaults to None.	None

Source code in deeplens-src/deeplens/geolens_pkg/ops.py

@torch.no_grad()
def prune_surf(self, mounting_margin=None):
    """Prune surface radii so all valid rays pass through, then enforce manufacturability.

    Traces 16 meridional fields from 0 to the full FoV to find the maximum
    ray height [mm] on each surface, expands it by a mounting margin, then
    caps the proposed radii to satisfy an edge-sag limit and edge-clearance
    (air-gap / edge-thickness) constraints with neighbouring surfaces.
    Aperture surfaces are not resized. The capped radii are committed via
    each surface's `update_r`.

    Args:
        mounting_margin (float or None, optional): Absolute mounting margin
            [mm] added to the ray-traced clear-aperture radius. If `None`,
            the margin is auto-selected per surface: 5% of the ray-traced
            radius when that radius is below 5 mm, otherwise 1 mm. Defaults
            to None.
    """
    surface_range = self.find_diff_surf()
    num_surfs = len(self.surfaces)

    # ------------------------------------------------------------------
    # 1. Temporarily remove radius limits so the trace is unclipped
    # ------------------------------------------------------------------
    saved_radii = [self.surfaces[i].r for i in range(num_surfs)]
    for i in surface_range:
        self.surfaces[i].r = self.surfaces[i].max_height()

    # ------------------------------------------------------------------
    # 2. Trace rays at full FoV to find maximum ray height per surface
    # ------------------------------------------------------------------
    assert self.rfov is not None, "prune_surf() requires self.rfov."
    fov_deg = self.rfov * 180 / torch.pi
    num_fov_samples = 16
    fov_y = torch.linspace(0.0, fov_deg, num_fov_samples, device=self.device)
    ray = self.sample_from_fov(fov_x=[0.0], fov_y=fov_y)
    _, ray_o_record = self.trace2sensor(ray=ray, record=True)

    # Ray record, shape [num_rays, num_surfaces + 2, 3]
    ray_o_record = torch.stack(ray_o_record, dim=-2)
    ray_o_record = torch.nan_to_num(ray_o_record, nan=0.0, posinf=0.0, neginf=0.0)
    ray_o_record = ray_o_record.reshape(-1, ray_o_record.shape[-2], 3)

    # Compute the maximum ray height for each surface
    ray_r_record = (ray_o_record[..., :2] ** 2).sum(-1).sqrt()
    surf_r_max = ray_r_record.max(dim=0)[0][1:-1]

    # ------------------------------------------------------------------
    # 3. Propose new radii (not yet committed to surfaces).
    # ------------------------------------------------------------------
    proposed_r = [float(self.surfaces[i].r) for i in range(num_surfs)]
    for i in surface_range:
        # Surface radius required by ray tracing
        if surf_r_max[i] > 0:
            base = float(surf_r_max[i].item())
        else:
            base = float(self.surfaces[i].r)

        # Expand the ray-traced radius by a mounting margin
        if mounting_margin is None:
            r_expand = 0.05 * base if base < 5.0 else 1.0
        else:
            r_expand = float(mounting_margin)

        # Propose the new radius, capped at the surface's physical maximum height
        proposed_r[i] = min(base + r_expand, float(self.surfaces[i].max_height()))

    # ------------------------------------------------------------------
    # 3b. Sag cap: edge sag must not exceed sag_factor * proposed radius.
    # Grid-search for the largest r in [r_min, proposed_r] where the
    # constraint holds. The grid is dense enough for typical aspheric sag
    # profiles; non-monotonic extremes are handled conservatively.
    # ------------------------------------------------------------------
    sag_factor = 0.4
    for i in surface_range:
        if not isinstance(self.surfaces[i], Aperture):
            r_prop = proposed_r[i]
            r_cands = torch.linspace(r_prop / 64, r_prop, 64, device=self.device)
            z0 = self.surfaces[i].surface_with_offset(
                torch.tensor(0.0, device=self.device), 0.0, valid_check=False
            )
            z_cands = self.surfaces[i].surface_with_offset(
                r_cands, torch.zeros_like(r_cands), valid_check=False
            )
            sag_valid = (z_cands - z0).abs() <= sag_factor * r_cands
            if sag_valid.any():
                proposed_r[i] = min(r_prop, float(r_cands[sag_valid].max().item()))
            else:
                proposed_r[i] = float(r_cands[0].item())

    # ------------------------------------------------------------------
    # 4. Edge-clearance pass — proactively cap adjacent pairs so the
    #    committed radii never produce self-intersection at the edge.
    #    Thresholds match loss_bound. The cap uses the common
    #    clear-aperture overlap between adjacent surfaces so one surface is
    #    not pruned against regions where the neighbour has already been
    #    apertured away. Aperture surfaces are skipped; the stop size is an
    #    optical specification and should not be changed by pruning. The cap
    #    is computed via a single vectorized grid search rather than a
    #    serial binary loop.
    #
    #    Each pruned surface is checked against both neighbours. The
    #    previous implementation only capped surface i against i + 1,
    #    which allowed surface i to expand into i - 1 and later crash
    #    tracing/optimization.
    # ------------------------------------------------------------------
    min_radius_floor = 0.1  # mm — guard against update_r(0) killing a surface
    n_cand = 64
    n_edge = 64
    r_frac = torch.linspace(0.5, 1.0, n_edge, device=self.device)
    cand_frac = torch.linspace(1.0 / n_cand, 1.0, n_cand, device=self.device)

    def cap_radius_against_pair(cap_idx, prev_idx, next_idx):
        prev_surf = self.surfaces[prev_idx]
        next_surf = self.surfaces[next_idx]
        if isinstance(prev_surf, Aperture) or isinstance(next_surf, Aperture):
            return
        if isinstance(self.surfaces[cap_idx], Aperture):
            return

        edge_min = 0.1  # mm
        r_check = proposed_r[cap_idx]

        other_idx = next_idx if cap_idx == prev_idx else prev_idx
        other_r = proposed_r[other_idx]
        z_off = {prev_idx: 0.0, next_idx: float(prev_surf.d_next)}

        required_r = max(
            float(surf_r_max[cap_idx].item()),
            min_radius_floor,
        )

        # Vectorized cap: evaluate gap for 64 candidate radii in one pass.
        cand_r = cand_frac * r_check
        cand_overlap_r = torch.minimum(
            cand_r, torch.tensor(other_r, device=self.device)
        )
        r_grid = cand_overlap_r.unsqueeze(1) * r_frac.unsqueeze(0)
        z_prev_grid = prev_surf.surface_with_offset(
            r_grid.reshape(-1),
            0.0,
            valid_check=False,
            d=z_off[prev_idx],
        ).reshape(n_cand, n_edge)
        z_next_grid = next_surf.surface_with_offset(
            r_grid.reshape(-1),
            0.0,
            valid_check=False,
            d=z_off[next_idx],
        ).reshape(n_cand, n_edge)
        per_cand_gap = (z_next_grid - z_prev_grid).min(dim=-1).values
        overlap_ok = per_cand_gap >= edge_min

        # Sag-bracket: the cap surface's edge z (at candidate r) must not
        # axially cross the other surface's edge z. Catches the case
        # where high-order aspheric terms blow up beyond the surface's
        # design r and drag its edge past the neighbour, while the
        # in-overlap gap above is still fine.
        cap_surf = self.surfaces[cap_idx]
        other_surf = self.surfaces[other_idx]
        z_other_edge = other_surf.surface_with_offset(
            torch.tensor(other_r, device=self.device),
            torch.tensor(0.0, device=self.device),
            valid_check=False,
            d=z_off[other_idx],
        )
        z_cap_at_cand = cap_surf.surface_with_offset(
            cand_r,
            torch.zeros_like(cand_r),
            valid_check=False,
            d=z_off[cap_idx],
        )
        if cap_idx > other_idx:
            # cap is later in light path — must stay axially after other
            bracket_ok = z_cap_at_cand > z_other_edge + edge_min
        else:
            # cap is earlier — must stay axially before other
            bracket_ok = z_cap_at_cand < z_other_edge - edge_min

        valid_mask = overlap_ok & bracket_ok
        if not bool(valid_mask.any()):
            logging.warning(
                f"Surf {prev_idx}-{next_idx} "
                f"({prev_surf.mat2.name}): no candidate "
                f"radius satisfies edge_min {edge_min:.3f} mm at "
                f"r_check {r_check:.3f} mm (possible sag crossing near "
                f"axis). Reducing surface {cap_idx} to the ray-required radius "
                f"{required_r:.3f} mm, but edge clearance may remain "
                f"violated."
            )
            proposed_r[cap_idx] = min(proposed_r[cap_idx], required_r)
            return

        r_safe = float((cand_frac[valid_mask].max() * r_check).item())
        if r_safe < required_r:
            logging.warning(
                f"Surf {prev_idx}-{next_idx} "
                f"({prev_surf.mat2.name}): ray-required "
                f"radius {required_r:.3f} mm exceeds edge-clearance-safe "
                f"radius {r_safe:.3f} mm for edge_min {edge_min:.3f} mm. "
                f"Reducing surface {cap_idx} to the ray-required radius; edge "
                f"clearance may remain violated."
            )
            proposed_r[cap_idx] = min(proposed_r[cap_idx], required_r)
            return

        r_safe = max(r_safe, min_radius_floor)
        if proposed_r[cap_idx] > r_safe:
            proposed_r[cap_idx] = r_safe

    for i in surface_range:
        if i > 0:
            cap_radius_against_pair(i, i - 1, i)
        if i < num_surfs - 1:
            cap_radius_against_pair(i, i, i + 1)

    # ------------------------------------------------------------------
    # 4b. Commit the capped proposed radii to the surfaces.
    # ------------------------------------------------------------------
    for i in surface_range:
        if proposed_r[i] > 0:
            self.surfaces[i].update_r(proposed_r[i])

correct_shape

correct_shape(mounting_margin=None)

Correct invalid lens shape during lens design optimization.

The first surface is always at z=0 under sequential geometry, so shape correction only needs to prune surfaces to let valid rays pass.

Parameters:

Name	Type	Description	Default
mounting_margin	float or None	Absolute mounting margin [mm] for surface pruning, passed through to prune_surf. Defaults to None.	None

Source code in deeplens-src/deeplens/geolens_pkg/ops.py

@torch.no_grad()
def correct_shape(self, mounting_margin=None):
    """Correct invalid lens shape during lens design optimization.

    The first surface is always at z=0 under sequential geometry, so shape
    correction only needs to prune surfaces to let valid rays pass.

    Args:
        mounting_margin (float or None, optional): Absolute mounting margin
            [mm] for surface pruning, passed through to `prune_surf`.
            Defaults to None.
    """
    self.prune_surf(mounting_margin=mounting_margin)

match_materials

match_materials(mat_table='CDGM')

Match each surface's material to the nearest entry in a glass catalog.

Replaces every surface's mat2 glass with the closest real catalog glass in-place, making an idealised design manufacturable.

Parameters:

Name	Type	Description	Default
mat_table	str	Glass catalog name. Supported values are 'CDGM' (default catalog) and 'PLASTIC'. Defaults to 'CDGM'.	'CDGM'

Raises:

Type	Description
NotImplementedError	If mat_table is an unrecognised catalog name.

Source code in deeplens-src/deeplens/geolens_pkg/ops.py

@torch.no_grad()
def match_materials(self, mat_table="CDGM"):
    """Match each surface's material to the nearest entry in a glass catalog.

    Replaces every surface's `mat2` glass with the closest real catalog
    glass in-place, making an idealised design manufacturable.

    Args:
        mat_table (str, optional): Glass catalog name. Supported values are
            'CDGM' (default catalog) and 'PLASTIC'. Defaults to 'CDGM'.

    Raises:
        NotImplementedError: If `mat_table` is an unrecognised catalog name.
    """
    for surf in self.surfaces:
        surf.mat2.match_material(mat_table=mat_table)

deeplens.geolens_pkg.render.GeoLensRender

Mixin providing image simulation for GeoLens.

Hosts render, the entry point that dispatches to reverse ray tracing or to the PSF-convolution methods inherited from Lens. The ray-tracing path is implemented here: rays are sampled at the sensor, traced backward through the lens to the object plane, and integrated against the object image. warp/unwarp apply and remove the corresponding geometric distortion, which the PSF path needs because its kernels are locally shift-invariant.

render

render(img_obj, depth=None, method=None, **kwargs)

Differentiable image simulation.

Image simulation methods

[1] PSF map block convolution. [2] PSF patch convolution. [3] Ray tracing rendering.

Parameters:

Name	Type	Description	Default
img_obj	Tensor	Input image object in raw space. Shape [N, C, H, W].	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
method	str	Image simulation method. One of 'psf_map', 'psf_patch', or 'ray_tracing'. When None (default), falls back to self._default_render_method ('ray_tracing' for GeoLens).	None
**kwargs		Additional arguments for different methods: - psf_grid (tuple): Grid size for PSF map method. Defaults to (10, 10). - psf_ks (int): Kernel size for PSF methods. Defaults to PSF_KS. - psf_spp (int): Rays per PSF for PSF map method. Defaults to SPP_PSF. - warp_grid (int): Inverse-distortion grid resolution for PSF map method. Defaults to 128. - patch_center (tuple): Center position for PSF patch method. Defaults to (0.0, 0.0). - spp (int): Samples per pixel for ray tracing. Defaults to SPP_RENDER.	{}

Returns:

Name	Type	Description
img_render	Tensor	Rendered image tensor. Shape of [N, C, H, W].

Source code in deeplens-src/deeplens/geolens_pkg/render.py

def render(self, img_obj, depth=None, method=None, **kwargs):
    """Differentiable image simulation.

    Image simulation methods:
        [1] PSF map block convolution.
        [2] PSF patch convolution.
        [3] Ray tracing rendering.

    Args:
        img_obj (torch.Tensor): Input image object in raw space. Shape [N, C, H, W].
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        method (str, optional): Image simulation method. One of 'psf_map', 'psf_patch',
            or 'ray_tracing'. When None (default), falls back to
            `self._default_render_method` ('ray_tracing' for `GeoLens`).
        **kwargs: Additional arguments for different methods:
            - psf_grid (tuple): Grid size for PSF map method. Defaults to (10, 10).
            - psf_ks (int): Kernel size for PSF methods. Defaults to PSF_KS.
            - psf_spp (int): Rays per PSF for PSF map method. Defaults to SPP_PSF.
            - warp_grid (int): Inverse-distortion grid resolution for PSF map method. Defaults to 128.
            - patch_center (tuple): Center position for PSF patch method. Defaults to (0.0, 0.0).
            - spp (int): Samples per pixel for ray tracing. Defaults to SPP_RENDER.

    Returns:
        img_render (torch.Tensor): Rendered image tensor. Shape of [N, C, H, W].
    """
    method = self._default_render_method if method is None else method
    depth = self.obj_depth if depth is None else depth
    B, C, Himg, Wimg = img_obj.shape
    Wsensor, Hsensor = self.sensor_res

    # Image simulation
    if method == "psf_map":
        # PSF rendering - uses PSF map to render image
        assert Wimg == Wsensor and Himg == Hsensor, (
            f"Sensor resolution {Wsensor}x{Hsensor} must match input image {Wimg}x{Himg}."
        )
        psf_grid = kwargs.get("psf_grid", (10, 10))
        psf_ks = kwargs.get("psf_ks", PSF_KS)
        psf_spp = kwargs.get("psf_spp", SPP_PSF)
        warp_grid = kwargs.get("warp_grid", 128)
        img_obj = self.warp(img_obj, depth=depth, num_grid=warp_grid)
        img_render = self.render_psf_map(
            img_obj,
            depth=depth,
            psf_grid=psf_grid,
            psf_ks=psf_ks,
            psf_spp=psf_spp,
        )

    elif method == "psf_patch":
        # PSF patch rendering - uses a single PSF to render a patch of the image
        patch_center = kwargs.get("patch_center", (0.0, 0.0))
        psf_ks = kwargs.get("psf_ks", PSF_KS)
        img_render = self.render_psf_patch(
            img_obj, depth=depth, patch_center=patch_center, psf_ks=psf_ks
        )

    elif method == "ray_tracing":
        # Ray tracing rendering
        assert Wimg == Wsensor and Himg == Hsensor, (
            f"Sensor resolution {Wsensor}x{Hsensor} must match input image {Wimg}x{Himg}."
        )
        spp = kwargs.get("spp", SPP_RENDER)
        img_render = self.render_raytracing(img_obj, depth=depth, spp=spp)

    else:
        raise Exception(f"Image simulation method {method} is not supported.")

    return img_render

sample_sensor

sample_sensor(spp=64, wvln=None, sub_pixel=False)

Sample rays from sensor pixels (backward rays). Used for ray-tracing based rendering.

Parameters:

Name	Type	Description	Default
spp	int	sample per pixel. Defaults to 64.	64
wvln	float	ray wvln in µm. When None (default), falls back to self.primary_wvln.	None
sub_pixel	bool	whether to sample multiple points inside the pixel. Defaults to False.	False

Returns:

Name	Type	Description
ray	Ray	Ray object. Shape [H, W, spp, 3]

Source code in deeplens-src/deeplens/geolens_pkg/render.py

@torch.no_grad()
def sample_sensor(self, spp=64, wvln=None, sub_pixel=False):
    """Sample rays from sensor pixels (backward rays). Used for ray-tracing based rendering.

    Args:
        spp (int, optional): sample per pixel. Defaults to 64.
        wvln (float, optional): ray wvln in µm. When ``None`` (default),
            falls back to ``self.primary_wvln``.
        sub_pixel (bool, optional): whether to sample multiple points inside the pixel. Defaults to False.

    Returns:
        ray (Ray): Ray object. Shape [H, W, spp, 3]
    """
    wvln = self.primary_wvln if wvln is None else wvln
    w, h = self.sensor_size
    W, H = self.sensor_res
    device = self.device

    # Sample points on sensor plane
    # Use top-left point as reference in rendering, so here we should sample bottom-right point
    x1, y1 = torch.meshgrid(
        torch.linspace(
            -w / 2,
            w / 2,
            W + 1,
            device=device,
            dtype=self.dtype,
        )[1:],
        torch.linspace(
            h / 2,
            -h / 2,
            H + 1,
            device=device,
            dtype=self.dtype,
        )[1:],
        indexing="xy",
    )
    z1 = torch.full_like(x1, self.d_sensor.item())

    # Sample second points on the pupil
    # sensor_res is (W, H) but meshgrid with indexing="xy" gives (H, W) arrays
    pupilz, pupilr = self.get_exit_pupil()
    ray_o2 = self.sample_circle(r=pupilr, z=pupilz, shape=(H, W, spp))

    # Form rays
    ray_o = torch.stack((x1, y1, z1), 2)
    ray_o = ray_o.unsqueeze(2).repeat(1, 1, spp, 1)  # [H, W, spp, 3]

    # Sub-pixel sampling for more realistic rendering
    if sub_pixel:
        delta_ox = (
            torch.rand(ray_o.shape[:-1], device=device, dtype=self.dtype)
            * self.pixel_size
        )
        delta_oy = (
            -torch.rand(ray_o.shape[:-1], device=device, dtype=self.dtype)
            * self.pixel_size
        )
        delta_oz = torch.zeros_like(delta_ox)
        delta_o = torch.stack((delta_ox, delta_oy, delta_oz), -1)
        ray_o = ray_o + delta_o

    # Form rays
    ray_d = ray_o2 - ray_o  # shape [H, W, spp, 3]
    ray = Ray(ray_o, ray_d, wvln, device=device)
    return ray

render_raytracing

render_raytracing(img, depth=None, spp=SPP_RENDER, vignetting=False)

Render an RGB image using ray-tracing rendering.

Parameters:

Name	Type	Description	Default
img	Tensor	RGB image tensor. Shape [N, 3, H, W].	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
spp	int	Samples per pixel. Defaults to SPP_RENDER.	SPP_RENDER
vignetting	bool	Whether to model the vignetting effect. Defaults to False.	False

Returns:

Name	Type	Description
img_render	Tensor	Rendered RGB image tensor. Shape [N, 3, H, W].

Source code in deeplens-src/deeplens/geolens_pkg/render.py

def render_raytracing(self, img, depth=None, spp=SPP_RENDER, vignetting=False):
    """Render an RGB image using ray-tracing rendering.

    Args:
        img (torch.Tensor): RGB image tensor. Shape [N, 3, H, W].
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        spp (int, optional): Samples per pixel. Defaults to SPP_RENDER.
        vignetting (bool, optional): Whether to model the vignetting effect. Defaults to False.

    Returns:
        img_render (torch.Tensor): Rendered RGB image tensor. Shape [N, 3, H, W].
    """
    depth = self.obj_depth if depth is None else depth
    img_render = torch.zeros_like(img)
    for i in range(3):
        img_render[:, i, :, :] = self.render_raytracing_mono(
            img=img[:, i, :, :],
            wvln=self.wvln_rgb[i],
            depth=depth,
            spp=spp,
            vignetting=vignetting,
        )
    return img_render

render_raytracing_mono

render_raytracing_mono(img, wvln, depth=None, spp=64, vignetting=False)

Render a monochrome image at a single wavelength using ray-tracing rendering.

Parameters:

Name	Type	Description	Default
img	Tensor	Monochrome image tensor. Shape [N, 1, H, W] or [N, H, W].	required
wvln	float	Wavelength in µm.	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
spp	int	Samples per pixel. Defaults to 64.	64
vignetting	bool	Whether to model the vignetting effect. Defaults to False.	False

Returns:

Name	Type	Description
img_mono	Tensor	Rendered monochrome image tensor. Shape [N, 1, H, W] or [N, H, W].

Source code in deeplens-src/deeplens/geolens_pkg/render.py

def render_raytracing_mono(self, img, wvln, depth=None, spp=64, vignetting=False):
    """Render a monochrome image at a single wavelength using ray-tracing rendering.

    Args:
        img (torch.Tensor): Monochrome image tensor. Shape [N, 1, H, W] or [N, H, W].
        wvln (float): Wavelength in µm.
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        spp (int, optional): Samples per pixel. Defaults to 64.
        vignetting (bool, optional): Whether to model the vignetting effect. Defaults to False.

    Returns:
        img_mono (torch.Tensor): Rendered monochrome image tensor. Shape [N, 1, H, W] or [N, H, W].
    """
    depth = self.obj_depth if depth is None else depth
    img = torch.flip(img, [-2, -1])
    scale = self.calc_scale(depth=depth)
    ray = self.sample_sensor(spp=spp, wvln=wvln)
    ray = self.trace2obj(ray)
    img_mono = self.render_compute_image(
        img, depth=depth, scale=scale, ray=ray, vignetting=vignetting
    )
    return img_mono

render_compute_image

render_compute_image(img, depth, scale, ray, vignetting=False)

Compute ray-image-plane intersections and integrate them into a rendered image.

Propagates the traced rays to the object plane, intersects them with the scaled object image, and accumulates radiance following the rendering equation. Back-propagation gradient flow: image -> w_i -> u -> p -> ray -> surface.

Parameters:

Name	Type	Description	Default
img	Tensor	Object image tensor. Shape [N, C, H, W] or [N, H, W].	required
depth	float	Object depth [mm].	required
scale	float	Object-to-image scale factor.	required
ray	Ray	Traced sensor rays. Shape [H, W, spp, 3].	required
vignetting	bool	Whether to model the vignetting effect. Defaults to False.	False

Returns:

Name	Type	Description
image	Tensor	Rendered image tensor. Shape [N, C, H, W] or [N, H, W].

Source code in deeplens-src/deeplens/geolens_pkg/render.py

def render_compute_image(self, img, depth, scale, ray, vignetting=False):
    """Compute ray-image-plane intersections and integrate them into a rendered image.

    Propagates the traced rays to the object plane, intersects them with the
    scaled object image, and accumulates radiance following the rendering
    equation. Back-propagation gradient flow: image -> w_i -> u -> p -> ray -> surface.

    Args:
        img (torch.Tensor): Object image tensor. Shape [N, C, H, W] or [N, H, W].
        depth (float): Object depth [mm].
        scale (float): Object-to-image scale factor.
        ray (Ray): Traced sensor rays. Shape [H, W, spp, 3].
        vignetting (bool): Whether to model the vignetting effect. Defaults to False.

    Returns:
        image (torch.Tensor): Rendered image tensor. Shape [N, C, H, W] or [N, H, W].
    """
    assert torch.is_tensor(img), "Input image should be Tensor."

    H, W = img.shape[-2:]
    squeeze_channel = False
    if len(img.shape) == 3:
        img = img.unsqueeze(1)
        squeeze_channel = True
    elif len(img.shape) == 4:
        pass
    else:
        raise ValueError("Input image should be [N, C, H, W] or [N, H, W] tensor.")

    # Scale object image physical size to get 1:1 pixel-pixel alignment with sensor image
    ray = ray.prop_to(depth)
    p = ray.o[..., :2]
    pixel_size = scale * self.pixel_size
    ray.is_valid = (
        ray.is_valid
        * (torch.abs(p[..., 0] / pixel_size) < (W / 2 + 1))
        * (torch.abs(p[..., 1] / pixel_size) < (H / 2 + 1))
    )

    image = backward_integral(
        ray=ray,
        img_obj=img,
        ps=pixel_size,
        vignetting=vignetting,
    )
    if squeeze_channel:
        image = image.squeeze(1)

    return image

warp

warp(img, depth=None, num_grid=128)

Apply lens distortion to an image using inverse distortion mapping.

Parameters:

Name	Type	Description	Default
img	Tensor	Undistorted image tensor, shape [B, C, H, W].	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
num_grid	int or tuple	Resolution of the inverse distortion grid.	128

Returns:

Name	Type	Description
img_warped	Tensor	Distorted image tensor, shape [B, C, H, W].

Source code in deeplens-src/deeplens/geolens_pkg/render.py

def warp(self, img, depth=None, num_grid=128):
    """Apply lens distortion to an image using inverse distortion mapping.

    Args:
        img (torch.Tensor): Undistorted image tensor, shape [B, C, H, W].
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        num_grid (int or tuple): Resolution of the inverse distortion grid.

    Returns:
        img_warped (torch.Tensor): Distorted image tensor, shape ``[B, C, H, W]``.
    """
    depth = self.obj_depth if depth is None else depth
    inv_distortion_map = self.calc_inv_distortion_map(
        depth=depth, num_grid=num_grid
    )
    inv_distortion_map = inv_distortion_map.permute(2, 0, 1).unsqueeze(0)
    inv_distortion_map = F.interpolate(
        inv_distortion_map, img.shape[-2:], mode="bilinear", align_corners=True
    )
    inv_distortion_map = inv_distortion_map.permute(0, 2, 3, 1).repeat(
        img.shape[0], 1, 1, 1
    )
    img_warped = F.grid_sample(img, inv_distortion_map, align_corners=True)
    return img_warped

unwarp

unwarp(img, depth=None, num_grid=128, crop=True, flip=True)

Unwarp (remove distortion from) a rendered image using the distortion map.

Parameters:

Name	Type	Description	Default
img	Tensor	Rendered image tensor. Shape [N, C, H, W].	required
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
num_grid	int	Resolution of the distortion grid. Defaults to 128.	128
crop	bool	Whether to crop the image. Defaults to True.	True
flip	bool	Whether to flip the distortion map. Defaults to True.	True

Returns:

Name	Type	Description
img_unwarpped	Tensor	Unwarped image tensor. Shape [N, C, H, W].

Source code in deeplens-src/deeplens/geolens_pkg/render.py

def unwarp(self, img, depth=None, num_grid=128, crop=True, flip=True):
    """Unwarp (remove distortion from) a rendered image using the distortion map.

    Args:
        img (torch.Tensor): Rendered image tensor. Shape [N, C, H, W].
        depth (float, optional): Object depth [mm]. When None (default),
            falls back to `self.obj_depth`.
        num_grid (int, optional): Resolution of the distortion grid. Defaults to 128.
        crop (bool, optional): Whether to crop the image. Defaults to True.
        flip (bool, optional): Whether to flip the distortion map. Defaults to True.

    Returns:
        img_unwarpped (torch.Tensor): Unwarped image tensor. Shape [N, C, H, W].
    """
    depth = self.obj_depth if depth is None else depth
    # Calculate distortion map, shape (num_grid, num_grid, 2)
    distortion_map = self.calc_distortion_map(depth=depth, num_grid=num_grid)

    # Interpolate distortion map to image resolution
    distortion_map = distortion_map.permute(2, 0, 1).unsqueeze(1)
    # distortion_map = torch.flip(distortion_map, [-2]) if flip else distortion_map
    distortion_map = F.interpolate(
        distortion_map, img.shape[-2:], mode="bilinear", align_corners=True
    )  # shape (B, 2, Himg, Wimg)
    distortion_map = distortion_map.permute(1, 2, 3, 0).repeat(
        img.shape[0], 1, 1, 1
    )  # shape (B, Himg, Wimg, 2)

    # Unwarp using grid_sample function
    img_unwarpped = F.grid_sample(
        img, distortion_map, align_corners=True
    )  # shape (B, C, Himg, Wimg)
    return img_unwarpped

deeplens.geolens_pkg.io.GeoLensIO

Mixin providing lens-file I/O for GeoLens.

Adds read/write methods for three lens prescription formats: DeepLens native JSON, Zemax sequential (.zmx), and Code V sequential (.seq). The JSON format is primary and human-readable, with parenthesised keys (e.g. "(d_sensor)") marking optimisable parameters. This class is not instantiated directly; it is mixed into GeoLens, and its methods read from and write to the host lens's state (surfaces, d_sensor, r_sensor, enpd, rfov_eff, etc.).

read_lens_zmx

read_lens_zmx(filename='./test.zmx')

Load the lens from a Zemax .zmx sequential lens file.

Parses STANDARD and EVENASPH surface types, glass materials, field definitions (YFLN, in degrees), and entrance pupil settings (ENPD/FLOA). Populates self.surfaces, self.d_sensor [mm], self.r_sensor [mm], self.enpd, self.float_enpd, and self.rfov_eff [rad].

Parameters:

Name	Type	Description	Default
filename	str	Path to the .zmx file. Both UTF-8 and UTF-16 encodings are accepted. Defaults to './test.zmx'.	'./test.zmx'

Returns:

Name	Type	Description
self	GeoLens	The updated lens (for chaining).

Source code in deeplens-src/deeplens/geolens_pkg/io.py

def read_lens_zmx(self, filename="./test.zmx"):
    """Load the lens from a Zemax .zmx sequential lens file.

    Parses STANDARD and EVENASPH surface types, glass materials, field
    definitions (YFLN, in degrees), and entrance pupil settings
    (ENPD/FLOA). Populates `self.surfaces`, `self.d_sensor` [mm],
    `self.r_sensor` [mm], `self.enpd`, `self.float_enpd`, and
    `self.rfov_eff` [rad].

    Args:
        filename (str, optional): Path to the .zmx file. Both UTF-8 and
            UTF-16 encodings are accepted. Defaults to './test.zmx'.

    Returns:
        self (GeoLens): The updated lens (for chaining).
    """
    # Read .zmx file
    try:
        with open(filename, "r", encoding="utf-8") as file:
            lines = file.readlines()
    except UnicodeDecodeError:
        with open(filename, "r", encoding="utf-16") as file:
            lines = file.readlines()

    # DeepLens stores every prescription length in millimetres. Silently
    # treating an inch/centimetre prescription as millimetres changes the
    # optical system, so reject unsupported declared units before creating
    # any surfaces.
    for line in lines:
        parts = line.strip().split()
        if parts and parts[0].upper() == "UNIT":
            if len(parts) < 2 or parts[1].upper() != "MM":
                declared = parts[1] if len(parts) >= 2 else "<missing>"
                raise ValueError(
                    f"Unsupported Zemax length unit {declared!r}; only MM is supported."
                )
            break

    # Iterate through the lines and extract SURF dict
    surfs_dict = {}
    current_surf = None
    for line in lines:
        # Strip leading/trailing whitespace for consistent parsing
        stripped_line = line.strip()

        if stripped_line.startswith("SURF"):
            current_surf = int(stripped_line.split()[1])
            surfs_dict[current_surf] = {}

        elif current_surf is not None and stripped_line != "":
            if len(stripped_line.split(maxsplit=1)) == 1:
                if stripped_line == "STOP":
                    surfs_dict[current_surf]["STOP"] = True
                continue
            else:
                key, value = stripped_line.split(maxsplit=1)
                if key == "PARM":
                    new_key = "PARM" + value.split()[0]
                    new_value = value.split()[1]
                    surfs_dict[current_surf][new_key] = new_value
                else:
                    surfs_dict[current_surf][key] = value

        elif stripped_line.startswith("FLOA") or stripped_line.startswith("ENPD"):
            if stripped_line.startswith("FLOA"):
                self.float_enpd = True
                self.enpd = None
            else:
                self.float_enpd = False
                self.enpd = float(stripped_line.split()[1])

        elif stripped_line.startswith("YFLN"):
            # Parse field of view from YFLN line (field coordinates in degrees)
            # YFLN format: YFLN 0.0 <0.707*rfov_deg> <0.99*rfov_deg>
            parts = stripped_line.split()
            if len(parts) > 1:
                field_values = [abs(float(x)) for x in parts[1:] if float(x) != 0.0]
                if field_values:
                    # The largest field value is typically 0.99 * rfov_deg
                    max_field_deg = max(field_values) / 0.99
                    self.rfov_eff = (
                        max_field_deg * math.pi / 180.0
                    )  # Convert to radians

    self.float_foclen = False
    self.float_rfov = False
    # Set default rfov_eff if not parsed from file
    if not hasattr(self, "rfov_eff"):
        self.rfov_eff = None

    lens_surface_types = {
        surf_dict.get("TYPE")
        for surf_idx, surf_dict in surfs_dict.items()
        if 0 < surf_idx < current_surf
    }
    unsupported_surface_types = sorted(
        "<missing>" if surface_type is None else surface_type
        for surface_type in lens_surface_types
        if surface_type not in {"STANDARD", "EVENASPH"}
    )
    if unsupported_surface_types:
        raise NotImplementedError(
            "Unsupported Zemax surface types: "
            + ", ".join(unsupported_surface_types)
        )

    # Reject coefficient records that the selected surface model cannot
    # represent instead of silently discarding prescription semantics.
    for surf_idx, surf_dict in surfs_dict.items():
        if not (0 < surf_idx < current_surf):
            continue
        surface_type = surf_dict.get("TYPE")
        for key, value in surf_dict.items():
            if not key.startswith("PARM"):
                continue
            try:
                coefficient = float(value)
                order = int(key.removeprefix("PARM"))
            except ValueError as exc:
                raise ValueError(
                    f"Invalid Zemax coefficient {key}={value!r} on surface {surf_idx}."
                ) from exc
            unsupported = surface_type == "STANDARD" or (
                surface_type == "EVENASPH" and order > 8
            )
            if unsupported and coefficient != 0.0:
                raise NotImplementedError(
                    f"Zemax {surface_type} surface {surf_idx} uses unsupported "
                    f"coefficient {key}={coefficient}."
                )

    # Read the extracted data from each SURF
    self.surfaces = []
    mat1_name = "air"
    for surf_idx, surf_dict in surfs_dict.items():
        if surf_idx > 0 and surf_idx < current_surf:
            # Lens surface parameters
            if "GLAS" in surf_dict:
                mat2_name = self._resolve_zmx_glass(surf_dict["GLAS"])
            else:
                mat2_name = "air"

            surf_r = (
                float(surf_dict["DIAM"].split()[0]) if "DIAM" in surf_dict else 1.0
            )
            surf_c = (
                float(surf_dict["CURV"].split()[0]) if "CURV" in surf_dict else 0.0
            )
            surf_d_next = (
                float(surf_dict["DISZ"].split()[0]) if "DISZ" in surf_dict else 0.0
            )
            surf_conic = float(surf_dict.get("CONI", 0.0))
            surf_param1 = float(surf_dict.get("PARM1", 0.0))
            surf_param2 = float(surf_dict.get("PARM2", 0.0))
            surf_param3 = float(surf_dict.get("PARM3", 0.0))
            surf_param4 = float(surf_dict.get("PARM4", 0.0))
            surf_param5 = float(surf_dict.get("PARM5", 0.0))
            surf_param6 = float(surf_dict.get("PARM6", 0.0))
            surf_param7 = float(surf_dict.get("PARM7", 0.0))
            surf_param8 = float(surf_dict.get("PARM8", 0.0))

            # Create surface object
            d_next_tensor = torch.as_tensor(surf_d_next, dtype=self.dtype)
            is_stop = bool(surf_dict.get("STOP", False))
            if surf_dict["TYPE"] == "STANDARD":
                if surf_conic != 0.0:
                    # Zemax permits a conic constant on STANDARD surfaces.
                    # Preserve it with the equivalent zero-polynomial asphere.
                    s = Aspheric(
                        c=surf_c,
                        r=surf_r,
                        d_next=d_next_tensor,
                        ai=[],
                        k=surf_conic,
                        mat2=mat2_name,
                    )
                elif surf_c == 0.0 and mat2_name == mat1_name == "air":
                    s = (
                        Aperture(r=surf_r, d_next=d_next_tensor)
                        if is_stop
                        else Plane(
                            r=surf_r,
                            d_next=d_next_tensor,
                            mat2="air",
                        )
                    )
                else:
                    s = Spheric(
                        c=surf_c,
                        r=surf_r,
                        d_next=d_next_tensor,
                        mat2=mat2_name,
                    )

            elif surf_dict["TYPE"] == "EVENASPH":
                # Aspherical surface
                s = Aspheric(
                    c=surf_c,
                    r=surf_r,
                    d_next=d_next_tensor,
                    ai2=surf_param1,
                    ai=[
                        surf_param2,
                        surf_param3,
                        surf_param4,
                        surf_param5,
                        surf_param6,
                        surf_param7,
                        surf_param8,
                    ],
                    k=surf_conic,
                    mat2=mat2_name,
                )

            else:
                print(f"Surface type {surf_dict['TYPE']} not implemented.")
                continue

            s.is_aperture = is_stop
            self.surfaces.append(s)
            mat1_name = mat2_name

        elif surf_idx == current_surf:
            # Image sensor
            self.r_sensor = float(surf_dict["DIAM"].split()[0])

        else:
            pass

    return self

write_lens_zmx

write_lens_zmx(filename='./test.zmx')

Write the lens to a Zemax .zmx sequential lens file.

Exports surfaces (STANDARD or EVENASPH), materials, field definitions (YFLN at 0, 0.707, and 0.99 of the effective half-FoV, in degrees), RGB wavelengths, and entrance-pupil settings in Zemax OpticStudio format. An extra image (sensor) surface is appended.

Parameters:

Name	Type	Description	Default
filename	str	Output file path. Defaults to './test.zmx'.	'./test.zmx'

Source code in deeplens-src/deeplens/geolens_pkg/io.py

    def write_lens_zmx(self, filename="./test.zmx"):
        """Write the lens to a Zemax .zmx sequential lens file.

        Exports surfaces (STANDARD or EVENASPH), materials, field definitions
        (YFLN at 0, 0.707, and 0.99 of the effective half-FoV, in degrees),
        RGB wavelengths, and entrance-pupil settings in Zemax OpticStudio
        format. An extra image (sensor) surface is appended.

        Args:
            filename (str, optional): Output file path. Defaults to './test.zmx'.
        """
        lens_zmx_str = ""
        if self.float_enpd:
            enpd_str = "FLOA"
        else:
            enpd_str = f"ENPD {self.enpd}"
        # Head string. Top-level directives are written at column 0 (not
        # indented to the surrounding Python block) so the emitted Zemax header
        # has no leading whitespace; the SURF 0 sub-keywords are indented to
        # match the per-surface blocks emitted by ``zmx_str``.
        head_str = f"""VERS 190513 80 123457 L123457
MODE SEQ
NAME
PFIL 0 0 0
LANG 0
UNIT MM X W X CM MR CPMM
{enpd_str}
ENVD 2.0E+1 1 0
GFAC 0 0
GCAT OSAKAGASCHEMICAL MISC
XFLN 0. 0. 0.
YFLN 0.0 {0.707 * self.rfov_eff * 57.3} {0.99 * self.rfov_eff * 57.3}
WAVL {self.wvln_rgb[2]:.7f} {self.wvln_rgb[1]:.7f} {self.wvln_rgb[0]:.7f}
RAIM 0 0 1 1 0 0 0 0 0
PUSH 0 0 0 0 0 0
SDMA 0 1 0
FTYP 0 0 3 3 0 0 0
ROPD 2
PICB 1
PWAV 2
POLS 1 0 1 0 0 1 0
GLRS 1 0
GSTD 0 100.000 100.000 100.000 100.000 100.000 100.000 0 1 1 0 0 1 1 1 1 1 1
NSCD 100 500 0 1.0E-3 5 1.0E-6 0 0 0 0 0 0 1000000 0 2
COFN QF "COATING.DAT" "SCATTER_PROFILE.DAT" "ABG_DATA.DAT" "PROFILE.GRD"
COFN COATING.DAT SCATTER_PROFILE.DAT ABG_DATA.DAT PROFILE.GRD
SURF 0
    TYPE STANDARD
    CURV 0.0
    DISZ INFINITY
"""
        lens_zmx_str += head_str

        # Surface string
        for i, s in enumerate(self.surfaces):
            surf_str = s.zmx_str(surf_idx=i + 1, d_next=s.d_next)
            if getattr(s, "is_aperture", False) and not isinstance(s, Aperture):
                lines = surf_str.splitlines(keepends=True)
                insert_at = min(2, len(lines))
                lines.insert(insert_at, "    STOP\n")
                surf_str = "".join(lines)
            lens_zmx_str += surf_str

        # Sensor (image) surface, formatted like the per-surface zmx_str blocks:
        # the SURF line at column 0 with its sub-keywords indented.
        sensor_str = f"""SURF {i + 2}
    TYPE STANDARD
    CURV 0.
    DISZ 0.0
    DIAM {self.r_sensor}
"""
        lens_zmx_str += sensor_str

        # Write lens zmx string into file
        with open(filename, "w") as f:
            f.writelines(lens_zmx_str)
        print(f"Lens written to {filename}")

read_lens_seq

read_lens_seq(filename='./test.seq')

Load the lens from a Code V .seq sequential file.

Parses standard and aspheric surfaces (conic K and polynomial coefficients A-I, mapped to even-aspheric terms ai[1]-ai[9]), entrance pupil diameter (EPD), field angles (YAN, in degrees), aperture stop (STO), and the image surface (SI). Populates self.surfaces, self.d_sensor [mm], self.r_sensor [mm], self.enpd, self.hfov [deg], and self.rfov_eff [rad]. Progress is printed to stdout.

Parameters:

Name	Type	Description	Default
filename	str	Path to the .seq file. Both UTF-8 and Latin-1 encodings are accepted. Defaults to './test.seq'.	'./test.seq'

Returns:

Name	Type	Description
self	GeoLens	The updated lens (for chaining).

Source code in deeplens-src/deeplens/geolens_pkg/io.py

def read_lens_seq(self, filename="./test.seq"):
    """Load the lens from a Code V .seq sequential file.

    Parses standard and aspheric surfaces (conic K and polynomial
    coefficients A-I, mapped to even-aspheric terms ai[1]-ai[9]), entrance
    pupil diameter (EPD), field angles (YAN, in degrees), aperture stop
    (STO), and the image surface (SI). Populates `self.surfaces`,
    `self.d_sensor` [mm], `self.r_sensor` [mm], `self.enpd`, `self.hfov`
    [deg], and `self.rfov_eff` [rad]. Progress is printed to stdout.

    Args:
        filename (str, optional): Path to the .seq file. Both UTF-8 and
            Latin-1 encodings are accepted. Defaults to './test.seq'.

    Returns:
        self (GeoLens): The updated lens (for chaining).
    """
    print(f"\n{'=' * 60}")
    print(f"Start reading CODE V file: {filename}")
    print(f"{'=' * 60}\n")

    # Read .seq file
    try:
        with open(filename, "r", encoding="utf-8") as file:
            lines = file.readlines()
        print(f"File read successfully (UTF-8)")
    except UnicodeDecodeError:
        try:
            with open(filename, "r", encoding="latin-1") as file:
                lines = file.readlines()
            print(f"File read successfully (Latin-1)")
        except Exception as e:
            print(f"Failed to read file: {e}")
            return self
    print(f"Total lines: {len(lines)}\n")

    # ============ Step 1: Parse file structure ============
    surfaces = []
    current_surface = {}
    surface_index = 0
    global_diameter = None

    print("Beginning to parse surface data...\n")

    for line_num, line in enumerate(lines, 1):
        line = line.strip()

        # Skip irrelevant lines
        if not line or line.startswith(
            (
                "RDM",
                "TITLE",
                "UID",
                "GO",
                "WL",
                "XAN",
                "REF",
                "WTW",
                "INI",
                "WTF",
                "VUY",
                "VLY",
                "DOR",
                "DIM",
                "THC",
            )
        ):
            continue
        # Read entrance pupil diameter
        if line.startswith("EPD"):
            self.enpd = float(line.split()[1])
            self.float_enpd = False
            global_diameter = self.enpd / 2.0
            print(
                f"[Line {line_num}] EPD={self.enpd} -> default radius={global_diameter}"
            )
            continue
        # Read field of view angle
        if line.startswith("YAN"):
            angles = [abs(float(x)) for x in line.split()[1:] if float(x) != 0.0]
            if angles:
                self.hfov = max(angles)
                # Also set rfov in radians for consistency with write functions
                self.rfov_eff = self.hfov * math.pi / 180.0
                print(f"[Line {line_num}] Max field of view={self.hfov} deg")
            continue
        # Object surface
        if line.startswith("SO"):
            parts = line.split()
            thickness = float(parts[2]) if len(parts) > 2 else 1e10

            current_surface = {
                "type": "OBJECT",
                "thickness": thickness,
                "index": surface_index,
            }
            surfaces.append(current_surface)
            print(f"[Line {line_num}] Object surface: T={thickness}")
            surface_index += 1
            current_surface = {}
            continue
        # Standard surface
        if line.startswith("S "):
            # Save the previous surface
            if current_surface:
                surfaces.append(current_surface)
                surface_index += 1

            parts = line.split()
            radius_value = float(parts[1]) if len(parts) > 1 else 0.0
            thickness = float(parts[2]) if len(parts) > 2 else 0.0
            material = parts[3].upper() if len(parts) > 3 else "AIR"

            # Key: compute curvature C = 1/R
            if abs(radius_value) > 1e-10:
                curvature = 1.0 / radius_value
            else:
                curvature = 0.0

            current_surface = {
                "type": "STANDARD",
                "radius": radius_value,
                "curvature": curvature,
                "thickness": thickness,
                "material": material,
                "index": surface_index,
                "diameter": global_diameter,
                "conic": 0.0,
                "asph_coeffs": {},
                "is_stop": False,
            }

            print(
                f"[Line {line_num}] Surface{surface_index}: R={radius_value:.4f} → C={curvature:.6f}, T={thickness}, Mat={material}"
            )
            continue
        # Image surface - do not append yet, wait for CIR
        if line.startswith("SI"):
            if current_surface:
                surfaces.append(current_surface)
                surface_index += 1

            parts = line.split()
            thickness = float(parts[1]) if len(parts) > 1 else 0.0

            current_surface = {
                "type": "IMAGE",
                "thickness": thickness,
                "diameter": None,  # Set to None first, wait for CIR line to update
                "index": surface_index,
            }
            print(f"[Line {line_num}] Image surface")
            continue
        # Handle surface attributes (CIR, STO, ASP, K, A~J, etc.)
        if current_surface:
            if line.startswith("CIR"):
                current_surface["diameter"] = float(
                    line.split()[1].replace(";", "")
                )
                print(f"[Line {line_num}]   → CIR={current_surface['diameter']}")

            elif line.startswith("STO"):
                current_surface["is_stop"] = True
                print(f"[Line {line_num}]   → Aperture stop flag")

            elif line.startswith("ASP"):
                current_surface["type"] = "ASPHERIC"
                print(f"[Line {line_num}]   → Aspheric surface")

            elif line.startswith("K "):
                current_surface["conic"] = float(line.split()[1].replace(";", ""))
                print(f"[Line {line_num}]   → K={current_surface['conic']}")

            # Only extract single-letter coefficients A-J
            elif any(
                line.startswith(p)
                for p in [
                    "A ",
                    "B ",
                    "C ",
                    "D ",
                    "E ",
                    "F ",
                    "G ",
                    "H ",
                    "I ",
                    "J ",
                ]
            ):
                parts = line.replace(";", "").split()
                i = 0
                while i < len(parts) - 1:
                    try:
                        key = parts[i]
                        # Only accept single letters within the range A-J
                        if len(key) == 1 and key in [
                            "A",
                            "B",
                            "C",
                            "D",
                            "E",
                            "F",
                            "G",
                            "H",
                            "I",
                            "J",
                        ]:
                            value = float(parts[i + 1])
                            current_surface["asph_coeffs"][key] = value
                            print(f"[Line {line_num}]   → {key}={value}")
                        i += 2
                    except:
                        i += 1

    # Save the last surface
    if current_surface:
        surfaces.append(current_surface)

    if not any(surface["type"] == "IMAGE" for surface in surfaces):
        raise ValueError("Code V lens file is missing its image surface (SI).")

    print(f"\nParsing complete, total {len(surfaces)} surfaces\n")

    # ============ Step 2: Create surface objects ============
    print(f"{'=' * 60}")
    print("Start creating surface objects:")
    print(f"{'=' * 60}\n")

    self.surfaces = []
    d = 0.0  # Cumulative distance from the first optical surface to the current surface
    previous_material = "air"
    pending_gap = 0.0

    for surf in surfaces:
        surf_idx = surf["index"]
        surf_type = surf["type"]

        print(f"{'=' * 50}")
        print(f"Processing surface{surf_idx} ({surf_type}), current d={d:.4f}")

        # Handle object surface
        if surf_type == "OBJECT":
            obj_thickness = surf["thickness"]
            if obj_thickness < 1e9:  # Finite object distance
                self.obj_depth = -float(obj_thickness)
                print(
                    f"   Object surface thickness={obj_thickness} → obj_depth={self.obj_depth:.4f}"
                )
            else:
                print("   Object surface at infinity")
            previous_material = "air"
            continue

        # Handle image surface
        if surf_type == "IMAGE":
            # Read diameter from surf dictionary (CIR value)
            self.r_sensor = (
                surf.get("diameter") if surf.get("diameter") is not None else 18.0
            )
            print(
                f"   Image plane position: d_sensor={d:.4f}, r_sensor={self.r_sensor:.4f}"
            )
            break

        # Get surface parameters
        current_material = surf.get("material", "AIR")
        if current_material in ["AIR", "0.0", "", None]:
            current_material = "air"
        else:
            current_material = current_material.lower()

        c = surf.get("curvature", 0.0)
        r = surf.get("diameter", 10.0)
        d_next = surf.get("thickness", 0.0)
        is_stop = surf.get("is_stop", False)

        print(f"   C={c:.6f}, R_aperture={r:.4f}, T={d_next:.4f}")
        print(f"   Material: {previous_material} → {current_material}")
        print(f"   is_stop={is_stop}")

        # Create surface object
        created = None
        try:
            # Case 1: pure aperture (air on both sides + STO flag)
            if is_stop and current_material == "air" and previous_material == "air":
                created = Aperture(r=r, d_next=d_next)
                print(f"   Created pure aperture: Aperture(r={r:.4f}, d={d:.4f})")

            # Case 2: refractive surface (material change)
            elif current_material != previous_material:
                if surf_type == "STANDARD":
                    created = Spheric(
                        c=c, r=r, d_next=d_next, mat2=current_material
                    )
                    status = " (stop surface)" if is_stop else ""
                    print(
                        f"   Created spherical surface{status}: Spheric(c={c:.6f}, r={r:.4f}, d={d:.4f}, mat2='{current_material}')"
                    )

                elif surf_type == "ASPHERIC":
                    k = surf.get("conic", 0.0)
                    asph_coeffs = surf.get("asph_coeffs", {})

                    # CODE V aspheric coefficient mapping (shift forward by one position):
                    # A → ai[1] (2nd term, ρ²)
                    # B → ai[2] (4th term, ρ⁴)
                    # C → ai[3] (6th term, ρ⁶)
                    # D → ai[4] (8th term, ρ⁸)
                    # E → ai[5] (10th term, ρ¹⁰)
                    # F → ai[6] (12th term, ρ¹²)
                    # G → ai[7] (14th term, ρ¹⁴)
                    # H → ai[8] (16th term, ρ¹⁶)
                    # I → ai[9] (18th term, ρ¹⁸)

                    # Initialize ai array (10 elements)
                    ai = [0.0] * 10
                    ai[0] = 0.0  # ρ⁰ term (unused)
                    ai[1] = asph_coeffs.get("A", 0.0)  # ρ²
                    ai[2] = asph_coeffs.get("B", 0.0)  # ρ⁴
                    ai[3] = asph_coeffs.get("C", 0.0)  # ρ⁶
                    ai[4] = asph_coeffs.get("D", 0.0)  # ρ⁸
                    ai[5] = asph_coeffs.get("E", 0.0)  # ρ¹⁰
                    ai[6] = asph_coeffs.get("F", 0.0)  # ρ¹²
                    ai[7] = asph_coeffs.get("G", 0.0)  # ρ¹⁴
                    ai[8] = asph_coeffs.get("H", 0.0)  # ρ¹⁶
                    ai[9] = asph_coeffs.get("I", 0.0)  # ρ¹⁸

                    created = Aspheric(
                        c=c,
                        r=r,
                        d_next=d_next,
                        ai=ai,
                        k=k,
                        mat2=current_material,
                    )
                    status = " (stop surface)" if is_stop else ""
                    print(
                        f"   Created aspheric surface{status}: Aspheric(c={c:.6f}, r={r:.4f}, d={d:.4f}, k={k}, mat2='{current_material}')"
                    )
                    if any(
                        ai[1:]
                    ):  # If there are non-zero higher-order terms (starting from ai[1])
                        print(
                            f"      Aspheric coefficients: A={ai[1]:.2e}, B={ai[2]:.2e}, C={ai[3]:.2e}, D={ai[4]:.2e}"
                        )

            else:
                print(f"   Skipped (same material on both sides and no stop flag)")

        except Exception as e:
            print(f"   Failed to create surface: {e}")
            import traceback

            traceback.print_exc()

        if created is not None:
            if pending_gap != 0.0 and self.surfaces:
                self.surfaces[-1].d_next += pending_gap
            pending_gap = 0.0
            self.surfaces.append(created)
        else:
            pending_gap += d_next

        # Key: accumulate distance at the end of the loop
        d += d_next
        print(f"   After accumulation: d={d:.4f}")
        previous_material = current_material

    if pending_gap != 0.0 and self.surfaces:
        self.surfaces[-1].d_next += pending_gap

    print(f"\n{'=' * 60}")
    print(f"   Done! Created {len(self.surfaces)} objects")
    print(f"   d_sensor={self.d_sensor:.4f}")
    print(f"   r_sensor={self.r_sensor:.4f}")
    print(f"   hfov={self.hfov:.4f}°")
    print(f"{'=' * 60}\n")

    return self

write_lens_seq

write_lens_seq(filename='./test.seq')

Write the lens to a Code V .seq sequential file.

Exports refractive surfaces (spheric and aspheric; pure apertures are skipped), materials, field angles (YAN at 0, 0.707, and 0.99 of the effective half-FoV, in degrees), entrance pupil diameter, and the image surface in Code V format.

Parameters:

Name	Type	Description	Default
filename	str	Output file path. Defaults to './test.seq'.	'./test.seq'

Returns:

Name	Type	Description
self	GeoLens	The updated lens (for chaining).

Source code in deeplens-src/deeplens/geolens_pkg/io.py

def write_lens_seq(self, filename="./test.seq"):
    """Write the lens to a Code V .seq sequential file.

    Exports refractive surfaces (spheric and aspheric; pure apertures are
    skipped), materials, field angles (YAN at 0, 0.707, and 0.99 of the
    effective half-FoV, in degrees), entrance pupil diameter, and the
    image surface in Code V format.

    Args:
        filename (str, optional): Output file path. Defaults to './test.seq'.

    Returns:
        self (GeoLens): The updated lens (for chaining).
    """

    import datetime

    current_date = datetime.datetime.now().strftime("%d-%b-%Y")

    head_str = f"""RDM;LEN       "VERSION: 2023.03       LENS VERSION: 89       Creation Date:  {current_date}"
TITLE 'Lens Design'
EPD   {self.enpd}
DIM   M
WL    650.0 550.0 480.0
REF   2
WTW   1 2 1
INI   '   '
XAN   0.0 0.0 0.0
YAN   0.0  {0.707 * self.rfov_eff * 57.3} {0.99 * self.rfov_eff * 57.3}
WTF   1.0 1.0 1.0
VUY   0.0 0.0 0.0
VLY   0.0 0.0 0.0
DOR   1.15 1.05
SO    0.0 0.1e14
"""

    lens_seq_str = head_str
    previous_material = "air"

    for i, surf in enumerate(self.surfaces):
        d_next = float(surf.d_next)

        current_material = getattr(surf, "mat2", "air")

        if current_material is None or current_material == "air":
            material_str = ""
            material_name = "air"
        elif isinstance(current_material, str):
            material_str = f" {current_material.upper()}"
            material_name = current_material
        else:
            material_name = getattr(current_material, "name", str(current_material))
            material_str = f" {material_name.upper()}"

        is_aperture = surf.__class__.__name__ == "Aperture"

        if is_aperture:
            surf_str = f"S     0.0 {d_next}\n"
            surf_str += "  CCY 0; THC 0\n"
            surf_str += "  STO\n"
            surf_str += f"  CIR {surf.r}\n"
            lens_seq_str += surf_str
            previous_material = "air"
            continue

        is_aspheric = surf.__class__.__name__ == "Aspheric"
        is_stop_surface = getattr(surf, "is_stop", False)

        if is_aspheric:
            if abs(surf.c) > 1e-10:
                radius = 1.0 / surf.c
            else:
                radius = 0.0

            k = surf.k if hasattr(surf, "k") else 0.0
            ai = surf.ai if hasattr(surf, "ai") else [0.0] * 10

            surf_str = f"S     {radius} {d_next}{material_str}\n"
            surf_str += f"  CCY 0; THC 0\n"
            surf_str += f"  CIR {surf.r}\n"
            if is_stop_surface:
                surf_str += f"  STO\n"
            surf_str += f"  ASP\n"
            surf_str += f"  K   {k}\n"

            if len(ai) > 4 and any(ai[1:5]):
                surf_str += f"  A   {ai[1]:.16e}; B {ai[2]:.16e}; C&\n"
                surf_str += f"   {ai[3]:.16e}; D {ai[4]:.16e}\n"

            if len(ai) > 8 and any(ai[5:9]):
                surf_str += f"  E   {ai[5]:.16e}; F {ai[6]:.16e}; G {ai[7]:.16e}; H {ai[8]:.16e}\n"

        else:
            if abs(surf.c) > 1e-10:
                radius = 1.0 / surf.c
            else:
                radius = 0.0

            surf_str = f"S     {radius} {d_next}{material_str}\n"
            surf_str += f"  CCY 0; THC 0\n"

            if is_stop_surface:
                surf_str += f"  STO\n"

            surf_str += f"  CIR {surf.r}\n"

        lens_seq_str += surf_str
        previous_material = material_name

    sensor_str = f"SI    0.0 0.0\n"
    sensor_str += f"  CIR {self.r_sensor}\n"
    lens_seq_str += sensor_str
    lens_seq_str += "GO \n"

    with open(filename, "w") as f:
        f.write(lens_seq_str)

    print(f"Lens written to CODE V file: {filename}")
    return self

read_lens_json

read_lens_json(filename='./test.json')

Read the lens from a DeepLens native JSON file.

Loads the surface list, sensor geometry, entrance pupil, and lens info, rebuilding each surface from its type field via init_from_dict. Absolute surface positions and self.d_sensor are derived from the per-surface d_next prefix sums. Sets self.r_sensor [mm], self.enpd, and self.float_enpd, then configures the sensor resolution from sensor_res (default 2000 x 2000).

Parameters:

Name	Type	Description	Default
filename	str	Path to the JSON lens file. Defaults to './test.json'.	'./test.json'

Raises:

Type	Description
Exception	If a surface type is not implemented in this loader.

Note

After loading, the lens is moved to self.device.

Source code in deeplens-src/deeplens/geolens_pkg/io.py

def read_lens_json(self, filename="./test.json"):
    """Read the lens from a DeepLens native JSON file.

    Loads the surface list, sensor geometry, entrance pupil, and lens info,
    rebuilding each surface from its `type` field via `init_from_dict`.
    Absolute surface positions and `self.d_sensor` are derived from the
    per-surface `d_next` prefix sums.
    Sets `self.r_sensor` [mm], `self.enpd`, and `self.float_enpd`, then
    configures the sensor resolution from `sensor_res` (default
    2000 x 2000).

    Args:
        filename (str, optional): Path to the JSON lens file. Defaults to './test.json'.

    Raises:
        Exception: If a surface `type` is not implemented in this loader.

    Note:
        After loading, the lens is moved to `self.device`.
    """
    self.surfaces = []
    with open(filename, "r") as f:
        data = json.load(f)
        for idx, surf_dict in enumerate(data["surfaces"]):
            surf_dict = dict(surf_dict)
            # Absolute `d`/`(d)` values are informational compatibility
            # metadata only. Runtime geometry comes solely from d_next.
            surf_dict["surf_idx"] = idx
            surf_dict["d_next"] = torch.as_tensor(
                surf_dict["d_next"], dtype=self.dtype
            )

            if surf_dict["type"] == "Aperture":
                s = Aperture.init_from_dict(surf_dict)

            elif surf_dict["type"] == "Aspheric":
                s = Aspheric.init_from_dict(surf_dict)

            elif surf_dict["type"] == "Cubic":
                s = Cubic.init_from_dict(surf_dict)

            # elif surf_dict["type"] == "GaussianRBF":
            #     s = GaussianRBF.init_from_dict(surf_dict)

            # elif surf_dict["type"] == "NURBS":
            #     s = NURBS.init_from_dict(surf_dict)

            elif surf_dict["type"] == "Phase":
                s = Phase.init_from_dict(surf_dict)

            elif surf_dict["type"] == "Binary2Phase":
                s = Binary2Phase.init_from_dict(surf_dict)

            elif surf_dict["type"] == "Plane":
                s = Plane.init_from_dict(surf_dict)

            # elif surf_dict["type"] == "PolyEven":
            #     s = PolyEven.init_from_dict(surf_dict)

            elif surf_dict["type"] == "Stop":
                s = Aperture.init_from_dict(surf_dict)

            elif surf_dict["type"] == "Spheric":
                s = Spheric.init_from_dict(surf_dict)

            elif surf_dict["type"] == "ThinLens":
                s = ThinLens.init_from_dict(surf_dict)

            else:
                raise Exception(
                    f"Surface type {surf_dict['type']} is not implemented in GeoLens.read_lens_json()."
                )

            s.is_aperture = bool(surf_dict.get("is_aperture", False))
            self.surfaces.append(s)

    self.lens_info = data.get("info", "None")
    primary_wvln = torch.as_tensor(
        data.get("primary_wvln", self.primary_wvln), dtype=torch.float64
    )
    wvln_rgb = torch.as_tensor(
        data.get("wvln_rgb", self.wvln_rgb), dtype=torch.float64
    )
    obj_depth = torch.as_tensor(
        data.get("obj_depth", self.obj_depth), dtype=torch.float64
    )
    if primary_wvln.numel() != 1 or not (0.1 < primary_wvln.item() < 10.0):
        raise ValueError(
            "primary_wvln must be a scalar satisfying 0.1 < wavelength < 10 µm."
        )
    if wvln_rgb.numel() != 3 or not bool(
        ((wvln_rgb > 0.1) & (wvln_rgb < 10.0)).all().item()
    ):
        raise ValueError("wvln_rgb must contain three wavelengths in (0.1, 10) µm.")
    if obj_depth.numel() != 1 or not (
        math.isfinite(obj_depth.item()) and obj_depth.item() < 0.0
    ):
        raise ValueError("obj_depth must be a finite negative distance [mm].")
    self.primary_wvln = primary_wvln.item()
    self.wvln_rgb = wvln_rgb.tolist()
    self.obj_depth = obj_depth.item()
    self.enpd = data.get("enpd", None)
    self.float_enpd = True if self.enpd is None else False
    self.float_foclen = False
    self.float_rfov = False
    self._set_sensor_from_data(data)

write_lens_json

write_lens_json(filename='./test.json')

Write the lens to a DeepLens native JSON file.

Saves lens info, focal length [mm], F-number, entrance pupil diameter, sensor radius/size [mm] and resolution, and all surfaces (each via surf_dict) with their per-surface spacing d_next [mm]. Numeric values are rounded to 4 decimal places.

Parameters:

Name	Type	Description	Default
filename	str	Path for the output JSON file. Defaults to './test.json'.	'./test.json'

Source code in deeplens-src/deeplens/geolens_pkg/io.py

def write_lens_json(self, filename="./test.json"):
    """Write the lens to a DeepLens native JSON file.

    Saves lens info, focal length [mm], F-number, entrance pupil diameter,
    sensor radius/size [mm] and resolution, and all surfaces (each via
    `surf_dict`) with their per-surface spacing `d_next` [mm]. Numeric
    values are rounded to 4 decimal places.

    Args:
        filename (str, optional): Path for the output JSON file. Defaults to './test.json'.
    """
    data = {}
    data["info"] = self.lens_info if hasattr(self, "lens_info") else "None"
    data["foclen"] = self.foclen
    data["fnum"] = self.fnum
    data["primary_wvln"] = self.primary_wvln
    data["wvln_rgb"] = list(self.wvln_rgb)
    data["obj_depth"] = self.obj_depth
    if self.float_enpd is False:
        data["enpd"] = self.enpd
    data["r_sensor"] = self.r_sensor
    data["(d_sensor)"] = self.d_sensor.item()
    data["(sensor_size)"] = list(self.sensor_size)
    data["sensor_res"] = list(self.sensor_res)
    data["surfaces"] = []
    for i, s in enumerate(self.surfaces):
        surf_dict = {"idx": i}
        surf_dict.update(s.surf_dict())
        if getattr(s, "is_aperture", False):
            surf_dict["is_aperture"] = True
        surf_dict["d_next"] = s.d_next.item()

        data["surfaces"].append(surf_dict)

    with open(filename, "w") as f:
        json.dump(data, f, indent=4)
    print(f"Lens written to {filename}")

deeplens.geolens_pkg.vis.GeoLensVis

Mixin providing 2D lens layout and ray visualization for GeoLens.

Generates publication-quality cross-section plots showing lens surfaces and traced ray bundles in either the meridional or sagittal plane.

This class is not instantiated directly; it is mixed into GeoLens.

sample_parallel_2D

sample_parallel_2D(
    fov=0.0, num_rays=7, wvln=None, plane="meridional", entrance_pupil=True, depth=0.0
)

Sample parallel rays (2D) in object space.

Used for (1) drawing lens setup, (2) 2D geometric optics calculation, for example, refocusing to infinity

Parameters:

Name	Type	Description	Default
fov	float	Incident angle [degree]. Defaults to 0.0.	0.0
num_rays	int	Number of rays. Defaults to 7.	7
wvln	float or None	Ray wavelength [µm]. When None, falls back to self.primary_wvln. Defaults to None.	None
plane	str	Sampling plane, "meridional" (y-z plane) or "sagittal" (x-z plane). Defaults to "meridional".	'meridional'
entrance_pupil	bool	If True, sample on the entrance pupil; otherwise sample on the first surface aperture. Defaults to True.	True
depth	float	Sampling depth [mm] to propagate rays to. Defaults to 0.0.	0.0

Returns:

Name	Type	Description
rays	Ray	Sampled rays with origin/direction tensors of shape [num_rays, 3].

Source code in deeplens-src/deeplens/geolens_pkg/vis.py

@torch.no_grad()
def sample_parallel_2D(
    self,
    fov=0.0,
    num_rays=7,
    wvln=None,
    plane="meridional",
    entrance_pupil=True,
    depth=0.0,
):
    """Sample parallel rays (2D) in object space.

    Used for (1) drawing lens setup, (2) 2D geometric optics calculation, for example, refocusing to infinity

    Args:
        fov (float, optional): Incident angle [degree]. Defaults to 0.0.
        num_rays (int, optional): Number of rays. Defaults to 7.
        wvln (float or None, optional): Ray wavelength [µm]. When None,
            falls back to `self.primary_wvln`. Defaults to None.
        plane (str, optional): Sampling plane, "meridional" (y-z plane) or
            "sagittal" (x-z plane). Defaults to "meridional".
        entrance_pupil (bool, optional): If True, sample on the entrance pupil;
            otherwise sample on the first surface aperture. Defaults to True.
        depth (float, optional): Sampling depth [mm] to propagate rays to.
            Defaults to 0.0.

    Returns:
        rays (Ray): Sampled rays with origin/direction tensors of shape [num_rays, 3].
    """
    wvln = self.primary_wvln if wvln is None else wvln
    # Sample points on the pupil
    if entrance_pupil:
        pupilz, pupilr = self.get_entrance_pupil()
    else:
        pupilz, pupilr = self.surf_d(0).item(), self.surfaces[0].r

    # Sample ray origins, shape [num_rays, 3]
    if plane == "sagittal":
        ray_o = torch.stack(
            (
                torch.linspace(-pupilr, pupilr, num_rays) * 0.99,
                torch.full((num_rays,), 0),
                torch.full((num_rays,), pupilz),
            ),
            axis=-1,
        )
    elif plane == "meridional":
        ray_o = torch.stack(
            (
                torch.full((num_rays,), 0),
                torch.linspace(-pupilr, pupilr, num_rays) * 0.99,
                torch.full((num_rays,), pupilz),
            ),
            axis=-1,
        )
    else:
        raise ValueError(f"Invalid plane: {plane}")

    # Sample ray directions, shape [num_rays, 3]
    if plane == "sagittal":
        ray_d = torch.stack(
            (
                torch.full((num_rays,), float(np.sin(np.deg2rad(fov)))),
                torch.zeros((num_rays,)),
                torch.full((num_rays,), float(np.cos(np.deg2rad(fov)))),
            ),
            axis=-1,
        )
    elif plane == "meridional":
        ray_d = torch.stack(
            (
                torch.zeros((num_rays,)),
                torch.full((num_rays,), float(np.sin(np.deg2rad(fov)))),
                torch.full((num_rays,), float(np.cos(np.deg2rad(fov)))),
            ),
            axis=-1,
        )
    else:
        raise ValueError(f"Invalid plane: {plane}")

    # Form rays and propagate to the target depth
    rays = Ray(ray_o, ray_d, wvln, device=self.device)
    rays.prop_to(depth)
    return rays

sample_point_source_2D

sample_point_source_2D(fov=0.0, depth=None, num_rays=7, wvln=None, entrance_pupil=True)

Sample point source rays (2D) in object space.

Used for (1) drawing lens setup.

Parameters:

Name	Type	Description	Default
fov	float	Incident angle [degree]. Defaults to 0.0.	0.0
depth	float or None	Object-plane depth [mm]. When None, falls back to self.obj_depth. Defaults to None.	None
num_rays	int	Number of rays. Defaults to 7.	7
wvln	float or None	Ray wavelength [µm]. When None, falls back to self.primary_wvln. Defaults to None.	None
entrance_pupil	bool	If True, aim rays at the entrance pupil; otherwise aim at the first surface aperture. Defaults to True.	True

Returns:

Name	Type	Description
ray	Ray	Sampled rays with origin/direction tensors of shape [num_rays, 3].

Source code in deeplens-src/deeplens/geolens_pkg/vis.py

@torch.no_grad()
def sample_point_source_2D(
    self,
    fov=0.0,
    depth=None,
    num_rays=7,
    wvln=None,
    entrance_pupil=True,
):
    """Sample point source rays (2D) in object space.

    Used for (1) drawing lens setup.

    Args:
        fov (float, optional): Incident angle [degree]. Defaults to 0.0.
        depth (float or None, optional): Object-plane depth [mm]. When None,
            falls back to `self.obj_depth`. Defaults to None.
        num_rays (int, optional): Number of rays. Defaults to 7.
        wvln (float or None, optional): Ray wavelength [µm]. When None,
            falls back to `self.primary_wvln`. Defaults to None.
        entrance_pupil (bool, optional): If True, aim rays at the entrance pupil;
            otherwise aim at the first surface aperture. Defaults to True.

    Returns:
        ray (Ray): Sampled rays with origin/direction tensors of shape [num_rays, 3].
    """
    wvln = self.primary_wvln if wvln is None else wvln
    depth = self.obj_depth if depth is None else depth
    # Sample point on the object plane
    ray_o = torch.tensor([depth * float(np.tan(np.deg2rad(fov))), 0.0, depth])
    ray_o = ray_o.unsqueeze(0).repeat(num_rays, 1)

    # Sample points (second point) on the pupil
    if entrance_pupil:
        pupilz, pupilr = self.calc_entrance_pupil_rayaiming()
    else:
        pupilz, pupilr = self.surf_d(0).item(), self.surfaces[0].r

    x2 = torch.linspace(-pupilr, pupilr, num_rays) * 0.99
    y2 = torch.zeros_like(x2)
    z2 = torch.full_like(x2, pupilz)
    ray_o2 = torch.stack((x2, y2, z2), axis=1)

    # Form the rays
    ray_d = ray_o2 - ray_o
    ray = Ray(ray_o, ray_d, wvln, device=self.device)

    # Propagate rays to the sampling depth
    ray.prop_to(depth)
    return ray

draw_layout

draw_layout(
    filename,
    depth=float("inf"),
    zmx_format=True,
    multi_plot=False,
    lens_title=None,
    show=False,
    return_fig=False,
)

Plot 2D lens layout with ray tracing.

The title is auto-generated when lens_title is None: it includes focal length, F-number, FoV, IMGH, RGB wavelengths, and a second line with per-FoV RMS spot radii from analysis_spot().

Parameters:

Name	Type	Description	Default
filename	str	Output filename.	required
depth	float	Object distance for ray tracing [mm]. Use float('inf') for collimated input. Defaults to float('inf').	float('inf')
zmx_format	bool	If True, draw surfaces in Zemax style. Defaults to True.	True
multi_plot	bool	If True, create one sub-plot per wavelength. Defaults to False.	False
lens_title	str or None	Title string. If None, auto-generated. Defaults to None.	None
show	bool	If True, display the figure interactively instead of saving. Defaults to False.	False
return_fig	bool	If True, return the axes and figure without saving or closing them (for overlay drawing by callers such as create_barrier). Defaults to False.	False

Returns:

Name	Type	Description
result	tuple or None	When return_fig is True, returns (ax, fig) matplotlib axes and figure objects; otherwise returns None.

Source code in deeplens-src/deeplens/geolens_pkg/vis.py

def draw_layout(
    self,
    filename,
    depth=float("inf"),
    zmx_format=True,
    multi_plot=False,
    lens_title=None,
    show=False,
    return_fig=False,
):
    """Plot 2D lens layout with ray tracing.

    The title is auto-generated when ``lens_title`` is None: it includes
    focal length, F-number, FoV, IMGH, RGB wavelengths, and a second line
    with per-FoV RMS spot radii from ``analysis_spot()``.

    Args:
        filename (str): Output filename.
        depth (float, optional): Object distance for ray tracing [mm]. Use ``float('inf')``
            for collimated input. Defaults to ``float('inf')``.
        zmx_format (bool, optional): If True, draw surfaces in Zemax style. Defaults to True.
        multi_plot (bool, optional): If True, create one sub-plot per wavelength.
            Defaults to False.
        lens_title (str or None, optional): Title string. If None, auto-generated. Defaults to None.
        show (bool, optional): If True, display the figure interactively instead of
            saving. Defaults to False.
        return_fig (bool, optional): If True, return the axes and figure without
            saving or closing them (for overlay drawing by callers such as
            `create_barrier`). Defaults to False.

    Returns:
        result (tuple or None): When `return_fig` is True, returns (ax, fig)
            matplotlib axes and figure objects; otherwise returns None.
    """
    num_rays = 11
    num_views = 3

    # Lens title
    if lens_title is None:
        eff_foclen = round(self.foclen, 2)
        fov_deg = round(2 * self.rfov * 180 / torch.pi, 1)
        imgh = round(self.r_sensor, 1)
        wvl_nm = [int(round(w * 1000)) for w in self.wvln_rgb]  # µm → nm

        if self.aper_idx is not None:
            _, pupil_r = self.calc_entrance_pupil_rayaiming()
            fnum = round(eff_foclen / pupil_r / 2, 2)
            line1 = (
                f"FocLen{eff_foclen}mm - F/{fnum} - FoV{fov_deg} - "
                f"IMGH{imgh}mm - RGB({wvl_nm[0]}/{wvl_nm[1]}/{wvl_nm[2]}nm)"
            )
        else:
            line1 = (
                f"FocLen{eff_foclen}mm - FoV{fov_deg} - "
                f"IMGH{imgh}mm - RGB({wvl_nm[0]}/{wvl_nm[1]}/{wvl_nm[2]}nm)"
            )

        spot = self.analysis_spot(num_field=3)
        rms0 = spot["fov0.0"]["rms"]
        rms5 = spot["fov0.5"]["rms"]
        rms10 = spot["fov1.0"]["rms"]
        line2 = f"RMS spot: 0.0FoV={rms0:.2f}\u03bcm  0.5FoV={rms5:.2f}\u03bcm  1.0FoV={rms10:.2f}\u03bcm"
        lens_title = f"{line1}\n{line2}"

    # Draw lens layout
    colors_list = ["#CC0000", "#006600", "#0066CC"]
    rfov_deg = float(np.rad2deg(self.rfov))
    fov_ls = np.linspace(0, rfov_deg * 0.99, num=num_views)

    if not multi_plot:
        ax, fig = self.draw_lens_2d(zmx_format=zmx_format)
        fig.suptitle(lens_title, fontsize=10, fontfamily="Nimbus Sans")
        for i, fov in enumerate(fov_ls):
            # Sample rays, shape (num_rays, 3)
            if depth == float("inf"):
                ray = self.sample_parallel_2D(
                    fov=fov,
                    wvln=self.wvln_rgb[2 - i],
                    num_rays=num_rays,
                    depth=-1.0,
                    plane="sagittal",
                )
            else:
                ray = self.sample_point_source_2D(
                    fov=fov,
                    depth=depth,
                    num_rays=num_rays,
                    wvln=self.wvln_rgb[2 - i],
                )
                ray.prop_to(-1.0)

            # Trace rays to sensor and plot ray paths
            _, ray_o_record = self.trace2sensor(ray=ray, record=True)
            ax, fig = self.draw_ray_2d(
                ray_o_record, ax=ax, fig=fig, color=colors_list[i]
            )

        ax.axis("off")

    else:
        fig, axs = plt.subplots(1, 3, figsize=(15, 5))
        fig.suptitle(lens_title, fontsize=10, fontfamily="Nimbus Sans")
        for i, wvln in enumerate(self.wvln_rgb):
            ax = axs[i]
            ax, fig = self.draw_lens_2d(ax=ax, fig=fig, zmx_format=zmx_format)
            for fov in fov_ls:
                # Sample rays, shape (num_rays, 3)
                if depth == float("inf"):
                    ray = self.sample_parallel_2D(
                        fov=fov,
                        num_rays=num_rays,
                        wvln=wvln,
                        plane="sagittal",
                    )
                else:
                    ray = self.sample_point_source_2D(
                        fov=fov,
                        depth=depth,
                        num_rays=num_rays,
                        wvln=wvln,
                    )

                # Trace rays to sensor and plot ray paths
                ray_out, ray_o_record = self.trace2sensor(ray=ray, record=True)
                ax, fig = self.draw_ray_2d(
                    ray_o_record, ax=ax, fig=fig, color=colors_list[i]
                )
                ax.axis("off")

    # Let an internal caller (e.g. create_barrier) keep drawing on the same
    # axes instead of saving and closing the figure here.
    if return_fig:
        return ax, fig

    if show:
        fig.show()
    else:
        fig.savefig(filename, format="png", dpi=300)
        # Close the specific figure to avoid leaking it.
        plt.close(fig)

draw_lens_2d

draw_lens_2d(
    ax=None, fig=None, color="k", linestyle="-", zmx_format=False, fix_bound=False
)

Draw lens cross-section layout in a 2D plot.

Renders each surface profile, connects lens elements with edge lines, and draws the sensor plane.

Parameters:

Name	Type	Description	Default
ax	Axes	Existing axes to draw on. If None, creates a new figure. Defaults to None.	None
fig	Figure	Existing figure. Defaults to None.	None
color	str	Line colour for lens outlines. Defaults to 'k'.	'k'
linestyle	str	Line style. Defaults to '-'.	'-'
zmx_format	bool	If True, draw stepped edge connections matching Zemax layout style. Defaults to False.	False
fix_bound	bool	If True, use fixed axis limits [-1,7]x[-4,4]. Defaults to False.	False

Returns:

Name	Type	Description
ax	Axes	The axes with the lens layout drawn.
fig	Figure	The figure.

Source code in deeplens-src/deeplens/geolens_pkg/vis.py

def draw_lens_2d(
    self,
    ax=None,
    fig=None,
    color="k",
    linestyle="-",
    zmx_format=False,
    fix_bound=False,
):
    """Draw lens cross-section layout in a 2D plot.

    Renders each surface profile, connects lens elements with edge lines,
    and draws the sensor plane.

    Args:
        ax (matplotlib.axes.Axes, optional): Existing axes to draw on. If None,
            creates a new figure. Defaults to None.
        fig (matplotlib.figure.Figure, optional): Existing figure. Defaults to None.
        color (str, optional): Line colour for lens outlines. Defaults to 'k'.
        linestyle (str, optional): Line style. Defaults to '-'.
        zmx_format (bool, optional): If True, draw stepped edge connections
            matching Zemax layout style. Defaults to False.
        fix_bound (bool, optional): If True, use fixed axis limits [-1,7]x[-4,4].
            Defaults to False.

    Returns:
        ax (matplotlib.axes.Axes): The axes with the lens layout drawn.
        fig (matplotlib.figure.Figure): The figure.
    """
    # If no ax is given, generate a new one.
    if ax is None and fig is None:
        # fig, ax = plt.subplots(figsize=(6, 6))
        fig, ax = plt.subplots()

    # Draw lens surfaces
    for i, s in enumerate(self.surfaces):
        s.draw_widget(ax, d=self.surf_d(i).item())

    # Connect two surfaces
    for i in range(len(self.surfaces)):
        if self.surfaces[i].mat2.n > 1.1:
            s_prev = self.surfaces[i]
            s = self.surfaces[i + 1]

            r_prev = float(s_prev.draw_r())
            r = float(s.draw_r())
            sag_prev = s_prev.surface_with_offset(
                r_prev, 0.0, valid_check=False, d=self.surf_d(i)
            ).item()
            sag = s.surface_with_offset(
                r, 0.0, valid_check=False, d=self.surf_d(i + 1)
            ).item()

            if r_prev >= r:
                # Front surface wider: go axially forward at r_prev, then step radially inward
                z = np.array([sag_prev, sag, sag])
                x = np.array([r_prev, r_prev, r])
            else:
                # Rear surface wider: step radially outward at z_prev, then go axially forward
                z = np.array([sag_prev, sag_prev, sag])
                x = np.array([r_prev, r, r])

            if not zmx_format:
                # In non-zmx mode use a direct diagonal between the two outer edges
                z = np.array([z[0], z[-1]])
                x = np.array([x[0], x[-1]])

            ax.plot(z, -x, color, linewidth=0.75)
            ax.plot(z, x, color, linewidth=0.75)
            s_prev = s

    # Draw sensor
    ax.plot(
        [self.d_sensor.item(), self.d_sensor.item()],
        [-self.r_sensor, self.r_sensor],
        color,
    )

    # Set figure size
    if fix_bound:
        ax.set_aspect("equal")
        ax.set_xlim(-1, 7)
        ax.set_ylim(-4, 4)
    else:
        ax.set_aspect("equal", adjustable="datalim", anchor="C")
        ax.minorticks_on()
        ax.set_xlim(-0.5, 7.5)
        ax.set_ylim(-4, 4)
        ax.autoscale()

    return ax, fig

draw_ray_2d

draw_ray_2d(ray_o_record, ax, fig, color='b')

Plot ray paths onto an existing 2D layout.

Each recorded ray origin is a [num_rays, 3] (or [num_view, num_rays, 3]) tensor; stacking them yields [num_view, num_rays, num_path, 3] where the last axis holds (x, y, z) in [mm]. The z (axial) and x (radial) components are drawn as polylines.

Parameters:

Name	Type	Description	Default
ray_o_record	list	List of ray-origin tensors, one per traced surface, each of shape [num_rays, 3] or [num_view, num_rays, 3].	required
ax	Axes	Matplotlib axes to draw on.	required
fig	Figure	Matplotlib figure.	required
color	str	Line colour for the ray paths. Defaults to 'b'.	'b'

Returns:

Name	Type	Description
ax	Axes	The axes with ray paths drawn.
fig	Figure	The figure.

Source code in deeplens-src/deeplens/geolens_pkg/vis.py

def draw_ray_2d(self, ray_o_record, ax, fig, color="b"):
    """Plot ray paths onto an existing 2D layout.

    Each recorded ray origin is a [num_rays, 3] (or [num_view, num_rays, 3])
    tensor; stacking them yields [num_view, num_rays, num_path, 3] where the
    last axis holds (x, y, z) in [mm]. The z (axial) and x (radial) components
    are drawn as polylines.

    Args:
        ray_o_record (list): List of ray-origin tensors, one per traced surface,
            each of shape [num_rays, 3] or [num_view, num_rays, 3].
        ax (matplotlib.axes.Axes): Matplotlib axes to draw on.
        fig (matplotlib.figure.Figure): Matplotlib figure.
        color (str, optional): Line colour for the ray paths. Defaults to 'b'.

    Returns:
        ax (matplotlib.axes.Axes): The axes with ray paths drawn.
        fig (matplotlib.figure.Figure): The figure.
    """
    # shape (num_view, num_rays, num_path, 2)
    ray_o_record = torch.stack(ray_o_record, dim=-2).cpu().numpy()
    if ray_o_record.ndim == 3:
        ray_o_record = ray_o_record[None, ...]

    for idx_view in range(ray_o_record.shape[0]):
        for idx_ray in range(ray_o_record.shape[1]):
            ax.plot(
                ray_o_record[idx_view, idx_ray, :, 2],
                ray_o_record[idx_view, idx_ray, :, 0],
                color,
                linewidth=0.8,
            )

            # ax.scatter(
            #     ray_o_record[idx_view, idx_ray, :, 2],
            #     ray_o_record[idx_view, idx_ray, :, 0],
            #     "b",
            #     marker="x",
            # )

    return ax, fig

create_barrier

create_barrier(filename, barrier_thickness=1.0, ring_height=0.5, ring_size=1.0)

Draw a lens barrel (barrier) overlay on the 2D lens layout and save it.

Computes barrier segments spanning each air gap (extending to the midpoint of the following air space, or to the sensor for the last segment), overlays them in green on the layout from draw_layout, and saves the figure as a PNG.

Parameters:

Name	Type	Description	Default
filename	str	Path to save the output PNG figure.	required
barrier_thickness	float	Barrier thickness [mm]. Defaults to 1.0.	1.0
ring_height	float	Annular ring height [mm]. Currently unused (ring drawing is not implemented). Defaults to 0.5.	0.5
ring_size	float	Annular ring size [mm]. Currently unused (ring drawing is not implemented). Defaults to 1.0.	1.0

Source code in deeplens-src/deeplens/geolens_pkg/vis.py

def create_barrier(
    self, filename, barrier_thickness=1.0, ring_height=0.5, ring_size=1.0
):
    """Draw a lens barrel (barrier) overlay on the 2D lens layout and save it.

    Computes barrier segments spanning each air gap (extending to the midpoint
    of the following air space, or to the sensor for the last segment), overlays
    them in green on the layout from `draw_layout`, and saves the figure as a PNG.

    Args:
        filename (str): Path to save the output PNG figure.
        barrier_thickness (float, optional): Barrier thickness [mm]. Defaults to 1.0.
        ring_height (float, optional): Annular ring height [mm]. Currently unused
            (ring drawing is not implemented). Defaults to 0.5.
        ring_size (float, optional): Annular ring size [mm]. Currently unused
            (ring drawing is not implemented). Defaults to 1.0.
    """
    barriers = []
    rings = []

    # Create barriers
    barrier_z = 0.0
    barrier_r = 0.0
    barrier_length = 0.0
    for i in range(len(self.surfaces)):
        barrier_r = max(self.surfaces[i].r, barrier_r)

        if self.surfaces[i].mat2.get_name() != "air":
            # Update the barrier radius
            # barrier_r = max(geolens.surfaces[i].r, barrier_r)
            pass
        else:
            # Extend the barrier till middle of the air space to the next surface
            max_curr_surf_d = self.surf_d(i).item() + max(
                self.surfaces[i].surface_sag(0.0, self.surfaces[i].r), 0.0
            )
            if i < len(self.surfaces) - 1:
                min_next_surf_d = self.surf_d(i + 1).item() + min(
                    self.surfaces[i + 1].surface_sag(0.0, self.surfaces[i + 1].r),
                    0.0,
                )
                extra_space = (min_next_surf_d - max_curr_surf_d) / 2
            else:
                min_next_surf_d = self.d_sensor.item()
                extra_space = min_next_surf_d - max_curr_surf_d

            barrier_length = max_curr_surf_d + extra_space - barrier_z

            # Create a barrier
            barrier = {
                "pos_z": barrier_z,
                "pos_r": barrier_r,
                "length": barrier_length,
                "thickness": barrier_thickness,
            }
            barriers.append(barrier)

            # Reset the barrier parameters
            barrier_z = barrier_length + barrier_z
            barrier_r = 0.0
            barrier_length = 0.0

    # # Create rings
    # for i in range(len(self.surfaces)):
    #     if self.surfaces[i].mat2.get_name() != "air":
    #         ring = {
    #             "pos_z": self.surf_d(i).item(),

    # Plot lens layout (keep the figure open so we can overlay the barrier)
    ax, fig = self.draw_layout(filename, return_fig=True)

    # Plot barrier
    barrier_z_ls = []
    barrier_r_ls = []
    for b in barriers:
        barrier_z_ls.append(b["pos_z"])
        barrier_z_ls.append(b["pos_z"] + b["length"])
        barrier_r_ls.append(b["pos_r"])
        barrier_r_ls.append(b["pos_r"])
    ax.plot(barrier_z_ls, barrier_r_ls, "green", linewidth=1.0)
    ax.plot(barrier_z_ls, [-i for i in barrier_r_ls], "green", linewidth=1.0)

    # Plot rings

    fig.savefig(filename, format="png", dpi=300)
    plt.close()

    pass

deeplens.geolens_pkg.vis3d.GeoLensVis3D

Mixin providing 3D mesh visualization for GeoLens.

Creates lens surface, aperture, barrier, sensor, and ray-path meshes as polygon data and optionally renders them with PyVista. All geometry is expressed in millimetres [mm] and stored as CrossPoly (vertex/face) objects that can be saved to .obj files for external renderers.

This class is not instantiated directly; it is mixed into GeoLens.

create_mesh

create_mesh(mesh_rings: int = 32, mesh_arms: int = 128, is_wrap: bool = False)

Build surface, bridge, and sensor meshes for the whole lens.

Surfaces are grouped into optical elements (split wherever a surface borders air). Adjacent surfaces within an element are joined by bridge face strips; with is_wrap the bridges are projected to form a cylindrical barrel between elements of differing radii.

Parameters:

Name	Type	Description	Default
mesh_rings	int	Number of rings per surface mesh. Defaults to 32.	32
mesh_arms	int	Number of arms per surface mesh. Defaults to 128.	128
is_wrap	bool	Whether to wrap the lens barrel around the elements as a cylinder. Defaults to False.	False

Returns:

Name	Type	Description
surf_meshes_cvt	List[FaceMesh]	Per-surface face meshes.
bridge_meshes	List[List[FaceMesh]]	Per-element lists of bridge face meshes (empty list for single-surface elements).
element_groups	List[List[int]]	Surface-index groups, one per optical element.
sensor_mesh	RectangleMesh	The rectangular sensor mesh.

Source code in deeplens-src/deeplens/geolens_pkg/vis3d.py

def create_mesh(
    self,
    mesh_rings: int = 32,
    mesh_arms: int = 128,
    is_wrap: bool = False,
):
    """Build surface, bridge, and sensor meshes for the whole lens.

    Surfaces are grouped into optical elements (split wherever a surface
    borders air). Adjacent surfaces within an element are joined by bridge
    face strips; with `is_wrap` the bridges are projected to form a
    cylindrical barrel between elements of differing radii.

    Args:
        mesh_rings (int, optional): Number of rings per surface mesh. Defaults to 32.
        mesh_arms (int, optional): Number of arms per surface mesh. Defaults to 128.
        is_wrap (bool, optional): Whether to wrap the lens barrel around the
            elements as a cylinder. Defaults to False.

    Returns:
        surf_meshes_cvt (List[FaceMesh]): Per-surface face meshes.
        bridge_meshes (List[List[FaceMesh]]): Per-element lists of bridge
            face meshes (empty list for single-surface elements).
        element_groups (List[List[int]]): Surface-index groups, one per
            optical element.
        sensor_mesh (RectangleMesh): The rectangular sensor mesh.
    """
    surf_meshes = []
    element_group = []
    element_groups = []
    bridge_meshes = []  # change to nested list for wrap around
    sensor_mesh = None

    # Create the surface meshes
    for i, surf in enumerate(self.surfaces):
        # Create the surface mesh (list of Surface objects)
        surf_meshes.append(
            surf.create_mesh(
                n_rings=mesh_rings,
                n_arms=mesh_arms,
                d=self.surf_d(i).item(),
            )
        )

        # Add the surface to the element group
        element_group.append(i)
        if surf.mat2.name == "air":
            element_groups.append(element_group)
            element_group = []

    # Create the bridge meshes (list of FaceMesh objects)
    for i, pair in enumerate(element_groups):
        if len(pair) == 1:
            bridge_meshes.append([])
            continue
        elif len(pair) == 2:
            a_idx, b_idx = pair
            a = surf_meshes[a_idx]
            b = surf_meshes[b_idx]
            bridge_mesh_group = []
            if not is_wrap:
                bridge_mesh = bridge(a.rim, b.rim)
                bridge_mesh_group.append(bridge_mesh)
            else:
                # create wrap by creating a new rim
                # from projecting the larger rim onto the smaller rim plane
                # assume the elements are always ordered on z-axis
                r_a = self.surfaces[a_idx].r
                r_b = self.surfaces[b_idx].r
                d_rim_a = np.mean(
                    a.rim.vertices[:, 2], keepdims=False
                )  # calc rim mean z
                d_rim_b = np.mean(b.rim.vertices[:, 2], keepdims=False)

                if r_a > r_b:
                    z = line_translate(a.rim, 0, 0, d_rim_b - d_rim_a)
                    bridge_mesh_wrap = bridge(z, b.rim)
                    bridge_mesh = bridge(a.rim, z)
                    bridge_mesh_group.append(bridge_mesh_wrap)
                elif r_a < r_b:
                    z = line_translate(b.rim, 0, 0, d_rim_a - d_rim_b)
                    bridge_mesh_wrap = bridge(a.rim, z)
                    bridge_mesh = bridge(z, b.rim)
                    bridge_mesh_group.append(bridge_mesh_wrap)
                else:
                    bridge_mesh = bridge(a.rim, b.rim)
                bridge_mesh_group.append(bridge_mesh)
            bridge_meshes.append(bridge_mesh_group)

        elif len(pair) == 3:
            a_idx, b_idx, c_idx = pair
            a = surf_meshes[a_idx]
            b = surf_meshes[b_idx]
            c = surf_meshes[c_idx]
            bridge_mesh_group = []
            if not is_wrap:
                bridge_mesh = bridge(a.rim, b.rim)
                bridge_mesh_group.append(bridge_mesh)
                bridge_mesh = bridge(b.rim, c.rim)
                bridge_mesh_group.append(bridge_mesh)
            else:
                # create wrap by creating a new rim
                # from projecting the larger rim onto the smaller rim plane
                # assume the elements are always ordered on z-axis
                r_a = self.surfaces[a_idx].r
                r_b = self.surfaces[b_idx].r
                r_c = self.surfaces[c_idx].r
                d_rim_a = np.mean(
                    a.rim.vertices[:, 2], keepdims=False
                )  # calc rim mean z
                d_rim_b = np.mean(b.rim.vertices[:, 2], keepdims=False)
                d_rim_c = np.mean(c.rim.vertices[:, 2], keepdims=False)

                rim_list = [a.rim, b.rim, c.rim]
                r_list = [r_a, r_b, r_c]
                d_rim_list = [d_rim_a, d_rim_b, d_rim_c]
                idx_wrap = r_list.index(max(r_list))
                r_wrap = r_list[idx_wrap]
                d_rim_wrap = d_rim_list[idx_wrap]

                for i in range(3):
                    if i != idx_wrap and r_list[i] != r_wrap:
                        # substitute the rim with the wrapped rim
                        d_diff = d_rim_list[i] - d_rim_wrap
                        z = line_translate(rim_list[idx_wrap], 0, 0, d_diff)
                        # add the wrap bridge between older rim and wrapped one
                        wrap_mesh = bridge(rim_list[i], z)
                        # update the rim
                        rim_list[i] = z
                        bridge_mesh_group.append(wrap_mesh)
                bridge_mesh = bridge(rim_list[0], rim_list[1])
                bridge_mesh_group.append(bridge_mesh)
                bridge_mesh = bridge(rim_list[1], rim_list[2])
                bridge_mesh_group.append(bridge_mesh)
            bridge_meshes.append(bridge_mesh_group)

        else:
            raise ValueError(f"Invalid bridge group length: {len(pair)}")

    # Create the sensor mesh (RectangleMesh object)
    sensor_d = self.d_sensor.item()
    sensor_r = self.r_sensor
    h, w = sensor_r * 1.4142, sensor_r * 1.4142
    sensor_mesh = RectangleMesh(
        np.array([0, 0, sensor_d]), np.array([1, 0, 0]), np.array([0, 1, 0]), w, h
    )

    # turn surf_meshes to list of FaceMesh
    surf_meshes_cvt = [surf_to_face_mesh(surf) for surf in surf_meshes]
    return surf_meshes_cvt, bridge_meshes, element_groups, sensor_mesh

draw_lens_3d

draw_lens_3d(
    plotter=None,
    save_dir: Optional[str] = None,
    mesh_rings: int = 32,
    mesh_arms: int = 128,
    surface_color: List[float] = [0.06, 0.3, 0.6],
    draw_rays: bool = True,
    fovs: List[float] = [0.0],
    fov_phis: List[float] = [0.0],
    ray_rings: int = 6,
    ray_arms: int = 8,
    is_wrap: bool = False,
)

Render the 3D lens layout (surfaces, sensor, and optional rays) with PyVista.

Parameters:

Name	Type	Description	Default
plotter	Plotter	Existing plotter to draw into. A new one is created when None. Defaults to None.	None
save_dir	str	Directory to save the rendered screenshot lens_layout3d.png. No image is saved when None. Defaults to None.	None
mesh_rings	int	Number of rings per surface mesh. Defaults to 32.	32
mesh_arms	int	Number of arms per surface mesh. Defaults to 128.	128
surface_color	List[float]	RGB surface color, each in [0, 1]. Defaults to [0.06, 0.3, 0.6].	[0.06, 0.3, 0.6]
draw_rays	bool	Whether to trace and draw rays. Defaults to True.	True
fovs	List[float]	Field-of-view angles to sample [degree]. Defaults to [0.0].	[0.0]
fov_phis	List[float]	Field azimuthal angles to sample [degree]. Defaults to [0.0].	[0.0]
ray_rings	int	Number of pupil rings to sample. Defaults to 6.	6
ray_arms	int	Number of pupil arms to sample. Defaults to 8.	8
is_wrap	bool	Whether to wrap the lens barrel as a cylinder. Defaults to False.	False

Returns:

Name	Type	Description
plotter	Plotter	The plotter with all meshes added.

Raises:

Type	Description
ImportError	If PyVista is not installed (imported lazily here).

Note

PyVista is imported lazily only when this method is called.

Source code in deeplens-src/deeplens/geolens_pkg/vis3d.py

def draw_lens_3d(
    self,
    plotter=None,
    save_dir: Optional[str] = None,
    mesh_rings: int = 32,
    mesh_arms: int = 128,
    surface_color: List[float] = [0.06, 0.3, 0.6],
    draw_rays: bool = True,
    fovs: List[float] = [0.0],
    fov_phis: List[float] = [0.0],
    ray_rings: int = 6,
    ray_arms: int = 8,
    is_wrap: bool = False,
):
    """Render the 3D lens layout (surfaces, sensor, and optional rays) with PyVista.

    Args:
        plotter (pyvista.Plotter, optional): Existing plotter to draw into. A
            new one is created when None. Defaults to None.
        save_dir (str, optional): Directory to save the rendered screenshot
            ``lens_layout3d.png``. No image is saved when None. Defaults to None.
        mesh_rings (int, optional): Number of rings per surface mesh. Defaults to 32.
        mesh_arms (int, optional): Number of arms per surface mesh. Defaults to 128.
        surface_color (List[float], optional): RGB surface color, each in [0, 1].
            Defaults to [0.06, 0.3, 0.6].
        draw_rays (bool, optional): Whether to trace and draw rays. Defaults to True.
        fovs (List[float], optional): Field-of-view angles to sample [degree].
            Defaults to [0.0].
        fov_phis (List[float], optional): Field azimuthal angles to sample [degree].
            Defaults to [0.0].
        ray_rings (int, optional): Number of pupil rings to sample. Defaults to 6.
        ray_arms (int, optional): Number of pupil arms to sample. Defaults to 8.
        is_wrap (bool, optional): Whether to wrap the lens barrel as a cylinder.
            Defaults to False.

    Returns:
        plotter (pyvista.Plotter): The plotter with all meshes added.

    Raises:
        ImportError: If PyVista is not installed (imported lazily here).

    Note:
        PyVista is imported lazily only when this method is called.
    """
    # Lazy import of pyvista
    try:
        import pyvista as pv
    except ImportError as e:
        raise ImportError(
            "PyVista is required for 3D GUI rendering. Install with `pip install pyvista`."
        ) from e

    # Create plotter if not provided
    if plotter is None:
        plotter = pv.Plotter()

    surf_color = surface_color
    sensor_color = [0.5, 0.5, 0.5]

    # Create meshes
    surf_meshes, bridge_meshes, _, sensor_mesh = self.create_mesh(
        mesh_rings, mesh_arms, is_wrap
    )

    # Draw meshes
    for surf in surf_meshes:
        if not isinstance(surf, Aperture):
            _draw_mesh_to_plotter(
                plotter, surf, color=surf_color, opacity=0.5, pv=pv
            )

    for bridge_group in bridge_meshes:
        for bridge_mesh in bridge_group:
            _draw_mesh_to_plotter(
                plotter, bridge_mesh, color=surf_color, opacity=0.5, pv=pv
            )

    _draw_mesh_to_plotter(
        plotter, sensor_mesh, color=sensor_color, opacity=1.0, pv=pv
    )

    # Draw rays
    if draw_rays:
        rays_curve = geolens_ray_poly(
            self, fovs, fov_phis, n_rings=ray_rings, n_arms=ray_arms
        )

        rays_poly_list = [curve_list_to_polydata(r) for r in rays_curve]
        rays_poly_fov = [merge(r) for r in rays_poly_list]
        rays_poly_fov = [_wrap_base_poly_to_pyvista(r, pv) for r in rays_poly_fov]
        for r in rays_poly_fov:
            plotter.add_mesh(r)

    # Save images
    if save_dir is not None:
        os.makedirs(save_dir, exist_ok=True)
        plotter.show(screenshot=os.path.join(save_dir, "lens_layout3d.png"))

    return plotter

save_lens_obj

save_lens_obj(
    save_dir: str,
    mesh_rings: int = 64,
    mesh_arms: int = 128,
    save_rays: bool = False,
    fovs: List[float] = [0.0],
    fov_phis: List[float] = [0.0],
    ray_rings: int = 6,
    ray_arms: int = 8,
    is_wrap: bool = False,
    save_elements: bool = True,
)

Save lens geometry, sensor, and optional rays as Wavefront .obj files.

Writes lens.obj (all surfaces and bridges merged, apertures excluded) and sensor.obj. When save_elements is True, also writes one element_{i}.obj per optical element; when save_rays is True, writes one lens_rays_fov_{i}.obj per traced field bundle.

Parameters:

Name	Type	Description	Default
save_dir	str	Directory to write the .obj files into.	required
mesh_rings	int	Number of rings per surface mesh. Defaults to 64.	64
mesh_arms	int	Number of arms per surface mesh. Defaults to 128.	128
save_rays	bool	Whether to trace and save rays. Defaults to False.	False
fovs	List[float]	Field-of-view angles to sample [degree]. Defaults to [0.0].	[0.0]
fov_phis	List[float]	Field azimuthal angles to sample [degree]. Defaults to [0.0].	[0.0]
ray_rings	int	Number of pupil rings to sample. Defaults to 6.	6
ray_arms	int	Number of pupil arms to sample. Defaults to 8.	8
is_wrap	bool	Whether to wrap the lens barrel as a cylinder. Defaults to False.	False
save_elements	bool	Whether to additionally save per-element .obj files. Defaults to True.	True

Note

Use #F2F7FFFF as the lens color when rendering in Blender. This routine writes .obj files directly and does not require PyVista.

Source code in deeplens-src/deeplens/geolens_pkg/vis3d.py

def save_lens_obj(
    self,
    save_dir: str,
    mesh_rings: int = 64,
    mesh_arms: int = 128,
    save_rays: bool = False,
    fovs: List[float] = [0.0],
    fov_phis: List[float] = [0.0],
    ray_rings: int = 6,
    ray_arms: int = 8,
    is_wrap: bool = False,
    save_elements: bool = True,
):
    """Save lens geometry, sensor, and optional rays as Wavefront ``.obj`` files.

    Writes ``lens.obj`` (all surfaces and bridges merged, apertures excluded)
    and ``sensor.obj``. When `save_elements` is True, also writes one
    ``element_{i}.obj`` per optical element; when `save_rays` is True, writes
    one ``lens_rays_fov_{i}.obj`` per traced field bundle.

    Args:
        save_dir (str): Directory to write the ``.obj`` files into.
        mesh_rings (int, optional): Number of rings per surface mesh. Defaults to 64.
        mesh_arms (int, optional): Number of arms per surface mesh. Defaults to 128.
        save_rays (bool, optional): Whether to trace and save rays. Defaults to False.
        fovs (List[float], optional): Field-of-view angles to sample [degree].
            Defaults to [0.0].
        fov_phis (List[float], optional): Field azimuthal angles to sample [degree].
            Defaults to [0.0].
        ray_rings (int, optional): Number of pupil rings to sample. Defaults to 6.
        ray_arms (int, optional): Number of pupil arms to sample. Defaults to 8.
        is_wrap (bool, optional): Whether to wrap the lens barrel as a cylinder.
            Defaults to False.
        save_elements (bool, optional): Whether to additionally save per-element
            ``.obj`` files. Defaults to True.

    Note:
        Use #F2F7FFFF as the lens color when rendering in Blender. This
        routine writes ``.obj`` files directly and does not require PyVista.
    """
    os.makedirs(save_dir, exist_ok=True)

    # Create surfaces & bridges meshes
    surf_meshes, bridge_meshes, element_groups, sensor_mesh = self.create_mesh(
        mesh_rings, mesh_arms, is_wrap
    )

    # Save individual lens elements (surfaces + bridges merged)
    if save_elements:
        for i, pair in enumerate(element_groups):
            print(f"Running in pair {i} with pair length {len(pair)}")
            # Collect surface polydata
            surf_polydata_list = [surf_meshes[idx].get_polydata() for idx in pair]

            # Collect bridge polydata if available
            bridge_polydata_list = []
            if i < len(bridge_meshes) and len(bridge_meshes[i]) > 0:
                print(f"Bridge mesh group number: {len(bridge_meshes[i])}")
                bridge_polydata_list = [b.get_polydata() for b in bridge_meshes[i]]

            # Merge surfaces and bridges together
            all_polydata = surf_polydata_list + bridge_polydata_list
            if len(all_polydata) == 1:
                element = all_polydata[0]
            else:
                element = merge(all_polydata)
            element.save(os.path.join(save_dir, f"element_{i}.obj"))

    # Merge all surfaces and bridges, and save as single lens.obj file
    surf_polydata = [
        surf.get_polydata()
        for surf in surf_meshes
        if not isinstance(surf, Aperture)
    ]
    bridge_polydata = [
        b.get_polydata() for group in bridge_meshes for b in group
    ]  # flatten the nested list
    lens_polydata = surf_polydata + bridge_polydata
    lens_polydata = merge(lens_polydata)
    lens_polydata.save(os.path.join(save_dir, "lens.obj"))

    # Save sensor
    sensor_polydata = sensor_mesh.get_polydata()
    sensor_polydata.save(os.path.join(save_dir, "sensor.obj"))

    # Save rays
    if save_rays:
        rays_curve = geolens_ray_poly(
            self, fovs, fov_phis, n_rings=ray_rings, n_arms=ray_arms
        )
        rays_poly_list = [curve_list_to_polydata(r) for r in rays_curve]
        rays_poly_fov = [merge(r) for r in rays_poly_list]
        for i, r in enumerate(rays_poly_fov):
            r.save(os.path.join(save_dir, f"lens_rays_fov_{i}.obj"))