This codebase is associated with the following published paper:
Lin, J., Yu, S., Peng, L., Beucler, T., Wong-Toi, E., Hu, Z., Gentine, P., Geleta, M., & Pritchard, M. (2025). Navigating the noise: Bringing clarity to ML parameterization design with O(100) ensembles. Journal of Advances in Modeling Earth Systems, 17(4), e2024MS004551. https://doi.org/10.1029/2024MS004551
Full text available at: https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2024MS004551
ClimScale is an end-to-end pipeline for large-scale testing of neural network parameterizations of Cloud Resolving Models in Super-Parameterized Community Atmosphere Model 3 (SPCAM3).
There are four main folders for this online-testing pipeline:
1.) preprocessing
2.) training
3.) coupling_folder
4.) coupled_results
Once all the results are created, the storytelling.ipynb notebook shows the results across configurations. The offline_evaluation folder generates results on offline test data. Since the "test set" of interest is online performance, the offline test data is relatively small and only used to cross-check validation results. That folder makes use of the tf2 env found in the envs folder.
The preprocessing folder contains one Python script and 6 jupyter notebooks. preprocessing.py is a Python script that defines all the preprocessing functions used by the 6 jupyter notebooks. The 6 jupyter notebooks create training, validation, and test input and target numpy arrays. The environment for this folder corresponds to preprocessing_env.txt in the envs folder. Code is designed such that one would just need to change functions in preprocessing.py and designated simulation file path in notebooks when creating new training, validation, and test data for a new configuration.
The training folder contains code for conducting large-scale testing on Bridges2; however, this code can be adapted for use on other HPCs.
To change the hyperparameter search space, modify the build_model function in training_functions.py. This code also assumes the use of an environment named tf2 that makes use of TensorFlow 2 and keras tuner. The environment I used corresponds to tf2_env.txt in the envs folder. To change the batch size, number of epochs, training objective, or early stopping callback, make changes to the tuner in 'tuning_template.py'.
Once those scripts have been adjusted for your purposes, you can execute a training run via the terminal command:
python train_models.py project_name max_trials partition clock_time
where project_name is replaced by the name of your project, max_trials is replaced by the number of trials you would like to run (I do 110 at a time), partition is replaced by the name of the partition (I use GPU-shared), and clock_time is replaced by the number of time desired on the partition.
An example would be:
python train_models.py relative_humidity 110 GPU-shared 24:00:00
Be sure to do a test run with a small number of trails and less time to debug any errors in training code before executing a large-scale training run. When adding more trials after a previously completed run, simply run the same line and update the number of trials. For example, if you are training 220 models and you have already trained 110, replace 110 with 220 in the above line.
The coupling folder contains the scripts necessary to couple models trained in the training folder. To use it, modify the build_model function in make_models.ipynb to match the hyperparameter search space used for the build_model function in training_functions.py and use make_models.ipynb to convert the trained models to text files that can be coupled thanks to FKB. This folder also makes use of the tf2 folder used in the training folder.
Running the coupled simulations is a one-line command in terminal:
python master_loop.py family start stop
where family is replaced by the name of the project, start is the model number you want to start with, and stop is the model number you want to end with. An example would be:
python master_loop.py relative_humidity 1 330
The coupled_results folder is a placeholder folder for runs generated by the coupling_folder.
The model is licensed under the Apache 2.0 license.
If you use ClimScale in your research, please use the following BibTeX entry.
@ARTICLE{Lin2025-ya,
title = "Navigating the noise: Bringing clarity to {ML} parameterization
design with {O} $\boldsymbol{\mathcal{O}}$(100) ensembles",
author = "Lin, Jerry and Yu, Sungduk and Peng, Liran and Beucler, Tom and
Wong-Toi, Eliot and Hu, Zeyuan and Gentine, Pierre and Geleta,
Margarita and Pritchard, Mike",
journal = "J. Adv. Model. Earth Syst.",
publisher = "American Geophysical Union (AGU)",
volume = 17,
number = 4,
pages = "e2024MS004551",
abstract = "AbstractMachine‐learning (ML) parameterizations of subgrid
processes (here of turbulence, convection, and radiation) may one
day replace conventional parameterizations by emulating
high‐resolution physics without the cost of explicit simulation.
However, uncertainty about the relationship between offline and
online performance (i.e., when integrated with a large‐scale
general circulation model) hinders their development. Much of
this uncertainty stems from limited sampling of the noisy,
emergent effects of upstream ML design decisions on downstream
online hybrid simulation. Our work rectifies the sampling issue
via the construction of a semi‐automated, end‐to‐end pipeline for
size ensembles of hybrid simulations, revealing important nuances
in how systematic reductions in offline error manifest in changes
to online error and online stability. For example, removing
dropout and switching from a Mean Squared Error to a Mean
Absolute Error loss both reduce offline error, but they have
opposite effects on online error and online stability. Other
design decisions, like incorporating memory, converting moisture
input from specific humidity to relative humidity, using batch
normalization, and training on multiple climates do not come with
any such compromises. Finally, we show that ensemble sizes of may
be necessary to reliably detect causally relevant differences
online. By enabling rapid online experimentation at scale, we can
empirically settle debates regarding subgrid ML parameterization
design that would have otherwise remained unresolved in the
noise.",
month = apr,
year = 2025,
keywords = "hybrid; parameterization; convective; sampling; machine learning;
climate",
language = "en"
}