4d-srda_sr-mixup

This repository contains the source code used in Spatio-Temporal Super-Resolution Data Assimilation (SRDA) Utilizing Deep Neural Networks with Domain Generalization Technique Toward Four-Dimensional SRDA (arXiv, james).

• Setup
  • Docker Containers
  • Singularity Containers
• How to Perform Experiments
  • CFD Simulations
  • Data Preparation
  • Training and Tuning
  • Evaluation
• Cite

Setup

• Basically, the Singularity containers were used for experiments.
• At least, 1 GPU board is required.
• Notes
  • The Docker containers have the same environments as in the Singularity containers.
  • tsubame means the super-computer at the Tokyo Institute of Technology (webpage).
    • Singularity containers can be used on TSUBAME.

Docker Containers

1. Install Docker.
2. Build docker containers: $ docker compose build
3. Start docker containers: $ docker compose up -d

Singularity Containers

1. Install Singularity.
2. Build Singularity containers:
   • $ singularity build -f pytorch_local.sif ./singularity/pytorch/pytorch.def
   • $ singularity build -f pytorch_tsubame.sif ./singularity/pytorch_tsubame/pytorch_tsubame.def
   • pytorch_tsubame.sif is for the super-computer, TSUBAME, at Tokyo Institute of Technology (webpage)
3. Start singularity containers:
   • The following command is for local environments

$ singularity exec --nv --env PYTHONPATH="$(pwd)/pytorch" \
    pytorch_local.sif jupyter lab \
    --no-browser --ip=0.0.0.0 --allow-root --LabApp.token='' --port=8888

How to Perform Experiments

• The Singularity container (on a local environment), pytorch_local.sif, is used in the following experiments.
  • We confirmed the following code works on an NVIDIA A100 40GB PCIe.
• Note
  • On TSUBAME, the same code was run using pytorch_tsubame.sif.
  • In deep learning, distributed data parallel (DDP) was used (train_ddp_ml_model.py) on TSUBAME, where four GPUs of TESLA P100 were used.

CFD Simulations

• In each script, a number of simulations must be specified.
• Run the following scripts:
  • ./pytorch/script/shell/simulate_cfd_jet_for_making_analysis_train_data.sh
  • ./pytorch/script/shell/simulate_cfd_jet_for_making_forecast_train_data.sh
  • ./pytorch/script/shell/simulate_cfd_jet_uhr_jet.sh

Data Preparation

• Run the following notebook on JupyterLab: split_npy.ipynb

Training and Tuning

• In each script, a configuration must be specified.
• ST-SRDA
  • ./pytorch/script/shell/train_ddp_ml_model.sh (for multiple GPUs)
  • ./pytorch/script/shell/train_ml_model.sh (for a single GPU)
• EnKF-SR
  • ./pytorch/script/shell/tune_enkf_sr.sh
• EnKF-HR
  • generate_additive_inflation_cov.ipynb
  • ./pytorch/script/shell/tune_enkf_hr.sh

Evaluation

• Run the notebooks in the following order.
  • ./pytorch/notebook/test_srda_using_uhr.ipynb
  • ./pytorch/notebook/test_enkf_sr_using_uhr.ipynb
  • ./pytorch/notebook/test_enkf_hr_using_uhr.ipynb

Cite

@article{
  author = {Yasuda, Yuki and Onishi, Ryo},
  title = {Spatio-Temporal Super-Resolution Data Assimilation (SRDA) Utilizing Deep Neural Networks With Domain Generalization},
  journal = {Journal of Advances in Modeling Earth Systems},
  volume = {15},
  number = {11},
  pages = {e2023MS003658},
  doi = {https://doi.org/10.1029/2023MS003658},
  url = {https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/2023MS003658},
  eprint = {https://agupubs.onlinelibrary.wiley.com/doi/pdf/10.1029/2023MS003658},
  year = {2023}
}