论文阅读：FV3GFS模式的Python封装

April 15, 2021 (最后修改: April 15, 2021)

fv3gfs-wrapper: a Python wrapper of the FV3GFS atmospheric model

McGibbon, J., Brenowitz, N. D., Cheeseman, M., Clark, S. K., Dahm, J., Davis, E., Elbert, O. D., George, R. C., Harris, L. M., Henn, B., Kwa, A., Perkins, W. A., Watt-Meyer, O., Wicky, T., Bretherton, C. S., and Fuhrer, O.: fv3gfs-wrapper: a Python wrapper of the FV3GFS atmospheric model, Geosci. Model Dev. Discuss. [preprint], https://doi.org/10.5194/gmd-2021-22, in review, 2021.

一篇投稿于 Geoscientific Model Development 的预印本论文，介绍使用 Python 对 FV3GFS 大气模式进行封装。文中将 Fortran 模式拆解为多个可以封装成 API 接口的步骤，比如动力框架和物理过程，使用 Python 控制整个模式积分过程，为在积分过程中整合其它代码提供基础。

这是一个非常值得借鉴的思路，即用 Python 等高层语言提供简化的顶层 API 接口，而对于需要优化性能的子模块，使用 C++、 Fortran 等底层语言进行封装。

注意：论文还在评审阶段，本文翻译的版本不是最终的发布版本，后续可能会有修改。

正文

因为需要满足严格的性能需求，地理科学仿真软件通常使用 Fortran 或 C++ 实现。这些代码通常难以理解，不易修改，很难与高效率语言进行交互，以进行探索性分析。

fv3gfs-wrapper 是 NOAA FV3GFS 全球大气模式的 Python 封装，提供简单接口，用于：

推进 Fortran 主循环
从 Fortran 模式中获取或设置状态

应用场景示例：

修改模式的行为
引入在线分析代码
直接从云存储读取或向云存储保存模式状态或强迫数据

拷贝数据的耗时在可接受范围内，可以通过修改 Fotran 源码避免额外增加的耗时。

上述封装方法同样也适用于其他 Fotran 模式，以提供更有生产力的科学工作流。

介绍

FV3GFS 模式使用 Fortran + OpenMP/MPI 实现。

使用用户基数较小的底层强类型编程语言需要权衡取舍。工具库不像在类似 Python 的高层语言中可用或广泛使用，例如：

与云存储交互
进行物理或统计学分析
机器学习

Fortran 编译代码的 Python 接口可以让有庞大基数的 Python 用户与代码进行交互，让 Python 工具生态系统可以与模式例程交互。

以 Fortran 为中心的工作流

以 Fortran 为中心的工作流，根据论文图片重新绘制

Python 通常作为后处理工具集成到 Fortran 模式工作流中，对 Fortran 模式保存到文件系统中的数据进行计算，如上图所示。这种方式有一些弊端：

每个时间步都保存全分辨率的模式状态基本不可行，所以通常使用一段时间的统计值代替
除非保存足够频率的快照，否则直接从全分辨率瞬时场计算新的统计量需要编写 Fortran 代码在模式中运行
需要将数据保存到文件系统中才能使用 Python 代码分析，如果数据不是必须的最终产品，可能会造成资源浪费
文件系统 IO 性能可能是显著的瓶颈 (译者注：这点非常关键)
无法使用 Python 修改 Fortran 模式的行为，修改模式行为的代码必须使用 Fortran 重新实现

以 Python 为中心的工作流

以 Python 为中心的工作流，根据论文图片重新绘制

FV3GFS 模式被编译成共享库，提供对模式进行控制和交互的 API 接口。

模式的核心是主积分循环 (main integration loop)，论文将其拆解为一系列子函数，并使用 Python 调用：

使用 Python 编写主循环，调用每个部分：step_dynamics, step_physics
支持从 Fortran 例程获取变量 (get_state)，也支持拷贝变量到 Fortran 例程 (set_state)，允许 Python 函数影响 Fortran 模式状态的积分

以 Python 为中心的工作流提供一种将 Python 生态系统集成到 Fortran 模式工作流中的方法：

诊断程序：在物理过程或动力框架步骤后添加在线诊断代码
定制IO：与云存储的 I/O 交互
物理过程参数化：定制逻辑，通过机器学习参数化增强 Fortran 模式

使用机器学习的参数化方案，需要在模式中在线运行。有以下几种方案：

从 Fortran 调用 Python (Brenowitz and Bretherton, 2019)
使用 Fortran 重新实现神经网络代码 (Ott et al., 2020; Curcic, 2019)
在大气模式中包含 Python 机器学习代码 (论文使用的方法)

当前 Python 代码只能在动力框架和物理方案例程外部集成，而不是在物理方案内部。在各个物理方案之间引入 Python 代码以增加灵活性仍然是未来的工作方向。

之前的工作

qtcm (Lin, 2009)，使用 f2py (Peterson, 2009)
PyCLES (Pressel et al., 2015)
CliMT (Monteiro et al., 2018)
nbodykit (Hand and Feng, 2017)

优先事项

