🔢 mathematics

TokaGrad: End-to-end differentiable tokamak simulator for L-to-H full scenario optimization

本文介绍了 TokaGrad，这是首个端到端可微的托卡马克模拟器，它能够自洽地模拟包括升流、L模态和H模态转换在内的全放电场景，从而实现对反应堆设计的高效梯度优化以及等离子体自主控制。

原作者： Jaemin Seo

发布于 2026-07-13

📖 1 分钟阅读🧠 深度阅读

原作者： Jaemin Seo

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

想象一下，你正试图烤出一个完美的蛋糕，但你手头没有食谱，只能通过向黑板投掷飞镖来猜测温度、搅拌速度和配料。如果蛋糕烤焦了，你就再试一次。如果太干了，你就再试一次。这基本上就是科学家们传统设计聚变反应堆的方式——这些被称为**托卡马克（tokamaks）**的巨大甜甜圈状机器，旨在地球上复制太阳的力量。它们依赖于昂贵的“试错法”，通过调节旋钮并寄希望于好运。

迎来 TokaGrad，这是一个由研究员 Jaemin Seo 开发的新型数字化工具。把 TokaGrad 想象成不只是一个让你盲目尝试的“黑箱”，而是一个超级聪明、具有自我感知能力的聚变视频游戏引擎。

“自我感知”模拟器的魔力

在过去，模拟聚变反应就像是在看一部电影，你无法倒退或暂停来观察角色为什么做出某个特定动作。你只能看到结果。如果等离子体（反应堆内部的超高温气体）表现异常，科学家们必须去猜哪里出了问题。

TokaGrad 改变了规则。它将整个模拟过程构建成一条相连的多米诺骨牌链。机器的每一个部分——从磁场的形状到加热功率——都是相互关联的。如果你推了一下其中一块多米诺骨牌（比如改变加热功率），该工具可以立即计算出这个微小的推力是如何层层传递，最终如何改变整个系统的结果。

这被称为可微编程（differentiable programming）。它不再是靠猜测，而是掌握了问题的“斜率”。这就像拥有一个 GPS，它不仅能告诉你你在哪里，还能瞬间计算出无论道路多么弯曲，到达目的地所需的精确转向角度。

重大挑战：“L 到 H”的跨越

聚变反应堆有着复杂的“个性”。它们可以处于一种“懒散模式”（称为 L-mode）或高性能模式（H-mode）。H-mode 是我们的目标，因为它能更好地锁住热量，但要达到这个状态，就像试图把一块沉重的巨石推过山丘。你需要恰到好处的力量，并在恰当的时机施加。

如果你推得太快太猛，可能无法翻过山丘；如果你推得太轻，则会滑落回来。在现实世界中，寻找这种完美的推力是非常困难的，因为物理过程既混乱又具有非线性特征。

TokaGrad 是第一个能够模拟整个旅程的模拟器——从冷启动、经过“懒散”的 L-mode，一直到高性能的 H-mode——同时保持这些“多米诺骨牌式连接”的工具。它不仅模拟最终状态，还模拟整个过渡过程。

这个工具实际做了什么（证明）

研究人员不仅开发了这个工具，还通过三种令人兴奋的方式对其进行了测试：

速度测试： 他们将 TokaGrad 与现有的工具 TORAX 进行了对比。如果说 TORAX 像是一个缓慢且谨慎的会计师，那么 TokaGrad 就是一名短跑运动员。在标准的笔记本电脑上，TokaGrad 的运行速度比 TORAX 快了十倍以上。它非常迅速，仅用 17 秒 的真实时间就能完成一个 150 秒聚变脉冲的模拟。
“智能”控制： 他们要求该工具计算如何从反应堆中获取最大能量。该工具并没有尝试随机设置，而是利用其对“斜率”的认知来调整加热和电流。它找到了一个方案：暂时增加加热功率以跳过 H-mode 的“山丘”，然后调低功率以节省能量，同时保持高性能。它是自动完成这一过程的，无需人类猜测。
“假设分析”设计师： 他们让该工具重新设计反应堆本身。通过微调机器的大小和形状（例如大半径和磁场强度），该工具提出了一个可以产生更多能量的新设计。它发现，通过让反应堆稍微变大并改变其形状，可以在保持机器稳定的同时提升性能。

这意味着什么（以及它并不意味着什么）

