The implementation of SIGGRAPH Asia 2025 paper (journal track) "Large-Area Fabrication-Aware Computational Diffractive Optics"
- 2025/11/19: First code/data release! We believe the proposed techinique in this work will serve as foundations to facilitate the mass-market implementation of computational diffractive optics technology.
Fabrication-aware, End-to-end Optimization for Large-area Diffractive Optics. We propose a fabrication-aware design method for diffractive optical elements fabricated by (left) direct-write grayscale lithography with nanoimprint replication (see the inset figures A-D for a step-by-step illustration) suited for inexpensive mass production. Enabled by tensor-parallel computing routines, our method jointly considers the fabrication 3-D geometry deformation and the downstream task-specific computational diffractive optics design. This combination of techniques allows for experimental findings with favorable quality to all tested existing methods, specifically closing the design-to-manufacturing gap in existing approaches
This implementation heavily relies on JAX, a differentiable programming library aimming for high performance array computing. You may want to install a complete jax AI ecosystem here. Besides, to ease the flexible implementation of tensor parallism, we adopt the generalized single program multiple data (GSPMD) programming model with the aid of scalax, scaling utilities for JAX. You may want to use my own customized version that adapted to the latest JAX version. Finally, our implementation follows the modern engineering practice of machine learning research & development, using hydra (for configuration management), wandb (for experiment tracking) and loguru (for logging).
Our neural lithography model checkpoints, lithography model calibration data (contrast curve and atomic force microscopy profiling data), refractive index curves of photoresist/resin and test images are available at Google Drive. After setting up the environment, you may want to quickly test the functionality of tensor-parallel FFT by running (assuming you have at least two GPUs in one node/machine)
pytest tests/test_dfft.py
Then, please take a look at scripts folder. I would recommend to start with holo2d_exp_paper.sh to see how we train and evaluate 2-D computer-generated holograms as in our paper. Neural lithography model training can be conducted with neural_litho.sh. The broadband imaging application can be performed in broadband_imaging.sh
@article{Wei2025LAFA,
author = {Kaixuan Wei and Hector A. Jimenez-Romero and Hadi Amata and Jipeng Sun and Qiang Fu and Felix Heide and Wolfgang Heidrich},
title = {Large-Area Fabrication-aware Computational Diffractive Optics},
journal = {ACM Transactions on Graphics (TOG)},
year = {2025},
volume = {44},
number = {6},
articleno = {243},
month = {dec},
doi = {10.1145/3763358},
publisher = {ACM}
}If you find any problem, please feel free to contact me (kaixuan.wei at kaust.edu.sa). A brief self-introduction (including your name, affiliation and position) is required, if you would like to get an in-depth help from me. I'd be glad to talk with you if more information (e.g. your personal website link) is attached.