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

Abstract

Many current microscopy methods incorporate computational modeling as an integral part of the imaging process, either to solve inverse problems or optimize the optical system design itself. These methods often depend on differentiable optics simulations, yet no standardized framework exists, forcing computational optics researchers to repeatedly and independently implement simulations with limited reusability and performance. These common problems limit the potential impact of computational optics as a field. Here we present Chromatix: an open-source, graphics processing unit (GPU)-accelerated, differentiable wave-optics simulation library. Chromatix builds on JAX to democratize fast, parallelized simulation of diverse optical systems and expand the design space in computational optics. Chromatix standardizes a growing collection of optical elements and propagation methods allowing a broad range of applications, which we demonstrate here for snapshot microscopy, holography and phase retrieval. We demonstrate speed improvements of 2–6 times on a single GPU and up to 22 times on 8 GPUs.