🔬 physics

TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

本文介绍了 TORAX，这是一个基于 JAX 构建的开源、可微托卡马克输运模拟器，它利用自动微分和 JIT 编译技术，为等离子体场景建模、脉冲设计和控制实现快速的基于梯度的优化。

原作者： Jonathan Citrin, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici, Philippe Hamel, Andrea Huber, Dmitry Nikulin, David Pfau, Brendan Tracey, Martin Riedmiller, Pus

发布于 2026-07-10

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jonathan Citrin, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici, Philippe Hamel, Andrea Huber, Dmitry Nikulin, David Pfau, Brendan Tracey, Martin Riedmiller, Pushmeet Kohli

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试烘焙出完美的面包，但你的原料不是面粉和水，而是超高温等离子体、磁场和聚变能。目标是？建造一台（托卡马克装置）能够为世界提供动力的机器。问题在于，预测这些等离子体在机器内部的行为，就像是在飓风发生时预测飓风中的天气一样。它既混乱又快速，而且涉及一套每毫秒都在变化的复杂物理规则。

于是，TORAX 登场了——这是由 Google DeepMind 研究人员创建的一个全新的开源计算机程序。可以将 TORAX 想象成这些聚变反应堆的“智能、超高速飞行模拟器”。但其中的玄机在于：不同于那些像笨重、移动缓慢的卡车一样的旧模拟器，TORAX 是基于一个名为 JAX 的特殊引擎构建的。这个引擎让模拟器不仅能运行得极其迅速，还能通过计算梯度（基本上就是弄清楚为了获得更好的结果应该如何旋转旋钮）来即时“从错误中学习”。

旧方法 vs. 新方法

多年来，科学家们一直使用其他模拟器（如 RAPTOR、JETTO 或 ASTRA）来模拟这些反应堆。这些工具就像是可靠、重型的计算器。它们工作良好，但通常是用较旧的语言（如 Fortran）编写的，这使得很难接入现代人工智能（AI）模型，或者难以快速调整数学模型以寻找完美设置。它们也很慢；在标准计算机上，运行几秒钟反应器时间的模拟可能需要数小时甚至数天。

本文作者认为，这种缓慢和缺乏灵活性的现状阻碍了聚变研究。他们明确指出，如果我们想要使用现代机器学习或进行复杂的优化，就不能固守这些陈旧、僵化的代码。他们表明，虽然旧代码表现出色，但它们无法轻松处理设计未来反应堆或实时控制所需的“如果……会怎样”的情景。

TORAX 如何运作： “飞行模拟器”的魔力

TORAX 通过求解一组主要的方程组来模拟等离子体的核心。你可以将这些方程视为管理以下内容的规则：

离子热量： 重粒子变得多热。
电子热量： 轻粒子变得多热。
粒子输运： 等离子体的密度如何变化。
电流扩散： 电流如何在等离子体中流动。

TORAX 并没有采用逐一解决这些方程的缓慢、步进式行进方式，而是使用了 有限体积法（Finite Volume Method）。想象一下将等离子体切成一叠薄薄的、看不见的“煎饼”（或单元格）。程序计算热量和粒子如何从一个“煎蛋饼”流向下一个。它使用一种被称为“幂律方案（power-law scheme）”的聪明数学技巧，根据等离子体运动的速度快慢，来决定是从左侧还是右侧观察流动。这确保了模拟过程不会变得“摇晃”或不准确。

真正的魔力在于 TORX 是使用 Python 和 JAX 编写的。这意味着：

速度： 它能将代码编译成超高效的版本，运行在 CPU、GPU 甚至 TPU 上。在测试中，它在标准计算机上仅用 22 秒 就完成了对一个 80 秒反应器启动过程的模拟。这比实时运行还要快！
可微性： 这是重头戏。TORAX 可以精确计算输入端的微小变化（例如将加热功率提高 1%）如何影响最终结果。这对于利用 AI 来自动设计完美的“脉冲”（一系列事件的序列）至型重要。

“代理模型”技巧：AI 作为捷径

TORAX 处理湍流的方式是其最令人兴奋的特性之一。等离子体中的湍流极其复杂，就像试图预测河流中每一个漩涡一样。通常，模拟这种现象需要耗费大量时间。TORAX 可以使用 机器学习（ML）代理模型 来代替。

