19th Workshop on high performance computing in meteorology

AI2 Climate Modeling is working to improve the FV3GFS weather and climate model using online machine learning parameterizations, and to add GPU capability and improve its useability by porting it to a Python-based domain specific language (DSL). Both these efforts have been greatly facilitated by wrapping the existing Fortran code as a library which can be imported in Python, called fv3gfs-wrapper (McGibbon et al. 2021). This package provides simple interfaces to progress the model main loop and get or set variables used by the Fortran model. This enables the main components of the time integration loop to be written in Python, allowing arbitrary Python code to interact with the model state. Furthermore, the base model is as performant when wrapped in Python as it is in pure Fortran.

To demonstrate the wrapper we will show executable examples, including simple MPI communication, adding machine learning to the model, and running interactively in a Jupyter notebook. Everything is written in Docker, allowing you to run these examples on your laptop. If the correct environment is installed, they can also be run on MPI-enabled supercomputing systems.

We will also discuss what it is like developing a weather and climate model based in Python. With this approach, we are able to independently and interactively develop and evaluate model components, and take advantage of a wide array of community libraries for data loading, data visualization, and machine learning. We will also showcase easy-to-use static typing, testing, and debugging tools available in Python which are particularly useful to a model developer.