通过对动力学核心的自动微分学习垂直坐标

Source: arXiv - 2512.17877v1

概述

本文提出了一种新方法，使天气模型的垂直网格能够 learn 自身的形状，而不再依赖人工调校的解析公式。通过在完全可微分的 dynamical core 中嵌入基于神经网络的 coordinate transformation，作者能够自动调整网格以最小化模拟误差，显著降低在陡峭地形上通常出现的 spurious motions。

关键贡献

• 可学习的随地形变化的坐标 – 作者提出了 NEUral Vertical Enhancement（NEUVE），这是一种单调的神经网络映射，将垂直坐标定义为地形高度的函数。
• 端到端可微分求解器 – 他们在 Arakawa C‑grid 上实现了 2‑D 非静力 Euler 方程的完全可微分实现，使得网格本身可以通过梯度优化。
• 通过 AD 精确计算度量 – 使用自动微分（AD）获得精确的几何度量项（例如 Jacobians、度量系数），消除了有限差分近似坐标导数时产生的截断误差。
• 耦合物理‑数值优化 – 将损失（例如相对于参考解的均方误差）通过时间积分方案反向传播到坐标参数，联合调优物理模型和数值方法。
• 实验性提升 – 在标准基准案例中，学习得到的坐标将均方误差降低了 1.4‑2 倍，并消除了传统混合/​SLEVE 网格在陡峭山脉上常见的“垂直速度条纹”。

方法论

1. 参数化垂直映射 – 垂直坐标 ( \eta ) 表示为神经网络输出 ( f_\theta(x,z) ) 的积分，该输出被约束为正值，保证从物理高度到模型层的单调映射。
2. 可微分动力核心 – 非静力学欧拉方程在 Arakawa C‑grid 上离散。所有算子（通量、压强梯度、度量项）均使用 JAX/PyTorch 风格的 AD‑兼容代码实现，使整个前向模拟成为计算图。
3. 度量项计算 – 不使用有限差分近似坐标变换的导数，而是直接从 AD 获得雅可比矩阵 ( \partial (x,z)/\partial (x,\eta) ) 及相关度量系数，得到机器精度的数值。
4. 训练循环 – 损失函数（通常是固定积分时间后模拟场与参考场差的 L2 范数）对神经网络参数 ( \theta ) 求导。使用随机梯度下降（或 Adam）更新 ( \theta ) 以最小化损失。
5. 评估 – 训练完成后，冻结学习得到的坐标并在传统（非梯度）模拟中使用，以评估误差降低和视觉伪影。

结果与发现

• 视觉质量： NEUVE 运行中，特征性的垂直速度“条纹”模式消失，导致近地表的场更平滑。
• 稳定性： 时间步限制与传统网格相当，表明学习得到的坐标系没有引入隐藏的 CFL 违规。
• 泛化能力： 在一种地形上学习的坐标能够相当好地迁移到相似但未见过的地形形状，暗示网络捕获了通用的平滑原理，而不是对单一案例过拟合。

Practical Implications

• Reduced manual tuning: Model developers can replace heuristic decay parameters (e.g., in hybrid or SLEVE coordinates) with a single training phase, freeing them from costly trial‑and‑error calibration.
• Improved forecast accuracy in complex terrain: Weather and climate models that operate over mountainous regions (e.g., alpine forecasting, wildfire spread modelling) can benefit from lower numerical noise without redesigning the whole dynamical core.
• Plug‑and‑play component: Because the NEUVE mapping is a thin wrapper around the vertical coordinate, it can be swapped into existing dynamical cores that already support differentiable programming (e.g., JAX‑based or TensorFlow‑based models).
• Potential for adaptive grids: The same framework could be extended to online adaptation, where the coordinate evolves during a simulation to follow moving fronts or convection, opening a path toward fully adaptive mesh refinement with minimal overhead.
• Cross‑disciplinary reuse: The idea of learning metric terms via AD is applicable to ocean models, plasma simulations, or any PDE solver that uses curvilinear coordinates.

限制与未来工作

• 训练的计算成本： 端到端可微求解器比手写的非自动微分版本慢，使得在高分辨率 3‑D 模型上的训练阶段成本高昂。
• 基准范围： 实验仅限于 2‑D 理想化测试；将其扩展到全物理、全球尺度模型仍是一个未解决的挑战。
• 学习映射的可解释性： 虽然单调性得到保证，但神经网络对 “最佳平滑” 的内部表征不透明，可能会影响在实际运行中的可信度。
• 稳定性保证： 论文展示了经验性的稳定性，但未提供形式化证明，说明学习得到的坐标始终满足 CFL 或能量守恒约束。
• 未来方向 包括将 NEUVE 扩展到三维、将其与物理参数化方案（例如微物理）集成，以及探索元学习方法，以在不同模型配置之间生成可迁移的坐标生成器。

作者

• Tim Whittaker
• Seth Taylor
• Elsa Cardoso‑Bihlo
• Alejandro Di Luca
• Alex Bihlo

论文信息

• arXiv ID: 2512.17877v1
• 分类: physics.ao-ph, cs.LG, physics.flu-dyn
• 出版日期: 2025年12月19日
• PDF: 下载 PDF