DiffRayFlow: A Differentiable Freeform Optical Design Engine Based on Discrete Optimal Transport

Freeform surfaces play a critical role in complex light-field modulation. However, traditional geometric mapping and standard optimization methods are limited by computational cost and convergence instability in large-scale ray tracing and complex surface modeling. This paper introduces DiffRayFlow, which integrates discrete optimal transport (OT), end-to-end differentiable ray tracing (DRT), and an adaptive multi-scale strategy. OT provides a global, energy-conserving geometric map. Differentiable tracing parameterizes the surface using the finite difference method (FDM) and constructs a differentiable link from height parameters to target landing points. The multi-scale approach, combined with early stopping, enhances efficiency and stability. For typical tasks involving over a million rays, the core heightmap optimization is usually completed within 20 s. The method can output standard Computer-Aided Design (CAD) data for rapid prototyping and physical validation. Ablation studies show that the multi-scale strategy is key to achieving high-precision convergence, while the early stopping mechanism can reduce optimization time by about 40% without sacrificing reconstruction quality. DiffRayFlow provides an efficient engineering path for interactive design and large-scale customization.