设计新的大气模式时要考虑的是，已经有大量的历史 Fortran 代码。因此，新的模式组件通常使用 Fortran 编写，可以与历史代码交互。论文中这种用其他语言重构或重写代码逐步替换现有组件的能力，能让重写现有 Fortran 模式 (例如在 GPU 架构上运行) 的工作受益。

优先事项：

保持 Fortran 模式的现有功能
最小程度地牺牲性能
易于理解和修改主时间步进循环
对 Fortran 代码修改最少

专注提高模式的易用性，面向对修改 Fortran 代码行为感兴趣的研究人员。

应用

论文介绍三种应用：

机器学习参数化
定制 MPI 通讯的在线诊断代码
在 Jupyter Notebook 中在线分析

验证

验证是否可以得到二进制一致的结果

import fv3gfs.wrapper

if __name__ == "__main__":
    fv3gfs.wrapper.initialize()
    for i in range(fv3gfs.wrapper.get_step_count()):
      fv3gfs.wrapper.step_dynamics()
      fv3gfs.wrapper.step_physics()
      fv3gfs.wrapper.save_intermediate_restart_if_enabled()
      fv3gfs.wrapper.cleanup()

流程被分为 5 个例程：

初始化模式
清理
动力框架
物理过程
输出中间数据，用于重启动

运行时间没有显著差异。进行 6 小时预报，C48 分辨率的耗时：

示例	运行耗时 (秒)
Fortran 基线	110
封装基线	110
随机森林	116
最小地面气压	110

代码：https://doi.org/10.5281/zenodo.4429297

应用

示例代码：https://doi.org/10.5281/zenodo.4429297

机器学习

通过机器学习增强模式，修改主循环流程，添加机器学习模型。

下面代码使用随机森林

import fv3gfs.wrapper
import fv3gfs.wrapper.examples
import f90nml
from datetime import timedelta

if __name__ == "__main__":
    # load timestep from the namelist
    namelist = f90nml.read("input.nml")
    timestep = timedelta(seconds=namelist["coupler_nml"]["dt_atmos"])

    # initialize the machine learning model
    rf_model = fv3gfs.wrapper.examples.get_random_forest()

    fv3gfs.wrapper.initialize()
    for i in range(fv3gfs.wrapper.get_step_count()):
        fv3gfs.wrapper.step_dynamics()
        fv3gfs.wrapper.step_physics()
        
        # apply an update from the machine learning model
        state = fv3gfs.wrapper.get_state(rf_model.inputs)
        rf_model.update(state, timestep=timestep)
        fv3gfs.wrapper.set_state(state)

        fv3gfs.wrapper.save_intermediate_restart_if_enabled()

    fv3gfs.wrapper.cleanup()

MPI 通讯

基于 MPI4py 实现 halo 更新，收集和分发操作。

下面示例计算全球地面温度最低值并在根进程中打印输出，展示如何在模式内使用 MPI4py 通过进程间通讯来计算诊断量。

import fv3gfs.wrapper
import numpy as np
from mpi4py import MPI

ROOT = 0

if __name__ == "__main__":
    fv3gfs.wrapper.initialize()
    # MPI4py requires a receive "buffer" array to store incoming data
    min_surface_temperature = np.array(0)
    for i in range(fv3gfs.wrapper.get_step_count()):
        fv3gfs.wrapper.step_dynamics()
        fv3gfs.wrapper.step_physics()

        # Retrieve model minimum surface temperature
        state = fv3gfs.wrapper.get_state(["surface_temperature"])
        MPI.COMM_WORLD.Reduce(
        state["surface_temperature"].view[:].min(),
        min_surface_temperature,
            root=ROOT,
            op=MPI.MIN,
        )
        if MPI.COMM_WORLD.Get_rank() == ROOT:
            units = state["surface_temperature"].units
            print(f"Minimum surface temperature: {min_surface_temperature} {units}")

        fv3gfs.wrapper.save_intermediate_restart_if_enabled()
    fv3gfs.wrapper.cleanup()

Jupyter Notebook

上述示例均可通过 ipyparallel 在 Jupyter Notebook 中进行交互分析。

图片来自论文，Jupyter Notebook 示例截图

实现

信息传递

读写模式状态的两种方式：

封装接口拷贝数据
共享内存

论文使用第一种方式，对 Fortran 代码改动最小，在 numpy 申请的 C 数组和 Fortran 数组之间拷贝数据。

用于传递数组的元编程

从 Python 拷贝数据到 Fortran 数组需要为每个 Fortran 变量提供 Python 变量到 Fortran 变量名之间的对应关系。

论文为每个变量提供两个 Fortran 封装例程：

C 封装层内的逻辑，用于调用 Fortran 封装的子例程
子例程的声明头

作者使用 Jinja 模板为每个变量生成封装代码，在 JSON 文件中声明的必要信息，例如：

Fortran name
Fortran container struct name
标准名称 (standard name)
单位 (units)
维度 (dimensions)

译者注：为什么不在代码中直接使用 JSON 文件而是要生成模板代码？

使用 Docker 封装和测试

