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

Λεπτομέρειες βιβλιογραφικής εγγραφής
Τίτλος:	Chromatix: a differentiable, GPU-accelerated wave-optics library.
Συγγραφείς:	Deb D; HHMI Janelia Research Campus, Ashburn, VA, USA. debd@janelia.hhmi.org., Both GJ; HHMI Janelia Research Campus, Ashburn, VA, USA. bothg@janelia.hhmi.org., Bezzam E; Audiovisual Communications Laboratory, Ecole Polytechnique Fédérale de Lausanne, Lausanne, Switzerland., Kohli A; Department of Electrical Engineering and Computer Sciences, UC Berkeley, Berkeley, CA, USA., Yang S; Chandra Family Department of Electrical and Computer Engineering, The University of Texas at Austin, Austin, TX, USA., Chaware A; Department of Biomedical Engineering, Duke University, Durham, NC, USA., Allier C; HHMI Janelia Research Campus, Ashburn, VA, USA., Cai C; Department of Applied Physical Sciences, The University of North Carolina at Chapel Hill, Chapel Hill, NC, USA.; Joint Department of Biomedical Engineering, The University of North Carolina at Chapel Hill, Chapel Hill, NC, USA., Anderberg G; Department of Radiology, University of Chicago, Chicago, IL, USA., Eybposh MH; Department of Applied Physical Sciences, The University of North Carolina at Chapel Hill, Chapel Hill, NC, USA.; Joint Department of Biomedical Engineering, The University of North Carolina at Chapel Hill, Chapel Hill, NC, USA., Schneider MC; HHMI Janelia Research Campus, Ashburn, VA, USA., Heintzmann R; Leibniz Institute of Photonic Technology, Jena, Germany.; Institute of Physical Chemistry and Abbe Center of Photonics, Friedrich-Schiller-Universität Jena, Jena, Germany., Rivera-Sanchez FA; Imaging Physics, Faculty of Applied Sciences, Technische Universiteit Delft, Delft, The Netherlands.; Advanced Research Center for Nanolithography (ARCNL), Amsterdam, The Netherlands., Simmerer C; Department of Biomedical Engineering, Johns Hopkins University, Baltimore, MD, USA., Meng G; Department of Electrical Engineering and Computer Sciences, UC Berkeley, Berkeley, CA, USA.; Herbert Wertheim School of Optometry and Vision Science, University of California, Berkeley, CA, USA., Tormes-Vaquerano J; Department of Applied Physical Sciences, The University of North Carolina at Chapel Hill, Chapel Hill, NC, USA., Han S; Department of Applied Physics, Yale University, New Haven, CT, USA., Shanmugavel SC; Chandra Family Department of Electrical and Computer Engineering, The University of Texas at Austin, Austin, TX, USA., Maruvada T; Department of Electrical and Computer Engineering, Johns Hopkins University, Baltimore, MD, USA., Yang X; Department of Biomedical Engineering, Duke University, Durham, NC, USA., Kim Y; Department of Mechanical Engineering, Yonsei University, Seoul, Republic of Korea., Diederich B; Leibniz Institute of Photonic Technology, Jena, Germany., Joo C; Department of Mechanical Engineering, Yonsei University, Seoul, Republic of Korea., Waller L; Department of Electrical Engineering and Computer Sciences, UC Berkeley, Berkeley, CA, USA., Durr NJ; Department of Biomedical Engineering, Johns Hopkins University, Baltimore, MD, USA., Pégard NC; Department of Applied Physical Sciences, The University of North Carolina at Chapel Hill, Chapel Hill, NC, USA.; Joint Department of Biomedical Engineering, The University of North Carolina at Chapel Hill, Chapel Hill, NC, USA.; Neuroscience Center, University of North Carolina at Chapel Hill, Chapel Hill, NC, USA.; Carolina Stress Initiative, University of North Carolina at Chapel Hill, Chapel Hill, NC, USA., La Rivière PJ; Department of Radiology, University of Chicago, Chicago, IL, USA., Horstmeyer R; Department of Biomedical Engineering, Duke University, Durham, NC, USA., Chowdhury S; Chandra Family Department of Electrical and Computer Engineering, The University of Texas at Austin, Austin, TX, USA., Turaga SC; HHMI Janelia Research Campus, Ashburn, VA, USA. turagas@janelia.hhmi.org.
Πηγή:	Nature methods [Nat Methods] 2026 Jul; Vol. 23 (7), pp. 1388-1398. Date of Electronic Publication: 2026 Jun 08.
Τύπος έκδοσης:	Journal Article
Γλώσσα:	English
Στοιχεία περιοδικού:	Publisher: Nature Pub. Group Country of Publication: United States NLM ID: 101215604 Publication Model: Print-Electronic Cited Medium: Internet ISSN: 1548-7105 (Electronic) Linking ISSN: 15487091 NLM ISO Abbreviation: Nat Methods Subsets: MEDLINE
Imprint Name(s):	Original Publication: New York, NY : Nature Pub. Group, c2004-
Ιατρικοί όροι (MeSH):	Optics and Photonics/methods , Microscopy/methods , Image Processing, Computer-Assisted/methods , Computer Graphics , Software*, Holography ; Computer Simulation ; Algorithms
Περίληψη:	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. (© 2026. The Author(s).)
Competing Interests:	Competing interests: The authors declare no competing interests.
Σχόλια:	Update of: bioRxiv. 2026 Mar 25:2025.04.29.651152. doi: 10.1101/2025.04.29.651152.. (PMID: 41929201)
References:	Betzig, E. et al. Imaging intracellular fluorescent proteins at nanometer resolution. Science 313, 1642–1645 (2006). (PMID: 1690209010.1126/science.1127344) Rust, M. J., Bates, M. & Zhuang, X. Sub-diffraction-limit imaging by stochastic optical reconstruction microscopy (STORM). Nat. Methods 3, 793–795 (2006). (PMID: 16896339270029610.1038/nmeth929) Deb, D. et al. FourierNets enable the design of highly non-local optical encoders for computational imaging. Neural Inf. Proc. Syst. 35, 25224–25236 (2022). (PMID: 10.52202/068431-1829) Cong, L. et al. Rapid whole brain imaging of neural activity in freely behaving larval zebrafish (Danio rerio). eLife 6, e28158 (2017). (PMID: 28930070564496110.7554/eLife.28158) Zhang, Z. et al. Imaging volumetric dynamics at high speed in mouse and zebrafish brain with confocal light field microscopy. Nat. Biotechnol. 39, 74–83 (2021). (PMID: 3277884010.1038/s41587-020-0628-7) Levoy, M., Ng, R., Adams, A., Footer, M. & Horowitz, M. Light field microscopy. ACM Trans. Graph. 25, 924–934 (2006). (PMID: 10.1145/1141911.1141976) Broxton, M. et al. Wave optics theory and 3-D deconvolution for the light field microscope. Opt. Express 21, 25418–25439 (2013). (PMID: 24150383386710310.1364/OE.21.025418) Andalman, A. S. et al. Neuronal dynamics regulating brain and behavioral state transitions. Cell 177, 970–985 (2019). (PMID: 31031000672613010.1016/j.cell.2019.02.037) Nehme, E. et al. DeepSTORM3D: dense 3D localization microscopy and PSF design by deep learning. Nat. Methods 17, 734–740 (2020). (PMID: 32541853761048610.1038/s41592-020-0853-5) Hulleman, C. N. et al. Simultaneous orientation and 3D localization microscopy with a Vortex point spread function. Nat. Commun. 12, 5934 (2021). (PMID: 34635658850543910.1038/s41467-021-26228-5) Yanny, K. et al. Miniscope3D: optimized single-shot miniature 3D fluorescence microscopy. Light Sci. Appl. 9, 171 (2020). (PMID: 33082940753214810.1038/s41377-020-00403-7) Hossein Eybposh, M., Caira, N. W., Atisa, M., Chakravarthula, P. & Pégard, N. C. DeepCGH: 3D computer-generated holography using deep learning. Opt. Express 28, 26636–26650 (2020). (PMID: 3290693310.1364/OE.399624) Oldenburg, I. A. et al. The logic of recurrent circuits in the primary visual cortex. Nat. Neurosci. 27, 137–147 (2024). (PMID: 381724371077414510.1038/s41593-023-01510-5) Faini, G. et al. Ultrafast light targeting for high-throughput precise control of neuronal networks. Nat. Commun. 14, 1888 (2023). (PMID: 3701989110.1038/s41467-023-37416-w) Pégard, N. C. et al. Three-dimensional scanless holographic optogenetics with temporal focusing (3D-SHOT). Nat. Commun. 8, 1228 (2017). (PMID: 29089483566371410.1038/s41467-017-01031-3) Chowdhury, S. et al. High-resolution 3D refractive index microscopy of multiple-scattering samples from intensity images. Optica 6, 1211–1219 (2019). (PMID: 385159601095670310.1364/OPTICA.6.001211) Yang, S., Kim, J., Swartz, M. E., Eberhart, J. K. & Chowdhury, S. DMD and microlens array as a switchable module for illumination angle scanning in optical diffraction tomography. Biomed. Opt. Express 15, 5932–5946 (2024). (PMID: 394217701148216910.1364/BOE.535123) Kellman, M. R., Bostan, E., Repina, N. A. & Waller, L. Physics-based learned design: optimized coded-illumination for quantitative phase imaging. IEEE Trans. Comput. Imaging 5, 344–353 (2019). (PMID: 10.1109/TCI.2019.2905434) Baek, S.-H. et al. Single-shot hyperspectral-depth imaging with learned diffractive optics. In Proc. IEEE/CVF International Conference on Computer Vision (ICCV) 2631–2640 (IEEE, 2021). Sitzmann, V. et al. End-to-end optimization of optics and image processing for achromatic extended depth of field and super-resolution imaging. ACM Trans. Graph. 37, 1–13 (2018). (PMID: 10.1145/3197517.3201333) Chang, J. & Wetzstein, G. Deep optics for monocular depth estimation and 3D object detection. In Proc. IEEE/CVF International Conference on Computer Vision (ICCV) 10192–10201 (IEEE, 2019). Ikoma, H., Nguyen, C. M., Metzler, C. A., Peng, Y. & Wetzstein, G. Depth from defocus with learned optics for imaging and occlusion-aware depth estimation. In Proc. IEEE International Conference on Computational Photography (ICCP) 1–12 (IEEE, 2021). Sinha, A., Lee, J., Li, S. & Barbastathis, G. Lensless computational imaging through deep learning. Optica 4, 1117 (2017). (PMID: 10.1364/OPTICA.4.001117) Paszke, A. et al. PyTorch: an imperative style, high-performance deep learning library. In Proc. 33rd International Conference on Neural Information Processing Systems, 721 (Curran Associates, 2019). Abadi, M. et al. TensorFlow: a system for large-scale machine learning. In Proc. 12th USENIX conference on Operating Systems Design and Implementation 265–283 (USENIX, 2016). Bradbury, J. et al. JAX: composable transformations of Python+NumPy programs. GitHub http://github.com/jax-ml/jax (2018). Ansys Zemax OpticStudio. Ansys http://www.ansys.com/products/optics/ansys-zemax-opticstudio (2023). CODE V Optical design software. Keysight http://www.synopsys.com/optical-solutions/codev.html (2026). Wang, C., Chen, N. & Heidrich, W. DO: a differentiable engine for deep lens design of computational imaging systems. IEEE Trans. Comput. Imaging 8, 905–916 (2022). (PMID: 10.1109/TCI.2022.3212837) Rodríguez, C., Arlt, S., Möckl, L. & Krenn, M. Automated discovery of experimental designs in super-resolution microscopy with XLuminA. Nat. Commun. 15, 10658 (2024). (PMID: 396585751163210010.1038/s41467-024-54696-y) Desdoigts, L., Pope, B. & Tuthill, P. Optical design, analysis, and calibration using ∂Lux. In Proc. Space Telescopes and Instrumentation: Optical, Infrared, and Millimeter Wave (eds Coyle, L. E. et al.) 1218032 (SPIE, 2022). Harrison, K. HarrisonKramer/optiland v.0.2.3. Zenodo http://doi.org/10.5281/ZENODO.14588962 (2025). Fienup, J. R. Iterative method applied to image reconstruction and to computer-generated holograms. In Proc. Applications of Digital Image Processing III (ed. Tescher, A. G.) (SPIE, 1979). Gerchberg, R. A practical algorithm for the determination of phase from image and diffraction plane pictures. Optik 35, 237–246 (1972). Beck, A. & Teboulle, M. A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM J. Imaging Sci. 2, 183–202 (2009). (PMID: 10.1137/080716542) Boyd, S. Distributed optimization and statistical learning via the alternating direction method of multipliers. Found. Trends Mach. Learn. 3, 1–122 (2010). (PMID: 10.1561/2200000016) Kingma, D. P. & Ba, J. Adam: a method for stochastic optimization. Preprint at http://arxiv.org/abs/1412.6980 (2014). Mu, S. et al. Multislice computational model for birefringent scattering. Optica 10, 81 (2023). (PMID: 10.1364/OPTICA.472077) Cao, R., Divekar, N. S., Nuñez, J. K., Upadhyayula, S. & Waller, L. Neural space-time model for dynamic multi-shot imaging. Nat. Methods 21, 2336–2341 (2024). (PMID: 393177291162102310.1038/s41592-024-02417-0) Cao, R., Kellman, M., Ren, D., Eckert, R. & Waller, L. Self-calibrated 3D differential phase contrast microscopy with optimized illumination. Biomed. Opt. Express 13, 1671–1684 (2022). (PMID: 35414990897319010.1364/BOE.450838) Rückert, D., Wang, Y., Li, R., Idoughi, R. & Heidrich, W. NeAT: neural adaptive tomography. ACM Trans. Graph. 41, 1–13 (2022). Peng, Y., Choi, S., Padmanaban, N. & Wetzstein, G. Neural holography with camera-in-the-loop training. ACM Trans. Graph. 39, 1–14 (2020). (PMID: 10.1145/3414685.3417802) Kang, I., Zhang, Q., Yu, S. X. & Ji, N. Coordinate-based neural representations for computational adaptive optics in widefield microscopy. Nat. Mach. Intell. 6, 714–725 (2024). (PMID: 10.1038/s42256-024-00853-3) Feng, B. Y. et al. NeuWS: neural wavefront shaping for guidestar-free imaging through static and dynamic scattering media. Sci. Adv. 9, eadg4671 (2023). (PMID: 373793861030629710.1126/sciadv.adg4671) Mildenhall, B. et al. NeRF: representing scenes as neural radiance fields for view synthesis. Commun. ACM 65, 99–106 (2022). (PMID: 10.1145/3503250) Liu, R., Sun, Y., Zhu, J., Tian, L. & Kamilov, U. S. Recovery of continuous 3D refractive index maps from discrete intensity-only measurements using neural fields. Nat. Mach. Intell. 4, 781–791 (2022). (PMID: 10.1038/s42256-022-00530-3) Harris et al. Array programming with NumPy. Nature 585, 357–362 (2020). (PMID: 32939066775946110.1038/s41586-020-2649-2) Kohli, A. et al. Ring deconvolution microscopy: an exact solution for spatially-varying aberration correction. Nat. Methods 22, 1311–1320 (2025). (PMID: 403016221216584610.1038/s41592-025-02684-5) Whang, E., McAllister, D., Reddy, A., Kohli, A. & Waller, L. SeidelNet: an aberration-informed deep learning model for spatially varying deblurring. In Proc. AI and Optical Data Sciences IV (eds Kitayama, K.-I. & Jalali, B.) Vol. 12438, 276–281 (SPIE, 2023). Ghosh, K. K. et al. Miniaturized integration of a fluorescence microscope. Nat. Methods 8, 871–878 (2011). (PMID: 21909102381031110.1038/nmeth.1694) Cai, D. J. et al. A shared neural ensemble links distinct contextual memories encoded close in time. Nature 534, 115–118 (2016). (PMID: 27251287506350010.1038/nature17955) Paganin, D. M. & Pelliccia, D. Tutorials on X-ray phase contrast imaging: some fundamentals and some conjectures on future developments. Preprint at http://arxiv.org/abs/1902.00364 (2019). Shechtman, Y., Sahl, S. J., Backer, A. S. & Moerner, W. E. Optimal point spread function design for 3D imaging. Phys. Rev. Lett. 113, 133902 (2014). (PMID: 25302889438186610.1103/PhysRevLett.113.133902) Gerchberg, R. W. & Saxton, W. O. Comment on ‘A method for the solution of the phase problem in electron microscopy’. J. Phys. D 6, L31–L32 (1973). (PMID: 10.1088/0022-3727/6/5/101) Chen, C.-Y., Li, W.-C., Chang, H.-T., Chuang, C.-H. & Chang, T.-J. 3-D modified Gerchberg–Saxton algorithm developed for panoramic computer-generated phase-only holographic display. J. Opt. Soc. Am. B 34, B42–B48 (2017). (PMID: 10.1364/JOSAB.34.000B42) Peng, Y., Choi, S., Padmanaban, N., Kim, J. & Wetzstein, G. Neural holography. In Proc. ACM SIGGRAPH Emerging Technologies (ACM, 2020). Brenner, B., Sun, C., Raymo, F. M. & Zhang, H. F. Spectroscopic single-molecule localization microscopy: applications and prospective. Nano Converg. 10, 14 (2023). (PMID: 369435411003075510.1186/s40580-023-00363-9) Owen, S. F., Liu, M. H. & Kreitzer, A. C. Thermal constraints on in vivo optogenetic manipulations. Nat. Neurosci. 22, 1061–1065 (2019). (PMID: 31209378659276910.1038/s41593-019-0422-3) Stujenske, J. M., Spellman, T. & Gordon, J. A. Modeling the spatiotemporal dynamics of light and heat propagation for in vivo optogenetics. Cell Rep. 12, 525–534 (2015). (PMID: 26166563451288110.1016/j.celrep.2015.06.036) Picot, A. et al. Temperature rise under two-photon optogenetic brain stimulation. Cell Rep. 24, 1243–1253 (2018). (PMID: 3006797910.1016/j.celrep.2018.06.119) Yang, X. et al. End-to-end hybrid refractive-diffractive lens design with differentiable ray-wave model. In Proc. SIGGRAPH Asia Conference Papers 1–11 (ACM, 2024). Osnabrugge, G., Leedumrongwatthanakun, S. & Vellekoop, I. M. A convergent Born series for solving the inhomogeneous Helmholtz equation in arbitrarily large media. J. Comput. Phys. 322, 113–124 (2016). (PMID: 10.1016/j.jcp.2016.06.034) Colburn, S. & Majumdar, A. Inverse design and flexible parameterization of meta-optics using algorithmic differentiation. Commun. Phys. 4, 65 (2021). (PMID: 10.1038/s42005-021-00568-6) Ansys Lumerical FDTD. Ansys http://www.ansys.com/products/optics/fdtd (2023). Hazineh, D. S. et al. D-Flat: a differentiable flat-optics framework for end-to-end metasurface visual sensor design. Preprint at http://arxiv.org/abs/2207.14780 (2022). Born, M. & Wolf, E. Principles of Optics: Electromagnetic Theory of Propagation, Interference and Diffraction of Light (Elsevier, 2013). Elsken, T., Metzen, J. H. & Hutter, F. Neural architecture search: a survey. J. Mach. Learn. Res. 20, 1997–2017 (2019). Edelstein, A. D. et al. Advanced methods of microscope control using µManager software. J. Biol. Methods 1, 1 (2014). (PMID: 10.14440/jbm.2014.36) Pinkard, H. et al. Pycro-Manager: open-source software for customized and reproducible microscope control. Nat. Methods 18, 226–228 (2021). (PMID: 33674797853217610.1038/s41592-021-01087-6) Heintzmann, R., Loetgering, L. & Wechsler, F. Scalable angular spectrum propagation. Optica 10, 1407–1416 (2023). (PMID: 10.1364/OPTICA.497809) Matsushima, K. & Shimobaba, T. Band-limited angular spectrum method for numerical simulation of free-space propagation in far and near fields. Opt. Express 17, 19662–19673 (2009). (PMID: 1999718610.1364/OE.17.019662) Goodman, J. W. Introduction to Fourier Optics (W.H. Freeman, 2005). Deb, D. et al. Zebrafish volumes dataset for ‘FourierNets enable the design of highly non-local optical encoders for computational imaging’. figshare http://doi.org/10.25378/janelia.25277269.v1 (2024). Ferdman, B., Saguy, A., Xiao, D. & Shechtman, Y. Diffractive optical system design by cascaded propagation. Opt. Express 30, 27509–27530 (2022). (PMID: 3623692110.1364/OE.465230) Collins, S. A. Jr. Lens-system diffraction integral written in terms of matrix optics. J. Opt. Soc. Am. 60, 1168–1177 (1970). (PMID: 10.1364/JOSA.60.001168)
Grant Information:	R35 GM155424 United States GM NIGMS NIH HHS; No.2023R1A2C3004040 National Research Foundation of Korea (NRF)
Entry Date(s):	Date Created: 20260608 Date Completed: 20260709 Latest Revision: 20260726
Update Code:	20260726
PubMed Central ID:	PMC13345901
DOI:	10.1038/s41592-026-03121-x
PMID:	42260310
Βάση Δεδομένων:	MEDLINE

Περιγραφή
ISSN:	1548-7105
DOI:	10.1038/s41592-026-03121-x