Scientific Reinforcement Learning for sub-grid scale modeling of 2D turbulence

[project website]

Table of contents

• Introduction
• Requirements
• Experiments
  • Case 1
• Citation
• References

Introduction

The checklist for refactor progress

• Check consistency with Py2D
  • RHS
  • Forcing: FK
  • CN for $\Pi$ formulation psiTemp = RHS/(1+dt*alpha+0.5*dt*(nu+ve)*Ksq)
  • Check operators and grid [ij] vs [xy] (all in [ij])
    • meshgrid Kx, Ky
    • meshgrid xx, yy
    • Checked $u,v$ calculations, and added a function: psi_2_uv
  • Change $\beta*v$ term to match the Py2D convention: changed with 'ij' indexing grid
• Option for calculation of $\Pi$
  • From $\sigma$
  • From $\tau$
• Optional CPU-GPU backend
• Consistent model action ($c_l^3$ and $c_s^2$)
• Options to save a list of parameters ($\omega$, $\psi$, $\nu_e$, $c_{model}$, $\Pi$, action list)
• Re-organize the state model, maybe have it as a list of options (accumalative):
  • Global: energy spectra
  • Global: enstrophy spectra
  • local: $\nabla u$
  • local: $\nabla \nabla u$
  • choice of invariants $\lambda_i$
• Update initial condition for cases (and the corresponding spectra)
  • Case 1: $\kappa_f=4$ , Re$=20\times10^3$, $\beta=0$
  • Case 2: $\kappa_f=4$ , Re$=20\times10^3$, $\beta=0$
  • Case 3: $\kappa_f=25$ , Re$=20\times10^3$, $\beta=20$
• Check consistency of IC mat files with the solver
• Case management system: Copy config file in the folder
• bring options to sh file:
  • ["Policy"]["Distribution"] choice
• Retrain with a [super-]gaussian spectra (to have a better behaving interscale transfers )
• Code to track interscale variation while training
• Double check the naming of time steps: n_init, ...

Requirements

• python 3.6

• scipy

• numpy

• Korali

• For GPU support:

  • jax 0.4.5
  • Korali + JAX docker
    • jax=0.4.5 - jaxlib=0.4.4+cuda11.cudnn86

Experiments

To run the training

bash run.sh

To run post process

bash runpost.sh

To delete all data files in the folder

bash clean.sh

Citation

• Mojgani, R., Waelchli, D., Guan, Y., Koumoutsakos, P., Hassanzadeh, P. "Extreme Event Prediction with Multi-agent Reinforcement Learning-based Parametrization of Atmospheric and Oceanic Turbulence", arXiv: 2312.00907, 2023.(url)
  BibTeX

  @article{Mojgani_arxiv_2023,
  title={Extreme Event Prediction with Multi-agent Reinforcement Learning-based Parametrization of Atmospheric and Oceanic Turbulence},
  author={Rambod Mojgani and Daniel Waelchli and Yifei Guan and Petros Koumoutsakos and Pedram Hassanzadeh},
  year={2023},
  eprint={2312.00907},
  archivePrefix={arXiv},
  }