Precipitation Prediction

Predicting precipitation over the continental United States by using convolutional recurrent networks (ConvLSTM/ConvGRU) for nowcasting
Author

Connor Fair

Published

May 16, 2026

Introduction

Humans have always had a need for predicting the weather to know when to plant food and to stay safe from severe weather events. Predictions have been made historically through observations and measurements of the atmosphere. This has led to an abundance of weather data being collected, especially since the 1960s, when radars and satellites could capture data over large areas.

In this project, I trained a spatiotemporal model on 2025 ERA5 data to predict precipitation over the continental US up to 3 hours ahead with 0.861 correlation at +1 hour and an overall correlation of 0.753 over the next 3 hours.

The code for this project

Dataset

Climate Data Store

The dataset used for this project was the ERA5 hourly data single level. The dataset contains information of weather conditions from 1940 to a few days behind real-time over the entire globe, with a resolution of 0.25° x 0.25° for atmosphere variables. To make the dataset more manageable, only a subset of the dataset was used:

  • Area:
    • Continental United States
  • Time:
    • Jan. 1 2025 to Dec. 31 2025
  • Variables:
    • Eastward (u) and Northward (v) wind components 10m off the ground (m/s)
    • Temperature 2 meters off the ground (K)
    • Surface pressure (Pa)
    • Total cloud coverage (ranged from 0 to 1)
    • Total precipitation (m)

Samples from the dataset:

The total size of the subset of the dataset was around 2.6 GB and was converted from the .grib file format to the .zarr format and chunked into 24 hour chunks to improve the performance of reading the dataset. All of the models trained read the previous 3 hours of all 6 of these variables to predict the next 3 hours of precipitation, so the dataset was split into 6-hour intervals.

The 6-hour intervals were split using a random stratified split across 3 groups so that all 3 groups had a mixture of samples from all 4 seasons in the data. 80% went into the training set which was used to train models, 10% went into the validation set which was used to detect over-fitting and 10% went into the test set which was used as the final evaluation to test for how well the models would generalize.

Note

In retrospect, a temporal split would have made more sense than the random split used here due to the autocorrelation between the time lagged steps, leading to optimistic results.

Preprocessing

Most of the variables in this dataset were very large or sparse and needed to be cleaned up before they are used for training. all scaling statistics were calculated from only the training set to prevent data leakage.

  • Wind (u & v components): Scaled by \(max|value|\)
  • Temperature: Transformed using a z-score standardization because the distribution was close to a bell-shape and symmetric, but the values were still much larger than most of the other variables
  • Total Cloud Coverage: None, the values were already between [0, 1]
  • Surface Pressure: Normalized to [0, 1] because the values here close to 100,000, which is significantly higher than the other variables, which would have hurt the stability of the training process.
  • Total Precipitation: This variable was very sparse, so it was scaled up to mm, then transformed with log1p to reduce the right-skew.

Before

Before

After

After

Evaluation

To measure the performance of each model, 5 metrics were used:

Note

All of the metrics were calculated after removing the log-scale preprocessing step.

The first 2 metrics are based on regression.

  • \(P\): Model Prediction
  • \(G\): Ground truth
  • \(i,j\): Grid Cells
  • Root Mean Squared Error (RMSE): The difference between the predicted rain vs the actual rain averaged.
    • \(\sqrt{\frac{\sum_{i,j}(G_{i,j}-P_{i,j})^2}{|G|}}\)
  • Correlation: How strongly correlated the model predictions are with the ground truth.
    • \(\frac{\sum_{i,j}P_{i,j}G_{i,j}}{\sqrt{(\sum_{i,j}P_{i,j}^2) (\sum_{i,j}G_{i,j}^2)} + \epsilon}\) where \(\epsilon = 10^{-9}\)

The other 3 metrics are from meteorology, where a threshold of ≥ 0.5mm/h of rain was applied to the model predictions and ground truth to see how well the model is able to predict the rain above the threshold. This creates a 2x2 contingency matrix out of the model predictions.

