TorchClim v1.0: A deep-learning framework for climate model physics

TorchClim is a framework that allows the introduction of ML/AI models that were trained using PyTorch into a climate model (aka GCM). It facilitates a fast turnover of the train-test-run workflow allowing for quick development of ML/AI-based parametrizations into parallel and distributed environments.

The framework consists of two components: first, a plugin that is in charge of the interactions with the underlying ML/AI framework. Currently, this plugin relies on LibTorch as the underlying implementation. This plugin is largely agnostic to the specifics of the GCM that is using it. Second, we provide a reference implementation of the framework into CESM version 1.0.6 where we replace moist and radiative parametrizations in the Community Atmospheric Model (CAM) version 4 with an ML surrogate. The reference implementation offers a range of innovative features, facilitating many aspects of the train-test-run workflow of ML models for GCMs. For example, it offers tools to extract data from CAM before and after the desired point of insertion of a surrogate parametrization. It also offers the ability to switch parametrizations at runtime based on time-space requirements, etc.