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

Area of Science:

• Computational optics
• Microscopy
• Wave optics

Background:

• Computational modeling is crucial in modern microscopy for solving inverse problems and optimizing optical systems.
• Current differentiable optics simulations lack standardization, leading to duplicated efforts and limited performance in computational optics research.

Purpose of the Study:

• To introduce Chromatix, an open-source, GPU-accelerated library for differentiable wave optics simulations.
• To provide a standardized framework for computational optics researchers, improving reusability and performance.
• To expand the design space in computational optics through fast, parallelized simulations.

Main Methods:

• Development of Chromatix, a library built on JAX for wave optics simulation.
• Implementation of GPU acceleration for parallelized computations.
• Standardization of optical elements and propagation methods within the library.

Main Results:

• Chromatix demonstrates significant speed improvements: 2-6x on a single GPU and up to 22x on 8 GPUs.
• The library supports diverse applications, including snapshot microscopy, holography, and phase retrieval.
• Standardization of simulations enhances reusability and performance for computational optics researchers.

Conclusions:

• Chromatix democratizes fast, parallelized simulation of optical systems.
• The library addresses the need for a standardized, high-performance framework in computational optics.
• Chromatix has the potential to significantly advance the field of computational optics.