notebooks/TFE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"assets/header_notebook.jpg\" />\n",
    "<hr style=\"color:#5A7D9F;\">\n",
    "<p align=\"center\">\n",
    "    <b style=\"font-size:2vw; color:#5A7D9F; font-weight:bold;\">\n",
    "    <center>Ocean subgrid parameterizations in an idealized model using machine learning</center>\n",
    "    </b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F;\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------\n",
    "#     Librairies\n",
    "# -----------------\n",
    "#\n",
    "# --------- Standard ---------\n",
    "import os\n",
    "import sys\n",
    "import json\n",
    "import glob\n",
    "import math\n",
    "import torch\n",
    "import random\n",
    "import fsspec\n",
    "import matplotlib\n",
    "import fourierflow\n",
    "import numpy                   as np\n",
    "import pandas                  as pd\n",
    "import xarray                  as xr\n",
    "import seaborn                 as sns\n",
    "import matplotlib.pyplot       as plt\n",
    "import matplotlib.gridspec     as gridspec\n",
    "\n",
    "from matplotlib.colorbar     import Colorbar\n",
    "from argparse                import ArgumentParser\n",
    "from scipy.stats             import gaussian_kde\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "# --------- PYQG ---------\n",
    "import pyqg\n",
    "import pyqg.diagnostic_tools\n",
    "from   pyqg.diagnostic_tools import calc_ispec         as _calc_ispec\n",
    "import pyqg_parameterization_benchmarks.coarsening_ops as coarsening\n",
    "\n",
    "calc_ispec = lambda *args, **kwargs: _calc_ispec(*args, averaging = False, truncate = False, **kwargs)\n",
    "\n",
    "# --------- PYQG Benchmark ---------\n",
    "from pyqg_parameterization_benchmarks.utils             import *\n",
    "from pyqg_parameterization_benchmarks.utils_TFE         import *\n",
    "from pyqg_parameterization_benchmarks.plots_TFE         import *\n",
    "from pyqg_parameterization_benchmarks.configurations    import *\n",
    "from pyqg_parameterization_benchmarks.online_metrics    import diagnostic_differences, diagnostic_similarities\n",
    "from pyqg_parameterization_benchmarks.neural_networks   import NN_Parameterization, NN_Parameterization_Handler\n",
    "from pyqg_parameterization_benchmarks.nn_analytical     import BackscatterBiharmonic, Smagorinsky, HybridSymbolic\n",
    "from pyqg_parameterization_benchmarks.nn_kaskade        import Kaskade\n",
    "from pyqg_parameterization_benchmarks.nn_fcnn           import FullyCNN\n",
    "from pyqg_parameterization_benchmarks.nn_fno            import FourierNO\n",
    "from pyqg_parameterization_benchmarks.nn_ffno           import FactorizedFNO\n",
    "from pyqg_parameterization_benchmarks.nn_unet           import UNet\n",
    "\n",
    "# --------- Jupyter ---------\n",
    "%matplotlib inline\n",
    "plt.rcParams.update({'font.size': 13})\n",
    "\n",
    "# Making sure modules are reloaded when modified\n",
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# Moving to correct folder\n",
    "%cd ../src/pyqg_parameterization_benchmarks/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<hr style=\"color:#5A7D9F;\">\n",
    "<p align=\"center\">\n",
    "    <b style=\"font-size:1.5vw; color:#5A7D9F;\">\n",
    "    <center>PYQG - Generating & Saving dataset<\\center>\n",
    "    </b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "<p align=\"justify\">\n",
    "    In this section, one will be able to generate:\n",
    "</p>\n",
    "\n",
    "- A **high resolution** (= HR) simulation from a quasi-geostrophic model (PYQG), \n",
    "        \n",
    "- A **low resolution** (= LR) simulation.\n",
    "        \n",
    "- An **augmented low resolution** (= ALR) simulations.\n",
    "        \n",
    "Furthermore, one will be able to:\n",
    "        \n",
    "- Observe the **state** and **subgrid variables** associated to the dataset in the corresponding folder.\n",
    "\n",
    "- **Save** the datasets on the hard drive. \n",
    "\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<p align=\"center\">\n",
    "\t<b style=\"font-size:1vw;\">\n",
    "\t<center>Simulation type</center>\n",
    "\t</b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<table style=\"width: 100%;\" border=\"1\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>INDEX</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">0</td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\">2</td>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\">3</td>\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\">4</td>\n",
    "\t\t\t<td style=\"width: 18%;\" align=\"center\">5</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\"align=\"center\"><b>TYPE</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Eddies</td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Jets  </td>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\">Eddies (Debug)</td>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\">Jets   (Debug)</td>\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\">Eddies (Random)</td>\n",
    "\t\t\t<td style=\"width: 18%;\" align=\"center\">Jets (Random)</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>\n",
    "\n",
    "<br>\n",
    "\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<p align=\"center\">\n",
    "\t<b style=\"font-size:1vw;\">\n",
    "\t<center>Parameters</center>\n",
    "\t</b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<table style=\"width: 100%;\" border=\"1\">\n",
    "\t<tbody>\n",
    "\t\t<tr style=\"height: 21px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>PARAMETERS</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">nx</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">dt</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">tmax</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">tavestart</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">rek</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">Δ</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">β</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr style=\"height: 21.5px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>DESCRIPTION</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">Number of real space grid points in the x directions</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">Numerical timestep (in hours)</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">Total time of integration (in years)</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">Start time for averaging (in years)</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">Linear drag in lower layer</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">Layer thickness ratio (H1/H2)</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">Gradient of coriolis parameter.</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr style=\"height: 21.5px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>EDDIES</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">256</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">10</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">5</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">5.789e-7</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">0.25</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">1.5 * 1e-11</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr style=\"height: 21px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>JETS</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">256</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">10</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">5</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">7e-08</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">0.1</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">1e-11</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr style=\"height: 21.5px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>EDDIES (Debug)</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">256</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">2</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">5.789e-7</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">0.25</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">1.5 * 1e-11</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr style=\"height: 21px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>JETS (Debug)</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">256</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">2</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">7e-08</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">0.1</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">1e-11</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr style=\"height: 21.5px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>EDDIES (Random)</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">256</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">10</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">5</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">[5.7, 5.9] * 1e-7</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">0.25</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">[1.45, 1.55] * 1e-11</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr style=\"height: 21px;\">\n",
    "\t\t\t<td style=\"width: 16%;\" align=\"center\"><b>JETS (Random)</b></td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">256</td>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\">1</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">10</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">5</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">[6.9, 7.1] * 1e-8</td>\n",
    "\t\t\t<td style=\"width: 12%;\" align=\"center\">0.1</td>\n",
    "\t\t\t<td style=\"width: 14%;\" align=\"center\">[0.95, 1.05] * 1e-11</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>\n",
    "\n",
    "<br>\n",
    "\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<p align=\"center\">\n",
    "\t<b style=\"font-size:1vw;\">\n",
    "\t<center>Coarsening operators</center>\n",
    "\t</b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<table style=\"width: 100%;\" border=\"1\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 10%;\" align=\"center\"><b>OPERATOR</b></td>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\">&nbsp;1</td>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\">&nbsp;2</td>\n",
    "\t\t\t<td style=\"width: 17%;\" align=\"center\">3</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 10%;\"align=\"center\"><b>DESCRIPTION</b></td>\n",
    "\t\t\t<td style=\"width: 15%;\"align=\"center\">Spectral Truncation, Sharp Filter</td>\n",
    "\t\t\t<td style=\"width: 15%;\"align=\"center\">Spectral Truncation, Gaussian Filter</td>\n",
    "\t\t\t<td style=\"width: 17%;\"align=\"center\">GCM Filter, Averaging and Coarsening</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<p align=\"center\">\n",
    "\t<b style=\"font-size:1vw;\">\n",
    "\t<center>Documentation</center>\n",
    "\t</b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<table style=\"width: 100%;\" border=\"1\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>save_folder</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Name of the folder used to save the datasets</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\"align=\"center\"><b>nb_threads</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Number of threads used to run the simulation</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>simulation_type</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Type of simulation used to generate the dataset</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\"align=\"center\"><b>memory</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Total number of memory allocated [GB] (used for security purpose)</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>skipped_time</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Time [year] at which the sampling of the simulation starts</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\"align=\"center\"><b>save_high_res</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Choose if the whole high resolution is saved or just the last sample (memory saving)</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>operator_cf</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Coarsening and filtering operator applied on the high resolution simulation</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\"align=\"center\"><b>target_sample_size</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Number of samples expected to be in the datasets (nb_sample >= target_sample_size)</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------- GENERATING DATASET -------------------------------\n",
    "%run generate_dataset.py --save_folder            NOTEBOOK_DATA                                                                                                                                   \\\n",
    "                         --simulation_type                    1                                                                                                                                   \\\n",
    "                         --target_sample_size                 5                                                                                                                                   \\\n",
    "                         --operator_cf                        1                                                                                                                                   \\\n",
    "                         --skipped_time                       1                                                                                                                                   \\\n",
    "                         --nb_threads                         2                                                                                                                                   \\\n",
    "                         --memory                            32                                                                                                                                   \\\n",
    "                         --save_high_res                   True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[comment]: <> (Section)\n",
    "<hr style=\"color:#5A7D9F;\">\n",
    "<p align=\"center\">\n",
    "    <b style=\"font-size:1.5vw; color:#5A7D9F;\">\n",
    "    <center>Ocean subgrid parameterization - Learning</center>\n",
    "    </b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "<p align=\"justify\">\n",
    "    In this section, one will be able to:\n",
    "</p>\n",
    " \n",
    "- **Create** and **train** a new parameterization using the datasets created previously.\n",
    "\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<p align=\"center\">\n",
    "\t<b style=\"font-size:1vw;\">\n",
    "\t<center>Documentation</center>\n",
    "\t</b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<table style=\"width: 100%;\" border=\"1\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>folder_training</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Folder used to load data as training data</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>folder_validation</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Folder used to load data as validation data</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>save_directory</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Folder used to load data as training data</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>inputs</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Type of inputs given to the parameterization for training</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>targets</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Parameterization ouptut</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>learning_rate </b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">The learning rate value used as starting point during training</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>batch_size</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">The number of samples flowing forward before performing the back-propagation</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>optimizer</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">The optimizer used during training</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>scheduler</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">The scheduler used during training</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>configuration</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Choose the default configuration or one of many in configurations.py</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>num_epochs</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Number of epochs made by the paremeterization while training</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>zero_mean</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Type of pre-processing made on the datasets</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>padding</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Type of padding used by the parameterization</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>memory</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Total number of memory allocated [GB] (used for security purpose)</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>param_type</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Choose the type of parameterization used to learn closure</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>sim_type</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Type of fluid simulation studied (used to order tensorboard folders)</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------- EDDIES -------------------------------------\n",
    "%run train_parameterization.py --folder_training        NOTEBOOK_EDDIES_TRAINING                                                                                                         \\\n",
    "                               --folder_validation    NOTEBOOK_EDDIES_VALIDATION                                                                                                         \\\n",
    "                               --param_name                                 FFNO                                                                                                         \\\n",
    "                               --inputs                                    q u v                                                                                                         \\\n",
    "                               --targets                                q_fluxes                                                                                                         \\\n",
    "                               --learning_rate                             0.001                                                                                                         \\\n",
    "                               --batch_size                                   32                                                                                                         \\\n",
    "                               --optimizer                                  adam                                                                                                         \\\n",
    "                               --scheduler                              constant                                                                                                         \\\n",
    "                               --configuration                           default                                                                                                         \\\n",
    "                               --num_epochs                                   20                                                                                                         \\\n",
    "                               --zero_mean                                 False                                                                                                         \\\n",
    "                               --padding                                circular                                                                                                         \\\n",
    "                               --memory                                       32                                                                                                         \\\n",
    "                               --param_type                      NOTEBOOK_EDDIES                                                                                                         \\\n",
    "                               --sim_type                        NOTEBOOK_EDDIES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---------------------------------- JETS --------------------------------------\n",
    "%run train_parameterization.py --folder_training          NOTEBOOK_JETS_TRAINING                                                                                                         \\\n",
    "                               --folder_validation      NOTEBOOK_JETS_VALIDATION                                                                                                         \\\n",
    "                               --param_name                                 FFNO                                                                                                         \\\n",
    "                               --inputs                                    q u v                                                                                                         \\\n",
    "                               --targets                                q_fluxes                                                                                                         \\\n",
    "                               --learning_rate                             0.001                                                                                                         \\\n",
    "                               --batch_size                                   32                                                                                                         \\\n",
    "                               --optimizer                                  adam                                                                                                         \\\n",
    "                               --scheduler                              constant                                                                                                         \\\n",
    "                               --configuration                           default                                                                                                         \\\n",
    "                               --num_epochs                                   20                                                                                                         \\\n",
    "                               --zero_mean                                 False                                                                                                         \\\n",
    "                               --padding                                circular                                                                                                         \\\n",
    "                               --memory                                       32                                                                                                         \\\n",
    "                               --param_type                        NOTEBOOK_JETS                                                                                                         \\\n",
    "                               --sim_type                          NOTEBOOK_JETS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---------------------------------- FULL --------------------------------------\n",
    "%run train_parameterization.py --folder_training          NOTEBOOK_FULL_TRAINING                                                                                                         \\\n",
    "                               --folder_validation               FULL_VALIDATION                                                                                                         \\\n",
    "                               --param_name                                 FFNO                                                                                                         \\\n",
    "                               --inputs                                    q u v                                                                                                         \\\n",
    "                               --targets                                q_fluxes                                                                                                         \\\n",
    "                               --learning_rate                             0.001                                                                                                         \\\n",
    "                               --batch_size                                   32                                                                                                         \\\n",
    "                               --optimizer                                  adam                                                                                                         \\\n",
    "                               --scheduler                              constant                                                                                                        \\\n",
    "                               --configuration                           default                                                                                                         \\\n",
    "                               --num_epochs                                   20                                                                                                         \\\n",
    "                               --zero_mean                                 False                                                                                                         \\\n",
    "                               --padding                                circular                                                                                                         \\\n",
    "                               --memory                                       32                                                                                                         \\\n",
    "                               --param_type                        NOTEBOOK_FULL                                                                                                         \\\n",
    "                               --sim_type                          NOTEBOOK_FULL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[comment]: <> (Section)\n",
    "<hr style=\"color:#5A7D9F;\">\n",
    "<p align=\"center\">\n",
    "    <b style=\"font-size:1.5vw; color:#5A7D9F;\">\n",
    "    <center>Ocean subgrid parameterization - Testing (Offline)</center>\n",
    "    </b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "<p align=\"justify\">\n",
    "    In this section, one will be able to:\n",
    "</p>\n",
    " \n",
    "- **Load** a trained FCNN parameterization;\n",
    "\n",
    "- Evaluate its **offline performances** on a test set, i.e. it's ability to **predict** accurately the **subgrid forcing terms**;\n",
    "\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<p align=\"center\">\n",
    "\t<b style=\"font-size:1vw;\">\n",
    "\t<center>Documentation</center>\n",
    "\t</b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<table style=\"width: 100%;\" border=\"1\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>folder_offline</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Folders used to load data as offline test data</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>folder_models</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Folder (inside the model folder) used to load all the different models to be tested</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>memory</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Total number of memory allocated [GB] (used for security purpose)</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>type_sim</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Type of fluid simulation studied (used to order tensorboard folders)</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ad0d46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----------------------------- EDDIES --------------------------------\n",
    "%run offline.py --folder_offline              NOTEBOOK_EDDIES_OFFLINE \\\n",
    "                --folder_models                                  hope \\\n",
    "                --memory                                           32 \\\n",
    "                --type_sim                                     EDDIES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "070aeaad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------ JETS ---------------------------------\n",
    "%run offline.py --folder_offline                         JETS_OFFLINE \\\n",
    "                --folder_models      PHASE_6_SENTITIVITY_ARCHITECTURE \\\n",
    "                --memory                                           32 \\\n",
    "                --type_sim                                       JETS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------ FULL ---------------------------------\n",
    "%run offline.py --folder_offline                NOTEBOOK_FULL_OFFLINE \\\n",
    "                --folder_models                        __P5__ONLINE__ \\\n",
    "                --memory                                           32 \\\n",
    "                --type_sim                                       FULL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[comment]: <> (Section)\n",
    "<hr style=\"color:#5A7D9F;\">\n",
    "<p align=\"center\">\n",
    "    <b style=\"font-size:1.5vw; color:#5A7D9F;\">\n",
    "    <center>Ocean subgrid parameterization - Testing (Online)</center>\n",
    "    </b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "<p align=\"justify\">\n",
    "    In this section, one will be able to:\n",
    "</p>\n",
    " \n",
    "- **Load** a trained FCNN parameterization;\n",
    "\n",
    "- Evaluate its **online performances** on a test set, i.e. it's ability to **perform** a **meaningfull** and **accurate** simulation.\n",
    "\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<p align=\"center\">\n",
    "\t<b style=\"font-size:1vw;\">\n",
    "\t<center>Documentation</center>\n",
    "\t</b>\n",
    "</p>\n",
    "<hr style=\"color:#5A7D9F; width: 100%;\" align=\"left\">\n",
    "<table style=\"width: 100%;\" border=\"1\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>folder_online</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Folders used to load data as offline test data</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>folder_models</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Folder (inside the model folder) used to load all the different models to be tested</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>memory</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Total number of memory allocated [GB] (used for security purpose)</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"width: 15%;\" align=\"center\"><b>type_sim</b></td>\n",
    "\t\t\t<td style=\"width: 13%;\" align=\"center\">Type of fluid simulation studied (used to order tensorboard folders)</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----------------------------- EDDIES --------------------------------\n",
    "%run online.py --folder_online                          EDDIES_ONLINE \\\n",
    "               --folder_models                         __P5__ONLINE__ \\\n",
    "               --memory                                            32 \\\n",
    "               --type_sim                                      EDDIES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----------------------------- JETS ----------------------------------\n",
    "%run online.py --folder_online                            JETS_ONLINE \\\n",
    "               --folder_models                        __P4__FFNO_FULL \\\n",
    "               --memory                                            32 \\\n",
    "               --type_sim                                        JETS"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TFE",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}