Differentiable wave propagation method for shape optimization of freeform optics beyond the paraxial approximation

Freeform optics provides precise control of light for beam shaping and mode conversion, but designing volumetric micro-optical elements requires efficient, low-runtime simulation methods. The wave propagation method (WPM) offers a fast alternative to full-wave solvers for high-index-contrast and non-paraxial systems but lacks differentiability, preventing its use with gradient descent methods for inverse design. We propose a differentiable formulation of WPM using smoothstep functions to create a soft partition of the spatial domain, allowing for gradient-based optimization. This crucial modification enables a general and computationally efficient framework for solving inverse design problems in micro-optics. We demonstrate its effectiveness through the simulation of a freeform hologram exhibiting high diffraction angles, highlighting the limitations of traditional paraxial methods.