Chromatix
Differentiable, GPU-accelerated wave-optics simulation
Chromatix is an open-source, GPU-accelerated, fully differentiable wave-optics simulation library built in JAX. It provides a standardized framework for computational optics — something that did not exist before. Labs previously had to independently implement their own simulations, leading to errors, limited reusability, and wasted effort.
Chromatix speeds up simulation and optimization 2–22× versus existing implementations depending on the system, scales seamlessly across multiple GPUs with a single-line code change, and is used in labs worldwide. It is now becoming the standard framework for computational optics research.
Co-developed with Diptodip Deb and Srinivas Turaga at HHMI Janelia Research Campus.
Key methods: JAX · Differentiable programming · Wave optics · Multi-GPU computing
Related publication: Deb*, Both* et al. Chromatix: a differentiable, GPU-accelerated wave-optics library. Nature Methods (2026).