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

Area of Science:

• Computational optics
• Microscopy
• Wave optics simulation

Background:

• Computational modeling is integral to modern microscopy for inverse problems and optical system design.
• Current methods rely on differentiable optics simulations but lack a standardized framework, hindering reusability and performance.
• This limitation restricts the potential impact of computational optics.

Purpose of the Study:

• Introduce Chromatix, an open-source, GPU-accelerated, differentiable wave-optics simulation library.
• Democratize fast, parallelized simulation of diverse optical systems.
• Expand the design space within computational optics.

Main Methods:

• Developed Chromatix, a library built on JAX for differentiable wave-optics simulations.
• Integrated a collection of standardized optical elements and propagation methods.
• Leveraged graphics processing unit (GPU) acceleration for parallelized computations.

Main Results:

• Chromatix provides significant speed improvements: 2-6x on a single GPU and up to 22x on 8 GPUs.
• Demonstrated applications in snapshot microscopy, holography, and phase retrieval.
• Enabled a broad range of applications through standardized components.

Conclusions:

• Chromatix offers a standardized, high-performance framework for computational optics research.
• The library facilitates faster and more accessible optical system design and simulation.
• Chromatix is poised to increase the impact and scope of computational optics.