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.

Based on JAX, Chromatix provides a standardized framework for developing, sharing and reusing wave-optics simulations. The capabilities are demonstrated in a number of computational microscopy applications.

Main

Many current microscopy methods increasingly rely on computation as an integral part of the imaging process. This model-based approach to optics—integrating optical system design with algorithmic reconstruction or optimization—has had major implications for biological discovery. Single-molecule localization microscopy can reveal individual protein complexes without requiring electron microscopy^1,2 and three-dimensional (3D) snapshot microscopy enables volumetric imaging at the frame rate of the camera, facilitating whole-brain imaging at high temporal resolution^3–8. For both techniques, using wave-optics models to engineer point spread functions (PSFs) allows improved resolution in 3D^3,9–11. Computer-generated holography uses wave-optics models of light propagation to design precise optogenetic stimulation of many neurons simultaneously^12–15. Integrating microscopy with wave-optics models also allows measurement of optical properties that are otherwise difficult or slow to obtain, such as using diffraction tomography through a strongly scattering sample to obtain 3D refractive index distributions of transparent tissues from only intensity images, allowing high-contrast, label-free imaging of optically transparent model organisms such as Caenorhabditis elegans or Danio rerio^16,17. Differentiable models of wave optics have also allowed for the end-to-end optimization of optical design and image processing algorithms for a variety of techniques such as quantitative phase imaging¹⁸, hyperspectral imaging¹⁹, extended depth-of-field²⁰, monocular depth estimation^21,22, localization microscopy⁹, lensless imaging²³ and 3D snapshot microscopy³.

However, each computational optics method typically requires the researcher to program an optics simulation from scratch. These de novo simulations have differing conventions from other simulations, can be difficult to reuse in other applications and are often computationally suboptimal due to the difficulty of programming fast optics simulations on current computer hardware such as graphics processing units (GPUs). Moreover, these simulations typically need to be differentiable to facilitate efficient optimization of the relevant optical parameters or to be easily combined with deep neural networks. A standard library providing differentiable wave-optics simulations that make efficient use of GPUs is therefore desirable.

Here we describe Chromatix, a high-performance differentiable wave-optics simulation library that could fill this gap in the field of computational optics. Drawing inspiration from current deep learning frameworks such as PyTorch²⁴ and TensorFlow²⁵ that construct neural networks from layers of mathematical operations, Chromatix enables researchers to describe optical systems as compositions of fundamental elements. This architectural parallel runs deeper than mere interface design: Chromatix shares the core requirements of differentiability for gradient-based optimization, scalability for tackling large-scale problems and composability for rapid method development. Chromatix is built on JAX (just after execution)²⁶, a numerical computation library for Python that provides GPU acceleration and automatic differentiation to automatically and efficiently calculate gradients with respect to any input of differentiable functions. By leveraging JAX’s capabilities, Chromatix simulations can be seamlessly accelerated on GPUs and parallelized across multiple devices with minimal code changes (Extended Data Fig. 1), enabling integration with existing deep learning models, optimization loops and even hardware control systems. Chromatix implements diverse optical models, from conventional lenses to diffractive elements such as liquid-crystal-on-silicon spatial light modulators (SLMs), and various models of scalar and vectorial wave propagation through both free space and scattering media, enabling simulation of optical systems interacting with multiple-scattering (potentially anisotropic) biological samples whose amplitude and refractive index vary in 3D.

Extended Data Fig. 1. Chromatix enables parallelization of optical simulations with minimal code changes.

For all methods, we show the optical model followed by its implementation using Chromatix in both an unparallelized single GPU simulation as well as the parallelized versions that run across multiple GPUs. Colors in the code denote the corresponding elements in the model diagrams on the left. Note that the code shown is not the complete code that is actually used in demonstrations across the rest of the paper; this code is meant to provide a high level depiction of Chromatix code and parallelization. Ellipses indicate code that has been intentionally omitted for brevity. Simulation results are identical for all of the parallelization scenarios demonstrated here. a, Optical model for quantitative phase imaging showing how a single simulation can be run independently on multiple GPUs. b, Optical model for 3D holography showing how a simulation can be run multiple times in parallel on one GPU, with this whole batch of simulations run in parallel across multiple GPUs. c, Optical model for 3D snapshot microscopy (the Holoscope) showing how a simulation can be run across one large 3D sample split across multiple GPUs to enable larger simulations/optimizations.

While there exist established tools for optical design such as Zemax²⁷ or CODE V²⁸ that do support wave optics, they are not efficiently implemented for the types of model we present here and are not differentiable or interoperable with deep learning models. Recently, a number of open-source differentiable optics simulation libraries have been released^29–32; however, these libraries are either primarily ray-based^29,32 (while the computational optics methods we are interested in are wave-based) or lack support for all the features that would be desirable for such a standard library in computational optics, for example multiple-scattering 3D samples or polarization. We provide a more detailed comparison of Chromatix versus other optics simulation software in Extended Data Table 1.

Extended Data Table 1.

Comparison of features against other simulation libraries/software

* Optiland does not provide a range of general propagation methods but does include a method to compute point spread functions with diffraction effects via the Fourier transform³².

** Chromatix uses a paraxial approximation to model the effect of thick lenses/ray transfer matrices via the Collins integral^74,75, which will not capture the effects of large systems of thick lenses at the edges of a field as accurately as ray tracing^27–29 or beamlet propagation^27,28.

Here we demonstrate Chromatix’s ability to simulate and optimize various optical systems for biological imaging, showing simulations of widefield and snapshot fluorescence microscopes, phase contrast microscopes and Fourier holography systems for optogenetics. In each case, we highlight Chromatix’s scalability, delivering results substantially faster than previous implementations—2–22× improvements through parallelization (depending on the problem)—while eliminating the challenges of correctly implementing wave-optics simulations from scratch. We anticipate that this open-source library will democratize high-performance, scalable wave-optics simulations, enabling exploration of a much richer design space in computational optics and accelerating innovation in biological imaging and beyond.

Results

Design and implementation

The design of Chromatix is strongly inspired by current deep learning frameworks. Here we detail the principles that informed our design choices and implementation decisions. We argue that effective frameworks for computational optics, like those for deep learning, must embody three key characteristics: differentiability, composability and scalability. We also discuss the high level implementation of Chromatix as it relates to these three characteristics.

Differentiability

