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

This paper introduces Chromatix, an open-source, GPU-accelerated, differentiable wave-optics library built on JAX that standardizes optical simulations to overcome reusability and performance limitations in computational microscopy, achieving significant speedups across applications like holography and phase retrieval.

The Gist

Imagine you are trying to build a complex machine, like a high-end camera or a medical scanner. In the past, if a scientist wanted to design a new type of microscope or figure out how to see through foggy tissue, they had to write their own computer code from scratch every single time. It was like every chef in the world having to invent their own recipe for chopping onions, mixing their own flour, and building their own ovens before they could even start cooking. It was slow, prone to mistakes, and nobody could easily share their "secret sauce" with others.

Chromatix is the solution to this chaos. Think of it as the "LEGO set" or the "App Store" for light and optics.

Here is a simple breakdown of what Chromatix does and why it's a big deal:

1. The Problem: Reinventing the Wheel

Currently, researchers in "computational optics" (using math and computers to improve how we see) have to program their own simulations of how light behaves.

• The Analogy: Imagine trying to drive a car, but every time you want to go from point A to point B, you have to build the engine, the wheels, and the steering column from raw metal. If you want to drive faster, you have to build a new engine. It's exhausting and inefficient.
• The Result: Because everyone builds their own "engines," the code is often slow, hard to fix, and impossible to combine with other tools.

2. The Solution: Chromatix (The Universal Toolkit)

Chromatix is a new, free software library that gives scientists a standard set of pre-built, high-quality "optical parts."

• The Analogy: Instead of building an engine from scratch, Chromatix gives you a shelf of perfectly engineered engines, wheels, and steering systems. You just snap them together to build your car (or microscope).
• The "LEGO" Aspect: You can take a "lens" block, a "scattering tissue" block, and a "camera sensor" block and click them together to simulate a complex system. If you want to swap a lens for a mirror, you just swap the block. No need to rebuild the whole car.

3. The Superpowers: Speed and "Self-Correcting"

Chromatix has two magical features that make it special:

A. It's Super Fast (The GPU Engine)

• The Analogy: Most old software runs on a single-lane road (a standard computer processor). Chromatix runs on a massive, 8-lane superhighway (Graphics Processing Units or GPUs).
• The Result: Tasks that used to take hours or days now take minutes. The paper shows it can be up to 22 times faster than previous methods. It's like going from a bicycle to a supersonic jet.

B. It's "Differentiable" (The Self-Correcting GPS)

• The Analogy: Imagine you are trying to find a hidden treasure. In the old days, you would guess a location, dig, find nothing, and then guess again blindly.
• Chromatix's Way: Chromatix is like a GPS that tells you exactly how to adjust your path to get closer to the treasure. If the image is blurry, the software automatically calculates exactly which part of the lens or the math needs to change to make it sharp. It learns and optimizes itself without a human needing to do the heavy lifting.

4. What Can You Do With It?

The paper shows scientists using Chromatix to solve real-world biological mysteries:

• Seeing Through Fog: They simulated light passing through a scattering sample (like a fish embryo or brain tissue) to reconstruct a clear 3D image, something that was previously very difficult.
• Fixing Bad Lenses: They used it to correct "aberrations" (blurry spots) in cheap, miniature microscopes, making them see as clearly as expensive ones.
• Holographic Light Sculpting: They used it to design patterns of light that can zap specific neurons in a brain to study how they work, even through the "foggy" tissue of the skull.
• Multicolor Vision: They engineered a microscope that can tell the difference between 25 different colors of light using just a black-and-white camera.

The Bottom Line

Chromatix is democratizing the science of light.

Just as libraries like TensorFlow and PyTorch allowed anyone to build AI without being a math genius, Chromatix allows biologists and engineers to build advanced optical systems without being coding wizards. It turns the field of "how to see better" from a solitary, slow struggle into a fast, collaborative, and creative playground.

In short: It's the tool that lets us stop building the car from scratch and start driving toward the future of seeing the invisible.

Technical Summary

1. Problem Statement

Modern computational microscopy increasingly relies on integrating optical system design with algorithmic reconstruction (inverse problems) or optimization. However, the field faces significant bottlenecks:

• Lack of Standardization: Researchers currently implement wave-optics simulations from scratch for every project, leading to non-standard conventions, limited reusability, and difficulty in comparing results.
• Computational Inefficiency: Many existing simulations are not optimized for modern hardware (GPUs) or lack parallelization capabilities, making large-scale reconstructions (e.g., whole-brain imaging) computationally prohibitive.
• Differentiability Gap: While deep learning frameworks (PyTorch, TensorFlow) offer automatic differentiation (AD), they are not natively designed for complex wave-optics physics. Conversely, traditional optical design software (Zemax, CODE V) supports wave optics but lacks differentiability and interoperability with deep learning models.
• Feature Limitations: Existing open-source differentiable optics libraries often rely on ray optics (unsuitable for diffraction-limited systems) or lack support for complex phenomena like 3D multiple scattering, polarization, and vectorial wave propagation.

2. Methodology: The Chromatix Framework

Chromatix is an open-source, GPU-accelerated, differentiable wave-optics simulation library built on JAX (Just After eXecution). It draws architectural inspiration from deep learning frameworks, treating optical systems as compositions of fundamental layers.

