examples/esn_hyperparameter_sweep_jallon.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "\n",
    "import torch\n",
    "import plotly.io as pio\n",
    "import plotly.express as px\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from dask.distributed import Client\n",
    "\n",
    "from time_series_prediction import settings, utils\n",
    "from biophysical_models import models\n",
    "\n",
    "pio.templates.default = 'plotly_white'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# settings.switch_device('cuda')\n",
    "settings.device == torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = [\n",
    "    ('leaking_rate', [0.05, 0.1, 0.3, 0.5, 0.7, 1]),\n",
    "    ('k_l2', 10**np.linspace(1, 3, 5)),\n",
    "    ('n_neurons', [400, 800, 1600, 2400, 3200]),\n",
    "    ('spectral_radius', [1, 2, 3, 4]),\n",
    "]\n",
    "# params = [\n",
    "#     ('k_l2', 10**torch.linspace(-0.5, -2.5, 5)),\n",
    "#     ('source', ['rossler', 'lorenz'])\n",
    "# ]\n",
    "keys, vals = zip(*params)\n",
    "param_dicts = []\n",
    "for val in product(*vals):\n",
    "    param_dicts.append({k: v for k, v in zip(keys, val)})\n",
    "print(f'{len(param_dicts)} simulations')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(n_workers=5, threads_per_worker=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_ensemble = 15\n",
    "\n",
    "kwargs_default = {\n",
    "}\n",
    "\n",
    "sweep = utils.SweepESN(\n",
    "    param_dicts,\n",
    "    kwargs_default,\n",
    "    n_ensemble,\n",
    "    noise=0.0,\n",
    "    source='jallon',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep.submit_jobs(client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = sweep.get_results()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(res)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_pickle('esn_sweep_jallon.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.sort_values(by='Median MAE').iloc[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = pio.read_json(df.sort_values(by='Median MAE').iloc[0]['_fig_path'])\n",
    "fig.update_layout(template='plotly_white', height=1200)\n",
    "for i, name in enumerate(models.JallonHeartLungs().state_names):\n",
    "    yname = 'yaxis' if i == 0 else f'yaxis{i+1}'\n",
    "    fig.update_layout({yname: {'title_text': name}})\n",
    "fig.update_traces(showlegend=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfq = df.query(\"leaking_rate==0.05 & n_neurons==3200 & spectral_radius==1\")\n",
    "dfq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = pio.read_json(df.query(\"leaking_rate==0.1 & n_neurons==3200 & spectral_radius==1\").iloc[2]['_fig_path'])\n",
    "fig.update_layout(template='plotly_white', height=1200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 'leaking_rate'\n",
    "color = 'n_neurons'\n",
    "y = 'Median MAE'\n",
    "\n",
    "dfq = df#.query('k_l2>=0.01')\n",
    "\n",
    "fig = px.line(\n",
    "    dfq,\n",
    "    x=x,\n",
    "    y=y,\n",
    "    color=color,\n",
    "    facet_row='k_l2',\n",
    "    line_dash='spectral_radius',\n",
    "    # facet_col='spectral_radius',\n",
    "    error_y=dfq['Max MAE'] - dfq[y],\n",
    "    error_y_minus=dfq[y] - dfq['Min MAE'],\n",
    "    # log_x=True,\n",
    "    # log_y=True,\n",
    "    color_discrete_sequence=px.colors.sample_colorscale(\n",
    "        px.colors.get_colorscale('Plasma'), \n",
    "        np.linspace(0, 0.9, len(df[color].unique())),\n",
    "    )\n",
    ")\n",
    "fig.update_traces(mode='markers+lines')\n",
    "fig.update_layout(height=800)\n",
    "fig.for_each_annotation(lambda x: x.update(text=x.text[:12]), lambda x: x.text[:4] == 'k_l2')\n",
    "fig.update_yaxes(title_text='median_err')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "57ab650755106788127778ca676b1d597fe3fd74372f745f5eeb1dc0c5eefa13"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}