Skip to content

Loss Functions API Reference

AutoLens provides two categories of loss functions: optimization losses (used during lens design) and PSF quality losses (neural network modules).

Optimization Losses

These are methods on GeoLensOptim (accessible via GeoLens). All return scalar tensors suitable for backpropagation.

loss_rms()

Primary optimization target. Computes RGB spot RMS error across sampled field positions, with optional error-adaptive field weighting.

loss_infocus()

Penalizes on-axis defocus when RMS exceeds a threshold:

\[\mathcal{L}_{\text{infocus}} = \text{softplus}(\text{rms}_{\text{on-axis}} - 0.005)\]

loss_reg()

Composite regularization loss:

\[\mathcal{L}_{\text{reg}} = \mathcal{L}_{\text{intersec}} + 0.1 \cdot \mathcal{L}_{\text{thickness}} + \mathcal{L}_{\text{surface}} + 2 \cdot \mathcal{L}_{\text{ray\_angle}}\]

loss_surface()

Penalizes extreme surface shapes via three sub-penalties:

Sub-penalty Constraint Description
Sag / diameter \(\text{sag} / \text{diam} < 0.1\) Prevents overly curved surfaces
Gradient \(\|\nabla z\| < \tan(30°)\) Limits surface steepness
D / T ratio \(\text{diam} / \text{thickness} < 15\) Prevents thin, wide elements

loss_intersec()

Prevents surfaces from intersecting each other:

\[\mathcal{L}_{\text{intersec}} = \sum \text{softplus}(\text{min\_gap} - \text{actual\_gap})\]

loss_thickness()

Enforces air gap, glass thickness, back focal length, and total track length constraints.

Cellphone lens constraints (\(r_{\text{sensor}} < 12\) mm):

Parameter Min (mm) Max (mm)
Air gap (center) 0.025 1.5
Air gap (edge) 0.025 3.0
Glass (center) 0.25 3.0
Glass (edge) 0.25 2.0
BFL 0.8 3.0
TTL — 15

Camera lens constraints (\(r_{\text{sensor}} \geq 12\) mm): significantly looser bounds (air gaps up to 100 mm, TTL up to 300 mm).

loss_ray_angle()

Constrains chief ray angle at the sensor to prevent extreme incidence:

\[\mathcal{L}_{\text{CRA}} = \text{softplus}(\cos\theta_{\text{ref}} - \cos\theta_{\text{CRA}})\]

Default maximum: \(30°\) for cellphone, \(40°\) for camera lenses.

loss_mat()

Constrains material parameters during optimization:

\[n \in [1.5, 1.9], \quad V \in [30, 70]\]

PSF Quality Losses

Neural network modules (nn.Module) for evaluating PSF quality.

PSFLoss

src.loss.PSFLoss 

PSFLoss(w_achromatic=1.0, w_psf_size=1.0)

Bases: Module

Source code in src/loss.py 
13
14
15
16
def __init__(self, w_achromatic=1.0, w_psf_size=1.0):
    super(PSFLoss, self).__init__()
    self.w_achromatic = w_achromatic
    self.w_psf_size = w_psf_size

Combines spatial concentration loss (variance of PSF spatial distribution) with achromatic loss (inter-channel PSF difference). Minimized during optimization.

from src.loss import PSFLoss

criterion = PSFLoss()
loss = criterion(psf)  # psf: (B, C, H, W)

PSFStrehlLoss

src.loss.PSFStrehlLoss 

PSFStrehlLoss()

Bases: Module

Source code in src/loss.py 
69
70
def __init__(self):
    super(PSFStrehlLoss, self).__init__()

forward 

forward(psf)

Compute Strehl-like score for PSFs with shape [B, 3, ks, ks].

The score is the center-pixel intensity after per-channel spatial normalization (sum over HxW equals 1), averaged over channels and batch. This value should be maximized during optimization.

Source code in src/loss.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def forward(self, psf):
    """Compute Strehl-like score for PSFs with shape [B, 3, ks, ks].

    The score is the center-pixel intensity after per-channel spatial
    normalization (sum over HxW equals 1), averaged over channels and batch.
    This value should be maximized during optimization.
    """
    # Ensure shape [B, 3, H, W]
    if psf.dim() == 3:
        psf = psf.unsqueeze(0)
    assert psf.dim() == 4 and psf.size(1) == 3, (
        f"Expected psf shape [B, 3, ks, ks], got {tuple(psf.shape)}"
    )

    eps = torch.finfo(psf.dtype).eps
    # Normalize per-sample, per-channel over spatial dims
    psf_sum = psf.sum(dim=(2, 3), keepdim=True)
    psf_norm = psf / (psf_sum + eps)

    # Center pixel indices
    h, w = psf.shape[-2:]
    cy, cx = h // 2, w // 2

    # Center intensity per sample and per channel
    center_vals = psf_norm[:, :, cy, cx]  # [B, 3]

    # Average across channels, then across batch
    strehl_per_sample = center_vals.mean(dim=1)  # [B]
    strehl = strehl_per_sample.mean()  # scalar

    return strehl

Strehl ratio proxy — center pixel intensity after normalization. Maximized during optimization (negate for gradient descent).

from src.loss import PSFStrehlLoss

criterion = PSFStrehlLoss()
strehl = criterion(psf)  # Higher is better