sosaa-data/experiments/model-comparison.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a2e664bc",
   "metadata": {
    "papermill": {
     "duration": 0.023998,
     "end_time": "2023-01-25T10:00:58.914605",
     "exception": false,
     "start_time": "2023-01-25T10:00:58.890607",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "import hashlib\n",
    "import itertools\n",
    "import json\n",
    "import os\n",
    "import pickle\n",
    "import re\n",
    "import warnings\n",
    "\n",
    "from collections import namedtuple\n",
    "from copy import copy\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5995374a",
   "metadata": {
    "papermill": {
     "duration": 11.586791,
     "end_time": "2023-01-25T10:01:10.517703",
     "exception": false,
     "start_time": "2023-01-25T10:00:58.930912",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import cartopy.feature as cfeature\n",
    "import cartopy.crs as ccrs\n",
    "\n",
    "import joblib\n",
    "\n",
    "import matplotlib as mpl\n",
    "from matplotlib import gridspec\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import netCDF4\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import scipy as sp\n",
    "\n",
    "import sklearn\n",
    "from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "841b3f34",
   "metadata": {},
   "source": [
    "Plots use `matplotlib-label-lines`:\n",
    "```bibtex\n",
    "@software{mpl-label-lines-2022,\n",
    "  author       = {Cadiou, Corentin},\n",
    "  title        = {Matplotlib label lines},\n",
    "  month        = {12},\n",
    "  year         = {2022},\n",
    "  publisher    = {Zenodo},\n",
    "  version      = {v0.5.1},\n",
    "  doi          = {10.5281/zenodo.7428071},\n",
    "  url          = {https://doi.org/10.5281/zenodo.7428071}\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "95a896b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from labellines import labelLines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "70eb71cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import colorcet as cc\n",
    "\n",
    "time_cmap = cc.m_CET_L19\n",
    "uncertainty_cmap = cc.m_CET_L8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7c88591e",
   "metadata": {},
   "outputs": [],
   "source": [
    "height_cmap = mpl.cm.rainbow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "68424caf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://stackoverflow.com/a/49601444\n",
    "def adjust_lightness(color, amount=0.5):\n",
    "    import matplotlib.colors as mc\n",
    "    import colorsys\n",
    "\n",
    "    try:\n",
    "        c = mc.cnames[color]\n",
    "    except:\n",
    "        c = color\n",
    "\n",
    "    c = colorsys.rgb_to_hls(*mc.to_rgb(c))\n",
    "\n",
    "    return colorsys.hls_to_rgb(c[0], max(0, min(1, amount * c[1])), c[2])\n",
    "\n",
    "def replace_lightness(color, lightness):\n",
    "    import matplotlib.colors as mc\n",
    "    import colorsys\n",
    "\n",
    "    try:\n",
    "        c = mc.cnames[color]\n",
    "    except:\n",
    "        c = color\n",
    "\n",
    "    c = colorsys.rgb_to_hls(*mc.to_rgb(c))\n",
    "\n",
    "    return colorsys.hls_to_rgb(c[0], max(0, min(1, lightness)), c[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "16c8e97b",
   "metadata": {
    "papermill": {
     "duration": 0.017674,
     "end_time": "2023-01-25T10:01:10.592952",
     "exception": false,
     "start_time": "2023-01-25T10:01:10.575278",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "TrajectoryPaths = namedtuple(\"TrajectoryPaths\", [\"date\", \"out\", \"aer\", \"ant\", \"bio\", \"met\"])\n",
    "TrajectoryDatasets = namedtuple(\"TrajectoryDatasets\", [\"date\", \"out\", \"aer\", \"ant\", \"bio\", \"met\"])\n",
    "MLDataset = namedtuple(\"MLDataset\", [\"date\", \"paths\", \"X_raw\", \"Y_raw\", \"X_train\", \"X_valid\", \"X_test\", \"Y_train\", \"Y_valid\", \"Y_test\", \"X_scaler\", \"Y_scaler\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7384f3e9",
   "metadata": {
    "papermill": {
     "duration": 0.01736,
     "end_time": "2023-01-25T10:01:10.625304",
     "exception": false,
     "start_time": "2023-01-25T10:01:10.607944",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "OUTDIR_PATTERN = re.compile(r\"(\\d{4})(\\d{2})(\\d{2})_T(\\d{2})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f04359d6",
   "metadata": {
    "papermill": {
     "duration": 2.78927,
     "end_time": "2023-01-25T10:01:13.502247",
     "exception": false,
     "start_time": "2023-01-25T10:01:10.712977",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "traj_datetimes = dict()\n",
    "\n",
    "base = Path.cwd().parent / \"trajectories\"\n",
    "\n",
    "for child in (base / \"outputs\" / \"baseline\").iterdir():\n",
    "    if not child.is_dir():\n",
    "        continue\n",
    "    \n",
    "    match = OUTDIR_PATTERN.match(child.name)\n",
    "    \n",
    "    if match is None:\n",
    "        continue\n",
    "        \n",
    "    date = datetime.datetime(\n",
    "        year=int(match.group(1)),\n",
    "        month=int(match.group(2)),\n",
    "        day=int(match.group(3)),\n",
    "        hour=int(match.group(4)),\n",
    "    )\n",
    "    \n",
    "    out_path = child / \"output.nc\"\n",
    "    aer_path = (\n",
    "        base / \"inputs\" / \"baseline\" / \"HYDE_BASE_Y2018\" /\n",
    "        f\"OUTPUT_bwd_{date.strftime('%Y%m%d')}\" /\n",
    "        \"EMISSIONS_0422\" /\n",
    "        f\"{date.strftime('%Y%m%d')}_7daybwd_Hyde_traj_AER_{24-date.hour:02}_L3.nc\"\n",
    "    )\n",
    "    ant_path = (\n",
    "        base / \"inputs\" / \"baseline\" / \"HYDE_BASE_Y2018\" /\n",
    "        f\"OUTPUT_bwd_{date.strftime('%Y%m%d')}\" /\n",
    "        \"EMISSIONS_0422\" /\n",
    "        f\"{date.strftime('%Y%m%d')}_7daybwd_Hyde_traj_ANT_{24-date.hour:02}_L3.nc\"\n",
    "    )\n",
    "    bio_path = (\n",
    "        base / \"inputs\" / \"baseline\" / \"HYDE_BASE_Y2018\" /\n",
    "        f\"OUTPUT_bwd_{date.strftime('%Y%m%d')}\" /\n",
    "        \"EMISSIONS_0422\" /\n",
    "        f\"{date.strftime('%Y%m%d')}_7daybwd_Hyde_traj_BIO_{24-date.hour:02}_L3.nc\"\n",
    "    )\n",
    "    met_path = (\n",
    "        base / \"inputs\" / \"baseline\" / \"HYDE_BASE_Y2018\" /\n",
    "        f\"OUTPUT_bwd_{date.strftime('%Y%m%d')}\" /\n",
    "        \"METEO\" /\n",
    "        f\"METEO_{date.strftime('%Y%m%d')}_R{24-date.hour:02}.nc\"\n",
    "    )\n",
    "    \n",
    "    if (\n",
    "        (not out_path.exists()) or (not aer_path.exists()) or\n",
    "        (not ant_path.exists()) or (not bio_path.exists()) or\n",
    "        (not met_path.exists())\n",
    "    ):\n",
    "        raise Exception(out_path, aer_path, ant_path, bio_path, met_path)\n",
    "    \n",
    "    traj_datetimes[date] = TrajectoryPaths(\n",
    "        date=date, out=out_path, aer=aer_path, ant=ant_path, bio=bio_path, met=met_path,\n",
    "    )\n",
    "\n",
    "traj_dates = sorted(set(d.date() for d in traj_datetimes.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5d551d7d",
   "metadata": {
    "papermill": {
     "duration": 0.018564,
     "end_time": "2023-01-25T10:01:13.535960",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.517396",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_trajectory_dataset(paths: TrajectoryPaths) -> TrajectoryDatasets:\n",
    "    outds = Dataset(paths.out, \"r\", format=\"NETCDF4\")\n",
    "    aerds = Dataset(paths.aer, \"r\", format=\"NETCDF4\")\n",
    "    antds = Dataset(paths.ant, \"r\", format=\"NETCDF4\")\n",
    "    biods = Dataset(paths.bio, \"r\", format=\"NETCDF4\")\n",
    "    metds = Dataset(paths.met, \"r\", format=\"NETCDF4\")\n",
    "    \n",
    "    return TrajectoryDatasets(\n",
    "        date=paths.date, out=outds, aer=aerds, ant=antds, bio=biods, met=metds,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c76ceb5a",
   "metadata": {
    "papermill": {
     "duration": 0.017935,
     "end_time": "2023-01-25T10:01:13.568160",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.550225",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_proj = ccrs.PlateCarree()\n",
    "projection = ccrs.LambertConformal(\n",
    "    central_latitude=50, central_longitude=20, standard_parallels=(25, 25)\n",
    ")\n",
    "extent = [-60, 60, 40, 80]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "004cbb83",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ccn_concentration(ds: TrajectoryDatasets):\n",
    "    ccn_bin_indices, = np.nonzero(ds.out[\"dp_dry_fs\"][:].data > 80e-9)\n",
    "    ccn_concentration = np.sum(ds.out[\"nconc_par\"][:].data[:,ccn_bin_indices,:], axis=1)\n",
    "    \n",
    "    return pd.DataFrame({\n",
    "        \"time\": np.repeat(get_output_time(ds), ds.out[\"lev\"].shape[0]),\n",
    "        \"level\": np.tile(ds.out[\"lev\"][:].data, ds.out[\"time\"].shape[0]),\n",
    "        \"ccn\": ccn_concentration.flatten(),\n",
    "    }).set_index([\"time\", \"level\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d6040303",
   "metadata": {
    "papermill": {
     "duration": 0.018385,
     "end_time": "2023-01-25T10:01:13.601713",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.583328",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_output_time(ds: TrajectoryDatasets):\n",
    "    fdom = datetime.datetime.strptime(\n",
    "        ds.out[\"time\"].__dict__[\"first_day_of_month\"], \"%Y-%m-%d %H:%M:%S\",\n",
    "    )\n",
    "    dt = (ds.date - fdom).total_seconds()\n",
    "    \n",
    "    out_t = ds.out[\"time\"][:].data\n",
    "    \n",
    "    return out_t - dt\n",
    "\n",
    "def interpolate_meteorology_values(ds: TrajectoryDatasets, key: str):\n",
    "    out_t = get_output_time(ds)\n",
    "    out_h = ds.out[\"lev\"][:].data\n",
    "    \n",
    "    met_t = ds.met[\"time\"][:].data\n",
    "    met_h = ds.met[\"lev\"][:].data\n",
    "    \n",
    "    met_t_h = ds.met[key][:]\n",
    "    \n",
    "    met_t_h_int = sp.interpolate.interp2d(\n",
    "        x=met_h, y=met_t, z=met_t_h, kind=\"linear\", bounds_error=False, fill_value=0.0,\n",
    "    )\n",
    "    \n",
    "    return met_t_h_int(x=out_h, y=out_t)\n",
    "\n",
    "def interpolate_meteorology_time_values(ds: TrajectoryDatasets, key: str):\n",
    "    out_t = get_output_time(ds)\n",
    "    out_h = ds.out[\"lev\"][:].data\n",
    "    \n",
    "    met_t = ds.met[\"time\"][:].data\n",
    "    \n",
    "    met_t_v = ds.met[key][:]\n",
    "    \n",
    "    met_t_int = sp.interpolate.interp1d(\n",
    "        x=met_t, y=met_t_v, kind=\"linear\", bounds_error=False, fill_value=0.0,\n",
    "    )\n",
    "    \n",
    "    return np.repeat(\n",
    "        met_t_int(x=out_t).reshape(-1, 1),\n",
    "        out_h.shape[0], axis=1,\n",
    "    )\n",
    "\n",
    "def interpolate_biogenic_emissions(ds: TrajectoryDatasets, key: str):\n",
    "    out_t = get_output_time(ds)\n",
    "    out_h = ds.out[\"lev\"][:].data\n",
    "    \n",
    "    # depth of each box layer, assuming level heights are midpoints and end points are clamped\n",
    "    out_d = (np.array(list(out_h[1:])+[out_h[-1]]) - np.array([out_h[0]]+list(out_h[:-1]))) / 2.0\n",
    "    \n",
    "    bio_t = ds.bio[\"time\"][:].data\n",
    "    \n",
    "    # Biogenic emissions are limited to boxes at <= 10m height\n",
    "    biogenic_emission_layers = np.nonzero(out_h <= 10.0)\n",
    "    biogenic_emission_layer_height_cumsum = np.cumsum(out_d[biogenic_emission_layers])\n",
    "    biogenic_emission_layer_proportion = biogenic_emission_layer_height_cumsum / biogenic_emission_layer_height_cumsum[-1]\n",
    "    num_biogenic_emission_layers = sum(out_h <= 10.0)\n",
    "    \n",
    "    bio_t_h = np.zeros(shape=(out_t.size, out_h.size))\n",
    "    \n",
    "    bio_t_int = sp.interpolate.interp1d(\n",
    "        x=bio_t, y=ds.bio[key][:], kind=\"linear\", bounds_error=False, fill_value=0.0,\n",
    "    )\n",
    "    \n",
    "    # Split up the biogenic emissions relative to the depth of the boxes\n",
    "    bio_t_h[:,biogenic_emission_layers] = (\n",
    "        np.tile(bio_t_int(x=out_t), (num_biogenic_emission_layers, 1, 1)) * biogenic_emission_layer_proportion.reshape(-1, 1, 1)\n",
    "    ).T\n",
    "    \n",
    "    return bio_t_h\n",
    "\n",
    "def interpolate_aerosol_emissions(ds: TrajectoryDatasets, key: str):\n",
    "    out_t = get_output_time(ds)\n",
    "    out_h = ds.out[\"lev\"][:].data\n",
    "    \n",
    "    aer_t = ds.aer[\"time\"][:].data\n",
    "    aer_h = ds.aer[\"mid_layer_height\"][:].data\n",
    "    \n",
    "    aer_t_h = ds.aer[key][:].T\n",
    "    \n",
    "    aer_t_h_int = sp.interpolate.interp2d(\n",
    "        x=aer_h, y=aer_t, z=aer_t_h, kind=\"linear\", bounds_error=False, fill_value=0.0,\n",
    "    )\n",
    "    \n",
    "    return aer_t_h_int(x=out_h, y=out_t)\n",
    "\n",
    "def interpolate_anthropogenic_emissions(ds: TrajectoryDatasets, key: str):\n",
    "    out_t = get_output_time(ds)\n",
    "    out_h = ds.out[\"lev\"][:].data\n",
    "    \n",
    "    ant_t = ds.ant[\"time\"][:].data\n",
    "    ant_h = ds.ant[\"mid_layer_height\"][:].data\n",
    "    \n",
    "    ant_t_h = ds.ant[key][:].T\n",
    "    \n",
    "    ant_t_h_int = sp.interpolate.interp2d(\n",
    "        x=ant_h, y=ant_t, z=ant_t_h, kind=\"linear\", bounds_error=False, fill_value=0.0,\n",
    "    )\n",
    "    \n",
    "    return ant_t_h_int(x=out_h, y=out_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "953273d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_meteorology_features(ds: TrajectoryDatasets):\n",
    "    return pd.DataFrame({\n",
    "        \"time\": np.repeat(get_output_time(ds), ds.out[\"lev\"].shape[0]),\n",
    "        \"level\": np.tile(ds.out[\"lev\"][:].data, ds.out[\"time\"].shape[0]),\n",
    "        \"met_t\": interpolate_meteorology_values(ds, \"t\").flatten(),\n",
    "        # \"met_u\": interpolate_meteorology_values(ds, \"u\").flatten(),\n",
    "        # \"met_v\": interpolate_meteorology_values(ds, \"v\").flatten(),\n",
    "        \"met_q\": interpolate_meteorology_values(ds, \"q\").flatten(),\n",
    "        # \"met_qc\": interpolate_meteorology_values(ds, \"qc\").flatten(),\n",
    "        # \"met_sp\": interpolate_meteorology_time_values(ds, \"sp\").flatten(),\n",
    "        # \"met_cp\": interpolate_meteorology_time_values(ds, \"cp\").flatten(),\n",
    "        # \"met_sshf\": interpolate_meteorology_time_values(ds, \"sshf\").flatten(),\n",
    "        \"met_ssr\": interpolate_meteorology_time_values(ds, \"ssr\").flatten(),\n",
    "        # \"met_lsp\": interpolate_meteorology_time_values(ds, \"lsp\").flatten(),\n",
    "        # \"met_ewss\": interpolate_meteorology_time_values(ds, \"ewss\").flatten(),\n",
    "        # \"met_nsss\": interpolate_meteorology_time_values(ds, \"nsss\").flatten(),\n",
    "        # \"met_tcc\": interpolate_meteorology_time_values(ds, \"tcc\").flatten(),\n",
    "        \"met_lsm\": interpolate_meteorology_time_values(ds, \"lsm\").flatten(),\n",
    "        # \"met_omega\": interpolate_meteorology_values(ds, \"omega\").flatten(),\n",
    "        # \"met_z\": interpolate_meteorology_time_values(ds, \"z\").flatten(),\n",
    "        # \"met_mla\": interpolate_meteorology_values(ds, \"mla\").flatten(),\n",
    "        # NOTE: lp is excluded because it allows the model to overfit\n",
    "        # \"met_lp\": interpolate_meteorology_values(ds, \"lp\").flatten(),\n",
    "        \"met_blh\": interpolate_meteorology_time_values(ds, \"blh\").flatten(),\n",
    "    }).set_index([\"time\", \"level\"])\n",
    "\n",
    "def get_bio_emissions_features(ds: TrajectoryDatasets):\n",
    "    return pd.DataFrame({\n",
    "        \"time\": np.repeat(get_output_time(ds), ds.out[\"lev\"].shape[0]),\n",
    "        \"level\": np.tile(ds.out[\"lev\"][:].data, ds.out[\"time\"].shape[0]),\n",
    "        \"bio_acetaldehyde\": interpolate_biogenic_emissions(ds, \"acetaldehyde\").flatten(),\n",
    "        \"bio_acetone\": interpolate_biogenic_emissions(ds, \"acetone\").flatten(),\n",
    "        \"bio_butanes_and_higher_alkanes\": interpolate_biogenic_emissions(ds, \"butanes-and-higher-alkanes\").flatten(),\n",
    "        \"bio_butanes_and_higher_alkenes\": interpolate_biogenic_emissions(ds, \"butenes-and-higher-alkenes\").flatten(),\n",
    "        \"bio_ch4\": interpolate_biogenic_emissions(ds, \"CH4\").flatten(),\n",
    "        \"bio_co\": interpolate_biogenic_emissions(ds, \"CO\").flatten(),\n",
    "        \"bio_ethane\": interpolate_biogenic_emissions(ds, \"ethane\").flatten(),\n",
    "        \"bio_ethanol\": interpolate_biogenic_emissions(ds, \"ethanol\").flatten(),\n",
    "        \"bio_ethene\": interpolate_biogenic_emissions(ds, \"ethene\").flatten(),\n",
    "        \"bio_formaldehyde\": interpolate_biogenic_emissions(ds, \"formaldehyde\").flatten(),\n",
    "        \"bio_hydrogen_cyanide\": interpolate_biogenic_emissions(ds, \"hydrogen-cyanide\").flatten(),\n",
    "        \"bio_iosprene\": interpolate_biogenic_emissions(ds, \"isoprene\").flatten(),\n",
    "        \"bio_mbo\": interpolate_biogenic_emissions(ds, \"MBO\").flatten(),\n",
    "        \"bio_methanol\": interpolate_biogenic_emissions(ds, \"methanol\").flatten(),\n",
    "        \"bio_methyl_bromide\": interpolate_biogenic_emissions(ds, \"methyl-bromide\").flatten(),\n",
    "        \"bio_methyl_chloride\": interpolate_biogenic_emissions(ds, \"methyl-chloride\").flatten(),\n",
    "        \"bio_methyl_iodide\": interpolate_biogenic_emissions(ds, \"methyl-iodide\").flatten(),\n",
    "        \"bio_other_aldehydes\": interpolate_biogenic_emissions(ds, \"other-aldehydes\").flatten(),\n",
    "        \"bio_other_ketones\": interpolate_biogenic_emissions(ds, \"other-ketones\").flatten(),\n",
    "        \"bio_other_monoterpenes\": interpolate_biogenic_emissions(ds, \"other-monoterpenes\").flatten(),\n",
    "        \"bio_pinene_a\": interpolate_biogenic_emissions(ds, \"pinene-a\").flatten(),\n",
    "        \"bio_pinene_b\": interpolate_biogenic_emissions(ds, \"pinene-b\").flatten(),\n",
    "        \"bio_propane\": interpolate_biogenic_emissions(ds, \"propane\").flatten(),\n",
    "        \"bio_propene\": interpolate_biogenic_emissions(ds, \"propene\").flatten(),\n",
    "        \"bio_sesquiterpenes\": interpolate_biogenic_emissions(ds, \"sesquiterpenes\").flatten(),\n",
    "        \"bio_toluene\": interpolate_biogenic_emissions(ds, \"toluene\").flatten(),\n",
    "        \"bio_ch2br2\": interpolate_biogenic_emissions(ds, \"CH2Br2\").flatten(),\n",
    "        \"bio_ch3i\": interpolate_biogenic_emissions(ds, \"CH3I\").flatten(),\n",
    "        \"bio_chbr3\": interpolate_biogenic_emissions(ds, \"CHBr3\").flatten(),\n",
    "        \"bio_dms\": interpolate_biogenic_emissions(ds, \"DMS\").flatten(),\n",
    "    }).set_index([\"time\", \"level\"])\n",
    "\n",
    "def get_aer_emissions_features(ds: TrajectoryDatasets):\n",
    "    return pd.DataFrame({\n",
    "        \"time\": np.repeat(get_output_time(ds), ds.out[\"lev\"].shape[0]),\n",
    "        \"level\": np.tile(ds.out[\"lev\"][:].data, ds.out[\"time\"].shape[0]),\n",
    "        \"aer_3_10_nm\": interpolate_aerosol_emissions(ds, \"3-10nm\").flatten(),\n",
    "        \"aer_10_20_nm\": interpolate_aerosol_emissions(ds, \"10-20nm\").flatten(),\n",
    "        \"aer_20_30_nm\": interpolate_aerosol_emissions(ds, \"20-30nm\").flatten(),\n",
    "        \"aer_30_50_nm\": interpolate_aerosol_emissions(ds, \"30-50nm\").flatten(),\n",
    "        \"aer_50_70_nm\": interpolate_aerosol_emissions(ds, \"50-70nm\").flatten(),\n",
    "        \"aer_70_100_nm\": interpolate_aerosol_emissions(ds, \"70-100nm\").flatten(),\n",
    "        \"aer_100_200_nm\": interpolate_aerosol_emissions(ds, \"100-200nm\").flatten(),\n",
    "        \"aer_200_400_nm\": interpolate_aerosol_emissions(ds, \"200-400nm\").flatten(),\n",
    "        \"aer_400_1000_nm\": interpolate_aerosol_emissions(ds, \"400-1000nm\").flatten(),\n",
    "    }).set_index([\"time\", \"level\"])\n",
    "\n",
    "def get_ant_emissions_features(ds: TrajectoryDatasets):\n",
    "    return pd.DataFrame({\n",
    "        \"time\": np.repeat(get_output_time(ds), ds.out[\"lev\"].shape[0]),\n",
    "        \"level\": np.tile(ds.out[\"lev\"][:].data, ds.out[\"time\"].shape[0]),\n",
    "        \"ant_co\": interpolate_anthropogenic_emissions(ds, \"co\").flatten(),\n",
    "        \"ant_nox\": interpolate_anthropogenic_emissions(ds, \"nox\").flatten(),\n",
    "        \"ant_co2\": interpolate_anthropogenic_emissions(ds, \"co2\").flatten(),\n",
    "        \"ant_nh3\": interpolate_anthropogenic_emissions(ds, \"nh3\").flatten(),\n",
    "        \"ant_ch4\": interpolate_anthropogenic_emissions(ds, \"ch4\").flatten(),\n",
    "        \"ant_so2\": interpolate_anthropogenic_emissions(ds, \"so2\").flatten(),\n",
    "        \"ant_nmvoc\": interpolate_anthropogenic_emissions(ds, \"nmvoc\").flatten(),\n",
    "        \"ant_alcohols\": interpolate_anthropogenic_emissions(ds, \"alcohols\").flatten(),\n",
    "        \"ant_ethane\": interpolate_anthropogenic_emissions(ds, \"ethane\").flatten(),\n",
    "        \"ant_propane\": interpolate_anthropogenic_emissions(ds, \"propane\").flatten(),\n",
    "        \"ant_butanes\": interpolate_anthropogenic_emissions(ds, \"butanes\").flatten(),\n",
    "        \"ant_pentanes\": interpolate_anthropogenic_emissions(ds, \"pentanes\").flatten(),\n",
    "        \"ant_hexanes\": interpolate_anthropogenic_emissions(ds, \"hexanes\").flatten(),\n",
    "        \"ant_ethene\": interpolate_anthropogenic_emissions(ds, \"ethene\").flatten(),\n",
    "        \"ant_propene\": interpolate_anthropogenic_emissions(ds, \"propene\").flatten(),\n",
    "        \"ant_acetylene\": interpolate_anthropogenic_emissions(ds, \"acetylene\").flatten(),\n",
    "        \"ant_isoprene\": interpolate_anthropogenic_emissions(ds, \"isoprene\").flatten(),\n",
    "        \"ant_monoterpenes\": interpolate_anthropogenic_emissions(ds, \"monoterpenes\").flatten(),\n",
    "        \"ant_other_alkenes_and_alkynes\": interpolate_anthropogenic_emissions(ds, \"other-alkenes-and-alkynes\").flatten(),\n",
    "        \"ant_benzene\": interpolate_anthropogenic_emissions(ds, \"benzene\").flatten(),\n",
    "        \"ant_toluene\": interpolate_anthropogenic_emissions(ds, \"toluene\").flatten(),\n",
    "        \"ant_xylene\": interpolate_anthropogenic_emissions(ds, \"xylene\").flatten(),\n",
    "        \"ant_trimethylbenzene\": interpolate_anthropogenic_emissions(ds, \"trimethylbenzene\").flatten(),\n",
    "        \"ant_other_aromatics\": interpolate_anthropogenic_emissions(ds, \"other-aromatics\").flatten(),\n",
    "        \"ant_esters\": interpolate_anthropogenic_emissions(ds, \"esters\").flatten(),\n",
    "        \"ant_ethers\": interpolate_anthropogenic_emissions(ds, \"ethers\").flatten(),\n",
    "        \"ant_formaldehyde\": interpolate_anthropogenic_emissions(ds, \"formaldehyde\").flatten(),\n",
    "        \"ant_other_aldehydes\": interpolate_anthropogenic_emissions(ds, \"other-aldehydes\").flatten(),\n",
    "        \"ant_total_ketones\": interpolate_anthropogenic_emissions(ds, \"total-ketones\").flatten(),\n",
    "        \"ant_total_acids\": interpolate_anthropogenic_emissions(ds, \"total-acids\").flatten(),\n",
    "        \"ant_other_vocs\": interpolate_anthropogenic_emissions(ds, \"other-VOCs\").flatten(),\n",
    "    }).set_index([\"time\", \"level\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7b4e2a63",
   "metadata": {
    "papermill": {
     "duration": 0.029729,
     "end_time": "2023-01-25T10:01:13.769583",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.739854",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# https://stackoverflow.com/a/67809235\n",
    "def df_to_numpy(df):\n",
    "    try:\n",
    "        shape = [len(level) for level in df.index.levels]\n",
    "    except AttributeError:\n",
    "        shape = [len(df.index)]\n",
    "    ncol = df.shape[-1]\n",
    "    if ncol > 1:\n",
    "        shape.append(ncol)\n",
    "    return df.to_numpy().reshape(shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b4a92854",
   "metadata": {
    "papermill": {
     "duration": 0.019778,
     "end_time": "2023-01-25T10:01:13.804285",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.784507",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_time_level_windows():\n",
    "    # -0.5h, -1.5h, -3h, -6h, -12h, -24h, -48h\n",
    "    # 0, -2, -5, -11, -23, -47, -95\n",
    "    time_windows = [(0, 0), (-2, -1), (-5, -3), (-11, -6), (-23, -12), (-47, -24), (-95, -48)]\n",
    "    \n",
    "    # +1l, +2l, +4l, +8l, +16l, +32l, +64\n",
    "    top_windows = [(1, 1), (1, 2), (1, 4), (2, 8), (2, 16), (3, 32), (3, 64)]\n",
    "    mid_windows = [(0, 0), (0, 0), (0, 0), (-1, 1), (-1, 1), (-2, 2), (-2, 2)]\n",
    "    bot_windows = [(-1, -1), (-2, -1), (-4, -1), (-8, -2), (-16, -2), (-32, -3), (-64, -3)]\n",
    "    \n",
    "    return list(itertools.chain(\n",
    "        zip(time_windows, top_windows), zip(time_windows, mid_windows), zip(time_windows, bot_windows),\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b162ef7f",
   "metadata": {
    "papermill": {
     "duration": 0.018687,
     "end_time": "2023-01-25T10:01:13.838147",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.819460",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_windowed_feature_names(columns):\n",
    "    time_windows = [\"-0.5h\", \"-1.5h\", \"-3h\", \"-6h\", \"-12h\", \"-24h\", \"-48h\"]\n",
    "    \n",
    "    top_windows = [\"+1l\", \"+2l\", \"+4l\", \"+8l\", \"+16l\", \"+32l\", \"+64l\"]\n",
    "    mid_windows = [\"+0l\", \"+0l\", \"+0l\", \"±1l\", \"±1l\", \"±2l\", \"±2l\"]\n",
    "    bot_windows = [\"-1l\", \"-2l\", \"-4l\", \"-8l\", \"-16l\", \"-32l\", \"-64l\"]\n",
    "    \n",
    "    names = []\n",
    "    \n",
    "    for (t, l) in itertools.chain(\n",
    "        zip(time_windows, top_windows), zip(time_windows, mid_windows), zip(time_windows, bot_windows),\n",
    "    ):\n",
    "        for c in columns:\n",
    "            names.append(f\"{c}{t}{l}\")\n",
    "    \n",
    "    return names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "07592d92",
   "metadata": {
    "papermill": {
     "duration": 0.023248,
     "end_time": "2023-01-25T10:01:13.876476",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.853228",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def time_level_window_mean_v1(input, t_range, l_range):\n",
    "    output = np.zeros(shape=input.shape)\n",
    "\n",
    "    for t in range(input.shape[0]):\n",
    "        for l in range(input.shape[1]):\n",
    "            for f in range(input.shape[2]):\n",
    "                window = input[\n",
    "                    min(max(0, t+t_range[0]), input.shape[0]):max(0, min(t+1+t_range[1], input.shape[0])),\n",
    "                    min(max(0, l+l_range[0]), input.shape[1]):max(0, min(l+1+l_range[1], input.shape[1])),\n",
    "                    f\n",
    "                ]\n",
    "\n",
    "                output[t,l,f] = np.mean(window) if window.size > 0 else 0.0\n",
    "    \n",
    "    return output\n",
    "\n",
    "def time_level_window_mean_v2(input, t_range, l_range):\n",
    "    output = np.zeros(shape=input.shape)\n",
    "\n",
    "    for t in range(input.shape[0]):\n",
    "        mint = min(max(0, t+t_range[0]), input.shape[0])\n",
    "        maxt = max(0, min(t+1+t_range[1], input.shape[0]))\n",
    "        \n",
    "        if mint == maxt:\n",
    "            continue\n",
    "        \n",
    "        for l in range(input.shape[1]):\n",
    "            minl = min(max(0, l+l_range[0]), input.shape[1])\n",
    "            maxl = max(0, min(l+1+l_range[1], input.shape[1]))\n",
    "            \n",
    "            if minl == maxl:\n",
    "                continue\n",
    "                \n",
    "            output[t,l,:] = np.mean(input[mint:maxt,minl:maxl,:], axis=(0,1))\n",
    "    \n",
    "    return output\n",
    "\n",
    "def time_level_window_mean_v3(input, t_range, l_range):\n",
    "    min_t = min(t_range[0], 0)\n",
    "    max_t = max(0, t_range[1])\n",
    "    abs_t = max(abs(min_t), abs(max_t))\n",
    "    \n",
    "    min_l = min(l_range[0], 0)\n",
    "    max_l = max(0, l_range[1])\n",
    "    abs_l = max(abs(min_l), abs(max_l))\n",
    "    \n",
    "    kernel = np.zeros(shape=(abs_t*2 + 1, abs_l*2 + 1, 1))\n",
    "    kernel[t_range[0]+abs_t:t_range[1]+abs_t+1,l_range[0]+abs_l:l_range[1]+abs_l+1,:] = 1.0\n",
    "    kernel = kernel[::-1,::-1]\n",
    "    \n",
    "    quot = sp.ndimage.convolve(np.ones_like(input), kernel, mode='constant', cval=0.0)\n",
    "    \n",
    "    result = np.zeros_like(input)\n",
    "    \n",
    "    np.divide(\n",
    "        sp.ndimage.convolve(input, kernel, mode='constant', cval=0.0),\n",
    "        quot, out=result, where=quot > 0,\n",
    "    )\n",
    "    \n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4243a406",
   "metadata": {
    "papermill": {
     "duration": 0.018056,
     "end_time": "2023-01-25T10:01:13.909591",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.891535",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_raw_features_for_dataset(ds: TrajectoryDatasets):\n",
    "    bio_features = get_bio_emissions_features(ds)\n",
    "    aer_features = get_aer_emissions_features(ds) * 1e21\n",
    "    ant_features = get_ant_emissions_features(ds)\n",
    "    met_features = get_meteorology_features(ds)\n",
    "    \n",
    "    return pd.concat([\n",
    "        bio_features, aer_features, ant_features, met_features,\n",
    "    ], axis=\"columns\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "344cf4a1",
   "metadata": {
    "papermill": {
     "duration": 0.019594,
     "end_time": "2023-01-25T10:01:13.943986",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.924392",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_features_from_raw_features(raw_features):\n",
    "    raw_features_np = df_to_numpy(raw_features)\n",
    "    \n",
    "    features_np = np.concatenate([\n",
    "        raw_features.index.get_level_values(0).to_numpy().reshape(\n",
    "            (raw_features.index.levels[0].size, raw_features.index.levels[1].size, 1)\n",
    "        ),\n",
    "        raw_features.index.get_level_values(1).to_numpy().reshape(\n",
    "            (raw_features.index.levels[0].size, raw_features.index.levels[1].size, 1)\n",
    "        )\n",
    "    ] + joblib.Parallel(n_jobs=-1)([\n",
    "        joblib.delayed(time_level_window_mean_v2)(raw_features_np, t, l) for t, l in generate_time_level_windows()\n",
    "    ]), axis=2)\n",
    "    \n",
    "    # Trim off the first two days, for which the time features are ill-defined\n",
    "    features_np_trimmed = features_np[95:-1,:,:]\n",
    "    \n",
    "    feature_names = [\"time\", \"level\"] + generate_windowed_feature_names(raw_features.columns)\n",
    "    \n",
    "    features = pd.DataFrame(features_np_trimmed.reshape(\n",
    "        features_np_trimmed.shape[0]*features_np_trimmed.shape[1], features_np_trimmed.shape[2],\n",
    "    ), columns=feature_names).set_index([\"time\", \"level\"])\n",
    "    \n",
    "    return features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "90da30b1",
   "metadata": {
    "papermill": {
     "duration": 0.01929,
     "end_time": "2023-01-25T10:01:13.978350",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.959060",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_labels_for_dataset(ds: TrajectoryDatasets):\n",
    "    ccn_concentration = get_ccn_concentration(ds)\n",
    "    \n",
    "    ccn_concentration_np = df_to_numpy(ccn_concentration)\n",
    "    \n",
    "    labels_np = np.concatenate([\n",
    "        ccn_concentration.index.get_level_values(0).to_numpy().reshape(\n",
    "            (ccn_concentration.index.levels[0].size, ccn_concentration.index.levels[1].size, 1)\n",
    "        ),\n",
    "        ccn_concentration.index.get_level_values(1).to_numpy().reshape(\n",
    "            (ccn_concentration.index.levels[0].size, ccn_concentration.index.levels[1].size, 1)\n",
    "        ),\n",
    "        ccn_concentration_np.reshape(\n",
    "            (ccn_concentration_np.shape[0], ccn_concentration_np.shape[1], 1)\n",
    "        ),\n",
    "    ], axis=2)\n",
    "    \n",
    "    # Trim off the first two days, for which the time features are ill-defined\n",
    "    labels_np_trimmed = labels_np[96:,:,:]\n",
    "    \n",
    "    label_names = [\"time\", \"level\", \"ccn\"]\n",
    "    \n",
    "    labels = pd.DataFrame(labels_np_trimmed.reshape(\n",
    "        labels_np_trimmed.shape[0]*labels_np_trimmed.shape[1], labels_np_trimmed.shape[2],\n",
    "    ), columns=label_names).set_index([\"time\", \"level\"])\n",
    "    \n",
    "    return labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4ef8315b",
   "metadata": {
    "papermill": {
     "duration": 0.018426,
     "end_time": "2023-01-25T10:01:14.013864",
     "exception": false,
     "start_time": "2023-01-25T10:01:13.995438",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def hash_for_dt(dt):\n",
    "    if not(isinstance(dt, tuple) or isinstance(dt, list)):\n",
    "        dt = [dt]\n",
    "    \n",
    "    dt_str = '.'.join(dtt.strftime('%d.%m.%Y-%H:00%z') for dtt in dt)\n",
    "    \n",
    "    h = hashlib.shake_256()\n",
    "    h.update(dt_str.encode('ascii'))\n",
    "    \n",
    "    return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "8f5e2e8a",
   "metadata": {
    "papermill": {
     "duration": 0.019423,
     "end_time": "2023-01-25T10:01:14.048460",
     "exception": false,
     "start_time": "2023-01-25T10:01:14.029037",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Clumped 0/1 sampler using a Markov Process\n",
    "\n",
    "P(0) = p and P(1) = 1-p\n",
    "clump = 0 => IID samples\n",
    "clump -> 1 => highly correlated samples\n",
    "\n",
    "\"\"\"\n",
    "class Clump:\n",
    "    def __init__(self, p=0.5, clump=0.0, rng=None):\n",
    "        a = 1 - (1-p)*(1-clump)\n",
    "        b = (1-a)*p/(1-p)\n",
    "        \n",
    "        self.C = np.array([[a, 1-a],[b, 1-b]])\n",
    "        \n",
    "        self.i = 0 if rng.random() < p else 1\n",
    "    \n",
    "    def sample(self, rng):\n",
    "        p = self.C[self.i,0]\n",
    "        u = rng.random()\n",
    "        \n",
    "        self.i = 0 if u < p else 1\n",
    "        \n",
    "        return self.i\n",
    "    \n",
    "    def steady(self, X):\n",
    "        return np.matmul(X, self.C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "607a98ef",
   "metadata": {
    "papermill": {
     "duration": 0.021247,
     "end_time": "2023-01-25T10:01:14.084943",
     "exception": false,
     "start_time": "2023-01-25T10:01:14.063696",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_test_split(X, Y, test_size=0.25, random_state=None, shuffle=True, clump=0.0):\n",
    "    assert len(X) == len(Y)\n",
    "    assert type(X) == type(Y)\n",
    "    assert test_size > 0.0\n",
    "    assert test_size < 1.0\n",
    "    assert random_state is not None\n",
    "    assert clump >= 0.0\n",
    "    assert clump < 1.0\n",
    "    \n",
    "    c = Clump(p=test_size, clump=clump, rng=random_state)\n",
    "    \n",
    "    if isinstance(X, pd.DataFrame):\n",
    "        assert X.index.values.shape == Y.index.values.shape\n",
    "        \n",
    "        # Split only based on the first-level index instead of flattening\n",
    "        n1 = len(X.index.levels[1])\n",
    "        n0 = len(X) // n1\n",
    "        \n",
    "        C = np.array([c.sample(random_state) for _ in range(n0)])\n",
    "        I_train, = np.nonzero(C)\n",
    "        I_train = np.repeat(I_train, n1) * n1 + np.tile(np.arange(n1), len(I_train))\n",
    "        I_test, = np.nonzero(1-C)\n",
    "        I_test = np.repeat(I_test, n1) * n1 + np.tile(np.arange(n1), len(I_test))\n",
    "    else:\n",
    "        C = np.array([c.sample(random_state) for _ in range(len(X))])\n",
    "        I_train, = np.nonzero(C)\n",
    "        I_test, = np.nonzero(1-C)\n",
    "    \n",
    "    if shuffle:\n",
    "        random_state.shuffle(I_train)\n",
    "        random_state.shuffle(I_test)\n",
    "    \n",
    "    if isinstance(X, pd.DataFrame):\n",
    "        X_train = X.iloc[I_train]\n",
    "        X_test = X.iloc[I_test]\n",
    "        \n",
    "        Y_train = Y.iloc[I_train]\n",
    "        Y_test = Y.iloc[I_test]\n",
    "    else:\n",
    "        X_train = X[I_train]\n",
    "        X_test = X[I_test]\n",
    "        \n",
    "        Y_train = Y[I_train]\n",
    "        Y_test = Y[I_test]\n",
    "    \n",
    "    return X_train, X_test, Y_train, Y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "0b92aa97",
   "metadata": {
    "papermill": {
     "duration": 0.024469,
     "end_time": "2023-01-25T10:01:14.124122",
     "exception": false,
     "start_time": "2023-01-25T10:01:14.099653",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_and_cache_dataset(dt: datetime.datetime, clump: float, datasets: dict) -> MLDataset:\n",
    "    if isinstance(dt, tuple) or isinstance(dt, list):\n",
    "        dt = tuple(sorted(dt))\n",
    "    \n",
    "    cached = datasets.get((dt, clump))\n",
    "    \n",
    "    if cached is not None:\n",
    "        return cached\n",
    "    \n",
    "    if isinstance(dt, tuple) or isinstance(dt, list):\n",
    "        mls = [load_and_cache_dataset(dtt, clump, datasets) for dtt in dt]\n",
    "\n",
    "        dp = tuple(ml.paths for ml in mls)\n",
    "        X_raw = pd.concat([ml.X_raw for ml in mls], axis='index')\n",
    "        Y = pd.concat([ml.Y_raw for ml in mls], axis='index')\n",
    "\n",
    "        train_features = np.concatenate([ml.X_scaler.inverse_transform(ml.X_train) for ml in mls], axis=0)\n",
    "        train_labels = np.concatenate([ml.Y_scaler.inverse_transform(ml.Y_train) for ml in mls], axis=0)\n",
    "        valid_features = np.concatenate([ml.X_scaler.inverse_transform(ml.X_valid) for ml in mls], axis=0)\n",
    "        valid_labels = np.concatenate([ml.Y_scaler.inverse_transform(ml.Y_valid) for ml in mls], axis=0)\n",
    "        test_features = np.concatenate([ml.X_scaler.inverse_transform(ml.X_test) for ml in mls], axis=0)\n",
    "        test_labels = np.concatenate([ml.Y_scaler.inverse_transform(ml.Y_test) for ml in mls], axis=0)\n",
    "    else:\n",
    "        dp = traj_datetimes[dt]\n",
    "        ds = load_trajectory_dataset(dp)\n",
    "\n",
    "        X_raw = get_raw_features_for_dataset(ds)\n",
    "\n",
    "        X = get_features_from_raw_features(X_raw)\n",
    "        Y = np.log10(get_labels_for_dataset(ds) + 1)\n",
    "\n",
    "        rng = np.random.RandomState(seed=int.from_bytes(hash_for_dt(dt).digest(4), 'little'))\n",
    "\n",
    "        train_features, test_features, train_labels, test_labels = train_test_split(\n",
    "            X, Y, test_size=0.25, random_state=rng, clump=clump,\n",
    "        )\n",
    "        train_features, valid_features, train_labels, valid_labels = train_test_split(\n",
    "            train_features, train_labels, test_size=1.0/3.0, random_state=rng, clump=clump,\n",
    "        )\n",
    "\n",
    "        # Close the NetCDF dataset\n",
    "        ds.out.close()\n",
    "\n",
    "    # Scale features to N(0,1)\n",
    "    # - only fit on training data\n",
    "    # - OOD inputs for constants at training time are blown up\n",
    "    feature_scaler = StandardScaler().fit(train_features)\n",
    "    feature_scaler.scale_[np.nonzero(feature_scaler.var_ == 0.0)] = np.nan_to_num(np.inf)\n",
    "\n",
    "    label_scaler = StandardScaler().fit(train_labels)\n",
    "\n",
    "    train_features = feature_scaler.transform(train_features)\n",
    "    train_labels = label_scaler.transform(train_labels)\n",
    "    valid_features = feature_scaler.transform(valid_features)\n",
    "    valid_labels = label_scaler.transform(valid_labels)\n",
    "    test_features = feature_scaler.transform(test_features)\n",
    "    test_labels = label_scaler.transform(test_labels)\n",
    "\n",
    "    dataset = MLDataset(\n",
    "        date=dt, paths=dp, X_raw=X_raw, Y_raw=Y,\n",
    "        X_train=train_features, X_valid=valid_features, X_test=test_features,\n",
    "        Y_train=train_labels, Y_valid=valid_labels, Y_test=test_labels,\n",
    "        X_scaler=feature_scaler, Y_scaler=label_scaler,\n",
    "    )\n",
    "\n",
    "    datasets[(dt, clump)] = dataset\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "72c8216f",
   "metadata": {
    "papermill": {
     "duration": 0.017777,
     "end_time": "2023-01-25T10:01:14.157251",
     "exception": false,
     "start_time": "2023-01-25T10:01:14.139474",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "DATASETS = dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "56f2a7dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import abc\n",
    "\n",
    "class TrainableModel(abc.ABC):\n",
    "    @abc.abstractmethod\n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        return self\n",
    "    \n",
    "    @abc.abstractmethod\n",
    "    def predict(self, X_test):\n",
    "        return None\n",
    "    \n",
    "    def predict_uncertainty(self, X_test):\n",
    "        return np.zeros_like(self.predict(X_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "5fe668a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "class LinearModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.model = LinearRegression(*args, **kwargs)\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        self.Y_shape = list(Y_train.shape[1:])\n",
    "        \n",
    "        self.model.fit(\n",
    "            np.concatenate([X_train, X_valid], axis=0),\n",
    "            np.concatenate([Y_train, Y_valid], axis=0),\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        return self.model.predict(X_test).reshape([len(X_test)] + self.Y_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "cc7288cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import ARDRegression\n",
    "\n",
    "class SparseBayesianModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.model = ARDRegression(*args, **kwargs)\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        self.Y_shape = list(Y_train.shape[1:])\n",
    "        \n",
    "        self.model.fit(\n",
    "            np.concatenate([X_train, X_valid], axis=0),\n",
    "            np.concatenate([Y_train, Y_valid], axis=0),\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        return self.model.predict(X_test).reshape([len(X_test)] + self.Y_shape)\n",
    "    \n",
    "    def predict_uncertainty(self, X_test):\n",
    "        return self.model.predict(X_test, return_std=True)[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "0dd80fe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.tree import DecisionTreeRegressor\n",
    "\n",
    "class DecisionTreeModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.model = DecisionTreeRegressor(*args, random_state=rng, **kwargs)\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        self.Y_shape = list(Y_train.shape[1:])\n",
    "        \n",
    "        self.model.fit(\n",
    "            np.concatenate([X_train, X_valid], axis=0),\n",
    "            np.concatenate([Y_train, Y_valid], axis=0),\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        return self.model.predict(X_test).reshape([len(X_test)] + self.Y_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "88825443",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor\n",
    "\n",
    "class RandomForestModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.model = RandomForestRegressor(*args, random_state=rng, **kwargs)\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        self.Y_shape = list(Y_train.shape[1:])\n",
    "        \n",
    "        self.model.fit(\n",
    "            np.concatenate([X_train, X_valid], axis=0),\n",
    "            np.concatenate([Y_train, Y_valid], axis=0),\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        return self.model.predict(X_test).reshape([len(X_test)] + self.Y_shape)\n",
    "    \n",
    "    def predict_uncertainty(self, X_test):\n",
    "        predictions = [\n",
    "            estimator.predict(X_test) for estimator in self.model.estimators_\n",
    "        ]\n",
    "        \n",
    "        return np.std(np.stack(predictions, axis=0), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "249d98c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor\n",
    "\n",
    "class PairwiseRandomForestModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.model = RandomForestRegressor(*args, random_state=rng, **kwargs)\n",
    "        self.rng = rng\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        self.Y_shape = list(Y_train.shape[1:])\n",
    "        \n",
    "        self.X = np.concatenate([X_train, X_valid], axis=0)\n",
    "        self.Y = np.concatenate([Y_train, Y_valid], axis=0)\n",
    "        \n",
    "        I1 = self.rng.choice(len(self.X), size=len(self.X))\n",
    "        I2 = self.rng.choice(len(self.X), size=len(self.X))\n",
    "        \n",
    "        self.model.fit(\n",
    "            np.concatenate([self.X[I1], self.X[I2]-self.X[I1]], axis=1),\n",
    "            self.Y[I2]-self.Y[I1],\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        predictions = []\n",
    "        \n",
    "        for estimator in self.model.estimators_:\n",
    "            I2 = self.rng.choice(len(self.X), size=len(X_test))\n",
    "            \n",
    "            predictions.append(self.Y[I2] - estimator.predict(\n",
    "                np.concatenate([X_test, self.X[I2]-X_test], axis=1),\n",
    "            ).reshape([len(X_test)] + self.Y_shape))\n",
    "        \n",
    "        return np.mean(np.stack(predictions, axis=0), axis=0)\n",
    "    \n",
    "    def predict_uncertainty(self, X_test):\n",
    "        predictions = []\n",
    "        \n",
    "        for estimator in self.model.estimators_:\n",
    "            I2 = self.rng.choice(len(self.X), size=len(X_test))\n",
    "            \n",
    "            predictions.append(self.Y[I2] - estimator.predict(\n",
    "                np.concatenate([X_test, self.X[I2]-X_test], axis=1),\n",
    "            ).reshape([len(X_test)] + self.Y_shape))\n",
    "        \n",
    "        return np.std(np.stack(predictions, axis=0), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "695f8f40",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.gaussian_process import GaussianProcessRegressor\n",
    "from sklearn.gaussian_process.kernels import RBF, WhiteKernel\n",
    "\n",
    "class GaussianProcessModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.model = GaussianProcessRegressor(*args, kernel=(\n",
    "            RBF() + WhiteKernel()\n",
    "        ), random_state=rng, **kwargs)\n",
    "        self.rng = rng\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        self.Y_shape = list(Y_train.shape[1:])\n",
    "        \n",
    "        num_X = len(X_train) + len(X_valid)\n",
    "        \n",
    "        # Safety hatch to limit computing time\n",
    "        if num_X < 10_000:\n",
    "            I = np.arange(num_X)\n",
    "        else:\n",
    "            I = self.rng.choice(num_X, size=num_X//6)\n",
    "        \n",
    "        self.model.fit(\n",
    "            np.concatenate([X_train, X_valid], axis=0)[I],\n",
    "            np.concatenate([Y_train, Y_valid], axis=0)[I],\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        return self.model.predict(X_test).reshape([len(X_test)] + self.Y_shape)\n",
    "    \n",
    "    def predict_uncertainty(self, X_test):\n",
    "        return self.model.predict(X_test, return_std=True)[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "0984faf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neural_network import MLPRegressor\n",
    "\n",
    "class MLPModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.model = MLPRegressor(*args, random_state=rng, **kwargs)\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        self.Y_shape = list(Y_train.shape[1:])\n",
    "        \n",
    "        self.model.fit(\n",
    "            np.concatenate([X_train, X_valid], axis=0),\n",
    "            np.concatenate([Y_train, Y_valid], axis=0),\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        return self.model.predict(X_test).reshape([len(X_test)] + self.Y_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "8954224a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "class NormalisedMLPModel(TrainableModel):\n",
    "    def __init__(self, *args, rng=None, **kwargs):\n",
    "        self.rng = rng\n",
    "    \n",
    "    def fit(self, X_train, Y_train, X_valid, Y_valid):\n",
    "        # Transform [-inf, +inf] into N(0, 1)\n",
    "        def force_normal(x):\n",
    "            return tf.math.erfinv(tf.math.softsign(x)) * tf.math.sqrt(2.0)\n",
    "        \n",
    "        tf.random.set_seed(int.from_bytes(self.rng.bytes(4), \"little\"))\n",
    "        \n",
    "        self.model = tf.keras.Sequential(\n",
    "            [\n",
    "                tf.keras.layers.InputLayer(input_shape=X_train.shape[1:]),\n",
    "                tf.keras.layers.Dense(128, activation=\"relu\", name=\"dense-1\"),\n",
    "                tf.keras.layers.Dense(128, activation=\"relu\", name=\"dense-2\"),\n",
    "                tf.keras.layers.Dense(128, activation=\"relu\", name=\"dense-3\"),\n",
    "                tf.keras.layers.Dense(Y_train.shape[1], activation=force_normal, name=\"output\"),\n",
    "            ]\n",
    "        )\n",
    "        self.model.compile(optimizer='adam', loss=\"mse\", metrics=[\"mse\", \"mae\"])\n",
    "        \n",
    "        self.model.fit(\n",
    "            X_train, Y_train, validation_data=(X_valid, Y_valid),\n",
    "            batch_size=200, epochs=200, verbose=1, callbacks=[\n",
    "                tf.keras.callbacks.EarlyStopping(\n",
    "                    monitor='val_loss',\n",
    "                    patience=10,\n",
    "                    restore_best_weights=True,\n",
    "                )\n",
    "            ],\n",
    "        )\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X_test):\n",
    "        return self.model.predict(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "f511a726",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_cache_model(dt: datetime.datetime, clump: float, datasets: dict, models: dict, cls, *args, **kwargs):\n",
    "    if isinstance(dt, tuple) or isinstance(dt, list):\n",
    "        dt = tuple(sorted(dt))\n",
    "    \n",
    "    model_key = (cls.__name__, dt, clump)\n",
    "    \n",
    "    cached = models.get(model_key)\n",
    "    \n",
    "    if cached is not None:\n",
    "        return cached\n",
    "    \n",
    "    model_path = f\"{cls.__name__.lower()}.{hash_for_dt(dt).hexdigest(8)}.{clump}.jl\"\n",
    "    \n",
    "    if Path(model_path).exists():    \n",
    "        model = joblib.load(model_path)\n",
    "        \n",
    "        models[model_key] = model\n",
    "        \n",
    "        return model\n",
    "    \n",
    "    dataset = load_and_cache_dataset(dt, clump, datasets)\n",
    "    \n",
    "    rng = np.random.RandomState(seed=int.from_bytes(hash_for_dt(dt).digest(4), 'little'))\n",
    "    \n",
    "    model = cls(*args, rng=rng, **kwargs).fit(\n",
    "        X_train=dataset.X_train, Y_train=dataset.Y_train,\n",
    "        X_valid=dataset.X_valid, Y_valid=dataset.Y_valid,\n",
    "    )\n",
    "    \n",
    "    joblib.dump(model, model_path)\n",
    "    \n",
    "    models[model_key] = model\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "da5b0801",
   "metadata": {},
   "outputs": [],
   "source": [
    "dt_train = datetime.datetime(year=2018, month=5, day=15, hour=19)\n",
    "dt_test = [\n",
    "    datetime.datetime(year=2018, month=5, day=14, hour=10),\n",
    "    datetime.datetime(year=2018, month=5, day=15, hour=19),\n",
    "    datetime.datetime(year=2018, month=5, day=17, hour=0),\n",
    "    datetime.datetime(year=2018, month=5, day=19, hour=4),\n",
    "    datetime.datetime(year=2018, month=5, day=21, hour=15),\n",
    "    datetime.datetime(year=2018, month=5, day=23, hour=13),\n",
    "]\n",
    "\n",
    "ds_train = load_and_cache_dataset(dt_train, 0.75, DATASETS)\n",
    "ds_test = load_and_cache_dataset(dt_test, 0.75, DATASETS)\n",
    "\n",
    "X_train_full = get_features_from_raw_features(ds_train.X_raw.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "9d6570be",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2204/4169729806.py:17: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel().\n",
      "  self.model.fit(\n",
      "/workspace/msc-tcm/miniconda3/envs/demo/lib/python3.10/site-packages/sklearn/base.py:450: UserWarning: X does not have valid feature names, but StandardScaler was fitted with feature names\n",
      "  warnings.warn(\n",
      "/workspace/msc-tcm/miniconda3/envs/demo/lib/python3.10/site-packages/sklearn/base.py:450: UserWarning: X does not have valid feature names, but StandardScaler was fitted with feature names\n",
      "  warnings.warn(\n",
      "/tmp/ipykernel_2204/4169729806.py:17: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel().\n",
      "  self.model.fit(\n",
      "/tmp/ipykernel_2204/161812582.py:10: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel().\n",
      "  self.model.fit(\n",
      "/workspace/msc-tcm/miniconda3/envs/demo/lib/python3.10/site-packages/sklearn/base.py:450: UserWarning: X does not have valid feature names, but StandardScaler was fitted with feature names\n",
      "  warnings.warn(\n",
      "/workspace/msc-tcm/miniconda3/envs/demo/lib/python3.10/site-packages/sklearn/base.py:450: UserWarning: X does not have valid feature names, but StandardScaler was fitted with feature names\n",
      "  warnings.warn(\n",
      "/tmp/ipykernel_2204/161812582.py:10: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel().\n",
      "  self.model.fit(\n"
     ]
    }
   ],
   "source": [
    "MODELS = dict()\n",
    "\n",
    "for cls, args, kwargs in [\n",
    "    (LinearModel, [], dict()),\n",
    "    (SparseBayesianModel, [], dict()),\n",
    "    (DecisionTreeModel, [], dict()),\n",
    "    (PairwiseRandomForestModel, [], dict(n_estimators=10, min_samples_leaf=5, max_features=1.0/3.0)),\n",
    "    (RandomForestModel, [], dict(n_estimators=10, min_samples_leaf=5, max_features=1.0/3.0)),\n",
    "    (GaussianProcessModel, [], dict()),\n",
    "    (MLPModel, [], dict(\n",
    "       hidden_layer_sizes=[128, 128, 128], activation=\"relu\", solver=\"adam\",\n",
    "    )),\n",
    "    (NormalisedMLPModel, [], dict()),\n",
    "]:      \n",
    "    model = train_and_cache_model(dt_train, 0.75, DATASETS, MODELS, cls, *args, **kwargs)\n",
    "    \n",
    "    Y_pred_np = np.concatenate([\n",
    "        ds_train.Y_raw.index.get_level_values(0).to_numpy().reshape(\n",
    "            (ds_train.Y_raw.index.levels[0].size, ds_train.Y_raw.index.levels[1].size, 1)\n",
    "        ),\n",
    "        ds_train.Y_raw.index.get_level_values(1).to_numpy().reshape(\n",
    "            (ds_train.Y_raw.index.levels[0].size, ds_train.Y_raw.index.levels[1].size, 1)\n",
    "        ),\n",
    "        ds_train.Y_scaler.inverse_transform(\n",
    "            model.predict(ds_train.X_scaler.transform(X_train_full))\n",
    "        ).reshape(\n",
    "            (ds_train.Y_raw.index.levels[0].size, ds_train.Y_raw.index.levels[1].size, 1)\n",
    "        ),\n",
    "        (model.predict_uncertainty(\n",
    "            ds_train.X_scaler.transform(X_train_full)\n",
    "        ) * ds_train.Y_scaler.scale_).reshape(\n",
    "            (ds_train.Y_raw.index.levels[0].size, ds_train.Y_raw.index.levels[1].size, 1)\n",
    "        ),\n",
    "    ], axis=2)\n",
    "    Y_pred = pd.DataFrame(Y_pred_np.reshape(\n",
    "        Y_pred_np.shape[0]*Y_pred_np.shape[1], Y_pred_np.shape[2],\n",
    "    ), columns=[\"time\", \"level\", \"ccn\", \"ccn-std\"]).set_index([\"time\", \"level\"])\n",
    "    \n",
    "    fig, axs = plt.subplots(5, 1, figsize=(6, 4))\n",
    "    \n",
    "    fig.tight_layout()\n",
    "    plt.subplots_adjust(wspace=None, hspace=0.0)\n",
    "    \n",
    "    for ax in axs[:-1]:\n",
    "        ax.get_xaxis().set_visible(False)\n",
    "    \n",
    "    ccn_concentration = ds_train.Y_raw\n",
    "\n",
    "    level_mask = ccn_concentration.index.get_level_values(1)\n",
    "    level_heights = ccn_concentration.index.levels[1]\n",
    "    \n",
    "    lis = [0]*5 + [1]*7 + [2]*8 + [3]*20 + [4]*10\n",
    "    \n",
    "    colours = height_cmap(np.linspace(0, 1, len(level_heights)))\n",
    "    \n",
    "    for l, h in enumerate(level_heights):\n",
    "        if l % 6 != 0:\n",
    "            continue\n",
    "        \n",
    "        axs[lis[l]].fill_between(\n",
    "            Y_pred[level_mask == h].index.get_level_values(0) / (60*60),\n",
    "            np.power(10.0, Y_pred[level_mask == h][\"ccn\"] - Y_pred[level_mask == h][\"ccn-std\"]) - 1,\n",
    "            np.power(10.0, Y_pred[level_mask == h][\"ccn\"] + Y_pred[level_mask == h][\"ccn-std\"]) - 1,\n",
    "            color=colours[l], alpha=0.35,\n",
    "        )\n",
    "    \n",
    "    ccn_lines = []\n",
    "    \n",
    "    for l, h in enumerate(level_heights):\n",
    "        if l % 6 != 0:\n",
    "            continue\n",
    "        \n",
    "        ccn_lines.append((axs[lis[l]].plot(\n",
    "            Y_pred[level_mask == h].index.get_level_values(0) / (60*60),\n",
    "            np.power(10.0, Y_pred[level_mask == h][\"ccn\"]) - 1,\n",
    "            color=adjust_lightness(colours[l], 0.4), label=f\"{int(level_heights[l])}m\"\n",
    "        ), colours[l]))\n",
    "        \n",
    "    true_lines = []\n",
    "        \n",
    "    for l, h in enumerate(level_heights):\n",
    "        if l % 6 != 0:\n",
    "            continue\n",
    "        \n",
    "        true_lines.append(axs[lis[l]].plot(\n",
    "            ccn_concentration[level_mask == h].index.get_level_values(0) / (60*60),\n",
    "            np.power(10.0, ccn_concentration[level_mask == h][\"ccn\"]) - 1,\n",
    "            c=replace_lightness(colours[l], 0.1), ls=(0, (1, 3)),\n",
    "        )[0])\n",
    "        axs[lis[l]].plot(\n",
    "            ccn_concentration[level_mask == h].index.get_level_values(0) / (60*60),\n",
    "            np.power(10.0, ccn_concentration[level_mask == h][\"ccn\"]) - 1,\n",
    "            c=replace_lightness(colours[l], 0.9), ls=(2, (1, 3)),\n",
    "        )\n",
    "        \n",
    "    axs[-1].set_xticks([\n",
    "        h for h in range(\n",
    "            int((ccn_concentration.index.levels[0][0] // (60*60))), 0, 24\n",
    "        )\n",
    "    ])\n",
    "    axs[-1].set_xticklabels([\n",
    "        (ds_train.date + datetime.timedelta(hours=h)).strftime(\"%d.%m\") for h in range(\n",
    "            int((ccn_concentration.index.levels[0][0] // (60*60))), 0, 24\n",
    "        )\n",
    "    ])\n",
    "    \n",
    "    for ax in axs:\n",
    "        ax.set_xlim((\n",
    "            ccn_concentration.index.levels[0][0] / (60*60) - 1.5,\n",
    "            ccn_concentration.index.levels[0][-1] / (60*60) + 1.5,\n",
    "        ))\n",
    "        ax.set_yscale(\"log\")\n",
    "        \n",
    "        with warnings.catch_warnings():\n",
    "            warnings.simplefilter(\"ignore\")\n",
    "            labelLines(ax.get_lines(), zorder=2.5, align=False, alpha=0.75)\n",
    "        \n",
    "    for ls, c in ccn_lines:\n",
    "        for l in ls:\n",
    "            l.set_color(c)\n",
    "        \n",
    "    ylim = list(axs[-1].get_ylim())\n",
    "    ylim[0] *= 0.75\n",
    "    axs[-1].scatter(\n",
    "        ccn_concentration.index.levels[0] / (60*60),\n",
    "        [ylim[0]] * ccn_concentration.index.levels[0].size,\n",
    "        c=time_cmap(np.linspace(0, 1, ccn_concentration.index.levels[0].size)),\n",
    "    )\n",
    "    axs[-1].set_ylim(ylim)\n",
    "    \n",
    "    _, _, train_indices, test_indices = train_test_split(\n",
    "        np.array(ccn_concentration.index.levels[0]),\n",
    "        np.arange(len(ccn_concentration.index.levels[0])),\n",
    "        test_size=0.25, clump=0.75, shuffle=False,\n",
    "        random_state=np.random.RandomState(\n",
    "            seed=int.from_bytes(hash_for_dt(dt_train).digest(4), 'little'),\n",
    "        ),\n",
    "    )\n",
    "    \n",
    "    train_mask = np.ones(len(ccn_concentration.index.levels[0]))\n",
    "    train_mask[test_indices] = 0\n",
    "\n",
    "    for ax in axs:\n",
    "        ylim = ax.get_ylim()\n",
    "        ax.fill_between(\n",
    "            ccn_concentration.index.levels[0] / (60*60),\n",
    "            ylim[0], ylim[1], where=train_mask, step=\"mid\",\n",
    "            alpha=0.1, color=\"grey\", zorder=-2,\n",
    "        )\n",
    "        \n",
    "    axs[-1].legend(\n",
    "        [ccn_lines[-1][0][0], true_lines[-1]], [\"predicted\", \"true CCN\"],\n",
    "        loc=\"lower left\", ncol=2, handlelength=1.0, columnspacing=1.0,\n",
    "        bbox_to_anchor=(0.0,0.05)\n",
    "    )\n",
    "    \n",
    "    axs[0].set_title(f\"{cls.__name__} predictions for the training trajectory\")\n",
    "    axs[2].set_ylabel(\"CCN concentration [m${}^{-3}$]\")\n",
    "    \n",
    "    plt.savefig(f\"{cls.__name__.lower()}-training-prediction.pdf\", dpi=100, transparent=True, bbox_inches='tight')\n",
    "    # plt.show()\n",
    "    plt.close(fig)\n",
    "    \n",
    "    fig, ax1 = plt.subplots(1, 1, figsize=(4, 4))\n",
    "    \n",
    "    Y_true = ds_test.Y_scaler.inverse_transform(ds_test.Y_test)\n",
    "    Y_pred = ds_train.Y_scaler.inverse_transform(model.predict(\n",
    "        ds_train.X_scaler.transform(ds_test.X_scaler.inverse_transform(ds_test.X_test))\n",
    "    ))\n",
    "    Y_stdv = model.predict_uncertainty(\n",
    "        ds_train.X_scaler.transform(ds_test.X_scaler.inverse_transform(ds_test.X_test))\n",
    "    )\n",
    "    \n",
    "    Y_valid = (Y_pred.flatten() > -250) & (Y_pred.flatten() < 250)\n",
    "    \n",
    "    umin = np.amin(Y_stdv[Y_valid])\n",
    "    umax = np.amax(Y_stdv[Y_valid])\n",
    "    \n",
    "    if np.amax(Y_stdv) > np.amin(Y_stdv):\n",
    "        vmin = umin-(umax-umin)*0.1\n",
    "        vmax = umax+(umax-umin)*0.1\n",
    "    else:\n",
    "        vmin = umin-0.1\n",
    "        vmax = umax+0.1\n",
    "\n",
    "    ax1.scatter(\n",
    "        np.power(10.0, Y_true[Y_valid]) - 1,\n",
    "        np.power(10.0, Y_pred[Y_valid]) - 1,\n",
    "        c=Y_stdv[Y_valid], cmap=uncertainty_cmap, s=1,\n",
    "        rasterized=True, vmin=vmin, vmax=vmax,\n",
    "    )\n",
    "    \n",
    "    ax1.set_xscale(\"log\")\n",
    "    ax1.set_yscale(\"log\")\n",
    "    \n",
    "    xlim = ax1.get_xlim()\n",
    "    ylim = ax1.get_ylim()\n",
    "    lim = (min(xlim[0], ylim[0]), max(xlim[1], ylim[1]))\n",
    "            \n",
    "    ax1.plot(lim, lim, c=\"black\")\n",
    "    ax1.plot(lim, lim, c=\"white\", ls=(0, (5, 5)))\n",
    "    \n",
    "    ax1.set_xlim(xlim)\n",
    "    ax1.set_ylim(ylim)\n",
    "    \n",
    "    ax2 = ax1.twinx()\n",
    "    \n",
    "    ax2.scatter(\n",
    "        [xlim[1]] * 100,\n",
    "        (\n",
    "            np.linspace(np.amin(Y_stdv), np.amax(Y_stdv), 100)\n",
    "            if np.amax(Y_stdv) > np.amin(Y_stdv) else\n",
    "            np.linspace(np.amin(Y_stdv) - 1, np.amax(Y_stdv) + 1, 100)\n",
    "        ),\n",
    "        c=uncertainty_cmap(\n",
    "            np.linspace(0.1, 0.9, 100) if np.amax(Y_stdv) > np.amin(Y_stdv)\n",
    "            else np.linspace(0.5, 0.5, 100)\n",
    "        ),\n",
    "    )\n",
    "    \n",
    "    if np.amax(Y_stdv) <= np.amin(Y_stdv):\n",
    "        ax2.set_yticks([np.amax(Y_stdv)])\n",
    "    \n",
    "    ax1.set_xlim(xlim)\n",
    "    ax2.set_xlim(xlim)\n",
    "    \n",
    "    mse = mean_squared_error(Y_true, Y_pred)\n",
    "    mae = mean_absolute_error(Y_true, Y_pred)\n",
    "    r2 = r2_score(Y_true, Y_pred)\n",
    "    \n",
    "    ax1.text(\n",
    "        0.037, 0.95,\n",
    "        f\"MSE(log$_{{10}}$) = {mse:.02}\\nMAE(log$_{{10}}$) = {mae:.02}\\nR$^2$(log$_{{10}}$) = {r2:.02}\",\n",
    "        ha=\"left\", va=\"top\", transform=ax1.transAxes, c=\"black\",\n",
    "        bbox=dict(facecolor='white', alpha=0.75, edgecolor='black'),\n",
    "    )\n",
    "    \n",
    "    ax1.set_title(\"Generalisation to all six trajectories\")\n",
    "    ax1.set_xlabel(\"true CCN concentration [m${}^{-3}$]\")\n",
    "    ax1.set_ylabel(\"predicted CCN concentration [m${}^{-3}$]\")\n",
    "    ax2.set_ylabel(\"log$_{10}$(CCN) prediction uncertainty [m${}^{-3}$]\")\n",
    "    \n",
    "    plt.savefig(f\"{cls.__name__.lower()}-test-generalisation.pdf\", dpi=100, transparent=True, bbox_inches='tight')\n",
    "    # plt.show()\n",
    "    plt.close(fig)\n",
    "    \n",
    "    model2 = train_and_cache_model(dt_test, 0.75, DATASETS, MODELS, cls, *args, **kwargs)\n",
    "    \n",
    "    fig, ax1 = plt.subplots(1, 1, figsize=(4, 4))\n",
    "    \n",
    "    Y_true = ds_test.Y_scaler.inverse_transform(ds_test.Y_test)\n",
    "    Y_pred = ds_test.Y_scaler.inverse_transform(model2.predict(ds_test.X_test))\n",
    "    Y_stdv = model2.predict_uncertainty(ds_test.X_test)\n",
    "    \n",
    "    Y_valid = (Y_pred.flatten() > -250) & (Y_pred.flatten() < 250)\n",
    "\n",
    "    umin = np.amin(Y_stdv[Y_valid])\n",
    "    umax = np.amax(Y_stdv[Y_valid])\n",
    "    \n",
    "    if np.amax(Y_stdv) > np.amin(Y_stdv):\n",
    "        vmin = umin-(umax-umin)*0.1\n",
    "        vmax = umax+(umax-umin)*0.1\n",
    "    else:\n",
    "        vmin = umin-0.1\n",
    "        vmax = umax+0.1\n",
    "    \n",
    "    ax1.scatter(\n",
    "        np.power(10.0, Y_true[Y_valid]) - 1,\n",
    "        np.power(10.0, Y_pred[Y_valid]) - 1,\n",
    "        c=Y_stdv[Y_valid], cmap=uncertainty_cmap, s=1,\n",
    "        rasterized=True, vmin=vmin, vmax=vmax,\n",
    "    )\n",
    "    \n",
    "    ax1.set_xscale(\"log\")\n",
    "    ax1.set_yscale(\"log\")\n",
    "    \n",
    "    xlim = ax1.get_xlim()\n",
    "    ylim = ax1.get_ylim()\n",
    "    lim = (min(xlim[0], ylim[0]), max(xlim[1], ylim[1]))\n",
    "            \n",
    "    ax1.plot(lim, lim, c=\"black\", ls=(5, (5, 5)))\n",
    "    ax1.plot(lim, lim, c=\"white\", ls=(0, (5, 5)))\n",
    "    \n",
    "    ax1.set_xlim(xlim)\n",
    "    ax1.set_ylim(ylim)\n",
    "    \n",
    "    ax2 = ax1.twinx()\n",
    "    \n",
    "    ax2.scatter(\n",
    "        [xlim[1]] * 100,\n",
    "        (\n",
    "            np.linspace(np.amin(Y_stdv), np.amax(Y_stdv), 100)\n",
    "            if np.amax(Y_stdv) > np.amin(Y_stdv) else\n",
    "            np.linspace(np.amin(Y_stdv) - 1, np.amax(Y_stdv) + 1, 100)\n",
    "        ),\n",
    "        c=uncertainty_cmap(\n",
    "            np.linspace(0.1, 0.9, 100) if np.amax(Y_stdv) > np.amin(Y_stdv)\n",
    "            else np.linspace(0.5, 0.5, 100)\n",
    "        ),\n",
    "    )\n",
    "    \n",
    "    if np.amax(Y_stdv) <= np.amin(Y_stdv):\n",
    "        ax2.set_yticks([np.amax(Y_stdv)])\n",
    "    \n",
    "    ax1.set_xlim(xlim)\n",
    "    ax2.set_xlim(xlim)\n",
    "    \n",
    "    mse = mean_squared_error(Y_true, Y_pred)\n",
    "    mae = mean_absolute_error(Y_true, Y_pred)\n",
    "    r2 = r2_score(Y_true, Y_pred)\n",
    "    \n",
    "    ax1.text(\n",
    "        0.037, 0.95,\n",
    "        f\"MSE(log$_{{10}}$) = {mse:.02}\\nMAE(log$_{{10}}$) = {mae:.02}\\nR$^2$(log$_{{10}}$) = {r2:.02}\",\n",
    "        ha=\"left\", va=\"top\", transform=ax1.transAxes, c=\"black\",\n",
    "        bbox=dict(facecolor='white', alpha=0.75, edgecolor='black'),\n",
    "    )\n",
    "    \n",
    "    ax1.set_title(\"... after training on all six trajectories\")\n",
    "    ax1.set_xlabel(\"true CCN concentration [m${}^{-3}$]\")\n",
    "    ax1.set_ylabel(\"predicted CCN concentration [m${}^{-3}$]\")\n",
    "    ax2.set_ylabel(\"log$_{10}$(CCN) prediction uncertainty [m${}^{-3}$]\")\n",
    "    \n",
    "    plt.savefig(f\"{cls.__name__.lower()}-test-prediction.pdf\", dpi=100, transparent=True, bbox_inches='tight')\n",
    "    # plt.show()\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97b97868",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 12737.072597,
   "end_time": "2023-01-25T13:33:11.655593",
   "environment_variables": {},
   "exception": null,
   "input_path": "generalisation.ipynb",
   "output_path": "generalisation.ipynb",
   "parameters": {},
   "start_time": "2023-01-25T10:00:54.582996",
   "version": "2.3.4"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "019031af7d8c48e69022c1cb57b2f9d3": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "19a7f13716ca4daf81a3c2f5568df16c": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_dfab174198f44685bc8e11b9912259a2",
       "max": 100,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_f8baa2bab0574d989b93569a13c54266",
       "tabbable": null,
       "tooltip": null,
       "value": 100
      }
     },
     "1b37107810414e828546b148557c66a8": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "2651cb61da234172af0b46e69103ba21": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "28f81896ed5b4f7ba6346541750a27e8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_e031ff83c380437c9c22bdd9b62903de",
       "placeholder": "​",
       "style": "IPY_MODEL_3a1ff08eb6694ad484a822b55b79690a",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "2bfbccfa44e34a0e9b839940f1f840f2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "2d2e7e129f644397849f21e014563416": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "2f1bf43a34ef41babf9771e9dc9a0fdf": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "31201e786cb0442f9d5d99eec85afe99": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_f2d8a2f02c5c435f8632d4f03f35748d",
       "placeholder": "​",
       "style": "IPY_MODEL_9e75d52a8e2a44eb82127a4d8a40cb60",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "3a1ff08eb6694ad484a822b55b79690a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "40c7c519d7f74843abbbeaf9c7355f12": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_7514425f5ee848c79abb969f383b4e46",
       "placeholder": "​",
       "style": "IPY_MODEL_8b3b4f97b0a84a2fb31812dd44317792",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "414f8ca5df0e43df8f4349e41101af28": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "41da05f10df242cc93d09f8417ff55b1": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "425d25c6f3cf446a8fb287f6f844525f": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_f69effc1db1a45b89beb8ec695ad323d",
       "placeholder": "​",
       "style": "IPY_MODEL_f08a94be5e7d40adaf14baa99f5d7039",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "466dea72255c4c269f0329f626c8c149": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_41da05f10df242cc93d09f8417ff55b1",
       "max": 100,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_2d2e7e129f644397849f21e014563416",
       "tabbable": null,
       "tooltip": null,
       "value": 100
      }
     },
     "48d1fdfb9fde452288455ecbd3ef8c80": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_31201e786cb0442f9d5d99eec85afe99",
        "IPY_MODEL_74e87566b9e543e9be174b313392bed6",
        "IPY_MODEL_e8170b5e76ca43079848afc04fb6dfc7"
       ],
       "layout": "IPY_MODEL_2651cb61da234172af0b46e69103ba21",
       "tabbable": null,
       "tooltip": null
      }
     },
     "5e20360724014f9b8413f64d373e98c4": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "5e98b9da1dd643f491fd33a46a0e43ec": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "63d75d770b674ccaa2049a10e3c45b0f": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "678c72d46baa47439867d3fec767c5cc": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "74e302337398458e97602a413293b495": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_40c7c519d7f74843abbbeaf9c7355f12",
        "IPY_MODEL_466dea72255c4c269f0329f626c8c149",
        "IPY_MODEL_c5a1cd58bace49c68461416790dda952"
       ],
       "layout": "IPY_MODEL_414f8ca5df0e43df8f4349e41101af28",
       "tabbable": null,
       "tooltip": null
      }
     },
     "74e87566b9e543e9be174b313392bed6": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_2bfbccfa44e34a0e9b839940f1f840f2",
       "max": 100,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_5e98b9da1dd643f491fd33a46a0e43ec",
       "tabbable": null,
       "tooltip": null,
       "value": 100
      }
     },
     "7514425f5ee848c79abb969f383b4e46": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "7bbbedf3e34d4aec98e0e046c109d508": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_425d25c6f3cf446a8fb287f6f844525f",
        "IPY_MODEL_19a7f13716ca4daf81a3c2f5568df16c",
        "IPY_MODEL_f4a034837d1847cf96d61a320106b3cc"
       ],
       "layout": "IPY_MODEL_dbbb78bea9f340f497eafc29b1fb1f0f",
       "tabbable": null,
       "tooltip": null
      }
     },
     "8723cebae357498ea0bbd3cb7c8dd979": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_fea27a5789d24ada95ccfd4cdb11147a",
       "max": 100,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_2f1bf43a34ef41babf9771e9dc9a0fdf",
       "tabbable": null,
       "tooltip": null,
       "value": 100
      }
     },
     "8b3b4f97b0a84a2fb31812dd44317792": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "9e75d52a8e2a44eb82127a4d8a40cb60": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "9ed0ca273f14498d855bb1641804db13": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "ab9bcf89cf9b4b789f5f1878c1ea6465": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_eee31f717c544b10b30da7ed0736803a",
       "placeholder": "​",
       "style": "IPY_MODEL_678c72d46baa47439867d3fec767c5cc",
       "tabbable": null,
       "tooltip": null,
       "value": " 100/100 [00:35&lt;00:00,  2.86it/s]"
      }
     },
     "b9365f68125c4b59ab67132e84334dda": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "c557ac89352a420085c5816bb05d3f77": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_28f81896ed5b4f7ba6346541750a27e8",
        "IPY_MODEL_8723cebae357498ea0bbd3cb7c8dd979",
        "IPY_MODEL_ab9bcf89cf9b4b789f5f1878c1ea6465"
       ],
       "layout": "IPY_MODEL_63d75d770b674ccaa2049a10e3c45b0f",
       "tabbable": null,
       "tooltip": null
      }
     },
     "c5a1cd58bace49c68461416790dda952": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_1b37107810414e828546b148557c66a8",
       "placeholder": "​",
       "style": "IPY_MODEL_df7b83b2aba54c0ebcdb31f4583b02bb",
       "tabbable": null,
       "tooltip": null,
       "value": " 100/100 [00:34&lt;00:00,  2.86it/s]"
      }
     },
     "dbbb78bea9f340f497eafc29b1fb1f0f": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "df7b83b2aba54c0ebcdb31f4583b02bb": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "dfab174198f44685bc8e11b9912259a2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "e031ff83c380437c9c22bdd9b62903de": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "e8170b5e76ca43079848afc04fb6dfc7": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_9ed0ca273f14498d855bb1641804db13",
       "placeholder": "​",
       "style": "IPY_MODEL_5e20360724014f9b8413f64d373e98c4",
       "tabbable": null,
       "tooltip": null,
       "value": " 100/100 [00:34&lt;00:00,  2.87it/s]"
      }
     },
     "eee31f717c544b10b30da7ed0736803a": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "f08a94be5e7d40adaf14baa99f5d7039": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "f2d8a2f02c5c435f8632d4f03f35748d": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "f4a034837d1847cf96d61a320106b3cc": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_b9365f68125c4b59ab67132e84334dda",
       "placeholder": "​",
       "style": "IPY_MODEL_019031af7d8c48e69022c1cb57b2f9d3",
       "tabbable": null,
       "tooltip": null,
       "value": " 100/100 [00:34&lt;00:00,  2.92it/s]"
      }
     },
     "f69effc1db1a45b89beb8ec695ad323d": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "f8baa2bab0574d989b93569a13c54266": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "fea27a5789d24ada95ccfd4cdb11147a": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}