Accurate future climate projections are crucial to a variety of sectors, including agricultural, energy, and public health. For instance, they can help predict shifts in growing seasons, optimize renewable energy deployment, and improve understanding of climate-sensitive diseases. Currently, physics-based earth system models, also referred to as general circulation models (GCMs), are responsible for generating these projections and require a massive amount of computing resources in order to produce one climate realization . These models need to balance accuracy with feasibility, and most opt to use coarse spatial resolution, typically around 100 km, allowing them to produce ensembles of climate simulations for decades or centuries. However, the downside of using coarse-resolution GCMs is their inability to resolve storms, clouds, and complex topography . The Simple Cloud-Resolving E3SM Atmosphere Model (SCREAM) addresses the resolution problem by using kilometer-scale resolution, but it is too computationally expensive to run more than a few years at a time . This project aims to apply a computationally efficient machine learning (ML) based emulator of high-resolution SCREAM nudging tendencies, thus replicating the high accuracy of the cloud-resolving GCM at the reduced computational cost of the coarse-resolution GCM.

ML has the potential of revolutionizing how weather forecasts and climate predictions are generated . Within this active area of research, there are two primary approaches. The first approach relies solely on ML models trained on historical observations, reanalysis datasets, or outputs from existing GCMs . The second approach involves a hybrid framework that combines machine learning with traditional GCMs. In this method, machine learning methods are often used to replace a specific physical process of GCMs or to apply column-wise correction to the coarse-grid GCM . This study builds on the success of a study using the second approach by the Allen Institute for Artificial Intelligence (Ai2) working with the Finite-Volume Cubed-Sphere Global Forecast System (FV3GFS) GCM (https://github.com/ai2cm/fv3gfs-fortran, last access: 9 July 2024) . Ai2 was able to improve the predictive accuracy of relatively coarse-resolution global simulations using a machine-learning based correction to the atmosphere state .

This study applies their corrective-ML approach to SCREAM. In addition, it developed software infrastructure that will be helpful for subsequent studies coupling ML to SCREAM. Section  provides the recipe for the development of the ML training data. Section  covers how the corrective-ML model was embedded into the SCREAM code base. Section  examines how well the ML-corrected coarse GCM performs with respect to the fine-resolution target, followed by discussion of the results and conclusions of the study in Sects.  and .

SCREAM is a configuration of the atmosphere component of the Energy Exascale Earth System Model (E3SM) targeting kilometer-scale global resolutions. In order to accomplish performant simulations at this scale, the E3SM Atmosphere Model (EAMxx) was rewritten from the original Fortran code to C++/Kokkos . The adoption of Kokkos in EAMxx unlocks the computational power of mixed CPU/GPU machines and ensures that EAMxx maintains good performance across a number of high performance computing systems . EAMxx employs a spectral-element, nonhydrostatic dynamical core for resolved-scale processes . Major subgrid physics parameterizations include P3 microphysics , SHOC turbulence and boundary layer physics , and RRTMG++ radiation . Notably, the SCREAM configuration omits a deep convection parameterization, regardless of grid resolution . For all simulations in this work, the model configuration and initial start date are identical. Initial conditions for all simulations are generated from ERA5 .

While EAMxx achieves computational performance of more than one simulated year per wallclock day for a km-scale GCM , its computational cost is still prohibitive for global km-scale simulations of a decade or more in length. This limitation makes it attractive to apply corrective-ML to EAMxx in order to represent effects of fine scale features at coarse resolution. We follow Ai2's approach with FV3GFS , using a km-scale EAMxx simulation to produce fine-scale training data for a corrective-ML model embedded in coarse-resolution EAMxx simulations. The approach is modified to handle technical and architectural differences between FV3GFS and EAMxx, see Sects.  and  for more details.

The corrective-ML approach described in this study can be represented by four steps (Fig. ). The first step involves the generation of a fine-grid reference simulation. This simulation serves two purposes: to provide the nudging targets for the training data set and to provide a validation data set for the ML-corrected solution. The second step involves nudging an EAMxx simulation at coarse resolution to the reference state. “Nudging tendencies” derived from this step are used for the third step, training of the corrective-ML models. Finally, a coarse-grid EAMxx simulation is run with inline corrections from these ML models; ideally the resulting ML-corrected coarse GCM will more closely approximate the time evolution and solution of the target fine-grid GCM. Each of these steps is described in more detail in the following subsections.

All tunable model parameters are held fixed between the coarse- and fine-resolution simulations, so in every case the parameter space is consistent with that of the target fine-resolution simulation. As a result, the climate states of the untuned coarse simulations would likely diverge from the fine resolution over long timescales. However, this study focuses on the shorter timescale of one year. We therefore expect the coarse simulations, when nudged toward the fine-resolution state, to agree more closely with the fine-resolution simulation.

We opted to utilize an existing framework, fv3net , built by Ai2 and used in . fv3net includes all aspects of the ML workflow, including pre-processing native GCM output, ML training, offline and online testing, and reporting. fv3net was originally developed for Geophysical Fluid Dynamics Laboratory's (GFDL) FV3GFS. We made necessary modifications to work with EAMxx's data, and all updates are available in the open-source fv3net repository.

One of the main challenges is that fv3net is written in Python whereas EAMxx is in C++. We use pybind11 to bridge between C++ and Python . The strategies for running on CPU and GPU are slightly different. On CPU, we utilize pybind11's numpy bindings to transfer EAMxx's Kokkos view data to Python and declare them as numpy arrays. We do this as a non-copy operation, making the operations relatively cheap. On GPU, the data passing is different since numpy does not support GPU. In this case we pass EAMxx's Kokkos view data as pointers to Python and reconstruct numpy-like arrays from the Cupy package through unmanaged memory access . As in the CPU case, we do not do any memory copies. Once Cupy arrays are constructed on the Python side, we are able to reuse the same workflow as in the CPU case. For both the CPU and GPU implementation the state variables are directly overwritten during the Python calls. The ML workflow adds 10 % overhead for the CPU runs and 6 % for GPU. There is also a one-time initialization cost for loading the necessary libraries in Python, which becomes negligible once the simulation is sufficiently long (e.g., 1 simulation month for the 100 km resolution case).

While it is possible to avoid Python and instead directly interface with a trained ML model in C++ (e.g. via Tensorflow's C++ interface), our approach allowed us to take advantage of the existing fv3net package's robustness and stability. For example, we are able to take advantage of safeguards built into fv3net to ensure ML has physically-consistent output. Since fv3net is an end-to-end ML workflow, we also utilize its reporting for data analysis and visualization. Overall, using fv3net drastically reduced development time in every step of the ML workflow.

We present results from a 1-year free-running coarse-grid simulation that includes “online” ML correction toward the fine-grid reference . We use an uncorrected 100 km run as a baseline comparison. While the ML model captures the overall vertical structure of the target nudging tendencies for both temperature and specific humidity, it consistently underestimates the strength. This is illustrated in Fig.  for the temperature and specific humidity fields. Figures  and  show the annual-mean bias and root mean square error (RMSE) of baseline and ML-corrected runs from the reference data (3.25 km coarsened to 100 km, also used as ML training data). For this experiment, we found the largest improvement in surface downwelling longwave flux, with 16 % improvement in global RMSE when compared to the baseline. Additionally, we found marginal improvement (5 %–10 %) for net shortwave flux at the surface, precipitation, and surface pressure. However, RMSE was worse for surface temperature (24 %) and total water vapor path (43 %) in the ML-corrected run. Not shown are the eastward and northward wind comparisons, which showed little to no change between the nudged and ML-corrected solutions.

As shown in Fig. , the ML model exhibits its largest bias in the tropics. This underestimation of the target nudging tendencies likely contributes to the observed weakening of the ITCZ, because the model provides insufficient local heating and moistening to sustain the expected convergence and convective intensity. Recall that EAMxx currently does not include a deep convection scheme. This omission may help explain the greater weakening observed in the tropics, where stronger ML corrections would be required to reduce the bias.

The pattern of variables that improve or degrade is broadly consistent with the underestimation of nudging tendencies shown in Fig. . Surface radiative fluxes show the greatest improvement, which is plausible because the surface forcing model learns these fluxes directly from the fine-grid reference rather than inferring them from predicted tendencies. Precipitation and surface pressure improve more modestly. By contrast, water vapor path and surface temperature are closely linked to the temperature and specific humidity profiles that the model is designed to correct. When these corrections are systematically too weak, particularly in the tropics, as evident in Fig. , the ML-corrected simulation may become less self-consistent than the uncorrected baseline, leading to larger annual-mean errors in these fields.

To evaluate the stability and robustness of the corrective-ML approach, we conducted a series of training experiments using three different random seeds in the neural networks. During the online prognostic runs, models trained with seeds 1 and 3 crashed due to the inputs being out-of-sample. To address this issue, we adopted the method outlined in , training a separate out-of-sample novelty detector to complement the existing neural networks for temperature and specific humidity modeling. This approach proved reliable, allowing all random seeds to successfully complete a full-year simulation. We found the ML-predicted tendencies to be similar across random seeds, although these results are not shown.

This study investigated the transferability of an ML parameterization from one GCM to another, a subject of growing importance as modeling groups plan to integrate ML machinery from other groups. Although we adopted the same strategy as Ai2, it quickly became clear that applying the workflow to a different atmospheric model would be challenging because of differences in GCM design and implementation. Each step in the corrective-ML workflow – nudging from coarse to fine-resolution simulation, prescribing surface fluxes for the coarse simulation, integrating ML, and diagnostics – is tightly coupled to the specifics of the underlying GCM, making generalization difficult. Even though this project was a collaboration between the team that developed EAMxx and the team that successfully applied this corrective-ML approach to FV3GFS, the end result fell short of the desired improvement from corrective-ML.

There are several possibilities for why corrective-ML was less successful for EAMxx than it was for FV3GFS. The most obvious difference between the two modeling frameworks is that the ML-corrected version of EAMxx does not include a deep convective parameterization and therefore produces a less realistic large-scale state. As a result, more is being asked of corrective-ML in EAMxx than in FV3GFS. Our initial expectation was that larger differences between the nudged and free-running low-resolution simulations would provide a stronger signal for training the ML. But in such an environment, imperfections in the corrective-ML will also stand out more. Potential improvements in the ML model from the stronger signal may be too small to offset the larger biases in EAMxx associated with the absence of a subgrid convection scheme. Additionally, this study was limited to 3-hourly averaged fields, rather than the instantaneous fields used in the FV3GFS study. Training on these smoother averaged fields may have reduced the ML model's skill in predicting the appropriate nudging tendencies for instantaneous states. Future work will address these hypotheses when EAMxx is enhanced to support a low-resolution configuration with parameterized deep convection.

As discussed in the Methods section, there are three independent ML models that make up the full corrective-ML implementation, each with its own set of parameters to train on. In this study, we focused on training the models following the same approach as with the FV3GFS test case . Given the differences between EAMxx and FV3GFS, it is possible that a successfully trained model may require an adjustment to the training hyperparameters. An example is the use of latitude as an input into the ML training data set. We found that the inclusion of latitude in the training input improved model accuracy, presumably as a proxy for physical variables that correlate strongly with latitude but are not well represented in the coarse GCM. A subsequent study could investigate what other parameters and input variables improve the performance of the ML models.

Adoption of the corrective-ML framework involved the development of several new features in EAMxx (https://github.com/E3SM-Project/E3SM, last access: 24 June 2024). EAMxx is heavily unit tested, but there is always a chance that a difficult bug or untested edge case is able to make it past the testing. There are three major developments that were added to EAMxx as part of this effort; Python to C++ bridging for ML, ability to nudge the EAMxx model state, and prescribed atmospheric surface fluxes in EAMxx. While significant effort went into developing and testing these new features, future work may want to revisit their implementation as a potential source of error in the ML implementation.

Finally, as outlined in the Methods section, the fine-resolution solution was coarsened to the target resolution both horizontally and vertically. The decision to map onto fixed vertical pressure coordinates resulted in situations where grid points over topography could lead to many or all data points being masked. This could create either masked data in the target state or grid points with disproportionate weighting based on only a few source columns. To investigate this hypothesis, a separate lightweight study was conducted using a 25 km fine-resolution simulation, comparing two corrective-ML models. One model utilized the approach described in this study, while the other did not apply vertical interpolation for the target state. Preliminary analysis showed that the latter model had slightly better performance. Unfortunately, due to the high cost of 3.25 km simulations, it was not feasible within the scope of this project to repeat this study using horizontal-only interpolation. A more detailed investigation of the impact from vertical interpolation should be included in any subsequent studies.

Although this work did not achieve the anticipated accuracy improvements from corrective-ML in EAMxx , it did have several beneficial effects. In particular, this study marked the first time Python packages were integrated into EAMxx, providing a blueprint for future hybrid ML without degrading EAMxx performance on either CPU or GPU architectures. This work also led to the development of new infrastructure within EAMxx to support prescribed surface fluxes and atmospheric state nudging. While the latter was already a planned milestone for the regionally-refined configuration of EAMxx, this effort expanded that capability. The introduction of prescribed surface fluxes is important for some doubly-periodic use cases and useful for debugging and hypothesis testing . Finally, the fact that EAMxx benefited less from corrective-ML than FV3GFS is a useful data point in the quest to understand generalizability of ML approaches in GCMs.

There is increasing interest in the earth system modeling community in utilizing differentiable algorithms and software. This strategy has already been successfully applied to NeuralGCM, which marries a spectral dycore coded in the Jax machine learning language with a learned ML representation of the combined physical parameterizations of a global model, trained to evolve following a gridded reanalysis. In this setting, the ML model can be trained end-to-end including feedbacks with other earth system components within this differentiable framework, dramatically improving the prognostic performance of this hybrid model such that it has highly accurate weather forecasts and strong climate skill . In principle, an analogous approach could be implemented to dramatically improve the prognostic skill of corrective-ML within EAMxx. However, refactoring EAMxx to enable such a capability while retaining leadership-class performance would be a substantial software engineering challenge.