使用 Docker 容器进行测试，无需单独再为不同环境进行单独配置。 Docker 容器编排文档 (Dockerfile) 也说明了如何为 FV3GFS 和 fv3gfs-wrapper 构建环境，可以当做指南。 Docker 容器可以让无法使用 HPC 的用户很方便地运行模式。

扩展到其它模式

论文中的封装方式可以扩展到其它模式。只要可以定义 Fortran 子例程来执行模式主循环的每个部分，就可以将这些子例程封装并从 Python 调用以进行模式积分。

挑战和限制

Python 载入问题

Python 程序在初始化加载包时会读取大量文件。

译者注：CMA-PI 上也有类似的问题。

解决方案：

并行文件系统
python-mpi-bcast
修改 CPython 二进制程序 (Enkovaara et al., 2011)

内部集成

当前封装将动力框架和物理过程看成单一的子例程，无法在物理框架中加入 Python 代码。

可能的解决方案：CCPP (Common Community Physics Package)

物理关联

FV3GFS 是一个复杂的并行模式，修改模式状态需要维持变量之间的物理关联。

结论

FV3GFS 大气模式的 Python 封装，允许用户控制用 Fortran 编写的大气模式并与之交互。提供以 Python 为中心的工作流，可以用于：

机器学习参数化开发
在线分析
交互运行模式

论文用示例说明 Python 和 Docker 可以用于重现和修改现有的 Fortran 模式，并介绍 Fortran 代码如何在交互 Jupyter 环境中被调用。

论文的封装方法可以推广到其它模式。用 Python 封装的 FV3GFS 大气模式显示了新一代天气和气候模式的方向，其中模式顶层控制流使用 Python 等高级语言实现，而性能关键部分则使用一种底层高性能的语言。该方法已应用在诸如 Numpy 和 Tensorflow 等流行的 Python 库中。

作者希望看到这种方法可以扩展到其它模式，从而在开发传统的 Fortran 模式时实现对 Python 工具的更广泛访问，并为有兴趣将在线分析代码引入模式的研究人员和学生减少访问障碍。

讨论

这篇论文使用 Python 控制数值天气预报模式积分循环的方式令我大开眼界。

我在几年前尝试使用 Python 替换控制 GRAPES MESO 模式运行的 Shell 脚本，另外也尝试使用 ecFlow 的 Python API 生成 SMS 的系统流程。相关代码放到两个项目中 (Private archived)：

不久后就发现这样做与主流方式差别太大，看不到明显的收益，没有推广价值。在我们决定采用某项“前沿”技术替换“现有”技术时，一定要从团队的角度去权衡，而不能仅凭个人喜好。

目前 Python 已在业务系统中逐步得到应用，我们使用 ecFlow 的 Python API 构建系统流程并开发监控工具，使用 MUSIC Python SDK 从 CMADaaS 检索观测资料，在数据查找、资料预处理、模式检验等工具中使用 Python 脚本，使用 Python 开发多种研发支撑软件平台，同时也在使用 Python 实现机器学习算法。

这篇论文让我惊喜地发现，Python 在气象领域中的应用还有这么“激进”的研究方向。虽然我并不清楚文中的方法是否适合当前的模式系统，但我认为这种高层和底层编程语言相结合的方式对研发模式支撑软件平台有很强的借鉴意义。比如，可以对于现有 API 接口进行二次开发，提供更方便使用、更匹配单位特定需求的新接口。我在去年有所尝试，对 NCL 绘图脚本进行封装，提供可以用于 Jupyter Notebook 的 Python 接口。

nwpc-oper/nwpc-graphics

去年我对 NCL 脚本的封装仅仅是为了验证 Jupyter Notebook 是否可以用于开发诊断工具。这篇论文则让我看到类似的方法有更广泛的应用空间，对于支撑工具开发人员来说是一个不错的方向。

另外，这篇文章给我带来非常有益的提醒：

要利用 Python 社区丰富的软件生态系统，但没有必要将所有程序都用 Python 重新实现

通过跨语言编程接口，可以将 Python 作为胶水语言，向用户提供灵活易用的 Python API，具体逻辑则由其他更底层的语言实现。 ECMWF 发布的大量 Python 库就是对原有 C++/Fortran 工具库的封装，比如 ecCodes，ecFlow 等等。

最近两年来我已很少编写 Python 之外的程序，其中一方面原因是 Python 语言确实上手容易，开发原型速度快。还有一方面原因是我熟悉的其它编程语言在单位几乎没有同行使用，想要推广应用必须封装成命令行程序，除非为个人工作需要而开发。不过，从程序员角度来说，在 Python 培训班铺天盖地的当前，仅使用 Python 一种语言编程不是一个很好的趋势。强类型编译型编程语言与脚本语言在编程方法上略有不同，还是应该多用几种编程语言来训练自己的编程能力 (C++？不，说的就是 GO) 。

参考

论文网址： https://gmd.copernicus.org/preprints/gmd-2021-22/

示例代码：

https://github.com/VulcanClimateModeling/fv3gfs-wrapper

https://github.com/VulcanClimateModeling/fv3gfs-fortran

https://github.com/VulcanClimateModeling/fv3gfs-util

Docker 镜像

gcr.io/vcm-ml/fv3gfs-wrapper:v0.6.0