Differentiability is the ability to calculate gradients, which can be used for gradient-based optimization (for example, of the parameters of an optical simulation). For small, low-dimensional inputs, numerical differentiation can be sufficient, but for high-dimensional inputs (for example, the pixels of an SLM) this becomes too computationally expensive to evaluate, and it is preferable to use backpropagation of the gradients of each step of the simulation. When combined with the wide variety of possible optical models or neural networks, automatic differentiation becomes a desirable property. Common programming languages such as MATLAB (https://www.mathworks.com) and C do not provide general-purpose automatic differentiation, requiring gradients to be manually derived and implemented: an error-prone, time-consuming and inflexible process. Current deep learning frameworks^24–26 provide automatic differentiation: given a function, they can automatically calculate the gradient with respect to any parameter of that function as long as the function is differentiable with respect to that parameter. Similarly, Chromatix can automatically calculate the gradient with respect to any parameter of a simulation as it has been written using JAX.

Differentiability has already found several uses in optical design, enabling end-to-end design of computational optics systems for a variety of problems^3,9,18–23. Differentiability also has the potential to improve solutions to inverse problems in optics. Traditional inverse problem approaches simplify both the sample and the optics^33–36, whereas differentiable models can handle more realistic complexity. Automatic differentiation opens up a new class of gradient-based optimizers such as Adam³⁷, which can improve reconstruction fidelity by allowing arbitrarily complex physics simulations (for example, scattering^16,38 or sample deformation³⁹) in the forward simulation. These benefits come at nearly zero programmer effort with automatic differentiation: once the forward model of the simulation has been defined, its gradients are automatically defined as well. Thus, automatic differentiation can also enable so-called self-calibrating algorithms⁴⁰. As hardware has a finite accuracy, optimizing certain physical parameters (such as angle of illumination in tomography) together with the sample has been shown to improve fidelity^41,42.

A different line of work replaces discrete voxel-based representations with neural network-based continuous representations, a concept known as implicit neural representations (INRs) or neural radiance fields^39,43–45. INRs have been applied to separation of motion artifacts from sample dynamics³⁹, estimation of dynamic aberrations⁴⁴, reconstruction of 3D quantitative phase of scattering samples⁴⁶ and aberration correction without wavefront sensors or calibration measurements⁴³. Here too, differentiable simulations are required to train these networks.

Composability

The principle underlying differentiability is composability: the gradient of a composition of two functions can be calculated from the gradient of each of those functions. Taking a broader and more practical view, we can interpret composability as being able to easily swap and replace components of a network, for example replacing the activation function in a multilayer perceptron, without requiring changes to the rest of the system. This composability is possible due to standardization in the field of machine learning, which enables machine learning researchers to conveniently incorporate their colleagues’ advances by quickly replacing a function rather than having to rewrite their code from scratch. The field of optics stands in stark contrast: implementations are often project-specific, each with their own conventions and quirks, and a baseline to compare these codes to with respect to accuracy and speed does not exist. This practice is time-consuming, error-prone and makes reproducing results challenging. Chromatix proposes a standard for wave-optics simulations to enable composition of a wide array of optical models (Fig. 1). The experiments presented in this paper all share a common, well-tested codebase and many more components are available in our documentation. We believe that the existence of both a standard library and baseline implementations can substantially speed up research and make it more reproducible.

Fig. 1. The design and components of Chromatix.

a, Chromatix combines wave-optics models, GPU acceleration and differentiability in a single library, providing a unified modeling framework to allow a wide range of applications. b, Chromatix implements a wide range of optical elements such as lenses, sensors, free-space propagation models for scalar and vectorial waves^70–72 and complex scattering samples^16,38. c, These elements can be combined to simulate a wide variety of experimental systems and solve a wide range of problems in computational optics. Green highlighted elements indicate the element or sample that would be optimized in each application. DMD, digital micromirror device; SLM, spatial light modulator; f, focal length; z_n, propagation distance.

Scalability

Optics is moving to ever larger fields of view (FOVs) and higher resolutions, sometimes requiring large compute clusters for sample reconstruction. A key requirement for new optics simulations is thus the ability to scale; researchers may want to run code on laptops for quick prototyping, but also easily scale up to GPU-clusters for large-scale sample reconstruction. Previous popular programming environments have made this difficult: NumPy runs only on central processing units⁴⁷; MATLAB requires specific code for GPU usage and does not support general-purpose automatic differentiation; PyTorch²⁴/TensorFlow²⁵ make writing GPU code with automatic differentiation relatively easy but it can be tricky to support multiple GPUs (both PyTorch and TensorFlow) or achieve good performance for typical operations in optical simulation that differ greatly from typical operations in neural networks (PyTorch). Writing device-specific programs for simulations requires substantial effort and calcifies their capabilities, which is not appropriate for the fast iteration demanded by scientific research. Chromatix instead relies on JAX²⁶ and its underlying XLA (accelerated linear algebra) compiler to support fast optical simulation on central processing units, GPUs and tensor processing units with only a single implementation (and without requiring custom lower level GPU code for fast operations as in PyTorch²⁴). JAX also offers several functions to automatically vectorize code (that is, parallelize a batch on a single GPU) or parallelize over multiple GPUs, independently of the description of the optical elements in an optical system²⁶. For example, with only a couple of lines of changes to the code we can scale a two-dimensional (2D), single-wavelength simulation to a 3D, multi-wavelength simulation running on multiple GPUs in parallel (see Extended Data Fig. 1 for code examples).

Implementation

In deep learning frameworks, these principles manifest themselves as models consisting of sequences of deep learning operations (layers). We observe a clear correspondence to optics, where optical systems consist of a sequence of optical elements and propagations. A key difference, however, is that the ‘hidden state’ of an optical system has a clear physical meaning: it is the complex light field moving through the system. To completely describe this field, and thus the state of the system at any time, additional information, such as the wavelength, polarization and spatial sampling, is required. The core idea behind Chromatix is that all this information can be encoded in a single, fundamental structure. Any optical element can then be written as a transformation of this structured field, and any optical system as a sequence of these elements. This allows Chromatix to model a wide variety of optical systems under a unified interface, which makes extending its capabilities straightforward.

Experiments

We present six computational experiments demonstrating four major features of Chromatix: solving inverse problems to reconstruct samples, accelerating reconstruction and optical design using deep learning, composing modular optical elements and models in arbitrary ways, and scaling optical simulation speed by an order of magnitude. To do this, we showcase both reproductions of existing computational methods in optics that rely on wave models as well as in silico demonstrations of solving inverse problems in underexplored combinations of optical phenomena.

Inverse problems for reconstructing samples

A microscope’s aberrations are usually assumed to be spatially invariant due to the calibration and simulation complexity of simulating and measuring field-varying PSFs, respectively. One research group^48,49 argued that most imaging systems are rotationally equivariant due to the rotational symmetry of many optical elements, and that many common aberrations therefore vary only along the radius away from the center of the FOV, up to a rotation. Measuring only the variation of aberrations along this radial dependence requires substantially less calibration time and is much more efficient to model (linear versus quadratic scaling with the number of rows of the camera sensor). This efficiency enables tractable deconvolution of spatially varying aberrations. The researchers introduced ‘ring deconvolution microscopy’⁴⁸, which models a spatially varying aberrated PSF efficiently by exploiting the rotational invariance present in many standard microscopes using incoherently illuminated samples or fluorescent samples emitting incoherent light. We implemented this ring deconvolution method in Chromatix (Fig. 2b) for UCLA (University of California, Los Angeles) Miniscope^50,51 data from a miniature widefield microscope. The microscope is modeled as a 4f optical system with rotationally invariant Seidel aberrations in the Fourier plane. After estimating Seidel coefficients from calibration images⁴⁸, the measured image of the sample is deconvolved using this rotationally invariant but spatially varying model of the PSF.

Fig. 2. Chromatix solves inverse problems in multiple types of sample and sample representation.

a, Sample type and representation for ring deconvolution microscopy. b, Implementation of ring deconvolution microscopy following ref. ⁴⁸. c–f, Highlighted regions show zoomed in cutouts from the center (red) and edge (blue) of the FOV on the bottom row. These images have been corrected to reduce vignetting for display purposes. Uncorrected images are shown in Extended Data Fig. 2. Measured image of incoherently illuminated rabbit liver c from a Miniscope with a gradient index lens without correction of spatially varying aberrations across the FOV. d, Spatially invariant (standard) deconvolution of the same FOV. e,f, Chromatix rotationally invariant deconvolution of the same FOV (e) and the original PyTorch rotationally invariant deconvolution that does not parallelize and cannot fit the whole FOV on 1 H100 GPU (80 GB of memory) (f). g, Sample type and representation for computational aberration correction using implicit neural networks. h, CoCoA self-supervised framework implementation using coordinate-based implicit neural networks for simultaneous aberration inference and 3D sample reconstruction from a single 3D measurement. i–k, Visual comparisons showing a maximum projection of a 4-µm slice from a raw aberrated measurement (i), original implementation reconstruction using PyTorch (j) and Chromatix reconstruction of the same dendrite (k). l–n, Validation of aberration recovery showing the true measured wavefront using direct wavefront sensing (l), inferred aberration from original PyTorch implementation (m) and Chromatix inference of the aberration (n). o, Sample type and representation for 3D refractive index microscopy. p, Refractive index microscopy to recover the 3D refractive index distribution of a strongly scattering sample from intensity measurements. q–s, Raw measurements (q) from multiple angles of coherent illumination are used to reconstruct the full 3D refractive index map of the tail of a D. rerio embryo at 24 hpf, demonstrated by the original MATLAB reconstruction of refractive index maximum projection (r) and the Chromatix reconstruction maximum projection (s). t,u, Slices through the full volume (colored borders denote slice regions) are shown to highlight the reduced grid artifacts in the Chromatix reconstruction (u(i), u(ii)) compared with the original MATLAB reconstruction (t(ii), t(ii)). For all panels, green highlights denote optimized parameters and dashed gray arrows denote propagation of gradients during iterative optimization.

We show, respectively, the measured image of incoherently illuminated rabbit liver, ring deconvolution using Chromatix that recovers detail across the whole FOV, standard deconvolution that only recovers detail in the center of the FOV and ring deconvolution from the original implementation^48,49 (Fig. 2c–f and Extended Data Fig. 2). Note that, compared with the original implementation, our reconstruction has a substantially larger FOV; our implementation was able to reconstruct a larger image by parallelizing across multiple GPUs. The original implementation fails to reconstruct the entire FOV of the camera (without suffering substantial degradation in reconstruction speed) due to overflowing the memory limitations of a single GPU (48 GB for an RTX 8000 or 80 GB for an H100). Chromatix’s implementation is also substantially faster, showing a 4.5× speedup versus the original PyTorch implementation on a single GPU and scaling up to almost 19× when using 8 GPUs (below).

Extended Data Fig. 2. Reconstruction performance of ring deconvolution without correction of vignetting.

Highlighted regions in a - d show zoomed in cutouts from the center (red) and edge (blue) of the field of view on the bottom row. Intensity values in a - d are normalized. Measured image of incoherently illuminated rabbit liver a from a Miniscope with a GRIN lens without correction of spatially varying aberrations across the field of view; b, Spatially invariant (standard) deconvolution of the same field of view; c, Chromatix rotationally invariant deconvolution of the same field of view; and d, original PyTorch rotationally invariant deconvolution which does not parallelize and cannot fit the whole field of view on 1 H100 GPU (80 GB memory).

In addition to voxel grids, Chromatix also enables INRs that in some cases improve the optimization loss landscape⁴⁴. CoCoA (coordinate-based neural representations for computational adaptive optics)⁴³ jointly reconstructs the sample (represented as an INR) and the aberrations (represented as Zernike coefficients) in a self-supervised fashion, as shown by the Chromatix implementation of the CoCoA method (Fig. 2g,h). Contrary to the previous section, the aberrations here are modeled as spatially invariant and the sample emits incoherent fluorescent light, rather than transmitting incoherent illumination. We show a maximum intensity projection of a 4-µm slice through a measured widefield mouse neuron volume, the reconstruction of the unaberrated sample from the original implementation and finally the reconstruction using Chromatix (Fig. 2i–k). We note that the Chromatix implementation retains more uniform dendrites that become spotted in the original implementation. Chromatix is also twice as fast at performing the reconstruction when using a single GPU, or almost 9× faster when using 8 GPUs (below). Further, on a fluorescent bead dataset with controlled, intentional aberrations provided in ref. ⁴³, the Chromatix implementation recovers the applied Zernike mode coefficients with a root-mean-square (r.m.s.) of 3.56 nm versus an r.m.s. of 6.97 nm for the original implementation⁴³ (r.m.s. is computed on the three nonzero Zernike modes that were used to intentionally aberrate the system; Extended Data Fig. 3).

Extended Data Fig. 3. Chromatix outperforms the original CoCOA implementation for estimating aberrations on fluorescent bead data.

a, Widefield measurement of fluorescent beads, without correction of aberrations and background fluorescence. b, Kang et al.⁴³ reconstruction of the fluorescent beads. c, Chromatix reconstruction of the fluorescent beads. d, Ground truth aberrated wavefront based on the imposed Zernike modes. e, Kang et al.⁴³ estimation of the aberrated wavefront. f, Chromatix estimation of the aberrated wavefront. g, Bar plot comparing the wavefront error (nm) between the true imposed coefficients and the estimated coefficients for each Zernike mode across both Chromatix (blue) and the original implementation⁴³ (red). Note that all Zernike modes aside from 5, 6, and 7 should have a coefficient of 0 nm in the imposed aberration. h, Bar plot of root mean square (RMS) of the difference between the true wavefront coefficients and the estimated coefficients across the three non-zero Zernike coefficients of the imposed aberration for both Chromatix and the original implementation⁴³. i, Bar plots of RMS of the pixelwise difference between the estimated aberrated wavefront and the true imposed wavefront for both Chromatix and the original implementation⁴³.

The loss of detail in the original implementation may be mitigated by increasing the number of layers of the INR as reported in ref. ⁴³, but here we compare reconstructions using identical network architectures. The Chromatix implementation uses a paraxial approximation for the field at the pupil plane while the original implementation in ref. ⁴³ mixes the exact model for the field at the pupil (requiring a higher sampling rate to avoid aliasing the propagation kernel) with a paraxial approximation of the second lens of the 4f system. In this case, the fully paraxial approximation used by Chromatix results in an improved reconstruction of the dendrite (Fig. 2i–k). This demonstration highlights not only our increased performance but also the utility provided by a standard set of models such as Chromatix when faced with the potential for model mismatch.

The research team in ref. ¹⁶ showed that computational imaging could also quantitatively recover the 3D refractive index distribution of strongly scattering samples (that is, beyond the first Born approximation) from intensity measurements, a quantity otherwise inaccessible to conventional widefield light microscopy. The sample (the tail of a D. rerio embryo 24 hours postfertilization (hpf)) is coherently illuminated at different angles¹⁷, and the refractive index (Fig. 2r) is recovered by matching the measurements (Fig. 2q) to a differentiable simulation (Fig. 2p) of imaging through the scattering sample. Typical samples are beyond the single-scattering regime and are hence modeled using a multislice approach⁵². While we use exactly the same forward model as the original MATLAB implementation¹⁷, our implementation is 3–13× faster than the original implementation on equivalent volumes (bringing reconstruction time from hours to minutes). Our implementation is also far fewer lines of code (approximately 25 lines for a differentiable simulation in Chromatix versus approximately 107 lines in the original implementation^16,17) and more flexible compared with the original implementation with respect to changes in the forward model or optimization parameters due to the use of automatic differentiation. This increase in speed allows us to choose more appropriate reconstruction settings so that the large grid artifacts in the original reconstruction (Fig. 2r) are removed in the Chromatix reconstruction (Fig. 2s).

Programmable optics and deep learning

The commercial availability of SLMs has made fine-grained control of light possible through millions of controllable pixels. These degrees of freedom are often used for holography but can also be used to design a PSF for a specific purpose⁵³, such as volumetric snapshot imaging for fluorescent 3D samples. The dimensionality of this optimization problem requires gradient-based optimization to be effectively solved and is well-suited to GPU acceleration. One research group³ introduced a deep learning method for engineering 3D snapshot PSFs that combined a programmable microscope design (the Holoscope) with a neural network to reconstruct fluorescent 3D volumes from 2D snapshot images taken by the microscope (Fig. 3a). The parameters of the neural network along with the pixels of the programmable phase mask implemented by an SLM are jointly optimized using a differentiable simulation of the microscope. In this snapshot microscope, the PSF essentially acts as a compression function from the 3D volume to the 2D image. The volume is then reconstructed from this image using a FourierNet neural network, so that structural priors of the sample type can be taken advantage of by both the optics and the computational reconstruction algorithm.

Fig. 3. Chromatix enables deep learning for optical design and acceleration of inverse problems.

a, Holoscope3 implementation in Chromatix showing a programmable 3D snapshot microscope, compressing volumetric information into a 2D representation with subsequent FourierNet reconstruction. b–e, Holoscope demonstration using Chromatix showing the sample-specific PSF (b), simulated 2D image of the simulated 3D sample (c) (approximately 0.01 s to capture), ground truth 3D volume (d) (approximately 7.0 s to capture via confocal microscopy assuming a scan speed of 100 ns per voxel) and Chromatix-enabled 3D reconstruction (e) from the single simulated 2D image. f, DeepCGH architecture using a UNet and a propagation step to directly generate a hologram in a single feedforward step from target 3D patterns. g–l, Demonstration of DeepCGH12 using Chromatix showing requested stimulation patterns at three planes spaced 10 mm apart around the focal plane (g,i,k) and their resulting simulated intensity distributions (h,j,l). Colored insets show detail of the 3D patterns at each plane. Intensity values in g–l are normalized. For all panels, green highlights denote optimized parameters of either neural networks or optical systems.

This microscope design can therefore be programmed to function as a snapshot microscope optimized for various sample types, while using exactly the same hardware. The microscope is modeled using a 4f system with an SLM (phase mask) in the Fourier plane and is optimized for whole-brain imaging of fluorescently labeled D. rerio larvae. The PSF of the 4f system is simulated with coherent propagation, and the image is simulated as the incoherent sum of these PSFs that is efficiently implemented as a convolution of the PSF and the sample intensity. We show the learned PSF (Fig. 3b), the simulated 2D measurement of a virtual zebrafish volume (Fig. 3c) and the ground truth volume and simulated reconstruction³ (Fig. 3d,e). Chromatix reproduces the original results³ nearly exactly: on a test set of 10 volumes and their simulated images, reconstruction networks trained with identical PSFs offer a structure similarity index measure on a test set of 10 volumes of 0.979 ± 0.003 (mean ± standard error; higher is better) for both Chromatix and the original implementation³ (not significantly different at P = 0.695 via two-sided t-test, Extended Data Fig. 4). Chromatix also outperforms the original implementation³ in training speed by a factor of approximately 7× (Fig. 5). Practically, this reduces the optimization time for a single PSF from weeks to days.

Extended Data Fig. 4. Chromatix matches the in silico Holoscope reconstruction performance of Deb et al.

a, Holoscope PSF optimized for zebrafish with a circular aperture in the original implementation by Deb et al.³. b, Simulated camera image showing a 3D zebrafish sample captured in a single 2D image. c, Ground truth confocal volume of the brain of a larval zebrafish (D. rerio), which will be over an order of magnitude slower to capture than the single snapshot image in (b). d, Reconstruction using the original implementation by Deb et al.³. e, Reconstruction using Chromatix. On a test set of 10 volumes, reconstruction networks trained with identical PSFs offer an SSIM of 0.979 ± 0.003 (mean ± standard error; higher is better; not significantly different at p = 0.695 via two-sided t-test) for both Chromatix and the original implementation³ and a peak signal to noise ratio (PSNR) of 35.91 ± 0.92 for Chromatix versus 36.99 ± 1.07 for the original implementation³ (not significantly different at p = 0.063 via two-sided t-test).

Fig. 5. Chromatix is the fastest implementation of existing computational optics methods.

Vertical axis shows relative speedup of Chromatix on 1–8 GPUs compared with the original implementations of each method as the baseline (represented via the gray dashed line at 1×). Points are centered on the mean relative speedup for each method. Error bars show standard error of speedups on individual iterations of the optimization for each method. The distribution for each method is visualized as a violin plot in a lighter shade. For all methods, speedups are relative to the original single GPU implementations. CGH, computer-generated holography.

SLMs also enable computer-generated holography systems for optogenetics, where 3D holographic stimulation patterns are used to perturb neural activity in the brain. Most holography systems rely on some form of iterative optimization (for example, refs. ^33,34,54,55) to find the phase to display on the SLM. While this produces accurate solutions, iterating does become problematic when speed is paramount. For optogenetics, point cloud holography can be used to stimulate multiple neurons without iterative optimization of phase patterns, but this only allows for placing copies of a single pattern at the desired locations¹⁵. Due to the interest in holography for displays, fast holography algorithms for arbitrary patterns have emerged that use neural networks to quickly generate a hologram given a target pattern^42,56. Applied to optogenetics, DeepCGH¹² also demonstrates fast computer-generation of holograms by training a neural network to generate phase patterns from intensity images of arbitrary 3D patterns in a single feedforward inference step. We implemented DeepCGH¹² (Fig. 3f). We show the desired target patterns and resulting simulated pattern at three different depth planes using the phase pattern produced by the DeepCGH method in Chromatix (Fig. 3g–l). We achieve nearly identical results to the original TensorFlow implementation: on a test set of 16 target patterns, Chromatix achieves a structure similarity index measure of 0.985 ± 0.001 (mean ± standard error; higher is better) versus 0.982 ± 0.001 for the original implementation (significantly different at P = 0.018 < 0.05 via two-sided t-test, Extended Data Fig. 5) and peak signal to noise ratio of 35.40 ± 0.37 (mean ± standard error; higher is better) for Chromatix versus 34.95 ± 0.16 for the original implementation¹² (not significantly different at P = 0.177 via two-sided t-test). Our implementation is approximately 17 lines of code for a differentiable hologram simulation versus 33 lines in the original work¹². While achieving the same quality, Chromatix provides a 2.5× performance improvement on a single GPU, which increases to over 10× when using 8 GPUs in parallel (Fig. 5).

Extended Data Fig. 5. Chromatix matches the in silico DeepCGH performance of Eybposh et al.

a,d,g, Target pattern of dots of random intensity at planes -10 mm, 0 mm, and 10 mm from the focal plane respectively. b,e,h, Simulated pattern using holograms generated by a DeepCGH¹² model trained using Chromatix at planes -10 mm, 0 mm, and 10 mm from the focal plane respectively. c,f,i, Simulated pattern using holograms generated by a DeepCGH model trained using the original implementation by Eybposh et al.¹² at planes -10 mm, 0 mm, and 10 mm from the focal plane respectively.

Flexible modeling with optical building blocks

Because Chromatix models are constructed from components that can be flexibly combined (Fig. 1), we can straightforwardly construct complex optical models and also optimize them with arbitrary objective functions. We show another programmable microscope modeled as a 4f system with an SLM in the Fourier plane, followed by a neural network-based reconstruction step (Fig. 4a–f). The objective in this demonstration is to optimize the PSF of this programmable microscope to perform spectroscopic single-molecule localization⁵⁷ from a single snapshot image: that is, to reconstruct multicolor point sources using only a single-channel image. The simulated samples consist of several point sources incoherently emitting fluorescence at 25 wavelengths from 400 nm to 650 nm that are simulated in parallel using Chromatix. We train the neural network to reconstruct the multicolor sample at the corresponding 10-nm intervals, giving us a hyperspectral cube from a single-channel 2D measurement. The optimized PSF allows visual classification of the color of these point sources on a monochrome simulated camera image (Fig. 4c,d) by taking advantage of different fringe patterns for different wavelengths. The reconstruction (Fig. 4f) reasonably matches the true colors of the points in the sample (Fig. 4e). We highlight that here we are optimizing the same programmable microscope model that was used for snapshot microscopy (Fig. 3a), but for an entirely new combination of sample type and objective.

Fig. 4. Chromatix enables arbitrary combinations of optical models.

a, Demonstration of PSF engineering for spectroscopic single-molecule localization microscopy using deep learning, where a neural network reconstructs both the structure and spectrum of a sparse 2D point sample from a single-channel image. b, Microscope model for multicolor PSF optimization with a single SLM in the Fourier plane. c, Optimized multicolor PSF for spectral imaging using a neural network, with different wavelengths and colors overlaid. d, Simulated single-channel image of multicolor fluorescent point sources. e, True simulated multicolor point sources, with different wavelengths and colors overlaid. f, Reconstructed multicolor point sources, with different wavelengths and colors overlaid. g, Iterative optimization workflow generating optimal phase masks for scattering-compensated holography. i, Chromatix model of holographic pattern formation through scattering media, which combines the holography model of Fig. 3f and the scattering sample model of Fig. 2p. h, Peak axial intensity distribution (normalized within the range 0–1) along the direction of propagation for the target pattern, an uncorrected pattern and the corrected pattern that are visualized below. j–m, Visual evolution of holographic pattern quality. Target pattern (j) is compared to the pattern constructed by an optimized hologram simulated through free space (k), the pattern constructed by the same hologram simulated with scattering through a sample with a known 3D refractive index distribution (l) and finally the pattern constructed by the Chromatix-corrected hologram simulated through the same scattering volume (m). For all panels, green highlights denote optimized parameters and dashed gray arrows denote propagation of gradients during iterative optimization.

Arbitrary combinations of wave-optics models can also open up further applications of differentiable simulations to biological research. Optogenetic experiments rely on complex, 3D patterns of light to selectively control the behavior of neurons, often using holography^12,14,15. 3D holography in itself is challenging, but this is compounded in optogenetics by the scattering nature of the biological tissue. As the light propagates through the tissue, it gets scattered, possibly activating the wrong neurons and increasing phototoxicity^58–60. We show how Chromatix can be used to optimize the desired holographic pattern in such a strongly scattering tissue by imaging the scattered pattern (Fig. 4g–m). We model this scenario as a plane wave incident on a phase mask (SLM), focusing through a thin lens and finally propagating through a scattering volume. The propagation through the volume is modeled using the multislice beam propagation method (using the same code to model the scattering sample as Fig. 2p), and we observe in simulation the intensity throughout the entire volume. Without correction, the stimulation is uneven due to the unaccounted-for scattering (Fig. 4h). Once we include the observed scattered intensity as feedback in the optimization, we obtain a near uniform stimulation (blue line in Fig. 4h) across the whole axial range. This serves as an in silico demonstration that Chromatix enables researchers to rapidly implement and iterate on their ideas, transforming intuition into tangible results.

High performance through parallelization

To demonstrate the computational performance and scalability of Chromatix, we benchmark iteration speed for all of training and optimization problems we have presented. Chromatix has superior performance across all reproduced optical methods, ranging from 2 to 6× on single GPUs to 22× faster on 8 GPUs in the best case (Fig. 5). Single GPU performance improvements are typically due to less overhead after compilation through JAX compared with other implementations in MATLAB, PyTorch or TensorFlow. In general, order-of-magnitude improvements are possible via parallelization with Chromatix. These parallelization schemes can be implemented on Chromatix models with virtually no changes to the code defining the models due to Chromatix being implemented in JAX²⁶. This acceleration allows Chromatix to scale to solving large problems, and also makes existing inverse problem solutions substantially more tractable as demonstrated by the larger FOV for the Chromatix reconstruction in ring deconvolution microscopy (Fig. 2a–f) as well as the substantial decrease in optimization time for refractive index microscopy (Fig. 2o–u) and snapshot PSF optimization with deep learning (Fig. 3a–e).

Discussion

We introduced Chromatix, an open-source differentiable wave-optics simulation library that standardizes a wide and growing range of wave-optics models. We anticipate that such an open-source standard library could address the reusability of computational optics methods, which is limited by a lack of standards and bespoke calibration procedures. However, open-source libraries depend on the ability of the community to contribute new features and address limitations.

In this open-source spirit, several of the major features in Chromatix were added by participants at hackathons we held to teach Chromatix. Across three hackathons at both HHMI Janelia Research Campus and Johns Hopkins University, a total of 32 participants with backgrounds in biology and optics—but not necessarily advanced numerical computing—were able to implement models and features required for their research in just a few days. Participants at subsequent hackathons and other researchers could then build on the work added to the library due to the common framework imposed by Chromatix. This highlights that Chromatix’s accessibility allows researchers to explore state-of-the-art computational optics, and also that standardization of models makes it easier to share and disseminate their work. We expect this trend of open-source contributions to continue, implementing an ever-increasing collection of models that would be necessary for use as a standard library in computational optics.

We are also continuing to develop Chromatix in several ways to address its limitations. Chromatix simulates some optical elements such as lenses in idealized ways (with experimental support for thicker lenses); the library of optical elements can be expanded with more realistic modeling of existing lens geometries with lens coatings²⁷, and sensors with noise patterns and pixel responses that are at present missing from Chromatix. Further, while Chromatix focuses on wave-based modeling of optics, it may be necessary to use ray-based modeling for computationally efficient simulation of complicated series of lenses or other surfaces. The research group in ref. ⁶¹ showed how to combine ray-based models with diffractive models, providing a path for Chromatix to incorporate ray-based models of lenses in the future. Computationally efficient modeling of forward and backward scattering through inhomogeneous samples can be implemented via the modified Born series⁶². Chromatix also does not support rigorous coupled-wave analysis, which is often used to efficiently simulate and design metasurfaces^63,64. Another study in ref. ⁶⁵ showed an efficient differentiable method for this inverse problem that could allow Chromatix to support metasurface design. Finally, Chromatix currently models completely coherent propagation, which can be expanded to include support for partially coherent propagation models⁶⁶ as well as changes in the spectrum of light due to propagation⁶⁶.

Looking further, we foresee a wide range of possible applications for fast simulations provided by Chromatix, such as architecture search for designing optical systems (analogous to neural architecture search for optimizing network architectures⁶⁷) or integration with hardware control software such as Micro-Manager⁶⁸ or Pycro-Manager⁶⁹ to enable automatic calibration or alignment with hardware-in-the-loop techniques⁴² (increasing reproducibility for computational optics techniques); both of these applications would benefit from being able to evaluate many simulations quickly and in parallel. Just as accessible deep learning libraries enabled deep learning researchers to explore a rich space of neural network architectures and methods, we expect the standardization of computationally efficient and scalable wave-optics models in Chromatix to enable the exploration of a rich space of computational optics methods that may have previously been impractical to implement.

Methods

Ring deconvolution microscopy

Visualized results

All procedures and data followed the original work in ref. ⁴⁸ (Fig. 2a–f). We modeled a microscope with spatially varying aberrations, namely the UCLA Miniscope v.3⁵⁰ with Ximea MU9PM-MBRD 12 bit, 2.2 micrometer pixel sensor⁴⁸. We showed deconvolution of rabbit liver tissue embedded on a glass slide illuminated with 550-nm incoherent light, also shown originally in ref. ⁴⁸. Imaging was described as the incoherent sum of coherently simulated PSFs with spatially varying aberrations. Ring convolution efficiently models spatially varying aberrations by convolving each ring of a sample (the pixels at a certain radius from the center of the FOV) with the aberrated PSF corresponding to that ring⁴⁸. Because Seidel aberrations are rotationally invariant, the PSF can be simulated at each pixel along a single diagonal path away from the center of the FOV, which defines up to a rotation the shape of the PSF at any other pixel⁴⁸.

For reconstruction, we normalized pupil coordinates (u, v) across the camera image to have magnitude less than or equal to 1 within the original cropped 1,024 × 1,024 pixel region used in ref. ⁴⁸, and scale coordinate magnitudes beyond 1 outside of this region. This allowed us to use the original estimated Seidel coefficients from the calibration performed in ref. ⁴⁸ to simulate the spatially varying aberrations. For the Miniscope data, these coefficients were 0.85, 0.56, 0.25, 0.29 and 0 waves for the five Seidel aberrations, respectively. The spatially varying PSF was simulated at all pixel coordinates along one radial line from the center of the image to one corner, and these PSFs were rotationally Fourier transformed for use in ring convolution, as outlined in ref. ⁴⁸. The forward model could then be defined by the ring convolution between the spatially varying PSFs and the sample. Deconvolution was performed by initializing the sample as the measured camera image, and then iteratively performing gradient descent with the Adam optimizer³⁷ to optimize the pixels of the reconstruction by comparing the loss between the simulated image from the forward model and the measured camera image⁴⁸.

Scalability results

We demonstrated reconstructions parallelized across 1, 2, 4 or 8 GPUs (NVIDIA H100s). Parallelization was performed by spreading different rings (different radii) of the ring convolution across different GPUs. To obtain the final image, the different rings were synchronized across the different GPUs after the convolution was performed. We reconstructed a 1,024 × 1,024 pixel image so that the original implementation (which is single GPU only) could complete a reconstruction. We reconstructed a simulated image of point sources rather than an experimental image, because we were only concerned with measuring the reconstruction iteration speed and not reconstruction quality. We measured the iteration time to perform the forward simulation and gradient descent update of the reconstruction with the Adam optimizer.

Computational aberration correction

Visualized results

All procedures and data followed the original work in ref. ⁴³ (Fig. 2g–n). We modeled a widefield microscope with Zernike aberrations from American National Standards Institute indices 3–14 (without normalization). The sample shown is a 3D volume of fixed mouse brain slices (Thy1-GFP line M) with a number 1.5 cover glass (0.16–0.19 mm thickness) tilted at 3° on top to emulate a glass cranial window⁴³. The sample was excited in widefield by an expanded 488-nm laser beam. The emitted fluorescence was collected by a 1.1 numerical aperture objective lens, relayed to the camera by a 4f system.

We first optimized an INR to recreate the measured widefield volume by minimizing the structured similarity loss between the output of the network and the measured widefield volume. The INR used to represent the volume took as input an x–y location (radially encoded) and output the pixel intensity values for that location across all 200 z-planes of the sample. We then performed a second stage of optimization of the INR using the forward model of the aberrated microscope. Here we simulated imaging the volume output by the INR using the forward model and minimized the structured similarity loss between the simulated widefield volume and the measured widefield volume. The simulation of the microscope required first simulating the aberrated coherent PSF of the system and then convolving the 3D PSF with the 3D sample volume. In this second stage, we also updated the Zernike coefficients of the forward model of the aberrated microscope. The final output was an INR of the deconvolved sample and the Zernike coefficients representing the aberration of the microscope.

Scalability results

We measured iteration time during the second stage of training, including the time to perform a forward pass of the neural representation, a forward simulation of the aberrated microscope and a backward pass to update the neural representation parameters as well as the Zernike coefficients of the aberrated microscope model using the Adam optimizer. We simulated a measured volume of point sources rather than the experimental data used in the visualization results to demonstrate the speed of optimization of Chromatix across 1–8 GPUs (NVIDIA H100s).

Refractive index microscopy

Visualized results

All procedures and data followed the original work in ref. ¹⁷ (Fig. 2o–u). We showed a reconstruction from the tail of a zebrafish (D. rerio) embryo (24 hpf). Images were measured by illuminating the transparent sample with coherent light from a 660-nm red laser (Thorlabs, catalogue number S4FC660)¹⁷. This light passed through the sample at a specific angle and then scattered through the transparent sample¹⁷. The light scattered through the sample at the focal plane of the objective (0.37 numerical aperture) was measured by the camera¹⁷. Many images were collected with different angles of illumination. Measured images were downsampled via cubic interpolation to 80% of their original size in both height and width, as performed in ref. ¹⁷. Certain images were manually omitted from use in the iterative reconstruction process for each FOV, as outlined in the data and reconstruction settings provided in ref. ¹⁷. Using the measured images, we reconstructed a volume of size 400 × 800 × 800 voxels (FOV 1) at a voxel size of (0.43125 µm z, 0.2875 µm y, 0.2875 µm x).

Reconstruction is an iterative optimization process over the voxels of the sample. Each voxel represented the change in refractive index at that location in the sample. The reconstruction was initialized to an all zero change in refractive index¹⁷. Images were simulated from each angle of illumination that was measured using multislice propagation through the scattering reconstruction, and the mean-squared error between each simulated image and each measured image was computed¹⁷. The gradient of this loss was used to update the reconstruction via gradient descent. Once all images had been used to update the reconstruction, the current reconstruction was smoothed via a proximal total variation operator to a relative error of 10⁻⁴. This smoothed reconstruction was updated from the previous smoothed reconstruction using the Nesterov accelerated gradient¹⁷. This whole process (forward simulating and updating from all the images) was repeated until the reconstruction converged¹⁷.

Scalability results

We reconstructed a volume of 400 × 800 × 800 from FOV 1 where we used all 168 images. We measured the iteration time to perform one forward simulation and backward pass to compute the gradient descent update, not including any regularization. We split these 168 images across 1, 2, 3, 4, 6, 7 or 8 GPUs (NVIDIA H100s). For a given number of GPUs, we forward simulated and computed gradients for that number of images in parallel. The gradients from this batch of images were averaged across all GPUs to compute the update for the given set of images. We reported iteration time as the time to process a single image by dividing the measured iteration time for a batch by the size (number of images) of the batch.

Snapshot PSF engineering using end-to-end deep learning

Visualized results

All procedures and data followed the original work in ref. ³. We followed the two-stage optimization procedure outlined in ref. ³, where we first optimized a weak FourierNet reconstruction network and phase mask to obtain a snapshot PSF for zebrafish larvae (D. rerio) at 6–8 days postfertilization³. We then optimized a larger FourierNet reconstruction network while fixing the phase mask to obtain higher quality reconstructions (Fig. 3a–e). This optimization process used 58 volumes of pan-neuronally labeled larval zebrafish (expressing nuclear-restricted GCaMP6s/7f) that were scanned at high resolution using a high-resolution confocal microscope. These volumes were all resampled to the same pixel size (1.0 µm z, 1.5 µm y, 1.5 µm x) for training. The simulated snapshot microscope design imaged a 386-µm diameter region of the whole zebrafish that was 250 µm tall, which simulated imaging zebrafish in 3D with an aperture in the system blocking light from the sample outside of the desired range. The simulated camera FOV was 1,132.8 µm.

Optimizing a PSF involved sampling different 386-µm diameter regions from the 58 total training volumes of larval zebrafish, applying affine transformation and brightness augmentations, simulating imaging the fluorescent volume at a wavelength of 513 nm using the snapshot microscope and then reconstructing the 3D image from the 2D simulated image. Simulating imaging in this stage of optimization required first simulating the PSF and then convolving each plane of the PSF with each corresponding plane of the sample after downsampling the PSF plane to the resolution of the sample plane. The final 2D image was the sum of these planewise convolutions. We then jointly updated both the weak neural network parameters for reconstruction and the phase mask parameters (pixels) using a loss that averaged the mean-squared error between the reconstruction and the true volume as well as the mean-squared error between the high pass filtered reconstruction and high pass filtered true volume. In the second stage of training, we used a more powerful neural network and fixed the phase mask parameters.

Scalability results

We reported iteration times for the first stage of training, where we optimized a phase mask (2,560 × 2,560 pixels) by jointly optimizing phase mask pixels and reconstruction neural network parameters. Simulating the PSF at 64 planes spanning 250 µm where each plane of the PSF had 2,560 × 2,560 pixels meant that we simulated a PSF containing 0.419 unique gigavoxels. The neural network reconstructed 64 planes spanning 250 µm where each plane of the volume had 512 × 512 pixels. Here for both the original implementation in ref. ³ using PyTorch and the Chromatix implementation, we reported iteration times across 1, 2, 4 or 8 GPUs (compared on NVIDIA RTX 8000s due to limitations of old PyTorch code). Parallelization for both the original implementation and the Chromatix implementation involved splitting the simulation of different planes (both for the PSF and for the imaging) across different GPUs, so that each GPU simulated its own chunk of the overall volume. Forming the final 2D image by summing across planes required synchronization across the GPUs. Iteration times included the forward simulation of the PSF, simulating the 2D image by convolving the PSF with the sample, reconstructing the 3D volume using the neural network and jointly updating the parameters of both the neural network and the phase mask using the rectified Adam optimizer.

Computer-generated holography using deep learning

Visualized results

All procedures and data followed the original work in ref. ¹². We generated a set of target patterns of 2,048 × 2,048 pixels for which we wanted to generate holograms (Fig. 3f–l). Each target pattern consisted of three planes of dots with varying intensities. The goal was to generate a hologram (phase mask pattern) that would reproduce these three target patterns at the three desired image planes. The system we modeled for holography consisted of a phase mask (quantized to 8 bits) illuminated by a plane wave, followed by a Fourier transforming lens to bring the generated pattern to the first image plane¹². We simulated the resulting pattern for these three planes at −10 mm, 0 mm or 10 mm from the focal plane of the lens.

We trained a neural network to generate these holograms by predicting the complex field (amplitude and phase) at the first image plane given the target pattern¹². We then propagated this field back to the phase mask plane to obtain the phase values for the pixels of the phase mask¹². This phase mask ess propagated to the three image planes to obtain the resulting pattern¹². The network parameters were updated according to the accuracy loss between the resulting pattern and the target pattern¹² This training was performed in batches of 16 patterns at a time.

Scalability results

We reported iteration times for training on a single batch of 16 patterns. Iteration time measured the time to perform a forward pass of the neural network, propagate the result to the phase mask plane to compute the phase mask pixel values, propagate to the three target image planes and finally update the parameters of the neural network based on the accuracy loss between the simulated patterns and the target patterns. Parallelization was performed by splitting up the batch of 16 patterns across 1, 2, 4 or 8 GPUs (NVIDIA H100s).

Multicolor PSF engineering

We demonstrated how a single phase mask could be engineered for spectral imaging of 25 different wavelengths ranging from 400 nm to 650 nm using a single-channel camera image (Fig. 4a–f). We simulated a programmable widefield fluorescence microscope and optimized its PSF to enable reconstruction of point sources emitting light of different wavelengths per point (at the focal plane). The reconstruction was performed using a FourierNet neural network³, which involved a global convolution (implemented as a Fourier convolution for efficiency) followed by direct convolution layers with smaller filters. We jointly trained both the reconstruction network and the pixels of the phase mask. The reconstruction network could exploit the sparsity of the sample as well as the PSF optimized to encode these specific wavelengths via fringe patterns to accurately reconstruct the points at their respective wavelengths.

Computer-generated holography through scattering media

We demonstrated the optimization of a hologram through scattering tissue rather than a homogeneous medium, as is typically assumed in computer-generated holography (Fig. 4g–m). This simulation followed a typical Fourier holography system, similar to the system used in our demonstration of computer-generated holography using deep learning. However, here we added the multislice scattering sample model that was used for refractive index microscopy and we also optimized the phase mask directly via gradient descent based on the Pearson correlation loss between the simulated pattern and the target pattern. Here we did not relay a focal plane to the camera but rather directly computed the intensity of the resulting pattern at each plane. The target pattern here was a 3D distribution of points of uniform intensity. For the scattering tissue, we used a volume 50 × 1,200 × 1,200 voxels large containing a reconstruction of a C. elegans head. We artificially scaled the refractive index values in the sample by a factor of five times to highlight the scattering for this demonstration, even though this was not a realistic sample. We first optimized a hologram through free space, using angular spectrum propagation through the same distance as the thickness of the scattering volume. We then simulated this hologram through both free space and through the scattering volume using the multislice beam propagation method. Finally, we optimized a hologram through the scattering volume and again simulated this hologram through the scattering volume to demonstrate that the intensity of the resulting points remained uniform throughout the volume, whereas the hologram optimized using the typical free-space model resulted in degradation of the resulting pattern.

Scalability results

For ring deconvolution microscopy, computational aberration correction, refractive index microscopy, snapshot PSF engineering and computer-generated holography we presented a comparison of timing between the original implementations and the Chromatix implementations (Fig. 5). Generally, we reported the average time to perform one iteration of training or optimization when solving these inverse problems. We show these results in Fig. 5 as relative performance versus number of GPUs used to parallelize each method. When measuring iteration times, outliers can arise due to compilation of the code by JAX or the GPU driver, scheduling of the process by the operating system or power management of the central processing unit and GPU(s); any and all of these can cause rare spikes in iteration time that can be ignored. For all methods, outlier iteration times (3 standard deviations higher than the mean) were discarded before plotting, which could lead to different numbers of samples for each point shown in Fig. 5. Different methods were also simulated on different numbers of GPUs in parallel depending on how many GPUs could cleanly parallelize the workload for a given problem. We collected the following numbers of iteration times for the various methods: snapshot PSF engineering 1 GPU: 499 iterations, 2 GPUs: 499 iterations, 4 GPUs: 499 iterations, 8 GPUs: 499 iterations; ring deconvolution 1 GPU: 4,980 iterations, 2 GPUs: 4,983 iterations, 4 GPUs: 4,981 iterations, 8 GPUs: 4,992 iterations; refractive index microscopy 1 GPU: 24,919 iterations, 2 GPUs: 12,530 iterations, 3 GPUs: 8,362 iterations, 4 GPUs: 6,284 iterations, 6 GPUs: 4,192 iterations, 7 GPUs: 3,494 iterations, 8 GPUs: 3,128 iterations; computer-generated holography 1 GPU: 99 iterations, 2 GPUs: 100 iterations, 4 GPUs: 99 iterations, 8 GPUs: 98 iterations; computational aberration correction 1 GPU: 998 iterations, 2 GPUs: 995 iterations, 3 GPUs: 996 iterations, 4 GPUs: 996 iterations, 5 GPUs: 991 iterations, 6 GPUs: 996 iterations, 7 GPUs: 992 iterations and 8 GPUs: 998 iterations. Note that refractive index microscopy performed fewer iterations scaling inversely with number of GPUs available due to greater numbers of GPUs supporting larger batch sizes to be optimized at a time. See the corresponding sections in Methods for more details on the scalability results for each application.

Reporting summary

Further information on research design is available in the Nature Portfolio Reporting Summary linked to this article.

Online content

Any methods, additional references, Nature Portfolio reporting summaries, source data, extended data, supplementary information, acknowledgements, peer review information; details of author contributions and competing interests; and statements of data and code availability are available at 10.1038/s41592-026-03121-x.

Author contributions

All authors contributed to the development of the Chromatix library. D.D. and G.-J.B. contributed equally by leading the development of the library and conducting the computational experiments for the paper. D.D., G.-J.B. and S.C.T. wrote the paper with contributions from all authors.

Peer review

Peer review information

Nature Methods thanks Christopher Metzler and the other, anonymous, reviewer(s) for their contribution to the peer review of this work. Peer reviewer reports are available. Primary Handling Editor: Nina Vogt, in collaboration with the Nature Methods team.

Funding

This work was funded by the Howard Hughes Medical Institute (S.C.T., D.D., G.-J.B., C.A., M.C.S.). B.D. was supported by a grant from the German-Israeli Foundation for Scientific Research and Development (GIF, grant number G-1566-413.13/2023). C.J. and the laboratory (Y.K.) are supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (grant number 2023R1A2C3004040) as well as The Ministry of Trade, Industry and Energy (MOTIE) and Korea Institute for Advancement of Technology (KIAT) through the ‘International Cooperative R&D program (Task Number P0028516)’. L.W. is a Chan Zuckerberg Biohub SF investigator; US Air Force Office of Scientific Research grant award number FA955-22-1-0521. S.C. and the laboratory (S.Y. and S.C.S.) are supported by the Chan Zuckerberg Initiative (grant numbers 2021-225666 and 2023-321173), the National Institutes of Health (NIBIB R21 EB033629, grant number NIGMS R35 GM155424) and the University of Texas at Austin; Cockrell School of Engineering; Chandra Family Department of Electrical and Computer Engineering. S.Y.H. is supported partly by grant number 2021-234544 from the Chan Zuckerberg Initiative DAF, an advised fund of the Silicon Valley Community Foundation. N.C.P. and the laboratory (C.C., M.H.E. and J.V.-T.) are supported by a Beckman Investigator grant from the Arnold and Mabel Beckman Foundation, a Sloan Fellowship in Neurosciences from the Alfred P. Sloan Foundation and a Kavli Innovation grant.

Data availability

We used publicly available data from previous work for all experiments shown. For ring deconvolution⁴⁸ (Fig. 2a–f), the data we used are available via GitHub at https://github.com/apsk14/rdmpy. For computational aberration correction⁴³ (Fig. 2g–n), the data we used are available via GitHub at https://github.com/iksungk/CoCoA. The measured wavefront for computational aberration correction (Fig. 2l) was obtained on request from in ref. ⁴³. For refractive index microscopy¹⁷ (Fig. 2o–u), the data we used are available via Dataverse at https://dataverse.tdl.org/dataverse/DMD-MLA-01. For snapshot PSF engineering³ (Fig. 3a–e), the data we used are available via figshare at 10.25378/janelia.25277269.v1 (ref. ⁷³). For holography through scattering (Fig. 4g–m), the data are available via GitHub at https://github.com/Waller-Lab/multi-slice (ref. ¹⁶). Any other data shown were completely programmatically generated in simulation.

Code availability

The Chromatix library is available via GitHub at https://github.com/chromatix-team/chromatix. The documentation and installation instructions are available via Chromatix at https://chromatix.readthedocs.io. This work used version 0.4, and Chromatix is distributed under the MIT License.

Competing interests

The authors declare no competing interests.

Extended data

is available for this paper at 10.1038/s41592-026-03121-x.

Supplementary information

The online version contains supplementary material available at 10.1038/s41592-026-03121-x.