Architecture
AutoLens models the automated lens design pipeline as a three-stage differentiable optimization process.
Flat Surfaces → [ Curriculum Learning ] → [ Fine-Tuning ] → [ Centroid Refinement ] → Optimized Lens
(Adam) (Adam) (Adam / L-BFGS)
Project Structure
AutoLens/
│
├── src/ (core optics library)
│ ├── geolens.py (GeoLens: main lens class)
│ ├── geolens_pkg/ (GeoLens methods: eval, optim, vis, io, psf)
│ ├── geometric_surface/ (surface types: Spheric, Aspheric, Aperture, ...)
│ ├── light/ (Ray and Wave classes)
│ ├── material/ (glass material database)
│ ├── imgsim/ (PSF computation and image simulation)
│ ├── phase_surface/ (diffractive optical elements)
│ ├── loss.py (loss functions)
│ └── config.py (global constants)
│
├── lens_zoo/ (pre-designed lens files)
├── 0_hello_deeplens.py (minimal demo)
├── 1_lens_optim.py (single-lens optimization)
├── 2_autolens_rms.py (automated lens design pipeline)
├── 3_random_autolens.py (batch design with random configs)
└── 9_autolens_aspheric.py (progressive aspheric introduction)
GeoLens Mixin Architecture
GeoLens uses a mixin architecture — functionality is split across specialized modules, composed via multiple inheritance:
GeoLens
├── GeoLensPSF (psf_compute.py — PSF computation)
├── GeoLensEval (eval.py — spot diagrams, MTF, distortion, wavefront)
├── GeoLensSeidel (eval_seidel.py — third-order aberration analysis)
├── GeoLensOptim (optim.py — loss functions and optimization loops)
├── GeoLensSurfOps (optim_ops.py — surface geometry operations)
├── GeoLensVis (vis.py — 2D layout visualization)
├── GeoLensVis3D (vis3d.py — 3D mesh visualization)
├── GeoLensIO (io.py — JSON / Zemax / Code V file I/O)
├── GeoLensTolerance (eval_tolerance.py — manufacturing tolerance analysis)
└── Lens (lens.py — abstract base with render interface)
All optical objects inherit from DeepObj, which provides to(device), clone(), and dtype conversion by introspecting instance tensors.
Differentiable Ray Tracing
Surface Intersection
Ray-surface intersection in geometric_surface/base.py uses Newton's method with a critical trick for differentiability:
- Non-differentiable loop (
torch.no_grad()) — iterates Newton's method to convergence - One differentiable Newton step — enables gradient flow through the intersection point back to surface parameters
This allows PyTorch autograd to compute \(\partial \text{loss} / \partial \text{surface\_params}\) while maintaining numerical stability.
Refraction
Surface refraction implements differentiable vector Snell's law:
where \(\mathbf{d}_i\), \(\mathbf{d}_t\) are incident and transmitted ray directions, and \(\mathbf{n}\) is the surface normal.
Three-Stage Optimization Pipeline
Stage 0: Lens Creation
create_lens() generates a random flat-surface starting point:
- Flat surfaces (curvature \(c \approx 0\)) with random glass materials
- Aperture radius: \(r_{\text{aper}} = f / (2 \cdot F_{\#})\)
- Surface radii sized by ray cone: \(r_i = r_{\text{aper}} + |d_i - d_{\text{stop}}| \cdot \tan(\theta/2)\)
Stage 1: Curriculum Learning
The key innovation. Starts with a small aperture (25% of target) and gradually opens it via a cosine schedule:
At each evaluation step, the optimizer also runs correct_shape() (clips surfaces to valid geometry) and match_materials() (snaps floating refractive indices to the nearest real glass in the catalog).
Optimizer: Adam with per-parameter learning rates [lr_d, lr_c, lr_k, lr_ai]
Loss: \(\mathcal{L} = \mathcal{L}_{\text{rms}} + 0.05 \cdot \mathcal{L}_{\text{infocus}} + 0.05 \cdot \mathcal{L}_{\text{reg}}\)
Stage 2: Fine-Tuning (Pinhole Center)
Full-aperture optimization with fixed materials and shape control. Uses pinhole model for PSF center reference. Error-adaptive field weighting gives higher loss weight to field positions with larger current RMS error.
Stage 3: Fine-Tuning (Centroid Center)
Same as Stage 2, but uses chief-ray centroid as the PSF center reference. This corrects for the difference between paraxial ideal and real image formation.
Parameter Reparametrization
For L-BFGS optimization, aspheric surface parameters are reparametrized to \(\sim\mathcal{O}(1)\) scale. This is critical because L-BFGS uses a single learning rate for all parameters.
Problem: Parameters span 8+ orders of magnitude — curvature \(c \sim 0.01\), conic \(k \sim 100\), aspheric coefficients \(a_4 \sim 10^{-3}\), \(a_{12} \sim 10^{-7}\).
Solution: Store normalized \(\phi\) tensors as optimizer targets. Physical values recovered via \(\theta = \phi \cdot s\):
| Parameter | Scale factor \(s\) | Effect |
|---|---|---|
| Curvature \(c\) | \(\max(\|c\|, 0.01)\) | \(c_\phi \approx 1\) |
| Conic \(k\) | \(\max(\|k\|, 1)\) | \(k_\phi \approx 1\) |
| Aspheric \(a_{2n}\) | \(1 / r^{2n}\) | Same insight as Adam lr scaling |
Activated via lens.optimize_lbfgs(reparam=True) or lens.enable_reparam() / lens.disable_reparam().
Loss Functions
| Loss | Purpose | Formula |
|---|---|---|
loss_rms |
RGB spot RMS vs reference | Primary optimization target |
loss_infocus |
On-axis focus penalty | \(\text{softplus}(\text{rms} - 0.005)\) |
loss_intersec |
Prevent surface self-intersection | \(\text{softplus}(\text{min\_gap} - \text{gap})\) |
loss_thickness |
Enforce thickness / TTL bounds | \(\text{softplus}(\text{actual} - \text{max})\) |
loss_surface |
Penalize extreme surface shapes | Sag/diam, gradient, D/T ratios |
loss_ray_angle |
Chief ray angle constraint | \(\text{softplus}(\cos_{\text{ref}} - \cos_{\text{CRA}})\) |
loss_mat |
Constrain \(n \in [1.5, 1.9]\), \(V \in [30, 70]\) | Linear penalty |
Lens File Format
Lenses are stored as JSON:
{
"foclen": 4.55,
"fnum": 2.0,
"r_sensor": 3.83,
"d_sensor": 5.2,
"sensor_size": [5.42, 5.42],
"surfaces": [
{
"type": "Aspheric",
"r": 3.5,
"(d)": 2.5,
"(c)": 0.08,
"k": 0.0,
"ai": [1e-4, 2e-5],
"mat2": "coc"
},
{
"type": "Aperture",
"r": 1.8,
"(d)": 5.0,
"mat2": "air"
}
]
}
Keys in parentheses (e.g., "(d)", "(c)") denote differentiable parameters. roc (radius of curvature) is accepted as an alias for c = 1/roc.