SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tung Nguyen*
tungnd@cs.ucla.edu
Jason Jewik*
jason.jewik@ucla.edu
Hritik Bansal
hbansal@ucla.edu
Prakhar Sharma
prakhar6sharma@gmail.com
Aditya Grover
adityag@cs.ucla.edu
Modeling weather and climate is an essential endeavor to understand the near- and long-term impacts of climate change, as well as inform technology and policymaking for adaptation and mitigation efforts. In recent years, there has been a surging interest in applying data-driven methods based on machine learning for solving core problems such as weather forecasting and climate downscaling. Despite promising results, much of this progress has been impaired due to the lack of large-scale, open-source efforts for reproducibility, resulting in the use of inconsistent or underspecified datasets, training setups, and evaluations by both domain scientists and artificial intelligence researchers. We introduce ClimateLearn, an open-source PyTorch library that vastly simplifies the training and evaluation of machine learning models for data-driven climate science. ClimateLearn consists of holistic pipelines for dataset processing (e.g., ERA5, CMIP6, PRISM), implementation of state-of-the-art deep learning models (e.g., Transformers, ResNets), and quantitative and qualitative evaluation for standard weather and climate modeling tasks. We supplement these functionalities with extensive documentation, contribution guides, and quickstart tutorials to expand access and promote community growth. We have also performed comprehensive forecasting and downscaling experiments to showcase the capabilities and key features of our library. To our knowledge, ClimateLearn is the first large-scale, open-source effort for bridging research in weather and climate modeling with modern machine learning systems. Our library is available publicly at https://github.com/aditya-grover/climate-learn.
The escalating extent, duration, and severity of extreme weather events such as droughts, floods, and heatwaves in recent decades are some of the most devastating outcomes of climate change. Moreover, as average surface temperature is anticipated to continue rising through the end of the century, such extreme weather events are likely to occur with even greater intensity and frequency in the future [12, 33, 43, 46, 67, 74]. The key devices used by scientists to understand historical trends and make such predictions about future weather and climate are the numerical weather prediction (NWP) models. These models represent Earth system components including the atmosphere, land surface, ocean, and sea ice as intricate dynamical systems, and they are the prevailing paradigm for weather and climate modeling today due to their established reliability, well-founded design, and extensive
study [3, 13, 29, 51]. However, they also suffer from notable limitations such as inadequate resolution of several subgrid processes, coarse representation of local geographical features, incapacity to utilize sources of observational data (e.g., weather stations, radar, satellites) in an automated manner, and substantial demand for computing resources [2, 39, 40, 64]. These deficiencies combined with the expanding availability of petabyte-scale climate data [18, 19, 50] and lowering compute requirements of machine learning (ML) models in recent years have motivated researchers from both the climate science and artificial intelligence (AI) communities to investigate the application of ML-based methods in weather and climate modeling [7, 14, 35, 36, 44, 47, 59, 77, 89].
In spite of this growing interest, the improvements have been marred by the lack of practically grounded data benchmarks, open-source model implementations, and transparency in evaluation. For example, many papers in weather forecasting [4, 8, 17, 22, 37, 54, 65, 75, 80, 86, 87, 90], climate projection [84], and climate downscaling [1, 41, 49, 68, 70, 82] choose to benchmark on different geographical regions, temporal ranges, evaluation metrics, and data augmentation strategies. These inconsistencies can confound the source of reported improvements and promotes a culture of irreproducible scientific practices [34]. Recently, there have been some leaderboard benchmarks, such as WeatherBench [63], ClimateBench [84], and FloodNet [61], that propose datasets and baselines for specific tasks in climate science, but a holistic software ecosystem that encompasses the entire data, modeling, and evaluation pipeline across several tasks is lacking.
To bridge this gap, we propose ClimateLearn, an open-source, user-friendly PyTorch library for data-driven climate science. To the best of our knowledge, it is the first software package to provide end-to-end ML pipelines for weather and climate modeling. ClimateLearn supports data preprocessing utilities, implements popular deep learning models along with traditional baseline methods, and enables easy quantification and visualization of data and model predictions for fundamental tasks in climate science, including weather forecasting, downscaling, and climate projections. One segment of ClimateLearn's target user demographic are weather and climate scientists, who possess expertise in the physical laws and phenomena relevant for constructing robust modeling priors, but might lack familiarity with optimal approaches for implementing, training, and evaluating machine learning models. Another segment of ClimateLearn's target user demographic are ML researchers, who might encounter difficulties framing weather and climate modeling problems in a scientifically sound and practically useful manner, working with climate datasets—which is quite heterogeneous and often exists in bespoke file formats uncommon in mainstream ML research (e.g., NetCDF), or appropriately quantifying and visualizing their results for interpretation and deployments.
To showcase the capabilities of ClimateLearn and establish benchmarks, we perform and report results from numerous experiments on the supported tasks with a variety of traditional methods and our own tuned implementations of deep learning models on weather and climate datasets. In addition to traditional evaluation setups, we have created novel dataset and benchmarking scenarios to test model robustness and applicability to forecasting extreme weather events. Further, the library is modular and easily extendable to include additional tasks, datasets, models, metrics, and visualizations. We have also provided extensive documentation and contribution guides for improving the ease of community adoption and accelerating open-source expansion. While our library is already public, we are releasing all our data, code, and model checkpoints for the benchmarked evaluation in this paper to aid reproducibility and broader interdisciplinary research efforts.
Recent works have proposed benchmark datasets for weather and climate modeling problems. Prominently, Rasp et al. [63] proposed WeatherBench, a dataset for weather forecasting based on ERA5, followed by an extension called WeatherBench Probability [21], which adds support for probabilistic forecasting. Mouatadid et al. [48] extend similar benchmarks to the subseasonal to seasonal timescale. For precipitation events such as rain specifically, there are prior datasets such as RainBench [15] and IowaRain [73]. There exist datasets such as ExtremeWeather [60], FloodNet [61], EarthNet [66], DroughtED [45] and ClimateNet [56] for detection and localization of extreme weather events, and NADBenchmarks [58] for natural disasters related tasks. Cachay et al. [9] recently proposed ClimART, a benchmark dataset for emulating atmospheric radiative transfer in weather and climate models. For identifying long-term, globally-averaged trends in climate, Watson-Parris et al. [84] proposed ClimateBench, a dataset for climate model emulation.
Figure 1: Key components of ClimateLearn. We support observational, simulated, and reanalysis datasets from a variety of sources. The currently supported tasks are weather forecasting, downscaling, and climate projection. ClimateLearn also provides a suite of standard baselines and deep learning architectures, along with common metrics, visualizations, and logging support.
ClimetLab [11] is a library which aims to simplify the process of downloading, preprocessing, and making climate data ML-friendly. While ClimetLab focuses purely on providing easy access to climate data, ClimateLearn is an end-to-end library for training and evaluating machine learning models on weather and climate problems. Moreover, while ClimetLab is intended to be used in Jupyter notebooks, ClimateLearn encompasses a wider range of usecases, from running quick starter code in Jupyter to large-scale benchmarking experiments.
Beyond plain datasets, libraries such as Scikit-downscale [24], CCdownscaling [55], and CMIP6-Downscaling [10] provide tools for post-processing of climate model outputs via statistical, nondeep-learning downscaling, or mapping low-resolution gridded, image-like inputs to high-resolution gridded outputs. In a slightly different approach, pyESD focuses on downscaling from gridded climate data to specific weather stations [6]. Pyrocast [79] proposes an integrated ML pipeline to forecast Pyrocumulonimbus (PyroCb) Clouds. Many of these works supply only individual components of an ML pipeline but do not always have an API for loading climate data into a ML-ready format, or standard model implementations and evaluation protocols across multiple climate science tasks. As an end-to-end ML pipeline, ClimateLearn holistically bridges the gap for applying ML to challenging weather and climate modeling tasks like forecasting, downscaling, and climate projection.
3 Key Components of ClimateLearn
ClimateLearn is a PyTorch library that implements a range of functionalities for benchmarking of ML models for weather and climate. Broadly, our library is comprised of four components: tasks, datasets, models, and evaluations. See Figure 1 for an illustration. Sample code snippets for configuring each component is provided in Appendix E.
3.1 Tasks
Weather forecasting is the task of predicting the weather at a future time step  $t + \Delta t$  given the weather conditions at the current step  $t$  and optionally steps preceding  $t$ . A ML model receives an input of shape  $C \times H \times W$  and predicts an output of shape  $C' \times H \times W$ .  $C$  and  $C'$  denote the number of input and output channels, respectively, which contain variables such as geopotential, temperature, and humidity.  $H$  and  $W$  denote the spatial coverage and resolution of each channel, which depend
on the region studied and how densely we grid it. In our benchmarking, we focus on forecasting all gird points globally, but ClimateLearn can be easily extended to regional forecasting.
Downscaling Due to their high computational cost, existing climate models often use large grid cells, leading to low-resolution predictions. While useful for understanding large-scale climate trends, these do not provide sufficient detail to analyze local phenomena and design regional policies. The process of correcting biases in climate model outputs and mapping them to higher resolutions is known as downscaling. ML models for downscaling are trained to map an input of shape  $C \times H \times W$  to a higher resolution output  $C' \times H' \times W'$ , where  $H' > H$  and  $W' > W$ . As in forecasting, in downscaling,  $H \times W$  and  $H' \times W'$  can span either the entire globe or a specific region.
Climate projection aims to obtain long-term predictions of the climate under different forcings, e.g., greenhouse gas emissions. We provide support to download data from ClimateBench [84], a recent benchmark designed for testing ML models for climate projections. Here, the task is to predict the annual mean distributions of 4 climate variables: surface temperature, diurnal temperature range, precipitation, and the 90th percentile of precipitation, given four anthropogenic forcing factors: carbon dioxide  $(\mathrm{CO}_{2})$ , sulfur dioxide  $(\mathrm{SO}_{2})$ , black carbon (BC), and methane  $(\mathrm{CH}_{4})$ .
3.2 Datasets
ERA5 is a commonly-used data source for training and benchmarking data-driven forecasting and downscaling methods [4, 38, 49, 52, 54, 63, 62]. It is maintained by the European Center for Medium-Range Weather Forecasting (ECMWF) [29]. ERA5 is a reanalysis dataset that provides the best guess of the state of the atmosphere and land-surface variables at any point in time by combining multiple sources of observational data with the forecasts of the current state-of-the-art forecasting model known as the Integrated Forecasting System (IFS) [85]. In its raw format, ERA5 contains hourly data from 1979 to the current time on a  $0.25^{\circ}$  grid of the Earth's sphere, with different climate variables at 37 different pressure levels plus the Earth's surface. This corresponds to nearly 400,000 data points with a resolution of  $721 \times 1440$ . As this data is too big for most deep learning models, ClimateLearn also supports downloading a smaller version of ERA5 from WeatherBench [63], which uses a subset of ERA5 climate variables and regrids the raw data to lower resolutions.
Extreme-ERA5 is a subset of ERA5 that we have constructed to evaluate forecasting performance in extreme weather situations. Specifically, we consider "simple extreme" events [83, 5], i.e., weather events that have individual climate variables exceeding critical values locally. Heat waves and cold spells are examples of extreme events that can be quantitatively captured by extreme localized surface-level temperatures over prolonged days. To mimic real-world scenarios, we calculate thresholds for each pixel of the grid using the 5th and 95th percentile of the 7-day localized mean surface temperature over the training period (1979-2015). We then select a subset of pixels from all the available pixels in the testing set (2017-18) that had a 7-day localized mean surface temperature beyond these thresholds. We refer to Appendix B.2.2 for more details.
CMIP6 is a collection of simulated data from the Coupled Model Intercomparison Project Phase 6 (CMIP6) [19], an international effort across different climate modeling groups to compare and evaluate their global climate models. While the main goal of CMIP6 is to improve the understanding of Earth's climate systems, the data from their experimental runs is freely accessible online. CMIP6 data covers a wide range of climate variables, including temperature and precipitation, from hundreds of climate models, providing a rich source of data. For our forecasting experiments, we specifically use the outputs of CMIP6's MPI-ESM1.2-HR model, as it contains similar climate variables to those represented in ERA5 and was also considered in previous works for pretraining deep learning models [62]. MPI-ESM1.2-HR provides data from 1850 to 2015 with a temporal resolution of 6 hours and spatial resolution of  $1^{\circ}$ . Since this again corresponds to a grid that is too big for most deep learning models, we provide lower resolution versions of this dataset for training and evaluation. Besides, we also perform experiments with ClimateBench, which contains data on a range of future emissions scenarios based on simulations by the Norwegian Earth System Model [71], another member of CMIP6. We refer to Appendix B.2.3 for time ranges and more details of the experiments.
PRISM is a dataset of various observed atmospheric variables like precipitation and temperature over the conterminous United States at varying spatial and temporal resolutions from 1895 to present day. It is maintained by the PRISM Climate Group at Oregon State University [57]. At the highest publicly available resolution, PRISM contains daily data on a grid of  $4\mathrm{km}$  by  $4\mathrm{km}$  cells (approximately
$0.03^{\circ}$ ), which corresponds to a matrix of shape  $621 \times 1405$ . For the same reason we regrid ERA5 and CMIP6, we also provide a regridded version of raw PRISM data to  $0.75^{\circ}$  resolution.
3.3 Models
Traditional baselines ClimateLearn provides the following traditional baseline methods for forecasting: climatology, persistence, and linear regression. The climatology method uses historical average values of the predictands as the forecast. In ClimateLearn, we consider two versions of the climatology, one in which we compute the average value over the entire training set, and the other keeps a mean for each of the 52 calendar weeks to account for the seasonal cycle of the climate. The persistence method uses the last observed values of the predictands as the forecast. For downscaling, ClimateLearn provides nearest and bilinear interpolation. Nearest interpolation estimates the value of an unknown pixel to be the value of the nearest known pixel. Bilinear interpolation estimates the value at an unknown pixel by taking the weighted average of neighboring pixels.
Deep learning models The data for gridded weather and climate variables is represented as a 3D matrix, where latitude, longitude, and the variables form the height, width, and channels, respectively. Hence, convolutional neural networks (CNNs) are commonly used for forecasting and downscaling, which can be viewed as instances of the image-to-image translation problem [17, 30, 49, 62, 68, 72, 75, 80, 81, 86, 87]. ClimateLearn supports ResNet [27] and U-Net [69]—two prominent variants of the commonly used CNN architectures. Additionally, ClimateLearn supports Vision Transformer (ViT) [4, 20, 52], a class of models that represent images as a sequence of pixel patches. ClimateLearn also supports loading benchmark models from the literature such as Rasp and Thuerey [62] in a single line of code and is built so that custom models can be added easily.
3.4 Evaluations
Forecasting metrics For deterministic forecasting, ClimateLearn provides metrics such as root mean square error (RMSE) and anomaly correlation coefficient (ACC), which measures how well model forecasts match ground truth anomalies. For probabilistic forecasting, ClimateLearn provides spread-skill ratio and continuous ranked probability score, as defined by Garg et al. [21]. ClimateLearn also provides latitude-weighted version of these metrics, which lends extra weight to pixels near the equator. This is needed because the curvature of the Earth means that grid cells at low latitudes cover less area than grid cells at high latitudes. We refer to Appendix B.4 for additional details, including equations.
Downscaling metrics For downscaling, ClimateLearn uses RMSE, mean bias, and Pearson's correlation coefficient, in which mean bias is the difference between the spatial mean of ground-truth values and the spatial mean of predictions. We refer to Appendix B.4 for additional details.
Climate projection metrics In addition to the standard RMSE metric, we provide two metrics suggested by ClimateBench: Normalized spatial root mean square error  $(\mathrm{NRMSE}_s)$  and Normalized global root mean square error  $(\mathrm{NRMSE}_g)$ . We refer to Appendix B.4 for more details.
Visualization Besides these quantitative evaluation procedures, ClimateLearn also provides ways for users to inspect model performance qualitatively through visualizations of data and model predictions. For instance, in a single line of code, users can visually inspect their forecasting model's per-pixel mean bias, or the expected values of forecast errors, over the testing period. Such a visualization can be useful for pinpointing the regions on which the model's predictions consistently deviate from the ground truth in a certain direction. For probabilistic forecasts, ClimateLearn can generate the corresponding rank histogram, which indicates the reliability and sharpness of the model. Sample visualizations of deterministic and probabilistic predictions are provided in Appendix D.
4 Benchmark Evaluation via ClimateLearn
In this section, we evaluate the performance of different deep learning methods supported by ClimateLearn on weather forecasting and climate downscaling. We refer to Appendix C.1 for experiments on the climate projection task. We conduct extensive experiments and analyses with different settings to showcase the features and flexibility of our library.
Figure 2: Performance on forecasting three variables at different lead times. Solid lines are deep learning methods, dashed lines are simple baselines, and the dotted line is the physics-based model. Lower RMSE and higher ACC indicate better performance.
4.1 Weather forecasting
We first benchmark on weather forecasting. In addition, we compare different approaches for training forecast models in Section 4.1.1, and investigate the robustness of these models to extreme weather events and data distribution shift in Section 4.1.2 and 4.1.3, respectively.
Task We consider the task of forecasting the geopotential at  $500\mathrm{hPa}$  (Z500), temperature at  $850\mathrm{hPa}$  (T850), and temperature at 2 meters from the ground  $(\mathrm{T}2\mathrm{m})$  at five different lead times: 6 hours, and  $\{1,3,5,10\}$  days. Z500 and T850 are often used for benchmarking in previous works [4, 38, 52, 54, 62, 63], while the surface variable T2m is relevant to human activities.
Baselines We consider ResNet [27], U-Net [69], and ViT [16] which are three common deep learning architectures in computer vision. We provide the architectural details of these networks in Appendix B.1. We perform direct forecasting, where we train one neural network for each lead time. In addition, we compare the deep learning methods with climatology, persistence, and IFS [85].
Data We use ERA5 [29] at  $5.625^{\circ}$  for training and evaluation, which is equivalent to having a  $32 \times 64$  grid for each climate variable. The input variables to the deep learning models include geopotential, temperature, zonal and meridional wind, relative humidity, and specific humidity at 7 pressure levels (50, 250, 500, 600, 700, 850, 925) hPa, 2-meter temperature, 10-meter zonal and meridional wind, incoming solar radiation, and finally 3 constant fields: the land-sea mask, orography, and the latitude, which together constitute 49 input variables. For non-constant variables, we use data at 3 timesteps  $t$ ,  $t - 6\mathrm{h}$ , and  $t - 12\mathrm{h}$  to predict the weather at  $t + \Delta t$ , resulting in  $46 \times 3 + 3 = 141$  input channels. Each channel is standardized to have a mean of 0 and a standard deviation of 1. The training period is from 1979 to 2015, validation in 2016, and test in 2017 and 2018.
Training and evaluation We use latitude-weighted mean squared error as the loss function. We use AdamW optimizer [42] with a learning rate of  $5 \times 10^{-4}$  and weight decay of  $1 \times 10^{-5}$ , a linear warmup schedule for 5 epochs, followed by cosine-annealing for 45 epochs. We train for 50 epochs with 128 batch size, and use early stopping with a patience of 5 epochs. We use latitude-weighted root mean squared error (RMSE) and anomaly correlation coefficient (ACC) as the test metrics.
Benchmark results Figure 2 shows the performance of different baselines. As expected, the forecast quality in terms of both RMSE and ACC of all baselines worsens with increasing lead times. The deep learning methods significantly outperform climatology and persistence but underperform IFS. ResNet is the best-performing deep learning model on most tasks in both metrics. We hypothesize that while being more powerful than ResNet in general, U-Net tends to perform better when trained on high-resolution data [69], and ViT often suffers from overfitting when trained from scratch [28, 52]. Our reported performance of ResNet closely matches that of previous work [62].
Figure 3: Comparison of direct, continuous, and iterative forecasting with ResNet architecture.
4.1.1 Should we perform direct, continuous, or iterative forecasting?