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 frameworks24,25,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 problems3,9,18,19,20,21,22,23. Differentiability also has the potential to improve solutions to inverse problems in optics. Traditional inverse problem approaches simplify both the sample and the optics33,34,35,36, whereas differentiable models can handle more realistic complexity. Automatic differentiation opens up a new class of gradient-based optimizers such as Adam37, which can improve reconstruction fidelity by allowing arbitrarily complex physics simulations (for example, scattering16,38 or sample deformation39) 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 algorithms40. 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 fidelity41,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 fields39,43,44,45. INRs have been applied to separation of motion artifacts from sample dynamics39, estimation of dynamic aberrations44, reconstruction of 3D quantitative phase of scattering samples46 and aberration correction without wavefront sensors or calibration measurements43. 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.
The alternative text for this image may have been generated using AI.
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 waves70,71,72 and complex scattering samples16,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; zn, 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 units47; MATLAB requires specific code for GPU usage and does not support general-purpose automatic differentiation; PyTorch24/TensorFlow25 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 JAX26 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 PyTorch24). 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 system26. 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 group48,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’48, 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) Miniscope50,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 images48, 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.
The alternative text for this image may have been generated using AI.
a, Sample type and representation for ring deconvolution microscopy. b, Implementation of ring deconvolution microscopy following ref. 48. 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 implementation48,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).
In addition to voxel grids, Chromatix also enables INRs that in some cases improve the optimization loss landscape44. CoCoA (coordinate-based neural representations for computational adaptive optics)43 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. 43, 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 implementation43 (r.m.s. is computed on the three nonzero Zernike modes that were used to intentionally aberrate the system; Extended Data Fig. 3).
The loss of detail in the original implementation may be mitigated by increasing the number of layers of the INR as reported in ref. 43, 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. 43 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. 16 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 angles17, 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 approach52. While we use exactly the same forward model as the original MATLAB implementation17, 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 implementation16,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 purpose53, 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 group3 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.
The alternative text for this image may have been generated using AI.
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 reconstruction3 (Fig. 3d,e). Chromatix reproduces the original results3 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 implementation3 (not significantly different at P = 0.695 via two-sided t-test, Extended Data Fig. 4). Chromatix also outperforms the original implementation3 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.
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 locations15. 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 pattern42,56. Applied to optogenetics, DeepCGH12 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 DeepCGH12 (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 implementation12 (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 work12. 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).
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 localization57 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.
The alternative text for this image may have been generated using AI.
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 holography12,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 phototoxicity58,59,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 JAX26. 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).
Fig. 5: Chromatix is the fastest implementation of existing computational optics methods.
The alternative text for this image may have been generated using AI.
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.