论文指出，这种方法可能会彻底改变我们设计聚变反应堆的方式。工程师们不再需要进行长达数年的试错实验，而是可以使用 TokaGrad 瞬间测试成千上万种场景，并找到通往实用化反应堆的最佳路径。

然而，必须记住这些都是模拟。该工具已经在计算机上展示了其有效性，预测了它可以找到更好的加热策略和反应堆形状。但它尚未建造出真实的物理反应堆，也尚未进行真实的聚变实验。这里的“突破”在于优化方法，而不是今天实验室里坐着一台真实的聚变反应堆。

作者还指出，尽管该工具快速且聪明，但它仍然依赖于输入其中的物理模型。如果模型本身是错误的，那么它的“智能”建议也可能是错误的。但就目前而言，TokaGrad 为我们提供了一种导航复杂且崎岖的道路、通往清洁、无限聚变能源的新途径。它将设计“地球上的恒星”的过程，从一场概率游戏变成了一场精准、有计划的策略博弈。

技术摘要：TokaGrad —— 端到端可微托卡马克模拟器

1. 问题陈述

随着磁约束核聚变研究从科学可行性向反应堆规模的商业化转型，新反应堆概念的设计和高性能等离子体场景的优化已变得至关重要。传统的场景开发方法过度依赖于高成本的试错法、基于专家经验的波形调优以及低维参数扫描。虽然集成建模工具（如 ASTRA、TRANSP）改进了虚拟实验，但场景设计仍是一个“人工参与”的过程。

当前方法的主要局限性包括：

黑盒优化： 许多机器学习方法（如贝叶斯优化、强化学习）将模拟器视为黑盒，需要重复的前向评估和间接的灵敏度推断，这对于高维、多执行器、全放电场景而言在计算上是难以承受的。
缺乏可微性： 传统的数值模拟作为逐步执行的过程，梯度无法在整个流水线中传播。这阻碍了利用内部灵敏度进行直接的基于梯度的控制和优化。
不完整的全场景建模： 现有的可微输运求解器（如 TORAX）主要侧重于核心输运和预设场景。它们缺乏对涉及 L-H 转换、边缘胎座（pedestal）形成以及随演化的等离子体形状与平衡重建耦合的动态全放电场景的自洽建模。

2. 方法论

作者引入了 TokaGrad，这是一个基于 JAX 框架构建的端到端可微托卡马克输运模拟器。该系统将整个模拟流程——从初始条件到最终等离子体性能——表示为一个连接的可微计算图。这使得自动微分（AD）能够计算整个模拟过程中的雅可比矩阵，将机器参数、执行器波形与等离子体响应联系起来。

核心架构

TokaGrad 在单个图中集成了以下可微模块：

输运框架： 核心由归一化径向通量坐标（ $\rho$ ）下电子/离子温度（ $T_{e,i}$ ）、电子密度（ $n_e$ ）和极向磁通（ $\psi$ ）的耦合一维输运方程组成。
- 数值方案： 使用半隐式有限体积法以保证稳定性，通过与 JAX 的即时编译（JIT）兼容的可微三对角矩阵求解器（Thomas 算法）进行求解。
- 方程： 包括粒子输运、保守能量输运（针对 $T_{e,i}$ ）以及与自举电流和电流驱动源耦合的感应通量演化。
几何与平衡：
- EQDSK 模式： 支持外部生成的平衡（如来自 EFIT），用于现实基准测试，尽管其对几何形状的梯度受限于参数化程度。
- Miller/Moment 模式： 一种使用解析通量面参数化（大半径、拉长比、三角化、Shafranov 位移）的简化固定边界模型。该模式是完全可微的，支持优化机器形状参数（ $R_0, a, \kappa$ ）。
输运闭合：
- 湍流输运： 使用神经网络（NN）代理模型（TGLF-NN, QLKNN）取代计算昂贵的回旋动力学求解器。这些代理模型提供了从等离子体轮廓到输运系数的可微映射。
- 新经典输运： 支持解析公式（Sauter-Angioni）和神经网络代理（NEO-NN）。
- 经验模型： 包括用于快速基准测试的 Bohm-gyroBohm (BgB) 模型。
源项与汇项：
- 加热： 包括欧姆加热、阿尔法粒子加热（使用 Bosch-Hale 拟合）、碰撞交换和辐射损失。辅助加热通过平滑的归一化高斯轮廓进行建模。
- 粒子源： 实现密度反馈控制（追踪 Greenwald 分数）或预设形状模型，以避免不可微的硬切换。
L-H 转换与胎座：
- 转换： 使用平滑的 Sigmoid 函数，根据跨越分隔面（separatrix）的功率（ $P_{sep}$ ）相对于经验阈值（ $P_{LH}$ ，例如 Martin 或 Delabie 标度）来控制 L 模与 H 模之间的转换。这避免了不可微的二元开关。
- 胎座： 使用 alpha-critical 简化模型或 EPED1-NN 代理构建胎座目标。胎座轮廓通过 smoothmax 函数与核心输运解进行平滑混合，确保梯度能通过转换过程传播。

