notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learning the underlying physics of a simulation model of the ocean's temperature (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/3286b92f-4fae-4cc6-a29e-e408bc844542/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/3286b92f-4fae-4cc6-a29e-e408bc844542/actions/workflows/render.yaml\">\n",
    "        <img alt=\"render\" src=\"https://github.com/eds-book-gallery/3286b92f-4fae-4cc6-a29e-e408bc844542/actions/workflows/render.yaml/badge.svg\">\n",
    "    </a>\n",
    "    <a href=\"https://github.com/alan-turing-institute/environmental-ds-book/issues/172\">\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=\"http://mybinder.org/v2/gh/eds-book-gallery/3286b92f-4fae-4cc6-a29e-e408bc844542/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/3286b92f-4fae-4cc6-a29e-e408bc844542/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/b078ae9c-2ae3-4c10-8a86-673b55d2b274\">\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/636822482\">\n",
    "        <img alt=\"doi\" src=\"https://zenodo.org/badge/636822482.svg\">\n",
    "    </a>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Context\n",
    "Can data-driven models for weather and climate predictions learn the underlying physics of the system against which they are trained? Or are they simply capable of identifying statistical patterns without any clear link to the underlying physics? Furner et al. (2022) run a sensitivity analysis of a regression-based ocean temperature model, trained against simulations from a 3D ocean model setup, demostrating that regression models are capable of learning much of the physics of the underlying system.\n",
    "    \n",
    "![Ocean Temperature](https://www.commondreams.org/media-library/ocean-surface-temperatures.jpg?id=33699188&width=1200&height=400&quality=90&coordinates=0%2C667%2C0%2C667)\n",
    "\n",
    "<p><center>\n",
    "Source: Image generated by the NASA and taken from Common Dreams.\n",
    "</center></p>\n",
    "\n",
    "### Purpose\n",
    "This notebook aims to complement the science and methodological development embedded within the original paper, using an open infrastructure that allows users to combine interactive code with text and graphical objects, translating research outputs into findable, accessible, interoperable and reusable outputs and code.\n",
    "\n",
    "### Description\n",
    "The notebook demonstrates the inputs, the training of the regression model of ocean temperature and its sensitivity analysis.\n",
    "\n",
    "### Highlights\n",
    "* Load the necessary libraries and dependencies to run this notebook, i.e., numpy, xarray, etc.\n",
    "* Fetch the original dataset used for training the regression models.\n",
    "* Visualize the simulator outputs and inspect the temperature distributions near the equator, the poles, and the deep ocean.\n",
    "* Showcase an example of model training for linear regression.\n",
    "* Evaluate the performance of the regressors.\n",
    "* Inspect the spatial distribution of the averaged errors.\n",
    "* Perform a sensitivity analysis and run a series of withholding experiments.\n",
    "* Summarize our findings.\n",
    "* Provide contact information.\n",
    "\n",
    "### Contributions\n",
    "#### Notebook\n",
    "* Garima Malhotra (author), University of Colorado Boulder, [@garimamalhotra](https://github.com/garimamalhotra) \n",
    "* Daniela Pinto Veizaga (author), University of California, Berkeley, [@dapivei](https://github.com/dapivei)\n",
    "* Jorge Eduardo Peña Velasco (author), Claremont McKenna College, [@jedpe](https://github.com/jedpe)\n",
    "* Rachel Furner (reviewer), University of Cambridge, [@RachelFurner](https://github.com/RachelFurner)\n",
    "* Oscar Bautista (reviewer), World Food Programme, [@oscarbau](https://github.com/oscarbau)\n",
    "* Ricardo Barros Lourenço (reviewer), McMaster University, [@ricardobarroslourenco](https://github.com/ricardobarroslourenco)\n",
    "\n",
    "#### Modelling codebase\n",
    "* Rachel Furner (author), University of Cambridge, [@RachelFurner](https://github.com/RachelFurner)\n",
    "* Garima Malhotra (contributor), University of Colorado Boulder, [@garimamalhotra](https://github.com/garimamalhotra)\n",
    "* Daniela Pinto Veizaga (contributor), University of California, Berkeley, [@dapivei](https://github.com/dapivei)\n",
    "* Jorge Eduardo Peña Velasco (contributor), Claremont McKenna College, [@jedpe](https://github.com/jedpe)\n",
    "\n",
    "#### Modelling publications\n",
    "```{bibliography}\n",
    "  :style: plain\n",
    "  :list: bullet\n",
    "  :filter: topic % \"3286b92f-4fae-4cc6-a29e-e408bc844542\"\n",
    "```\n",
    "\n",
    "### Data\n",
    "* Furner, Rachel. (2021). MITgcm Dataset for paper: Sensitivity analysis of a data-driven model of ocean temperature (v1.1) [Data set]. Zenodo. https://doi.org/10.5281/zenodo.7919172\n",
    "* Malhotra, Garima, Pinto-Veizaga, Daniela, & Pena, Jorge. (2023). Reproducible Challenge - Team 3 - Sensitivity analysis- Models. https://doi.org/10.5281/zenodo.7954232\n",
    "\n",
    "<div class=\"alert alert-block alert-info\"> <b>Note: </b> We provide the input feature arrays truncated to first 15000 points for large arrays (in trunc_io_arrays/-> data/interim/) because of space limitation. These are provided only for the withholding experiments that are required to run this notebook. For input arrays of other witholding experiments, please contact the authors. These can be provided if the users are interested in understanding the sensitivity of the regressor for other withholding experiments (not shown in this notebook). </div>\n",
    "\n",
    "### Source code\n",
    "The authors of the notebook acknowledge the original creator for providing public code available at [Regression Model OceanTemp](https://github.com/RachelFurner/RegressionModelOceanTemp). An updated and enhaced version of the code is available at [Reproducibility of \"A sensitivity analysis of a regression model of ocean temperature\"](https://github.com/eds-book-gallery/repro-challenge-team-3).\n",
    "\n",
    "<div class=\"alert alert-block alert-info\"> <b>Note: </b> Significant parts from source code were adapted or written from scratch for this notebook and the reproducible challenge. </div>\n",
    "\n",
    "```{admonition} Project Resources\n",
    ":class: dropdown\n",
    "\n",
    "*The submitted notebook was implemented and tested within a [cloud computing infrastructure, kindly provided by the European Open Science Cloud (EOSC)](https://github.com/pangeo-data/pangeo-eosc). The following lines thoroughly described the characteristics of this environment.*\n",
    "\n",
    "| | |\n",
    "|:------- | ------:| \n",
    "|**Operating System** | Ubuntu 22.04.1|\n",
    "|**Architecture**          | x86_64|\n",
    "|  CPU op-mode(s)       | 32-bit, 64-bit|\n",
    "|  Address sizes        | 48 bits physical, 48 bits virtual|\n",
    "|  Byte Order           | Little Endian|\n",
    "|**CPU(s)**                 | 128|\n",
    "|  On-line CPU(s) list  | 0-127|\n",
    "|**Vendor ID**              | AuthenticAMD|\n",
    "|  Model name         | AMD EPYC 7543 32-Core Processor|\n",
    "|    CPU family         | 25|\n",
    "|    Model             | 1|\n",
    "|    Thread(s) per core | 2|\n",
    "|    Core(s) per socket | 32|\n",
    "|    Socket(s)          | 2|\n",
    "|    Stepping           | 1|\n",
    "|    Frequency boost    | enabled|\n",
    "|    CPU max MHz       | 2800.0000|\n",
    "|    CPU min MHz        | 1500.0000|\n",
    "|    BogoMIPS           | 5589.72|\n",
    "|    Flags            | fpu vme de pse|\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# System imports\n",
    "import sys\n",
    "import warnings\n",
    "import pandas as pd\n",
    "\n",
    "# Data handling\n",
    "import intake\n",
    "\n",
    "# Modelling\n",
    "from sklearn import linear_model\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "\n",
    "# Custom functions from the authors' repo\n",
    "sys.path.append(\"src/general/\")\n",
    "sys.path.append(\"src/pipelines/\")\n",
    "import CreateDataName as cn\n",
    "import ReadRoutinesModifiedCloud as rr\n",
    "from make_scatter_data_plots_cloud import make_scatter_plots\n",
    "from make_error_data_plots_cloud import make_error_plots\n",
    "\n",
    "# Custom functions of the reproducibility team\n",
    "from constants_cloud import *\n",
    "\n",
    "warnings.filterwarnings(action=\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "The training and validation data derive from running the Massachusetts Institute of Technology general circulation model ([MITgcm](https://zenodo.org/record/4672260#.ZFgb_S-B2pY)) —a physically based model capable of simulating the ocean or the atmosphere dynamics due to isomorphisms in the governing equations. This research project relies on a simple 2° sector configuration that captures the fundamental dynamics of the ocean, including a realistic overturning circulation. This configuration features a single ocean basin, with:\n",
    "+ limited topography\n",
    "+ simplified coastlines \n",
    "+ constant idealized forcing \n",
    "\n",
    "The domain runs from 60°S to 60° N, and is just over 20° wide in longitude. The domain is bounded by land along its northern and southern edges, and a strip of land runs along the eastern/western boundary from 60° N to 40° S.  Below this, in the southernmost 20°, the simulator has a periodic boundary condition, allowing flow that exits to the east and west to return to the domain at the western and eastern boundary. The domain has flat-bottom bathymetry of 5,000 m over most of the domain, with a 2° region of 2,500-m depth at the southernmost 20° of the eastern edge.\n",
    "\n",
    "The simulator has:\n",
    "+ 42 depth levels, with the surface layer being the thinnest at 10 m, and the bottom 10 levels being the maximum at 250 m. + 11 cells in the longitudinal direction\n",
    "+ 78 cells in the latitudinal direction\n",
    "\n",
    "The simulator has a 12-hr time step, with fields output daily. We focus on daily-mean outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download input material\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "**Note:** \n",
    "The submitted notebook (located at `notebooks/notebook_submitted.ipynb`) downloads input datasets and saved models from Zenodo repositories. This notebook version skips the download step by fetching input files directly from a cloud storage kindly provided by the [European Open Science Cloud (EOSC)](https://github.com/pangeo-data/pangeo-eosc).\n",
    "    \n",
    "</div>\n",
    "\n",
    "The original dataset is available [here](https://zenodo.org/record/7919172). This dataset is used to train the regression model described in the paper.  To download this data locally, run the following line of code. Note that this process will take $\\sim 1$ hour and will require at least 52 GB of available space.\n",
    "\n",
    "```python\n",
    "%%time\n",
    "!bash \"infrastructure/downloader_datasets.sh\" > \"../outputs/logs/download_datasets.txt\"\n",
    "```\n",
    "\n",
    "Training this model and performing all the experiments presented in the original paper requires plenty of computing power, so, to facilitate running this notebook, we have publicly released few trained models, their predictions, and input arrays available [here](https://doi.org/10.5281/zenodo.7954232). To download this data locally, run the following line of code. This process will take $\\sim 15$ minutes and requires at least $\\sim 20$ GB available space.\n",
    "\n",
    "```python\n",
    "%%time\n",
    "!bash \"infrastructure/downloader_experiments.sh\" > \"../outputs/logs/download_experiments.txt\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reproduce paper\n",
    "### Ocean dynamics (Figure 1)\n",
    "The paper uses a  2° sector configuration to simulate ocean dynamics. This configuration features a single ocean basin, with limited topography, simplified coastlines, and constant idealized forcing. This has been used in a number of idealized simulations of Southern Ocean processes and their impacts on the global circulation. This configuration, while relatively simple, captures the fundamental dynamics of the ocean, including a realistic overturning circulation.\n",
    "\n",
    "In the following script, we load the MITGCM DATASET and visualize the configuration used to simulate the ocean dynamics. Focusing particuarly on:\n",
    "\n",
    "<ol style=\"list-style-type: lower-alpha\">\n",
    "<li>Temperature (°C) at 25 m below the surface for one particular day.</li>\n",
    "<li>Change in temperature between over 1 day at 25 m below the surface.</li>\n",
    "<li>Standard deviation in temperature at 25 m below the surface.</li>\n",
    "<li>Temperature (°C) at 25 m below the surface and 13° E for one particular day.</li>\n",
    "<li>Change in temperature between over 1 day in temperature at 25 m below the surface & 13° E.</li>\n",
    "<li>Standard deviation in temperature at 25 m below the surface & 13° E.</li>\n",
    "<li>Time series at 57° N, 17° E, and −25 m.</li>\n",
    "<li>Time series at 55° S, 9° E, and −25 m.</li>\n",
    "</ol>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%run \"src/pipelines/graphs_section_cloud.py\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the domain:\n",
    "+ runs from  60°𝑆 to  60°𝑁, and is just over  20° wide in longitude. \n",
    "+ is bounded by land along the northern (and southern) edge, and a strip of land runs along the eastern (and western) boundary from 60°.\n",
    "  N to  40°𝑆 (see Figure 1a). Below this, in the southernmost 20°, the simulator has a periodic boundary condition, allowing flow that exits to the east (west) to return to the domain at the western (eastern) boundary. \n",
    "+ has flat-bottom bathymetry of 5,000 m over most of the domain, with a 2° region of 2,500-m depth at the southernmost 20° of the eastern edge (i.e., the spit of land forming the eastern boundary continues to the southern boundary as a 2,500-m high subsurface ridge).\n",
    "\n",
    "Importantly, note from the figures presented, that the depth axis is scaled to give each GCM grid cell equal spacing. The simulator shows a realistic temperature distribution with warm surface water near the equator, and cooler water near the poles and in the deep ocean. Temperature changes are largest in the very north of the domain and throughout the southern region. Though changes per day are small, they accumulate over time to give cycles of around 0.2° in some regions of the domain."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model training\n",
    "Model training is performed for the purely linear regression model (only linear terms) and the nolinear regression model (linear and multiplicative terms). Run parameters for the control and withholding experiments were set through the `run_vars` variable in the script `src/pipelines/TrainLinearRegressionModel.py`.\n",
    "\n",
    "### Linear regression (A toy example)\n",
    "For illustrative purposes, in the following section we showcase how to set parameters for a linear regression model that takes as input all the relevant variables, that is,:\n",
    "> Temperature, salinity, U (East–West) and V (North–South) current components, density, U, V, and W (vertical) components of the Gent-McWilliams (GM) bolus velocities, sea surface height (SSH), latitude, longitude, and depth. The GM bolus velocities are a parameterization of the velocities resulting from ocean eddies and are used in the GM scheme to calculate the advective effects of eddy action on tracers.\n",
    "\n",
    "First, we define a dictionary named `run_vars`, which has all the input variables set to `True` and, `poly_degree` set to `1` -- the latter parameter defines we will be running a purely linear regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run parameters\n",
    "run_vars = {\n",
    "    \"dimension\": 3,\n",
    "    \"lat\": True,\n",
    "    \"lon\": True,\n",
    "    \"dep\": True,\n",
    "    \"current\": True,\n",
    "    \"bolus_vel\": True,\n",
    "    \"sal\": True,\n",
    "    \"eta\": True,\n",
    "    \"density\": True,\n",
    "    \"poly_degree\": 1,\n",
    "    \"StepSize\": 1,\n",
    "    \"predict\": \"DelT\",\n",
    "}\n",
    "data_name = data_prefix + cn.create_dataname(run_vars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following chunk of code performs the following tasks:\n",
    "+ loads the MITGCM data into input and output arrays.\n",
    "+ splits the dataset into training, validation and testing samples.\n",
    "+ normalises the dataset.\n",
    "+ saves the datasets into 1) norm_outputs_tr, array containing the normalized outputs that will be used for training, 2) norm_outputs_val, containing the normalized outputs that would be used for validation.\n",
    "\n",
    "The training-validation split ratio is set by default to $0.7$, while the validation-testing split ratio is set to $0.9$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data\n",
    "cat = intake.open_catalog(\"data/inputs_paper.yml\")\n",
    "\n",
    "mitgcm_filename = cat[\"MITGCM_model\"]\n",
    "density_file = cat[\"density_data\"]\n",
    "clim_filename = cat[\"clim_data\"]\n",
    "trainval_ratio = 0.7\n",
    "valtest_ratio = 0.9\n",
    "\n",
    "(\n",
    "    norm_inputs_tr,\n",
    "    _,\n",
    "    _,\n",
    "    norm_outputs_tr,\n",
    "    norm_outputs_val,\n",
    "    _,\n",
    "    _,\n",
    "    _,\n",
    "    _,\n",
    "    _,\n",
    "    _,\n",
    "    _,\n",
    "    _,\n",
    ") = rr.ReadMITGCM(mitgcm_filename, clim_filename, density_file, data_name, run_vars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model fit example\n",
    "\n",
    "After splitting our dataset into the corresponding subsamples (training, validation and testing), we proceed to train the model setting the following  hyperparameter:\n",
    "+ $\\alpha = 0.001$\n",
    "\n",
    "Our cross-validation splitting strategy is set to $3$ folds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model hyperparameters for the grid search\n",
    "parameters = [{\"alpha\": [0.001]}]\n",
    "n_folds = 3\n",
    "lr = linear_model.Ridge(fit_intercept=False)\n",
    "lr = GridSearchCV(\n",
    "    lr, param_grid=parameters, cv=n_folds, scoring=\"neg_mean_squared_error\", refit=True\n",
    ")\n",
    "\n",
    "# Fit the model\n",
    "lr.fit(norm_inputs_tr, norm_outputs_tr)\n",
    "# Parameters for the estimator.\n",
    "pd.DataFrame.from_dict(lr.get_params(), orient=\"index\", columns=[\"values\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performance of the regressor\n",
    "#### Scatter plot for the control regressor (Figure 2) - Different from above\n",
    "We update the dictionary `run_vars`, such that `poly_degree` is set to `2` -- This helps us in defining the model so that it has second-order pairwise polynomial terms, in order to capture a limited amount of nonlinear behavior through interaction between terms. This means that as well as the above inputs, we have multiplicative pairs of features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up the parameters for the run\n",
    "run_vars = {\n",
    "    \"dimension\": 3,\n",
    "    \"lat\": True,\n",
    "    \"lon\": True,\n",
    "    \"dep\": True,\n",
    "    \"current\": True,\n",
    "    \"bolus_vel\": True,\n",
    "    \"sal\": True,\n",
    "    \"eta\": True,\n",
    "    \"density\": True,\n",
    "    \"poly_degree\": 2,\n",
    "    \"StepSize\": 1,\n",
    "    \"predict\": \"DelT\",\n",
    "}\n",
    "data_name = data_prefix + cn.create_dataname(run_vars)\n",
    "model_name = model_prefix + data_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "**Note:** \n",
    "Some models were trained using large input datasets which will not work in the memory available of the standard Binder. Therefore, below we display the code needed for loading the input data in a markdown cell, rather than an executable cell. Instead, in the Sensitivity analysis (Figure 5) section we demonstrate how steps below work on another model which fits the memory and run the prediction function.\n",
    "</div>\n",
    "\n",
    "```python\n",
    "# Open data\n",
    "cat = intake.open_catalog('data/inputs_models.yml')\n",
    "\n",
    "#Loading Input Feature Arrays\n",
    "norm_inputs_tr=cat[model_name].to_dask()['norm_inputs_tr']\n",
    "norm_inputs_val=cat[model_name].to_dask()['norm_inputs_val']\n",
    "```\n",
    "\n",
    "The following script performs the following steps to make the scatter plot of predicted and true temperature change, allowing us to evaluate the performance of the control regressor.\n",
    "\n",
    "<ol style=\"list-style-type: lower-alpha\">\n",
    "<li>Load the true temperature change arrays, persistence arrays to calculate climatology.</li>\n",
    "<li>Load the trained model.</li>\n",
    "<li>Calculate the predicted temperature changes using the input feature arrays and the trained model.</li>\n",
    "<li>Denormalise the predictions and true outputs using the mean and standard deviation of the arrays.</li>\n",
    "<li>Calculate the mean squared errors and make the scatter plot using the denormalized arrays.</li>\n",
    "</ol>\n",
    "\n",
    "```python\n",
    "make_scatter_plots(run_vars,data_name,model_name,norm_inputs_tr,norm_inputs_val,lim=15000,fig_prefix='fig2')\n",
    "```\n",
    "\n",
    "The model captures the central part of the distribution well. While the majority of the temperature change is dominated by small near-zero changes, capturing these is key to producing a good forecast system. To a lesser extent, the regressor also captures the tails of the distribution, where temperature changes are larger, although the underprediction is more significant here. However, it is noteworthy that the model still shows some skill for these points, given that the model used is very simple and there are a relatively limited number of training samples in the tails.\n",
    "\n",
    "#### Spatial patterns of errors (Figure 3)\n",
    "To understand how the regressor performs spatially, temporally averaged absolute errors are plotted. \n",
    "These averaged errors are shown in Figure 3. Note that the regressor is only applied away from boundary and land points (in its current form, it cannot deal with spatial locations that are not surrounded on all sides by ocean points); hence, points close to land are not included in these plots. The script below performs the following steps:\n",
    "\n",
    "<ol style=\"list-style-type: lower-alpha\">\n",
    "<li>It loads the the MITgcm state at 500 different times from the 20-year dataset, masks the locations that are not surrounded on all sides by ocean points, and uses these fields as inputs to the regressor to forecast a single time step ahead.</li>\n",
    "<li>The set of forecasts created from these 500 input states is compared to the truth from the GCM run, and the absolute errors between the truth and the predictions are then temporally averaged.</li>\n",
    "</ol>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"src/pipelines/spatial_patterns_errors_cloud.py\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The regressor shows the largest errors are located in the north of the domain and in the Southern Ocean. We see that the errors in the north of the domain are co-located with regions of high vertical advective temperature fluxes, and regions of high convective fluxes. These results imply the regression model struggles to fully capture the vertical processes, and the associated heat flux, in the north of the domain. The high errors in the Southern Ocean are again co-located with regions of high vertical diffusive fluxes, this time both explicit and implicit, and vertical advection.  Throughout the ocean interior where temperature changes and the fluxes associated with these are small, errors are also small as would be expected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sensitivity analysis\n",
    "#### Coefficient analysis (Figure 4)\n",
    "Figure 4a shows coefficients averaged over all input locations for each variable type (i.e., for most variables, there are 27 inputs, corresponding to the 27 neighboring cells; we average over these to give a single value for each variable (temperature, salinity, etc.) and for each polynomial combination of variables). \n",
    "\n",
    "Figure 4b shows the coefficients related to polynomial interactions of temperature with temperature. These are the raw coefficients, without any averaging applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"src/pipelines/coecien_analysis_cloud.py\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "High-weighted inputs (those with a large magnitude coefficient) are variables which are heavily used in the predictions and therefore considered important for the predictive skill. From the above figures we infer: \n",
    "\n",
    "<ol style=\"list-style-type: lower-alpha\">\n",
    "<li>From Figure 4a: Density (as a linear term, not in combination with other variables) is by far the most highly weighted variable in this model. This agrees well with the current physical understanding of ocean dynamics.</li>\n",
    "<li>From Figure 4a: The location information is low weighted, particularly when interacting with other variables. This indicates that the regressor is not basing its predictions predominantly on the location of points, but on the physical variables themselves.</li>\n",
    "<li>From Figure 4b: Multiplicative interaction between temperatures at the grid point we are predicting for and the temperature at all surrounding points is very highly weighted. This fits well with our physical expectation of the system—as diffusive and advective fluxes of temperature are dominated by local gradients in temperature.</li>\n",
    "</ol>\n",
    "    \n",
    "#### Witholding experiments\n",
    "In this section, we run a series of withholding experiments. For each of the variables described with the exception of temperature, we train a new regressor leaving out that one variable group, for example, we train a new regressor with all the existing inputs except for salinity at all surrounding points and any multiplicative terms including salinity.\n",
    "\n",
    "##### Withholding multiplicative terms: Figure 5; Importance of Non-Linearity\n",
    "We update the dictionary `run_vars`, such that `poly_degree` is set to `1` -- This helps us in evaluating the performance of the purely linear model, that is, the model trained without any multiplicative terms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up the parameters for the run - Setting Polynomial Degree=1\n",
    "run_vars = {\n",
    "    \"dimension\": 3,\n",
    "    \"lat\": True,\n",
    "    \"lon\": True,\n",
    "    \"dep\": True,\n",
    "    \"current\": True,\n",
    "    \"bolus_vel\": True,\n",
    "    \"sal\": True,\n",
    "    \"eta\": True,\n",
    "    \"density\": True,\n",
    "    \"poly_degree\": 1,\n",
    "    \"StepSize\": 1,\n",
    "    \"predict\": \"DelT\",\n",
    "}\n",
    "data_name = data_prefix + cn.create_dataname(run_vars)\n",
    "model_name = model_prefix + data_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to the input arrays from an intake catalog\n",
    "cat = intake.open_catalog(\"data/inputs_models.yml\")\n",
    "\n",
    "inputs_tr_filename = cat[model_name](partition=\"Tr\")\n",
    "inputs_val_filename = cat[model_name](partition=\"Val\")\n",
    "\n",
    "# Loading Input Arrays\n",
    "norm_inputs_tr = inputs_tr_filename.to_dask()[\"norm_inputs_tr\"]\n",
    "norm_inputs_val = inputs_val_filename.to_dask()[\"norm_inputs_val\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_scatter_plots(\n",
    "    run_vars,\n",
    "    data_name,\n",
    "    model_name,\n",
    "    norm_inputs_tr,\n",
    "    norm_inputs_val,\n",
    "    plot_val=False,\n",
    "    fig_prefix=\"fig5\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that, without multiplicative terms, the model can capture the mean behavior of the system (zero change in temperature) but is unable to capture any of the variability. This mean behavior alone does not provide useful forecasts, as can be seen from the statistics for this experiment. Nonlinearity is shown to be critical to modeling the variability of temperature change.\n",
    "\n",
    "##### Withholding 3D structure: Figure 6; Impact of vertical structure\n",
    "\n",
    "To assess how information about the vertical structure of the ocean impacts predictions, we look at spatially averaged errors from the model trained with only a 2D neighborhood of inputs, along with the difference in error between this and the control regressor from previous Section.\n",
    "\n",
    "We update the dictionary `run_vars`, such that `dimension` is set to `2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading the 2D model\n",
    "run_vars = {\n",
    "    \"dimension\": 2,\n",
    "    \"lat\": True,\n",
    "    \"lon\": True,\n",
    "    \"dep\": True,\n",
    "    \"current\": True,\n",
    "    \"bolus_vel\": True,\n",
    "    \"sal\": True,\n",
    "    \"eta\": True,\n",
    "    \"density\": True,\n",
    "    \"poly_degree\": 2,\n",
    "    \"StepSize\": 1,\n",
    "    \"predict\": \"DelT\",\n",
    "}\n",
    "data_name = data_prefix + cn.create_dataname(run_vars)\n",
    "model_name = model_prefix + data_name\n",
    "exp_name = exp_prefix + model_name\n",
    "cntrl_name = (\n",
    "    exp_prefix\n",
    "    + model_prefix\n",
    "    + data_prefix\n",
    "    + \"3dLatLonDepUVBolSalEtaDnsPolyDeg2_Step1_PredictDelT\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_error_plots(mitgcm_filename, exp_name, cntrl_name, fig_prefix=\"fig6\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interestingly, this regressor shows some regions (the deep water in the south of the domain) where the errors are notably improved in a regressor using only 2D information. In this work, we have developed a regressor which learns one equation to be applied across all grid boxes in the domain. We optimize for best performance averaged over all relevant grid cells, but this does not enforce the best possible performance over each individual grid point/region, and so some of the resultant models will favor certain types of dynamics more than others. Given this, it is not unexpected that the new equations discovered for the withholding experiments (which again optimize for best performance averaged over the entire domain interior) may outperform the control in some locations, despite being poorer overall. Here, we see that the control model is able to perform well across the domain, and optimizes for good performance overall, rather than the much more varied performance seen in the withholding experiments. It seems that as the model which withholds vertical information is not capable of performing well in many regions of the domain, a solution is found which highly optimizes performance in other regions to minimize error overall.\n",
    "\n",
    "##### Withholding currents: Figure 7; Impact of currents\n",
    "We update the dictionary `run_vars`, such that `current` is set to `False` -- This helps us in evaluating the impact of currents on the performance of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting the run parameters\n",
    "run_vars = {\n",
    "    \"dimension\": 3,\n",
    "    \"lat\": True,\n",
    "    \"lon\": True,\n",
    "    \"dep\": True,\n",
    "    \"current\": False,\n",
    "    \"bolus_vel\": True,\n",
    "    \"sal\": True,\n",
    "    \"eta\": True,\n",
    "    \"density\": True,\n",
    "    \"poly_degree\": 2,\n",
    "    \"StepSize\": 1,\n",
    "    \"predict\": \"DelT\",\n",
    "}\n",
    "data_name = data_prefix + cn.create_dataname(run_vars)\n",
    "model_name = model_prefix + data_name\n",
    "exp_name = exp_prefix + model_name\n",
    "cntrl_name = (\n",
    "    exp_prefix\n",
    "    + model_prefix\n",
    "    + data_prefix\n",
    "    + \"3dLatLonDepUVBolSalEtaDnsPolyDeg2_Step1_PredictDelT\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_error_plots(mitgcm_filename, exp_name, cntrl_name, fig_prefix=\"fig7\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The horizontal (U and V) components of the currents directly drive horizontal advection of temperature. They are also indirectly related to horizontal diffusion, as this is increased in regions of high currents and steep gradients. As such, we would expect that suppressing information about the horizontal currents would cause increases in error in regions where horizontal advection and horizontal diffusion is high. However, again, we note that this region of increased error is one where many processes are present, and the increased errors seen also coincide, to a lesser extent, with regions of high vertical processes (advection, diffusion, and convection), which is less in line with our physical understanding. Here, errors appear more closely matched to the horizontal processes, and so a reasonable interpretation is that the model here is behaving as expected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "* There has been much recent interest in developing data-driven models for weather and climate predictions. However, there are open questions regarding their generalizability and robustness, highlighting a need to better understand how they make their predictions. In particular, it is important to understand whether data-driven models learn the underlying physics of the system against which they are trained, or simply identify statistical patterns without any clear link to the underlying physics. \n",
    "* In this notebook, we run a sensitivity analysis of a regression-based model of ocean temperature, trained against simulations from a 3D ocean model setup in a very simple configuration. \n",
    "* We show that the regressor heavily bases its forecasts on, and is dependent on, variables known to be key to the physics such as currents and density. By contrast, the regressor does not make heavy use of inputs such as location, which have limited direct physical impacts. \n",
    "* The model requires nonlinear interactions between inputs in order to show any meaningful skill—in line with the highly nonlinear dynamics of the ocean. Further analysis interprets the ways certain variables are used by the regression model. \n",
    "* We see that information about the vertical profile of the water column reduces errors in regions of convective activity, and information about the currents reduces errors in regions dominated by advective processes. \n",
    "* Our results demonstrate that even a simple regression model is capable of learning much of the physics of the system being modeled. We expect that a similar sensitivity analysis could be usefully applied to more complex ocean configurations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Citing this Notebook\n",
    "\n",
    "Please see [CITATION.cff](https://github.com/eds-book-gallery/3286b92f-4fae-4cc6-a29e-e408bc844542/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**: Malhotra, Garima, Pinto-Veizaga, Daniela, & Pena, Jorge. (2023). Reproducible Challenge - Team 3 - Sensitivity analysis- Models. https://doi.org/10.5281/zenodo.7954232 - Version 1\n",
    "\n",
    "Furner, Rachel. (2021). MITgcm Dataset for paper: Sensitivity analysis of a data-driven model of ocean temperature (v1.1) [Data set]. Zenodo. https://doi.org/10.5281/zenodo.7919172 - Version 1.1\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 or send a direct message to environmental.ds.book@gmail.com."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "\n",
    "print('Notebook repository version: v1.0.19')\n",
    "print(f'Last tested: {date.today()}')"
   ]
  }
 ],
 "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.8.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}