Systematic LLM Translation of Legacy Scientific Code to Differentiable Frameworks: Application to a Land Surface Model

技术摘要：将遗留科学代码系统性地翻译为可微框架的 LLM 方法

问题陈述
遗留的地表系统模型（ESMs），特别是陆面模型（LSMs），对于气候政策和预测至关重要，但由于其代码结构，目前仍难以利用现代可微编程技术。这些模型通常由数十万行 Fortran 代码组成，存在模块化程度低、组件过时以及开发者后继乏力等问题。虽然可微编程能够实现高效的基于梯度的参数估计、敏感性分析和数据同化，但将遗留代码库迁移到可微框架（如 JAX）仍然是一个尚未解决的挑战。现有的代码现代化尝试要么涉及丢弃遗留物理机制的彻底重写，要么因跨文件依赖关系以及大语言模型（LLM）翻译缺乏数值保真度而无法扩展到函数之外。此外，目前的 LLM 基准测试未能充分衡量物理模型所需的科学有效性和数值对等性。

方法论：五阶段智能体流水线
作者提出了一种系统性的五阶段智能体流水线，旨在将拥有 19,000 行代码、102 个模块的 Fortran 陆面模型 CLM-ml-v2 翻译为名为 clm-ml-jax 的全可微 JAX 实现。该流水线依赖于三个核心设计原则：静态依赖分析、模块化数值对等性测试以及全面的梯度验证。

1. 第一阶段：范围界定与依赖分析： 为了防止 LLM 在处理耦合代码库时产生关于数据类型的幻觉，作者对完整的模块依赖图进行了静态分析。通过拓扑排序确定一个遵循依赖关系的翻译顺序，确保所有被导入的模块在依赖它们的模块之前完成翻译。
2. 第二阶段：状态文档设置： 流水线利用持久化状态文档（CLAUDE.md、plan.md、CHANGELOG.md）作为外部智能体记忆。这些文档记录了项目规范、跟踪完成标志并记录失败尝试，从而实现跨上下文窗口边界的连贯会话恢复，并防止智能体陷入死循环。
3. 第三阶段：Fortran Oracle 构建： 在翻译之前，一个智能体团队构建了一个功能测试套件。子程序根据可测试性被分为三个等级。系统在固定输入下执行原始 Fortran 代码以生成“黄金”输入-输出对，从而建立一个用于数值验证的参考 Oracle（神谕）。
4. 第四阶段：迭代翻译、测试与修复： 模块按照第一阶段确定的拓扑顺序进行处理。翻译智能体生成 JAX 代码，随后由测试智能体对照 Fortran Oracle 进行验证（要求相对容差在 $10^{-4}$ 以内）。如果测试失败，则启动“Ralph Loop”（一种编排模式），调用修复智能体诊断差异、编辑代码并自主重新测试，直到模块通过测试或被标记为需要人工审查。
5. 第五阶段：集成与可微性验证： 模块翻译完成后，集成智能体组装完整的柱体流水线并验证端到端行为。随后，可微性智能体迭代解决 JAX 特有的问题，例如：
   • 静态循环展开： 将 Python 层级的 for 循环替换为 jax.lax.scan，以便在编译时对循环体进行一次性追踪，避免为 Runge-Kutta 解算器生成难以管理的计算图。
   • 静态与动态类型： 将常量转换为 JAX 标量，以防止因更改缓存键而导致的完整重新编译。
   • 分支安全保护： 应用下界（例如 jnp.maximum(x, 1e-30)），以防止在 jnp.where 表达式中出现 $0 \times \infty$ 等操作导致 NaN 梯度的产生。
   • 隐函数定理 (IFT) 伴随： 使用 IFT 伴随而非展开循环来对迭代求解器（如气孔导度双分法、Monin–Obukhov 长度）进行微分，因为展开循环会导致错误的梯度。

关键结果
生成的 clm-ml-jax 模型在使用了 2007 年 5 月气象数据的 CHATS7 核桃果园 AmeriFlux 站点上进行了评估。

• 数值等价性： 在 1488 个时间步长中，翻译后的模型与 Fortran 参考模型实现了近乎完美的数值对等，所有变量（如感热、潜热、GPP）的相关系数均为 1.000，且偏差极小。
• 梯度正确性： 反向模式微分已针对中心有限差分进行了验证。所有活跃参数的相对误差均小于 $1 \times 10^{-4}$。该流水线成功处理了复杂的结构，包括迭代求解器和 Runge-Kutta 子步。
• 参数校准效率： 在对三个参数进行的证明性概念校准中，基于梯度的 L-BFGS-B 优化器在少于 50 次评估内就恢复了接近机器精度的真值（$\sim 10^{-19}$）。这比无梯度算法 Nelder-Mead 优化器快了 8 倍（以评估次数计）。
• 性能加速： 在 GPU（Quadro RTX 8000）上使用 jax.vmap，在系综规模为 $N=2,048$ 时，模型实现了比顺序执行 Fortran 快 24 倍 的墙钟时间加速。随着系综规模增加，每个样本的摊销成本从 24.9 ms 降至 11.4 ms，而 Fortran 则稳定在 $\approx 54$ ms/样本。
• 智能体效率： 在 73 个模块翻译任务和 46 个智能体会话中，81.2% 的 Bug 在第一次尝试时即被自主解决。总的人工监督时间估计仅为 10–15 小时。

意义与主张
本文声称，可以系统地将遗留科学代码库翻译为完全可微的框架，而不必进行手动重构，从而保留科学真实性。

• 科学收益： 这项工作使得此前对于 CLM-ml-v2 而言难以实现的基于梯度的参数优化和敏感性分析成为可能。雅可比矩阵分析显示，空气温度和短波辐射在总初级生产力（GPP）和能量通量对敏感性中占据主导地位，而 $V_{c,max}$ 则发挥着独特的控制作用。
• 方法论贡献： 本文提出了一个包含自主修复循环和领域特定 Oracle 验证的可重复五步工作流。它量化了人类与 AI 智能体之间进行“共同创作”进行科学代码现代化的可行性，表明高层级的人类监督足以应对复杂的、持续数日的翻译任务。
• 可重用性： 翻译后的模型及流水线基础设施均已发布，作为一个可重用的框架，旨在支持社区对其他地球系统模型组件进行可微化处理。

作者将这项工作定位为并非针对所有 ESMs 的最终解决方案，而是一个经过验证的概念验证，它弥合了遗留 Fortran 物理学与现代可微编程之间的鸿沟，为物理-机器学习混合架构以及更高效的不确定性量化铺平了道路。