利用大型语言模型将地球系统模型从 Fortran 迁移至 Python/JAX
分布式、并行与集群计算
2024-05-02 v1 大气与海洋物理
摘要
地球系统模型（ESM）对于理解过去、现在和未来的气候至关重要，但它们受制于遗留的技术基础设施。ESM 主要用 Fortran 语言实现，这种语言对早期职业科学家来说入门门槛高，并且缺乏 GPU 运行时环境，而随着 GPU 算力的提升和 CPU 扩展速度的放缓，GPU 运行时已成为持续发展的必要条件。Fortran 还缺乏可微性——即对数值代码进行微分的能力——而这种能力正是集成机器学习方法的混合模型所需要的。将 ESM 从 Fortran 迁移到 Python/JAX 可以解决这些问题。本工作提出了一种利用大型语言模型（GPT-4）将单个模型组件从 Fortran 半自动翻译为 Python/JAX 的方法。通过翻译社区地球系统模型（CESM）中的光合作用模型，我们证明了 Python/JAX 版本利用 GPU 并行化可实现高达 100 倍的运行速度提升，并能通过自动微分进行参数估计。Python 代码也易于阅读和运行，可供教师在课堂上使用。这项工作为实现快速、包容且可微的气候模型这一最终目标指明了道路。
关键词
引用
@article{arxiv.2405.00018,
title = {Proof-of-concept: Using ChatGPT to Translate and Modernize an Earth System Model from Fortran to Python/JAX},
author = {Anthony Zhou and Linnia Hawkins and Pierre Gentine},
journal= {arXiv preprint arXiv:2405.00018},
year = {2024}
}