3. 关键贡献

首个自洽的可微全放电模拟器： TokaGrad 是第一个能够模拟动态全放电场景的框架，其中执行器、等离子体平衡、胎座形成和约束机制转换（L-to-H）在可微图中共同演化。
端到端可微性： 通过将高保真物理闭合模型（TGLF, EPED1, NEO）转换为兼容 JAX 的神经网络代理，该系统实现了梯度从最终性能指标向后传播，穿过输运、湍流及平衡模型，直至设计和控制输入。
交互式实时模拟： 该框架支持快于实时（faster-than-real-time）的模拟（在 JIT 编译后，对于长脉冲模拟，其速度比 TORAX 快 10 倍以上），从而实现对因果关系的交互式探索。
基于梯度的优化： 系统支持直接优化执行器波形、场景轨迹和反应堆设计参数，而无需进行有限差分扫描。

4. 结果与演示

论文通过基准测试和优化案例研究验证了 TokaGrad：

基准测试：
- 对比 TORAX： 在圆形几何结构的简化 L 模场景中，TokaGrad 重现的核输运轮廓（温度、密度、电流）与成熟的基于 JAX 的 TORAX 模拟器的误差在 ~5% 以内。
- 性能： 一旦摊销了 JIT 编译开销，TokaGrad 在长脉冲模拟中的墙钟时间（wall-clock time）比 TORAX 快一个数量级以上。
- H 模模拟： 使用 TGLF-NN/EPED1-NN 和 QLKNN/alpha-critical 模型的 ITER 基准场景模拟，在约 7 秒内实现了融合增益 $Q \sim 11$ 和融合功率 $\sim 550$ MW（快于实时）。
全放电 L-to-H 模拟：
- 150 秒的 ITER 基准模拟成功模拟了升流阶段（ramp-up）、L 模阶段以及向 H 模的转换。
- 模拟捕捉到了 L-H 转换的迟滞现象：一次瞬态加热过冲（73 MW）触发了转换，此后即使辅助功率降至 50 MW，阿尔法加热仍维持了 H 模。
优化应用：
- 执行器控制： 基于梯度的控制（Adam 优化器）成功调整了等离子体电流（ $I_p$ ）、辅助加热（ $P_{aux}$ ）和 Greenwald 分数（ $f_{GW}$ ），以最大化融合指标（ $Q + 0.01 P_{fus}$ ），同时遵守稳定性约束（ $q_{95} > 3, \beta_N < 3$ ）。
- 场景波形优化： 从一个无法进入 H 模的原始升流过程开始，优化器自动发现了瞬态加热过冲策略。这使得等离子体能够跨越 L-H 阈值，并在相同的最终执行器条件下维持高性能 H 模。
- 反应堆设计优化： 框架优化了机器设计参数（ $R_0, a, B_T, I_p$ ）以最大化融合性能。优化器正确识别了趋势，如增加大半径和磁场，同时调整参数以满足稳定性约束。

5. 重要性与主张

作者声称，TokaGrad 为自动化托卡马克场景设计提供了一条切实可行的路径。通过用可微物理实现的梯度方法取代黑盒优化，该框架允许：

直接灵敏度分析： 计算融合增益如何取决于加热轨迹，或者胎座高度如何取决于几何形状，而无需外部扰动。
路径依赖型优化： 不仅优化最终状态，还优化达到高性能机制（如 H 模）所需的整个动态路径（例如升流波形）。
可扩展性： 能够在笔记本级硬件（Apple M1 Pro）上进行重复的前向和后向评估，这使得优化复杂的、以往在计算上难以实现的多元目标反应堆场景变得可行。

论文总结道，尽管目前的演示侧重于类 ITER 基准条件，但该框架具有通用性，并为燃烧等离子体场景和下一代反应堆概念的自动化、基于物理的优化提供了新途径。