Chromatix：一种可微分、GPU加速的波动光学库

许多当前的显微学方法将计算建模作为成像过程的组成部分，要么用于求解逆问题，要么用于优化光学系统设计本身。这些方法往往依赖于可微分的光学模拟，然而目前缺乏标准化框架，迫使计算光学研究人员反复独立地实现模拟，导致可重用性和性能有限。这些普遍问题限制了计算光学作为一个领域的潜在影响。在此，研究人员提出Chromatix：一种开源、图形处理器（GPU）加速、可微分的波动光学模拟库。Chromatix基于JAX构建，旨在普及快速、并行化的多样光学系统模拟，并拓展计算光学中的设计空间。Chromatix标准化了日益丰富的光学元件与传播方法集合，可支持广泛应用，研究人员在此展示了其在快照显微术、全息术和相位恢复中的应用。研究人员演示了在单GPU上2–6倍的速度提升，以及在8个GPU上最高达22倍的速度提升。

论文发表于《Nature Methods》，针对计算光学领域长期缺乏标准化、高性能、可微分波动光学模拟框架的问题，现有商业软件（如Zemax、CODE V）不支持高效可微分与深度学习集成，其他开源库多为光线追迹或功能不完整，而研究人员从零实现波动光学模拟存在约定不一致、难复用、难在GPU上高效实现等瓶颈，制约了该领域的发展。为此，研究人员开发了基于JAX的开源库Chromatix，将光学系统建模为序列化的光学元件与传播变换，以统一结构表示复光场（Field对象，含波长、偏振、空间采样等信息），实现了可微分、可组合、可扩展的波动光学模拟。结论表明Chromatix在单GPU上可比原有实现快2–6倍，8GPU并行下最高达22倍，且能正确复现多种典型计算光学实验的结果，为计算光学提供了标准化基础设施，有助于提升研究可重复性与设计空间探索效率。

为开展研究，作者采用的主要关键技术方法包括：基于Python的JAX框架实现自动微分与JIT编译以支持GPU/TPU加速与多GPU并行（jax.pmap等）；将光学系统建模为光学元件（透镜、相位掩模、SLM、DMD、传感器、自由空间传播等）和波传播模块（标量/矢量角谱法、带限角谱法、多切片光束传播法等）的函数式组合；引入统一Field数据结构封装复振幅及元数据；利用现有公开数据集（如UCLA Miniscope数据、固定小鼠脑切片、斑马鱼胚胎、荧光微球等）和已有方法（环去卷积显微术、CoCoA、折射率显微术、FourierNets/Holoscope、DeepCGH等）在Chromatix中重构其正向模型与优化流程；通过JAX原生并行原语实现跨GPU批处理与多波长/多角度并行，比较单GPU与多GPU（NVIDIA H100/RTX 8000）下单次迭代耗时；所有基准均相对各原实现（MATLAB、PyTorch、TensorFlow）在同等问题规模下进行。

研究结果如下：

Design and implementation（设计与实现）：研究人员提出Chromatix的三个核心特征为可微分性（借助JAX自动微分对任意参数求梯度，支持基于梯度的优化如Adam用于光学参数、PSF、神经网络联合优化，也支持自校准算法与隐式神经表示（INR）训练）、可组合性（光学元件与传播模块标准化为可组合单元，类似深度学习层，便于复用与替换，提供透镜、传感器、自由空间标量/矢量传播、复杂散射样品等多组件）、可扩展性（基于JAX与XLA编译器，同一代码可在CPU/GPU/TPU运行，支持jax.vmap自动向量化与jax.pmap多设备并行，无需改写光学模型即可从单GPU扩至8GPU）。实现上，Chromatix用统一Field结构表示复光场及其波长、偏振、空间采样，光学系统定义为Field到Field的变换序列。

Inverse problems for reconstructing samples（用于样品重建的逆问题）：研究人员用Chromatix实现了环去卷积显微术（ring deconvolution microscopy），对UCLA Miniscope拍摄的兔肝组织像进行旋转不变PSF建模与空变去卷积，结果表明Chromatix能在多GPU并行下重建更大FOV（单GPU快4.5倍，8GPU快近19倍），而原PyTorch实现因单GPU显存限制无法完整重建。用Chromatix实现CoCoA（coordinate-based neural representations for computational adaptive optics），联合优化INR表示的样品与Zernike像差系数，在小鼠神经元数据上重建质量优于原PyTorch实现（Zernike系数恢复RMSE为3.56 nm vs 6.97 nm），单GPU快2倍，8GPU快近9倍。用Chromatix实现3D折射率显微术（refractive index microscopy），对斑马鱼胚胎尾部的相干照明多角度强度测量进行多切片传播正向模拟与梯度下降优化，代码量从原MATLAB约107行减至约25行，速度提升3–13倍，且可选用更大规模设置消除原重建中的网格伪影。

Programmable optics and deep learning（可编程光学与深度学习）：研究人员在Chromatix中重构了Holoscope（可编程4f系统含傅里叶面SLM的快照3D显微术），联合优化相位掩模像素与FourierNet重建网络，PSF设计压缩3D荧光体积到2D快照，重建SSIM达0.979±0.003，与原PyTorch结果无显著差异（P=0.695），训练速度快约7倍。实现DeepCGH（3D计算机生成全息术），用UNet加传播步直接从目标3D图案前馈生成相位全息图，在测试集上SSIM为0.985±0.001（原TensorFlow为0.982±0.001，P=0.018），PSNR为35.40±0.37（原34.95±0.16，P=0.177），代码从33行减至约17行，单GPU快2.5倍，8GPU快超10倍。

Flexible modeling with optical building blocks（利用光学积木的灵活建模）：研究人员演示同一可编程4f SLM模型可重新用于多色单分子定位显微术，优化多波长（400–650 nm）PSF使单通道快照图像可通过神经网络重建光谱立方体；还演示将全息模型（相位掩模+透镜）与多切片散射样品模型结合，通过迭代优化相位掩模补偿散射介质影响，在模拟中实现了均匀轴向刺激强度分布，显示Chromatix可灵活组合不同波动光学模块以解决新生物问题（如散射介质中全息光遗传学）。

High performance through parallelization（通过并行化的高性能）：对全部复现方法（环去卷积、CoCoA、折射率显微、Holoscope PSF优化、DeepCGH）进行1–8 GPU基准，Chromatix单GPU提速2–6倍，最佳情况8GPU达22倍；提速来自JAX编译降低开销与极少代码改动即可多GPU并行（jax.pmap分割批次/角度/径向环等），使大FOV重建、长时间优化（折射率显微、PSF优化）从数小时至数周缩减至分钟至数天。

讨论部分总结：研究人员指出Chromatix作为开源可微分波动光学模拟标准库可解决计算光学方法因缺乏标准与定制校准导致的复用受限问题；其开源社区已通过hackathon贡献多个功能（如更厚透镜实验支持、矢量传播等），证明可及性与标准化促进共享。作者承认当前局限：透镜为理想/有限厚模型，未来可加入实际镀膜透镜、噪声传感器模型；可融合光线追迹与波动模型以高效模拟复杂透镜组；需加入修正Born级数以更高效正反散射、支持严格耦合波分析（RCWA）用于超表面、部分相干传播模型、光谱变化效应等。展望未来，Chromatix可支持光学系统架构搜索、与硬件控制软件（Micro-Manager、Pycro-Manager）集成实现硬件在环自动校准，就像深度学习库推动神经网络探索一样，Chromatix的标准化高效波动光学模型将推动计算光学方法空间的探索。