2x2 Contingency Matrix
Event Observed: Yes Event Observed: No
Event Forecast Yes Hits False Alarms
Event Forecast No Misses Correct Rejections

\(hits\): number of times the model predicted rainfall correctly

\(false\ alarms\): number of times the model predicts rain but it doesn’t rain

\(misses\): number of times the model doesn’t predict rain, but it does rain

  • Critical Success Index (CSI): The fraction of model predictions that correctly forecast precipitation (1 is the best value)
    • \(\frac{hits}{hits+misses+false\ alarms}\)
  • False Alarm Ratio (FAR): The fraction of model predictions that are false alarms (0 is the best value)
    • \(\frac{false\ alarms}{false\ alarms+hits}\)
  • Probability Of Detection (POD): The fraction of observed events that were forecast (1 is the best value)
    • \(\frac{hits}{hits+misses}\)

Baseline

To ensure my model results are meaningful, a baseline model was created that just assumes the precipitation over the next 3 hours is the same as it was over the most recent hour. This model was used to ensure other models were actually learning meaningful patterns in the data.

Baseline Results
RMSE Correlation CSI FAR POD
0.394 0.621 0.411 0.411 0.575

Initial Design

The initial model design was the ConvLSTM from this paper. The model architecture uses a standard LSTM, but replaces the fully connected layers with convolutions and stores its hidden states as 3D tensors, where the last 2 dimensions are spatial. This helps the model keep track of both the spatial and temporal dependencies needed to make accurate predictions.

A diagram of the ConvLSTM Cell over a timestep

A diagram of the ConvLSTM Cell over a timestep

Because the model is taking all 6 variables for only the first 3 hours, and is predicting only the precipitation on the last 3 hours without looking at the other variables, an encoder-decoder structure was used, where the encoder ran through the first 3 hours and the decoder predicted the last 3 hours.

A diagram of the overall model architecture

A diagram of the overall model architecture

The parameters used to train the network:

  • Hidden State / Cell State Channels: 20
  • Kernel Size: 5x5
  • Optimizer: AdamW
  • Learning Rate: 1e-4
  • Epochs: 30
  • Loss Function: Mean Squared Error
Initial Results
RMSE Correlation CSI FAR POD
0.303 0.743 0.491 0.307 0.627

Improvements

The network was trained on a MacBook Air with the M1 chip and the graph that PyTorch uses to calculate back-propagation was so large that it was almost maxing out the swap memory on the laptop, due to the recurrent nature of the model and the large size of the hidden states leading to very large computation graphs. To reduce the strain on the laptop’s swap memory while keeping the performance as similar as possible to the ConvLSTM, a new architecture, ConvGRU, was tested.

ConvGRU

The GRU (Gated Recurrent Unit) is a special version of an RNN, which uses gates in the same way as the LSTM to prevent exploding or vanishing gradients, however it drops the cell state that the LSTM used for long-term dependencies. This reduces the number of parameters significantly, but loses some performance in long sequences. This use case only had short input and output sequence lengths so the model was expected to work well here. The GRU was converted into a ConvGRU by replacing the fully connected layers with convolutions, just like what was done for the ConvLSTM.

A diagram of the ConvGRU Cell over a timestep

A diagram of the ConvGRU Cell over a timestep

The basic idea for the rest of the model’s architecture stayed the same, but was simplified due to the removal of the cell states. This change reduced the number of parameters by around 25%.

A diagram of the overall model architecture

A diagram of the overall model architecture

When training this model, all other parameters stayed the same. The change in architecture led to a improvement in CSI, FAR and POD, but changes in the other metrics are likely noise.

ConvGRU Results
RMSE Correlation CSI FAR POD
0.305 0.743 0.500 0.303 0.639

Sharpening Predictions with a Spatial Loss