可以将 ML 代理模型想象成一个经过数百万次复杂物理模拟训练出来的“水晶球”。与其进行繁重的计算来追踪每一个漩涡，不如直接利用这个“水晶球”（AI 模型）瞬间预测结果。TORAX 构建得非常巧妙，可以无缝接入这些 AI 模型。例如，它使用了一个名为 QLKNN 的模型，该模型是基于更复杂的代码 QuaLiKiz 的数据训练而成的。这使得 TORAX 能够以极短的时间运行带有高保真湍流模型的模拟。

它奏效了吗？事实胜于雄辩

作者不仅构建了它，还对其进行了测试。他们将 TORAX 与成熟且受信任的模拟器 RAPTOR 进行了并排对比实验。他们模拟了一个“类 ITER”的情景（ITER 是正在法国建造的大型国际聚变实验装置）。

结果令人印象深刻：

一致性： TORAX 生成的温度、密度和磁场剖面与 RAPTOR 几乎完美匹配。差异极小——在稳态下小于 1%，在快速变化期间约为 2.5%。
速度： 虽然 RAPTOR 已经很快，但 TORAX 在使用 AI 代理模型时，能够比它更快地运行相同的场景。
验证： 文中明确指出 TORAX 已针对 RAPTOR 进行了“验证”，这意味着在测试场景内，其结果是值得信赖的。

未来如何？（以及目前还缺什么）

论文保持了谨慎的态度，没有过度吹捧。作者承认 TORAX 是一个“新”工具，且有些内容仍在开发路线图中。

缺失的部分： 它目前还无法动态处理复杂的 3D 磁场形状（它使用的是预先计算好的形状），并且尚未包含所有类型的辐射损失或重杂质输运。
即将到来的功能： 作者计划加入更多物理模型，例如处理“锯齿波”（一种磁不稳定性）的模型，以及更好地处理等离子体边缘（“台阶/pedestal”）的方法。

总结

TORAX 并不是一个能一夜之间解决聚变问题的魔杖。它是一个强大、灵活且极其快速的工具，让设计聚变反应堆的工作变得更加容易。通过将现代计算的速度与机器学习的智慧相结合，它允许科学家在过去运行一次模拟所需的时间内，运行数千个“如果……会怎样”的情景。这就像是从一台手摇计算器升级到了能够自我学习如何更好地解决问题的超级计算机。

论文结论认为，该工具已准备好助力加速托卡马克情景建模、脉冲设计和控制的研究，但同时也强调，它是迈向终点旅程中的一步，而非终点本身。聚变的未来是光明的，而像 TORAX 这样的工具正在帮助照亮前行的道路。

技术摘要：TORAX —— 基于 JAX 的快速且可微的托卡马克输运模拟器

问题陈述

准确的集成等离子体模拟对于诸如托卡马克等环形磁约束装置的设计与运行至关重要，其应用范围涵盖了从实验解释到控制器设计。然而，现有的核心输运代码（如 JETTO、ASTRA、PTRANSP）在物理保真度与模拟速度之间面临权衡；高保真度模型通常需要数小时甚至数天才能模拟几个能量约束时间。此外，传统的代码通常由 Fortran 或 MATLAB 编写，缺乏原生可微性，这阻碍了基于梯度的优化、敏感性分析以及与机器学习（ML）代理模型的无缝集成。虽然像 RAPTOR 这样的一些专用代码提供了速度优势，但它们依赖于手动编写的雅可比矩阵（Jacobian），限制了向复杂 ML 模型的扩展性。目前迫切需要一种既能结合快速运行时间，又具备自动可微性并原生支持 ML 集成的模拟器。

方法论

作者提出了 TORAX，这是一个使用 JAX 框架并在 Python 中实现的开源托卡马克核心输运模拟器。

控制方程与离散化

TORAX 在归一化的环向磁通坐标（ $\hat{\rho}$ ）下求解一组耦合的一维输运偏微分方程（PDEs）。该系统控制着：

离子热输运（ $T_i$ ）
电子热输运（ $T_e$ ）
电子粒子输运（ $n_e$ ）
电流扩散（极向磁通 $\psi$ ）

这些方程源自动力学方程和法拉第定律的通量面平均矩。空间离散化采用均匀网格上的有限体积法（FVM）。该方案采用幂律加权处理对流-扩散通量，以根据 Peclet 数在中心差分与上风格式之间进行平滑过渡。

