notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variational data assimilation with deep prior (CIRC23)\n",
    "\n",
    "{bdg-primary}`Ocean`\n",
    "{bdg-secondary}`Modelling`\n",
    "{bdg-warning}`Special Issue`\n",
    "{bdg-info}`Python`\n",
    "\n",
    "<p align=\"left\">\n",
    "    <a href=\"https://github.com/eds-book-gallery/39d9c177-11da-41b2-9b64-63f4c1c834b3/blob/main/LICENSE\">\n",
    "        <img alt=\"license\" src=\"https://img.shields.io/badge/license-MIT-yellow.svg\">\n",
    "    </a>\n",
    "    <a href=\"https://github.com/eds-book-gallery/39d9c177-11da-41b2-9b64-63f4c1c834b3/actions/workflows/render.yaml\">\n",
    "        <img alt=\"render\" src=\"https://github.com/eds-book-gallery/39d9c177-11da-41b2-9b64-63f4c1c834b3/actions/workflows/render.yaml/badge.svg\">\n",
    "    </a>\n",
    "    <a href=\"https://github.com/alan-turing-institute/environmental-ds-book/issues/178\">\n",
    "        <img alt=\"review\" src=\"https://img.shields.io/badge/view-review-purple\">\n",
    "    </a>\n",
    "    <br/>\n",
    "</p>\n",
    "\n",
    "<p align=\"left\">\n",
    "    <a href=\"https://mybinder.org/v2/gh/eds-book-gallery/39d9c177-11da-41b2-9b64-63f4c1c834b3/main?labpath=notebook.ipynb\">\n",
    "        <img alt=\"binder\" src=\"https://mybinder.org/badge_logo.svg\">\n",
    "    </a>\n",
    "    <a href=\"https://replay.notebooks.egi.eu/v2/gh/eds-book-gallery/39d9c177-11da-41b2-9b64-63f4c1c834b3/main?labpath=notebook.ipynb\">\n",
    "        <img alt=\"binder\" src=\"https://img.shields.io/badge/launch-EGI%20binder-F5A252.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAFkAAABZCAMAAABi1XidAAAB8lBMVEX///9XmsrmZYH1olJXmsr1olJXmsrmZYH1olJXmsr1olJXmsrmZYH1olL1olJXmsr1olJXmsrmZYH1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olJXmsrmZYH1olL1olL0nFf1olJXmsrmZYH1olJXmsq8dZb1olJXmsrmZYH1olJXmspXmspXmsr1olL1olJXmsrmZYH1olJXmsr1olL1olJXmsrmZYH1olL1olLeaIVXmsrmZYH1olL1olL1olJXmsrmZYH1olLna31Xmsr1olJXmsr1olJXmsrmZYH1olLqoVr1olJXmsr1olJXmsrmZYH1olL1olKkfaPobXvviGabgadXmsqThKuofKHmZ4Dobnr1olJXmsr1olJXmspXmsr1olJXmsrfZ4TuhWn1olL1olJXmsqBi7X1olJXmspZmslbmMhbmsdemsVfl8ZgmsNim8Jpk8F0m7R4m7F5nLB6jbh7jbiDirOEibOGnKaMhq+PnaCVg6qWg6qegKaff6WhnpKofKGtnomxeZy3noG6dZi+n3vCcpPDcpPGn3bLb4/Mb47UbIrVa4rYoGjdaIbeaIXhoWHmZYHobXvpcHjqdHXreHLroVrsfG/uhGnuh2bwj2Hxk17yl1vzmljzm1j0nlX1olL3AJXWAAAAbXRSTlMAEBAQHx8gICAuLjAwMDw9PUBAQEpQUFBXV1hgYGBkcHBwcXl8gICAgoiIkJCQlJicnJ2goKCmqK+wsLC4usDAwMjP0NDQ1NbW3Nzg4ODi5+3v8PDw8/T09PX29vb39/f5+fr7+/z8/Pz9/v7+zczCxgAABC5JREFUeAHN1ul3k0UUBvCb1CTVpmpaitAGSLSpSuKCLWpbTKNJFGlcSMAFF63iUmRccNG6gLbuxkXU66JAUef/9LSpmXnyLr3T5AO/rzl5zj137p136BISy44fKJXuGN/d19PUfYeO67Znqtf2KH33Id1psXoFdW30sPZ1sMvs2D060AHqws4FHeJojLZqnw53cmfvg+XR8mC0OEjuxrXEkX5ydeVJLVIlV0e10PXk5k7dYeHu7Cj1j+49uKg7uLU61tGLw1lq27ugQYlclHC4bgv7VQ+TAyj5Zc/UjsPvs1sd5cWryWObtvWT2EPa4rtnWW3JkpjggEpbOsPr7F7EyNewtpBIslA7p43HCsnwooXTEc3UmPmCNn5lrqTJxy6nRmcavGZVt/3Da2pD5NHvsOHJCrdc1G2r3DITpU7yic7w/7Rxnjc0kt5GC4djiv2Sz3Fb2iEZg41/ddsFDoyuYrIkmFehz0HR2thPgQqMyQYb2OtB0WxsZ3BeG3+wpRb1vzl2UYBog8FfGhttFKjtAclnZYrRo9ryG9uG/FZQU4AEg8ZE9LjGMzTmqKXPLnlWVnIlQQTvxJf8ip7VgjZjyVPrjw1te5otM7RmP7xm+sK2Gv9I8Gi++BRbEkR9EBw8zRUcKxwp73xkaLiqQb+kGduJTNHG72zcW9LoJgqQxpP3/Tj//c3yB0tqzaml05/+orHLksVO+95kX7/7qgJvnjlrfr2Ggsyx0eoy9uPzN5SPd86aXggOsEKW2Prz7du3VID3/tzs/sSRs2w7ovVHKtjrX2pd7ZMlTxAYfBAL9jiDwfLkq55Tm7ifhMlTGPyCAs7RFRhn47JnlcB9RM5T97ASuZXIcVNuUDIndpDbdsfrqsOppeXl5Y+XVKdjFCTh+zGaVuj0d9zy05PPK3QzBamxdwtTCrzyg/2Rvf2EstUjordGwa/kx9mSJLr8mLLtCW8HHGJc2R5hS219IiF6PnTusOqcMl57gm0Z8kanKMAQg0qSyuZfn7zItsbGyO9QlnxY0eCuD1XL2ys/MsrQhltE7Ug0uFOzufJFE2PxBo/YAx8XPPdDwWN0MrDRYIZF0mSMKCNHgaIVFoBbNoLJ7tEQDKxGF0kcLQimojCZopv0OkNOyWCCg9XMVAi7ARJzQdM2QUh0gmBozjc3Skg6dSBRqDGYSUOu66Zg+I2fNZs/M3/f/Grl/XnyF1Gw3VKCez0PN5IUfFLqvgUN4C0qNqYs5YhPL+aVZYDE4IpUk57oSFnJm4FyCqqOE0jhY2SMyLFoo56zyo6becOS5UVDdj7Vih0zp+tcMhwRpBeLyqtIjlJKAIZSbI8SGSF3k0pA3mR5tHuwPFoa7N7reoq2bqCsAk1HqCu5uvI1n6JuRXI+S1Mco54YmYTwcn6Aeic+kssXi8XpXC4V3t7/ADuTNKaQJdScAAAAAElFTkSuQmCC\">\n",
    "    </a>\n",
    "    <br/>\n",
    "</p>\n",
    "\n",
    "<p align=\"left\">\n",
    "    <a href=\"https://w3id.org/ro-id/a6376631-ef64-4c7a-b242-893d096fd33b\">\n",
    "        <img alt=\"RoHub\" src=\"https://img.shields.io/badge/RoHub-FAIR_Executable_Research_Object-2ea44f?logo=Open+Access&logoColor=blue\">\n",
    "    </a>\n",
    "    <a href=\"https://zenodo.org/badge/latestdoi/643572395\">\n",
    "        <img alt=\"doi\" src=\"https://zenodo.org/badge/643572395.svg\">\n",
    "    </a>\n",
    "</p>\n",
    "\n",
    "## Context\n",
    "### Purpose\n",
    "Solve data assimilation problems with implicit regularization with deep prior.\n",
    "\n",
    "### Description\n",
    "Deep image prior, a neural network with proven results in standard inverse problems, is used to determine the state of a physical system. It is a fully unsupervised model which acts as implicit regularization in the variational data assimilation method. The algorithm is demonstrated with a shallow-water toy model and the simulated observations are compared with variational assimilation without regularization, with Tikhonov regularization and with deep prior.\n",
    "\n",
    "### Highlights\n",
    "* Fetch the modelling codebase from the GitHub repo provided with the paper.\n",
    "* Run the main Python script to generate sample observations from a shallow-water model and variational assimilation solutions.\n",
    "* Recreate visualizations from the paper.\n",
    "* Compare assimilation scores and smoothness statistics between the different models.\n",
    "\n",
    "### Contributions\n",
    "\n",
    "#### Notebook\n",
    "* Mukulika Pahari (author), University of Mumbai, [@Mukulikaa](https://github.com/Mukulikaa)\n",
    "* Rutika Bhoir (author), University of Mumbai, [@Rutika-16](https://github.com/Rutika-16)\n",
    "* Tina Odaka (reviewer), Ifremer, [@tinaok](https://github.com/tinaok)\n",
    "* Caroline Arnold (reviewer), German Climate Computing Center, [@crlna16](https://github.com/crlna16)\n",
    "* Paolo Pelucchi (reviewer), Universitat de València, [@polpel](https://github.com/polpel)\n",
    "\n",
    "#### Modelling codebase\n",
    "* Arthur Filoche (author), Sorbonne Université, [@ArFiloche](https://github.com/ArFiloche)\n",
    "\n",
    "#### Modelling publications\n",
    "```{bibliography}\n",
    "  :style: plain\n",
    "  :list: bullet\n",
    "  :filter: topic % \"39d9c177-11da-41b2-9b64-63f4c1c834b3\"\n",
    "```\n",
    "\n",
    "### Source code\n",
    "The authors of the notebook acknowledge the original creator for providing public code available at [Deepprior4DVar_CI22](https://github.com/ArFiloche/Deepprior4DVar_CI22)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Clone the paper's GitHub repository\n",
    "\n",
    "The cloned repository is moved to the current directory to easily import the required Python modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!git clone -q https://github.com/ArFiloche/Deepprior4DVar_CI22"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# system\n",
    "import sys\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "sys.path.insert(0, os.path.join(os.getcwd(), \"Deepprior4DVar_CI22\"))\n",
    "\n",
    "# data\n",
    "import numpy as np\n",
    "import pooch\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "# modelling\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# custom functions from the Deepprior4DVar_CI22 repo\n",
    "import dynamics\n",
    "import utils\n",
    "import _4DVar\n",
    "import _DeepPrior4DVar\n",
    "\n",
    "# plotting and tables\n",
    "import matplotlib.pyplot as plt\n",
    "import tabulate\n",
    "\n",
    "# utils\n",
    "import time\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(action=\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set project structure\n",
    "\n",
    "The following directory structure is the same as given in the project repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "root_dir = \"data/generated/\"\n",
    "save_dir = \"data/estimated/\"\n",
    "\n",
    "if os.path.exists(\"data\"):\n",
    "    shutil.rmtree(\"data\")\n",
    "\n",
    "os.makedirs(root_dir)\n",
    "os.makedirs(\"data/results/\")\n",
    "os.makedirs(save_dir + \"4DVar/\")\n",
    "os.makedirs(save_dir + \"4DVar_reg/\")\n",
    "os.makedirs(save_dir + \"4DVar_deep/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate data and run assimilations\n",
    "\n",
    "The  `Deepprior4DVar_CI22` repo has the following main modules:\n",
    "* `dynamics`: For simulating the shallow-water model.\n",
    "* `_4DVar`: For assimilation with and without regularization.\n",
    "* `_DeepPrior4DVar`: For assimilation with deep prior.\n",
    "\n",
    "The repo has a script- `main.py`, to perform all the steps required to reproduce the paper. In this notebook, this script has been run with a smaller sample size to reproduce the results. The repo also contains Jupyter notebooks for demonstration in the `notebooks_demo` directory.\n",
    "\n",
    "The first step is to generate observations from a shallow-water model initialized with initial state variables. Keeping in mind the computational limitations, the sample size has been reduced to 20 from 100, as taken in the original script. The output is stored in `data/generated`.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\"> <b>Note: </b> You can change to the original value of n_sample=100 according to the availability of GPU resources </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_sample = 20\n",
    "dx = utils.data_generator.dx\n",
    "dy = utils.data_generator.dy\n",
    "dynamics_obj = dynamics.SW(dx, dy)\n",
    "generator = utils.DataGenerator(dynamics_obj, T=10, subsample=3, sigma=0.05)\n",
    "generator.dataset(n_sample, root_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The rest of the script is run as is. The next steps include:\n",
    "* extracting observations from the simulated model,\n",
    "* performing 4D variational assimilation with and without regularization,\n",
    "* performing deep prior 4D variational assimilation,\n",
    "* computing performance metrics,\n",
    "  * Endpoint Error - EPE\n",
    "  * Angular Error -AE\n",
    "  * ||grad|| - $\\parallel \\nabla\\mathbf{w}_0\\parallel_2$\n",
    "  * ||div|| - $\\parallel \\nabla . \\mathbf{w}_0\\parallel_2$\n",
    "  * ||lap|| - $\\parallel\\Delta \\mathbf{w}_0\\parallel_2$\n",
    "* storing computed metrics in a NumPy array."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "**Note:**\n",
    "The steps indicated above require access to GPU hardware which is not available in the standard Binder. Therefore, below we display the code in a markdown cell, rather than an executable cell. Users can copy and paste the code below to a new executable cell if GPU resources are available e.g. [Colab](https://colab.google/).\n",
    "</div>\n",
    "\n",
    "```python\n",
    "results=np.zeros((n_sample,4,5)) # EPE, AE, ||grad||, ||div||, ||lap||\n",
    "\n",
    "for i in range(n_sample):\n",
    "    #load data\n",
    "    initial_condition=np.load(root_dir+'initial_conditions/'+'{0:04}'.format(int(i))+'.npy')\n",
    "    w0_truth=torch.Tensor(initial_condition[1:,:,:])\n",
    "\n",
    "    # results\n",
    "    results[i,0,0]=utils.EPE(w0_truth, w0_truth) #0\n",
    "    results[i,0,1]=utils.Angular_error(w0_truth, w0_truth) #0\n",
    "    results[i,0,2]=utils.norm_gradw(w0_truth)\n",
    "    results[i,0,3]=utils.norm_divw(w0_truth)\n",
    "    results[i,0,4]=utils.norm_lapw(w0_truth)\n",
    "\n",
    "    Obs=np.load(root_dir+'Obs/'+'{0:04}'.format(int(i))+'.npy')\n",
    "    Obs=torch.Tensor(Obs)\n",
    "\n",
    "    Rm1=torch.ones(Obs.shape)*(Obs!=0)\n",
    "    Rm1=torch.Tensor(Rm1)\n",
    "\n",
    "    ## 4D-Var #####\n",
    "    assim = _4DVar.strong_4DVar(dynamics=dynamics_obj,\n",
    "                         optimizer=optim.LBFGS([torch.zeros(0)],lr=0.75, max_iter=250))\n",
    "    assim.fit(Obs=Obs,Rm1=Rm1)\n",
    "\n",
    "    w0_4dvar= assim.initial_condition[1:,:,:]\n",
    "    np.save(save_dir+'4DVar/'+'{0:04}'.format(int(i))+'.npy', w0_4dvar)\n",
    "\n",
    "    # results\n",
    "    results[i,1,0]=utils.EPE(w0_truth, w0_4dvar)\n",
    "    results[i,1,1]=utils.Angular_error(w0_truth, w0_4dvar)\n",
    "    results[i,1,2]=utils.norm_gradw(w0_4dvar)\n",
    "    results[i,1,3]=utils.norm_divw(w0_4dvar)\n",
    "    results[i,1,4]=utils.norm_lapw(w0_4dvar)\n",
    "\n",
    "    ## Thikonov 4D-Var #####\n",
    "    smoothreg= _4DVar.smooth_regul(alpha=5e3, beta=5e2,dx=dx,dy=dy)\n",
    "    assim = _4DVar.strong_4DVar(dynamics=dynamics_obj, regul=smoothreg,\n",
    "                         optimizer=optim.LBFGS([torch.zeros(0)],lr=0.75, max_iter=250))\n",
    "    assim.fit(Obs=Obs,Rm1=Rm1)\n",
    "\n",
    "    w0_4dvar_reg=assim.initial_condition[1:,:,:]\n",
    "    np.save(save_dir+'4DVar_reg/'+'{0:04}'.format(int(i))+'.npy', w0_4dvar_reg)\n",
    "\n",
    "    # results\n",
    "    results[i,2,0]=utils.EPE(w0_truth, w0_4dvar_reg)\n",
    "    results[i,2,1]=utils.Angular_error(w0_truth, w0_4dvar_reg)\n",
    "    results[i,2,2]=utils.norm_gradw(w0_4dvar_reg)\n",
    "    results[i,2,3]=utils.norm_divw(w0_4dvar_reg)\n",
    "    results[i,2,4]=utils.norm_lapw(w0_4dvar_reg)\n",
    "\n",
    "    ## Deep prior 4D-Var #####\n",
    "    netG = _DeepPrior4DVar.Generator()\n",
    "    netG.apply(_DeepPrior4DVar.weights_init)\n",
    "\n",
    "    noise = torch.randn(netG.nz, 1, 1)\n",
    "\n",
    "    assim=_DeepPrior4DVar.deeprior_strong_4DVar(generator=netG, dynamics=dynamics_obj,\n",
    "                               lr = 0.002, beta1 = 0.5, n_epoch=2000)\n",
    "    assim.fit(noise, Obs, Rm1)\n",
    "\n",
    "    w0_4dvar_deep=assim.initial_condition[1:,:,:]\n",
    "    np.save(save_dir+'4DVar_deep/'+'{0:04}'.format(int(i))+'.npy', w0_4dvar_deep)\n",
    "\n",
    "    # results\n",
    "    results[i,3,0]=utils.EPE(w0_truth, w0_4dvar_deep)\n",
    "    results[i,3,1]=utils.Angular_error(w0_truth, w0_4dvar_deep)\n",
    "    results[i,3,2]=utils.norm_gradw(w0_4dvar_deep)\n",
    "    results[i,3,3]=utils.norm_divw(w0_4dvar_deep)\n",
    "    results[i,3,4]=utils.norm_lapw(w0_4dvar_deep)\n",
    "\n",
    "    elapsed = time.perf_counter() - time_start\n",
    "    print(f'Sample {i} successfully run and stored in {elapsed //60} min')\n",
    "\n",
    "np.save('./data/results/main.npy',results)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fetch data\n",
    "We use `pooch` to fetch outputs of the above cell from Zenodo. This will allow us to run the rest of cells for virtual environments without GPU availability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "pooch.retrieve(\n",
    "    url=\"doi:10.5281/zenodo.8339319/data.zip\",\n",
    "    known_hash=\"md5:945c6f0768147a15448fb40c1e587462\",\n",
    "    processor=pooch.Unzip(extract_dir=f\"data\"),\n",
    "    path=f\".\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualisation\n",
    "\n",
    "The data generated from the script is loaded into variables for visualization. The plots are made with the scripts provided in `utils`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "root_dir = \"data/generated/\"\n",
    "\n",
    "initial_condition = np.load(\n",
    "    root_dir + \"initial_conditions/\" + \"{0:04}\".format(int(i)) + \".npy\"\n",
    ")\n",
    "Obs = np.load(root_dir + \"Obs/\" + \"{0:04}\".format(int(i)) + \".npy\")\n",
    "\n",
    "results = np.load(\"data/results/main.npy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Initial conditions (Figure 2)\n",
    "\n",
    "The plot below shows the evolution of the simulated system. State variables of the considered system are $\\eta$, the height deviation of the horizontal pressure surface from its mean height, and $w$, the associated velocity field. The considered temporal window $T$ has a fixed size, in this case equal to 10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "T = 10  # Assimilation window size\n",
    "\n",
    "X = torch.Tensor(initial_condition)\n",
    "fig2 = plt.figure(figsize=(20, 4))\n",
    "for t in range(T):\n",
    "    plt.subplot(2, T, t + 1)\n",
    "    plt.title(\"$\\eta_{\" + str(t) + \"}$\", fontsize=17)\n",
    "    plt.yticks([])\n",
    "    plt.xticks([])\n",
    "\n",
    "    plt.imshow(X[0, :, :])\n",
    "    plt.subplot(2, T, T + t + 1)\n",
    "    utils.plot_w(X[1, :, :], X[2, :, :], quiver=True, title=\"w_\" + str(t), q_scale=2.5)\n",
    "    plt.title(\"$w_{\" + str(t) + \"}$\", fontsize=17)\n",
    "\n",
    "    X = dynamics_obj.forward(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Observations with noise (Figure 3)\n",
    "\n",
    "The plot below shows the observations extracted from the simulated system and infused with noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "subsample = 3  # Subsampling frequency\n",
    "\n",
    "fig3 = plt.figure(figsize=(20, 4))\n",
    "\n",
    "for t in range(T):\n",
    "    plt.subplot(1, T, t + 1)\n",
    "    plt.title(\"$Y_{\" + str(t) + \"}$\", fontsize=17)\n",
    "    if t % subsample == 0:\n",
    "        plt.imshow(Obs[t, 0, :])\n",
    "    else:\n",
    "        plt.imshow(float(np.nan) * Obs[t, 0, :])\n",
    "    plt.yticks([])\n",
    "    plt.xticks([])\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Results (Figure 4)\n",
    "\n",
    "The following plots represent the dynamical state of the system through the different assimilation approaches and the observed state. It is a reconstruction of [Figure 4](https://static.cambridge.org/binary/version/id/urn:cambridge.org:id:binary:20230110163855547-0883:S2634460222000310:S2634460222000310_fig4.png) in the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "save_dir = \"data/estimated/\"\n",
    "\n",
    "w0_truth = initial_condition[1:, :, :]\n",
    "\n",
    "w0_4dvar = np.load(save_dir + \"4DVar/\" + \"{0:04}\".format(int(i)) + \".npy\")\n",
    "w0_4dvar_reg = np.load(save_dir + \"4DVar_reg/\" + \"{0:04}\".format(int(i)) + \".npy\")\n",
    "w0_4dvar_deep = np.load(save_dir + \"4DVar_deep/\" + \"{0:04}\".format(int(i)) + \".npy\")\n",
    "\n",
    "u0_truth = torch.Tensor(w0_truth[0, :, :])\n",
    "v0_truth = torch.Tensor(w0_truth[1, :, :])\n",
    "\n",
    "u0_4dvar = torch.Tensor(w0_4dvar[0, :, :])\n",
    "v0_4dvar = torch.Tensor(w0_4dvar[1, :, :])\n",
    "\n",
    "u0_4dvar_reg = torch.Tensor(w0_4dvar_reg[0, :, :])\n",
    "v0_4dvar_reg = torch.Tensor(w0_4dvar_reg[1, :, :])\n",
    "\n",
    "u0_4dvar_deep = torch.Tensor(w0_4dvar_deep[0, :, :])\n",
    "v0_4dvar_deep = torch.Tensor(w0_4dvar_deep[1, :, :])\n",
    "\n",
    "normalization = True\n",
    "quiver = True\n",
    "q_alpha = 0.5\n",
    "q_scale = 2.5\n",
    "sub = 4\n",
    "\n",
    "fig4 = plt.figure(figsize=(20, 10))\n",
    "\n",
    "plt.subplot(1, 4, 1)\n",
    "utils.plot_w(\n",
    "    u0_4dvar,\n",
    "    v0_4dvar,\n",
    "    normalization=normalization,\n",
    "    quiver=quiver,\n",
    "    q_alpha=q_alpha,\n",
    "    q_scale=q_scale,\n",
    "    title=\"4D-Var\",\n",
    ")\n",
    "plt.title(\"4D-Var\", fontsize=20)\n",
    "\n",
    "plt.subplot(1, 4, 2)\n",
    "utils.plot_w(\n",
    "    u0_4dvar_deep,\n",
    "    v0_4dvar_deep,\n",
    "    normalization=normalization,\n",
    "    quiver=quiver,\n",
    "    q_alpha=q_alpha,\n",
    "    q_scale=q_scale,\n",
    "    title=\"Deep prior 4D-Var\",\n",
    ")\n",
    "plt.title(\"Deep prior 4D-Var\", fontsize=20)\n",
    "\n",
    "plt.subplot(1, 4, 3)\n",
    "utils.plot_w(\n",
    "    u0_4dvar_reg,\n",
    "    v0_4dvar_reg,\n",
    "    normalization=normalization,\n",
    "    quiver=quiver,\n",
    "    q_alpha=q_alpha,\n",
    "    q_scale=q_scale,\n",
    "    title=\"4D-Var + Thikhonov\",\n",
    ")\n",
    "plt.title('\"Thikhonov\" 4D-Var ', fontsize=20)\n",
    "\n",
    "plt.subplot(1, 4, 4)\n",
    "utils.plot_w(\n",
    "    u0_truth,\n",
    "    v0_truth,\n",
    "    normalization=normalization,\n",
    "    quiver=quiver,\n",
    "    q_alpha=q_alpha,\n",
    "    q_scale=q_scale,\n",
    "    title=\"\",\n",
    ")\n",
    "plt.title(\"Ground truth\", fontsize=20)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Metrics (Figure 5)\n",
    "\n",
    "The following plots compare the different algorithms with respect to the five metrics as mentioned in the paper:\n",
    "* Endpoint error\n",
    "* Angular error\n",
    "* Smoothness statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "bins = 10\n",
    "figextra = plt.figure(figsize=(20, 5))\n",
    "\n",
    "titles = [\"Endpoint error\", \"Angular error\"]\n",
    "for i in range(0, 2):\n",
    "    plt.subplot(1, 2, i + 1)\n",
    "\n",
    "    plt.hist(\n",
    "        [results[:, j, i] for j in range(1, 4)],\n",
    "        bins,\n",
    "        label=[\"4D-Var\", '\"Thikonov\" 4D-Var', \"Deep prior 4D-Var\"],\n",
    "    )\n",
    "    if i == 1:\n",
    "        plt.legend(loc=\"upper right\")\n",
    "    plt.title(titles[i - 1], fontsize=20)\n",
    "    plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "bins = 15\n",
    "fig5 = plt.figure(figsize=(20, 5))\n",
    "\n",
    "titles = [\n",
    "    r\"$\\parallel \\nabla\\mathbf{w}_0\\parallel_2$\",\n",
    "    r\"$\\parallel \\nabla . \\mathbf{w}_0\\parallel_2$\",\n",
    "    \"$\\parallel\\Delta \\mathbf{w}_0\\parallel_2$\",\n",
    "]\n",
    "for i in range(2, results.shape[2]):\n",
    "    plt.subplot(1, 3, i - 1)\n",
    "\n",
    "    plt.hist(\n",
    "        [results[:, j, i] for j in range(4)],\n",
    "        bins,\n",
    "        label=[\"Ground truth\", \"4D-Var\", '\"Thikonov\" 4D-Var', \"Deep prior 4D-Var\"],\n",
    "    )\n",
    "    if i == 4:\n",
    "        plt.legend(loc=\"upper right\")\n",
    "    plt.title(titles[i - 2], fontsize=20)\n",
    "    plt.grid()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following table is a direct comparison between the metrics. They are within the limits of the error rates as mentioned in [Table 1](https://static.cambridge.org/binary/version/id/urn:cambridge.org:id:binary:20230110163855547-0883:S2634460222000310:S2634460222000310_tab1.png)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "table_format = \"grid\"\n",
    "table = [\n",
    "    [\"Metrics\", \"Ground truth\", \"4D-Var\", \"'Thikonov' 4D-Var\", \"Deep prior 4D-Var\"],\n",
    "    [\n",
    "        \"EPE\",\n",
    "        results.mean(axis=0)[0, 0],\n",
    "        results.mean(axis=0)[1, 0],\n",
    "        results.mean(axis=0)[2, 0],\n",
    "        results.mean(axis=0)[3, 0],\n",
    "    ],\n",
    "    [\n",
    "        \"Angular error\",\n",
    "        results.mean(axis=0)[0, 1],\n",
    "        results.mean(axis=0)[1, 1],\n",
    "        results.mean(axis=0)[2, 1],\n",
    "        results.mean(axis=0)[3, 1],\n",
    "    ],\n",
    "    [\n",
    "        \"||grad||\",\n",
    "        results.mean(axis=0)[0, 2],\n",
    "        results.mean(axis=0)[1, 2],\n",
    "        results.mean(axis=0)[2, 2],\n",
    "        results.mean(axis=0)[3, 2],\n",
    "    ],\n",
    "    [\n",
    "        \"||div||\",\n",
    "        results.mean(axis=0)[0, 3],\n",
    "        results.mean(axis=0)[1, 3],\n",
    "        results.mean(axis=0)[2, 3],\n",
    "        results.mean(axis=0)[3, 3],\n",
    "    ],\n",
    "    [\n",
    "        \"||lap||\",\n",
    "        results.mean(axis=0)[0, 4],\n",
    "        results.mean(axis=0)[1, 4],\n",
    "        results.mean(axis=0)[2, 4],\n",
    "        results.mean(axis=0)[3, 4],\n",
    "    ],\n",
    "]\n",
    "\n",
    "\n",
    "formatted_table = tabulate.tabulate(table, headers=\"firstrow\", tablefmt=table_format)\n",
    "\n",
    "print(formatted_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Challenges and improvements\n",
    "\n",
    "* While most cloud-based notebook servers come with pre-installed packages, it will be helpful to include an environment file in the codebase to aid local runs.\n",
    "* The codebase is modularized and easy to understand but a brief outline in the Readme file would be useful."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* The results in the paper \"Deep prior in variational assimilation to estimate an ocean circulation without explicit regularization\" can be correctly reproduced with the help of the materials provided with the paper. \n",
    "* This paper encourages the use of deep learning techniques in traditional inverse problems which are a basis for many geoscience modelling experiments.\n",
    "* PyTorch is used for running the assimilation models and NumPy is used to store the results.\n",
    "* Matplotlib is used for all visualizations.\n",
    "* Possible improvements include add an environment file and a comprehensive README file to support the modelling codebase."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Citing this Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please see [CITATION.cff](https://github.com/eds-book-gallery/39d9c177-11da-41b2-9b64-63f4c1c834b3/blob/main/CITATION.cff) for the full citation information. The citation file can be exported to APA or BibTex formats (learn more [here](https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-citation-files))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Additional information\n",
    "**Dataset**: All data are generated at runtime.\n",
    "\n",
    "**License**: The code in this notebook is licensed under the MIT License. The Environmental Data Science book is licensed under the Creative Commons by Attribution 4.0 license. See further details [here](https://github.com/alan-turing-institute/environmental-ds-book/blob/master/LICENSE.md).\n",
    "\n",
    "**Contact**: If you have any suggestion or report an issue with this notebook, feel free to [create an issue](https://github.com/alan-turing-institute/environmental-ds-book/issues/new/choose) or send a direct message to [environmental.ds.book@gmail.com](mailto:environmental.ds.book@gmail.com)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "\n",
    "print('Notebook repository version: v1.0.2')\n",
    "print(f'Last tested: {date.today()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "source": [
    "## Outputs registration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "outputs = {\n",
    "    \"static_figures\": {\n",
    "        \"filenames\": [\n",
    "            \"fig2_simulatedtrajectory\",\n",
    "            \"fig3_simulatedsystem\",\n",
    "            \"fig4_estimatedfields\",\n",
    "            \"fig5_smoothness statistics\",\n",
    "            \"figextra_smoothness statistics\",\n",
    "        ],\n",
    "        \"data\": [\n",
    "            fig2,\n",
    "            fig3,\n",
    "            fig4,\n",
    "            fig5,\n",
    "            figextra,\n",
    "        ],\n",
    "    }\n",
    "}\n",
    "\n",
    "fig_folder = \"outputs/figures\"\n",
    "os.makedirs(fig_folder, exist_ok=True)\n",
    "\n",
    "# save static figures\n",
    "if len(outputs[\"static_figures\"][\"filenames\"]) > 0:\n",
    "    [\n",
    "        data.savefig(\n",
    "            os.path.join(\n",
    "                fig_folder, outputs[\"static_figures\"][\"filenames\"][x] + \".png\"\n",
    "            ),\n",
    "            bbox_inches=\"tight\",\n",
    "        )\n",
    "        for x, data in enumerate(outputs[\"static_figures\"][\"data\"])\n",
    "    ]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}