One of the main issues with the model’s predictions was the fuzziness in the predictions made by the model. To address this issue, a new loss function was added to the existing mean squared error that was used to train the model. The new loss function is called the spatial loss which uses the mean squared error between the image gradients of the predictions and ground truths. The image gradient is used in image processing to find the edges of objects in an image, however here we are using it to penalize the model for having softer edges than the ground truth.

An example of applying the image gradient to an image of a square

An example of applying the image gradient to an image of a square

To combine the 2 loss functions, there were 3 approaches considered. The first was weighting each component separately, but this would require finding the optimal weights for both, which would have been time-consuming. The second approach considered was using gradient surgery which removes the conflicting component of the gradients, however it had an expensive overhead swap memory usage. The approach that was settled on was to normalize the values of the error functions while preserving the directions of the gradients. This reduces the competition between the error functions (as shown by the diagram below) improving the stability of the training process.

An example of how normalization can help the gradients from 2 error functions balance out

The final results from combining this new loss function with the ConvGRU architecture.

Best Model Results
RMSE Correlation CSI FAR POD
0.299 0.753 0.507 0.299 0.647

Final Results

All three models beat the persistence baseline on every metric, where the best cut RMSE by about 24%.

Overall results
Model RMSE Correlation CSI FAR POD
Baseline 0.394 0.621 0.411 0.411 0.575
ConvLSTM 0.303 0.743 0.491 0.307 0.627
ConvGRU 0.305 0.743 0.500 0.303 0.639
ConvGRU + Spatial 0.299 0.753 0.507 0.299 0.647

Conclusion

Overall, this was a fun and challenging project that taught me a lot about the challenges in weather forecasting, recurrent neural networks and working with the limitations of my hardware. There are several issues with this project that I hope to address in future work. The first of these is the fuzziness in my model’s outputs. even though adding another loss function that targeted the edges of predicted rainfall led to overall improvements to model performance, there is still a lot that could be improved here. The best solution I want to explore next is the GAN (Generative Adversarial Network) which penalizes the model’s fuzzy outputs using a discriminator network to address this problem.

Another area of future work is the network’s performance over time. As can be seen in the table below, the best model’s performance drops rapidly over the course of the 3 hours. This may be due to the small number of layers in my model, but I would also like to explore more modern approaches to this problem, such as introducing attention or swapping over to a transformer-based architecture to help the model with recognizing time-dependent relationships better.

Performance of best model over time
Metric +1 hour +2 hours +3 hours
RMSE 0.232 0.305 0.347
Correlation 0.861 0.741 0.641
CSI 0.637 0.502 0.397
FAR 0.200 0.306 0.398
POD 0.757 0.645 0.537

References

Better weather warnings can save lives and billions of dollars globally: Published by Journal of Health Economics and Outcomes Research. Journal of Health Economics and Outcomes Research. (n.d.). https://jheor.org/post/2024-better-weather-warnings-can-save-lives-and-billions-of-dollars-globally

Ecmwf. (2026, March 16). ERA5 hourly data on single levels from 1940 to present. https://cds.climate.copernicus.eu/datasets/reanalysis-era5-single-levels?tab=overview

Shi, X., Chen, Z., Wang, H., Yeung, D.-Y., Wong, W., & Woo, W. (2015, September 19). Convolutional LSTM network: A machine learning approach for precipitation nowcasting. arXiv.org. https://arxiv.org/abs/1506.04214

Understanding LSTM networks. Understanding LSTM Networks – colah’s blog. (n.d.). https://colah.github.io/posts/2015-08-Understanding-LSTMs/

Weather forecasting | methods, importance, & history | britannica. (n.d.-c). https://www.britannica.com/science/weather-forecasting

Forecast verification glossary a accuracy. (n.d.-b). https://www.swpc.noaa.gov/sites/default/files/images/u30/Forecast%20Verification%20Glossary.pdf

Strategies for balancing multiple loss functions in deep learning | by Baicen Xiao | Medium. (n.d.-d). https://medium.com/@baicenxiao/strategies-for-balancing-multiple-loss-functions-in-deep-learning-e1a641e0bcc0