Chromatix: a differentiable, GPU-accelerated wave-optics library

Modern microscopy methods incorporate computational modeling of optical systems as an integral part of the imaging process, either to solve inverse problems or enable optimization of the optical system design. These methods often depend on differentiable simulations of optical systems, yet no standardized framework exists—forcing computational optics researchers to repeatedly and independently implement simulations that are prone to errors, difficult to reuse in other applications, and often computationally suboptimal. These common problems limit the potential impact of computational optics as a field. We present Chromatix: an open-source, GPU-accelerated differentiable wave optics library. Chromatix builds on JAX to enable fast simulation of diverse optical systems and inverse problem solving, scaling these simulations from single-CPU laptops to multi-GPU servers. The library implements various optical elements (e.g., lenses, polarizers and spatial light modulators) and multiple light propagation models (e.g., Fresnel approximation, angular spectrum and off-axis propagation) that can be flexibly combined to model various computational optics applications such as snapshot microscopy, holography, and phase retrieval of multiple scattering samples. These simulations can be automatically parallelized to scale across multiple GPUs with a single-line change to the modeling code, enabling simulation and optimization of previously impractical optical system designs. We demonstrate Chromatix’s capacity to substantially accelerate optics simulation and optimization on existing methods in computational optics, speeding up optical simulation and optimization from 2-6× on a single GPU to up to 22× on 8 GPUs (depending on the particular system being modeled) compared to the original implementations. Chromatix establishes a standard for wave optics simulations, democratizing access to and expanding the design space of computational optics.i