求解器实现

时间积分利用 theta 方法，它是显式与隐式 Euler 方案之间的加权平均。TORAX 提供三种求解器选项来处理由此产生的非线性代数系统：

线性求解器（预测-校正法）： 使用不动点迭代。为了增强在面对具有刚性系数（例如来自 ML 模型）时的稳定性，它可选地实现了 Pereverzev-Corrigan 方法，通过添加与对流平衡的人工扩散来增强稳定性。
Newton-Raphson 求解器： 一种基于梯度的求根方法。至关重要的是，它利用 JAX 的自动微分技术自动计算残差向量的雅可比矩阵，从而消除了手动编写雅可比矩阵的需求。
优化求解器： 将残差重新构建为通过迭代优化算法（使用 jaxopt）进行最小化的损失函数。

物理模型

TORAX 拥有一个模块化的物理组件框架：

几何结构： 支持 CHEASE 平衡数据和简化的圆圈几何结构。
输运： 包括常数模型、临界梯度模型（CGM）以及 QLKNN——一种基于 QuaLiKiz gyrokinetic 代码训练的机器学习代理模型。QLKNN 模型使用 Flax 库在 JAX 中原生实现。
源项： 处理离子-电子热交换、聚变功率（Bosch-Hale 参数化）、欧姆加热以及粒子源（气体注入、小球注入、中性束注入/NBI）。源项可以被视为隐式或显式。
新经典物理： 使用 Sauter 模型处理 Bootstrap 电流和新经典电导率。

软件架构

代码库被组织成不同的模块（几何、输运模型、源项、步进器、时间步长计算器、状态）。它利用 JAX 的**即时编译（JIT）**来加速计算密集型组件（源模型、输运系数、线性代数），同时保持控制流处于标准 Python 中以保证灵活性。这种设计允许在 CPU、GPU 和 TPU 上执行。

关键结果

作者通过已建立的 RAPTOR 代码对 TORAX 进行了验证，而 RAPTOR 本身已通过其他输运求解器进行了验证。

稳态基准测试（ITER L-mode）：
- 模拟了具有恒定电流和加热条件的 10 秒 ITER 类场景。
- 结果显示 TORAX 与 RAPTOR 在离子/电子温度、密度、极向磁通和安全因子剖面方面具有极佳的一致性。
- 稳态剖面的归一化均方根偏差（NRMSD）约为 0.7% 至 1.1%。
- 瞬态温度偏差限制在 ~2.5% 以内。
动态场景基准测试：
- 模拟了一个具有随时间变化的等离子体电流（上升期、平台期、下降期）和调制加热功率的 300 秒场景。
- TORAX 和 RAPTOR 在能量约束和电流扩散时间尺度上紧密追踪动态过程，温度剖面的 NRMSD 值在 ~3.5% 至 4.0% 之间。
性能：
- 在一个 48 核 CPU 节点上，使用 QLKNN 和 Newton-Raphson 求解器的 80 秒 ITER 上升期模拟在 22 秒 运行时间内完成，编译时间为 15.6 秒。
- 这证明了针对特定配置具有**快于实时（faster-than-realtime）**的能力，尽管作者指出这对于所有瞬态现象并非普遍结论。

意义与主张

本文将 TORAX 定位为一个多功能工具，旨在加速托卡马克场景建模、脉冲设计和控制的研究。其主要意义体现在三个方面：

可微性： 通过利用 JAX，TORAX 提供了一个可微的 PDE 残差。这使得无需像之前的可微代码（如 RAPTOR）那样进行手动雅可比编码，即可实现基于梯度的优化和敏感性分析（通过前向敏感性分析）。
ML 集成： 对 JAX 生态系统中神经网络的原生支持，促进了将机器学习代理（如 QLKNN）作为物理模型进行无缝耦合，使其能够在与模拟器相同的语言中进行开发。
速度与灵活性： JIT 编译确保了快速运行，而基于 Python 的模块化架构简化了代码的可扩展性、维护性以及与外部工作流的耦合。

作者总结道，TORAX 为融合领域提供了一个“强大且多功能的工具”，特别针对需要速度和可微性的应用，例如不确定性量化、场景优化和基于模型的控制器设计。未来的工作计划扩展物理保真度（例如 MHD、辐射、多离子组分）并增强可微性特性（例如持久化编译缓存、前向敏感性计算）。