Code for "Tropical Basin Interactions Reduce Spring Predictability Barrier of ENSO in a Deep Learning Model"

This is the public data and code for the manuscript “Tropical Basin Interactions Reduce Spring Predictability Barrier of ENSO in a Deep Learning Model” by Zhou and Zhang.

README

Due to data size limitations, both training and model output data are stored on Google Drive. You can access them through the following link: https://drive.google.com/drive/folders/1DvASC6Jdtlw-fWl6lkwdrO0Au-S22oOX?usp=sharing. The "train_validation_test_dataset" folder contains the model's training and testing data; the "model_predictions" folder contains prediction datasets made by the GL-Geoformer with monthly predictions initiated from Oct 1980 to Jan 2025 (each prediction predicts 1-24 month ahead).

-> Code Description:
The "pretrain" folder contains code for model pretraining. You can modify the training data path in "config.py", then run "trainer.py" to start training the model. The trained model parameters are saved in the "pretrain/model" folder and will be used for subsequent transfer learning.

The "transfer" folder contains code for transfer learning optimization of the pretrained model. Before performing transfer learning, you need to set "configs.TFtrain = True" in "config.py", modify the data path, and set "configs.output_length" to an appropriate value based on your GPU memory. Next, you need to modify the pretrained model parameter path in the "Trainer class" in "transfer_trainer.py". After completing these modifications, you can run "transfer_trainer.py" to perform model transfer learning.

After completing transfer learning and before starting model testing, you need to set "configs.TFtrain = False" in "config.py" and set "configs.output_length" to an appropriate value (e.g., configs.output_length=24). Then modify the model path and test data path in "test_model.py", and you can run this file to execute model testing.

Currently, 10 sets of model parameters have already been saved in the "transfer\model" folder. You can also directly use these parameters for predictions or sensitivity experiments.