Core Design Principles:

• Differentiability: Built on JAX, Chromatix provides automatic differentiation for any parameter in the simulation (e.g., sample refractive index, SLM phase masks, lens aberrations). This enables gradient-based optimization and end-to-end training of optical systems with neural networks.
• Composability: Optical systems are constructed by chaining modular components (lenses, sensors, propagation kernels, scattering samples). This allows researchers to swap components easily and build complex, custom optical setups without rewriting core physics engines.
• Scalability: Leveraging JAX's XLA compiler, Chromatix supports seamless acceleration on CPUs, single GPUs, and multi-GPU clusters (TPUs supported). It allows for batched processing and parallelization across devices with minimal code changes.

Key Technical Components:

• Unified Field Representation: The library encodes the "hidden state" of an optical system (complex light field, wavelength, polarization, spatial sampling) into a single structured object.
• Optical Elements: Implements a wide range of elements including thin/thick lenses, spatial light modulators (SLMs), digital micromirror devices (DMDs), and sensors.
• Propagation Models: Supports scalar and vectorial wave propagation through free space (Fresnel, Angular Spectrum, Band-limited Angular Spectrum) and through complex media.
• Scattering Models: Features a Multislice Beam Propagation (MSBP) method to simulate light propagation through strongly scattering, anisotropic 3D biological samples (beyond the single-scattering Born approximation).

3. Key Contributions

1. First Standardized Differentiable Wave-Optics Library: Chromatix fills the gap between deep learning frameworks and rigorous wave-optics physics, offering a unified interface for diverse applications.
2. Support for Complex Physics: Unlike previous differentiable tools, it natively supports 3D multiple scattering, polarization, and vectorial wave propagation, essential for modeling realistic biological tissues.
3. Hardware Acceleration: It achieves significant speedups (2–6× on single GPUs, up to 22× on 8 GPUs) compared to original implementations in MATLAB, PyTorch, or TensorFlow, enabling previously intractable large-scale simulations.
4. End-to-End Optimization: The library facilitates joint optimization of optical hardware (e.g., phase masks) and computational reconstruction algorithms (neural networks) in a single differentiable pipeline.

4. Results and Demonstrations

The authors validated Chromatix through six computational experiments spanning four major application areas:

• Inverse Problems & Sample Reconstruction:

  • Ring Deconvolution Microscopy: Successfully reconstructed spatially varying aberrations across a large field of view (2048×2048 pixels) for a Miniscope. Chromatix achieved a 4.5× speedup on a single GPU and ~19× on 8 GPUs compared to the original PyTorch implementation, which failed to fit the full field of view on a single GPU due to memory limits.
  • Computational Aberration Correction (COCOA): Used implicit neural representations (INRs) to jointly reconstruct a mouse brain volume and infer aberrations. Chromatix was 2× faster (single GPU) and ~9× faster (8 GPUs) than the original implementation, while producing higher-quality dendrite continuity and more accurate aberration recovery (RMS error of 3.56 nm vs. 6.97 nm).
  • 3D Refractive Index Microscopy: Recovered the 3D refractive index of a scattering D. rerio embryo. The Chromatix implementation was 3–13× faster than the original MATLAB code, reducing reconstruction time from hours to minutes and eliminating grid artifacts.

• Programmable Optics & Deep Learning:

  • 3D Snapshot Microscopy (Holoscope): Optimized a programmable phase mask and a neural network (FourierNet) to reconstruct 3D volumes from single 2D snapshots. Chromatix reproduced the original results (SSIM 0.979) but reduced training time by ~7×.
  • DeepCGH (Computer-Generated Holography): Trained a UNet to generate holograms for optogenetic stimulation. Chromatix achieved nearly identical results (SSIM 0.985) with a 2.5× speedup on a single GPU and >10× on 8 GPUs.

• Flexible Modeling & Novel Applications:

  • Spectroscopic Single-Molecule Localization: Engineered a single-phase mask to enable spectral imaging of 25 wavelengths from a single-channel camera, demonstrating the library's ability to optimize for novel, multi-objective tasks.
  • Scattering-Compensated Holography: Combined holography and scattering models to optimize phase masks for uniform light delivery through scattering tissue, correcting for distortions that standard free-space models miss.

5. Significance and Impact

• Democratization of Computational Optics: By providing a standardized, easy-to-use library, Chromatix lowers the barrier to entry for researchers, allowing them to focus on scientific discovery rather than re-implementing physics engines.
• Accelerated Innovation: The massive speedups enable the exploration of larger design spaces (e.g., optimizing entire optical systems end-to-end) and larger datasets (e.g., whole-brain imaging) that were previously computationally impossible.
• Reproducibility: The common codebase ensures that different research groups can compare methods fairly and reproduce results without the "black box" of custom, non-standard implementations.
• Future Directions: The authors envision Chromatix enabling neural architecture search for optical systems, hardware-in-the-loop calibration, and the integration of ray-based models for hybrid refractive-diffractive design.

In summary, Chromatix represents a paradigm shift in computational optics, bridging the gap between rigorous wave physics and modern differentiable programming to accelerate the design and analysis of next-generation imaging systems.