AI 中文总结
FESOM2-JAX v1.0是FESOM2基于JAX的可微GPU版本，与原Fortran版本结果一致，具备端到端可微性，是首个CMIP级非结构化网格可微全球海冰-海洋模型。
AI 中文摘要
我们提出FESOM2-JAX，它是有限体积海冰-海洋模型（Finite-volumE Sea ice-Ocean Model，FESOM2）基于JAX框架的Python重实现版本。该模型保留了原始模型的非结构化网格、单元顶点有限体积格式，可在从笔记本电脑CPU到256个GPU的不同硬件上无修改运行，且具备端到端可微性。FESOM2-JAX是Fortran模型的代码镜像：它被投射到Python生态系统中，通过大语言模型翻译，并逐内核与原始模型进行验证。其设计目的是降低实验门槛，涵盖从新数值方法、参数化方案到基于梯度的校准以及混合物理-机器学习组件等方向，同时与原始模型保持足够的一致性，使得在镜像版本上开发的成果可反向迁移回原始模型。在1958-2019年的后报模拟中，采用1°等效分辨率、相同的物理过程和强迫条件，JAX版本与Fortran版本的平均状态差异，远小于两者与观测结果的差异（差异幅度小两个数量级），且两个版本在六十年间的全球温度、盐度、热含量和海冰变化趋势一致。完整的1°配置可在单个GPU上运行，搭载4个GH200超级芯片的节点可实现约113个模拟年/每个挂钟日，表面顶点数达740万（约5km）的网格可扩展至128个GPU运行。模型的性能瓶颈在于通信而非计算。该镜像版本为原始模型新增了梯度功能：通过完整时间循环的单次反向模式传播，可返回模型诊断量对每个网格顶点参数的敏感性，且已通过有限差分法验证。据我们所知，FESOM2-JAX是首个原生基于可微框架编写的CMIP级复杂度全球海冰-海洋模型，也是首个采用非结构化网格的此类模型。
英文摘要
We present FESOM2-JAX, a Python re-implementation of the Finite-volumE Sea ice-Ocean Model (FESOM2) in JAX. The model retains the unstructured-mesh, cell-vertex finite-volume formulation of the original, runs unchanged from a laptop CPU to 256 GPUs, and is end-to-end differentiable. FESOM2-JAX is a code shadow of the Fortran model: a projection onto the Python ecosystem, translated with large language models and verified kernel by kernel against the original. It is built to lower the barrier to experimentation, from new numerics and parameterizations to gradient-based calibration and hybrid physics-machine-learning components, while remaining close enough to the original so that what is developed in the shadow can be transferred back. In a 1958-2019 hindcast at 1$^{\circ}$ equivalent resolution with identical physics and forcing, the mean states of the JAX and Fortran versions differ from each other by two orders of magnitude less than either differs from observations, and the two runs agree for six decades in global temperature, salinity, heat content, and sea ice. The complete 1$^{\circ}$ configuration fits on a single GPU, a node of four GH200 superchips integrates $\sim$113 simulated years per wall-clock day, and meshes of up to 7.4 million surface vertices ($\sim$5 km) scale to 128 GPUs. What limits the model is communication rather than arithmetic. What the shadow adds to the original is the gradient: a single reverse-mode pass through the full time loop returns the sensitivity of a model diagnostic to a parameter at every mesh vertex, verified against finite differences. To our knowledge, FESOM2-JAX is the first global ocean-sea-ice model of CMIP-class complexity written natively in a differentiable framework, and the first on an unstructured mesh.
