{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook contains all experiment results exhibited in our paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "sns.set(style='white')\n",
    "matplotlib.rcParams['pdf.fonttype'] = 42\n",
    "matplotlib.rcParams['ps.fonttype'] = 42\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "def func(x, N=80):\n",
    "    ret = x.ret.copy()\n",
    "    x = x.rank(pct=True)\n",
    "    x['ret'] = ret\n",
    "    diff = x.score.sub(x.label)\n",
    "    r = x.nlargest(N, columns='score').ret.mean()\n",
    "    r -= x.nsmallest(N, columns='score').ret.mean()\n",
    "    return pd.Series({\n",
    "        'MSE': diff.pow(2).mean(), \n",
    "        'MAE': diff.abs().mean(), \n",
    "        'IC': x.score.corr(x.label),\n",
    "        'R': r\n",
    "    })\n",
    "    \n",
    "ret = pd.read_pickle(\"data/ret.pkl\").clip(-0.1, 0.1)\n",
    "def backtest(fname, **kwargs):\n",
    "    pred = pd.read_pickle(fname).loc['2018-09-21':'2020-06-30']  # test period\n",
    "    pred['ret'] = ret\n",
    "    dates = pred.index.unique(level=0)\n",
    "    res = Parallel(n_jobs=-1)(delayed(func)(pred.loc[d], **kwargs) for d in dates)\n",
    "    res = {\n",
    "       dates[i]: res[i]\n",
    "       for i in range(len(dates))\n",
    "    }\n",
    "    res = pd.DataFrame(res).T\n",
    "    r = res['R'].copy()\n",
    "    r.index = pd.to_datetime(r.index)\n",
    "    r = r.reindex(pd.date_range(r.index[0], r.index[-1])).fillna(0)  # paper use 365 days\n",
    "    return {\n",
    "        'MSE': res['MSE'].mean(),\n",
    "        'MAE': res['MAE'].mean(),\n",
    "        'IC': res['IC'].mean(),\n",
    "        'ICIR': res['IC'].mean()/res['IC'].std(),\n",
    "        'AR': r.mean()*365,\n",
    "        'AV': r.std()*365**0.5,\n",
    "        'SR': r.mean()/r.std()*365**0.5,\n",
    "        'MDD': (r.cumsum().cummax() - r.cumsum()).max()\n",
    "    }, r\n",
    "\n",
    "def fmt(x, p=3, scale=1, std=False):\n",
    "    _fmt = '{:.%df}'%p\n",
    "    string = _fmt.format((x.mean() if not isinstance(x, (float, np.floating)) else x) * scale)\n",
    "    if std and len(x) > 1:\n",
    "        string += ' ('+_fmt.format(x.std()*scale)+')'\n",
    "    return string\n",
    "\n",
    "def backtest_multi(files, **kwargs):\n",
    "    res = []\n",
    "    pnl = []\n",
    "    for fname in files:\n",
    "        metric, r = backtest(fname, **kwargs)\n",
    "        res.append(metric)\n",
    "        pnl.append(r)\n",
    "    res = pd.DataFrame(res)\n",
    "    pnl = pd.concat(pnl, axis=1)\n",
    "    return {\n",
    "        'MSE': fmt(res['MSE'], std=True),\n",
    "        'MAE': fmt(res['MAE'], std=True),\n",
    "        'IC': fmt(res['IC']),\n",
    "        'ICIR': fmt(res['ICIR']),\n",
    "        'AR': fmt(res['AR'], scale=100, p=1)+'%',\n",
    "        'VR': fmt(res['AV'], scale=100, p=1)+'%',\n",
    "        'SR': fmt(res['SR']),\n",
    "        'MDD': fmt(res['MDD'], scale=100, p=1)+'%'\n",
    "    }, pnl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation\n",
    "\n",
    "\n",
    "You could prepare the source data as below for the backtest code:\n",
    "1. Linear: see Qlib examples\n",
    "2. LightGBM: see Qlib examples\n",
    "3. MLP: see Qlib examples\n",
    "4. SFM: see Qlib examples\n",
    "5. ALSTM: `qrun` configs/config_alstm.yaml\n",
    "6. Transformer: `qrun` configs/config_transformer.yaml\n",
    "7. ALSTM+TRA: `qrun` configs/config_alstm_tra_init.yaml && `qrun` configs/config_alstm_tra.yaml\n",
    "8. Tranformer+TRA: `qrun` configs/config_transformer_tra_init.yaml && `qrun` configs/config_transformer_tra.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "exps = {\n",
    "    'Linear': ['output/Linear/pred.pkl'],\n",
    "    'LightGBM': ['output/GBDT/lr0.05_leaves128/pred.pkl'],\n",
    "    'MLP': glob.glob('output/search/MLP/hs128_bs512_do0.3_lr0.001_seed*/pred.pkl'),\n",
    "    'SFM': glob.glob('output/search/SFM/hs32_bs512_do0.5_lr0.001_seed*/pred.pkl'),\n",
    "    'ALSTM': glob.glob('output/search/LSTM_Attn/hs256_bs1024_do0.1_lr0.0002_seed*/pred.pkl'),\n",
    "    'Trans.': glob.glob('output/search/Transformer/head4_hs64_bs1024_do0.1_lr0.0002_seed*/pred.pkl'),\n",
    "    'ALSTM+TS':glob.glob('output/LSTM_Attn_TS/hs256_bs1024_do0.1_lr0.0002_seed*/pred.pkl'),\n",
    "    'Trans.+TS':glob.glob('output/Transformer_TS/head4_hs64_bs1024_do0.1_lr0.0002_seed*/pred.pkl'),\n",
    "    'ALSTM+TRA(Ours)': glob.glob('output/search/finetune/LSTM_Attn_tra/K10_traHs16_traSrcLR_TPE_traLamb2.0_hs256_bs1024_do0.1_lr0.0001_seed*/pred.pkl'),\n",
    "    'Trans.+TRA(Ours)': glob.glob('output/search/finetune/Transformer_tra/K3_traHs16_traSrcLR_TPE_traLamb1.0_head4_hs64_bs512_do0.1_lr0.0005_seed*/pred.pkl')\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0acd535e05944e539fd001009ed0748d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "res = {\n",
    "    name: backtest_multi(exps[name])\n",
    "    for name in tqdm(exps)\n",
    "}\n",
    "report = pd.DataFrame({\n",
    "    k: v[0]\n",
    "    for k, v in res.items()\n",
    "}).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>MSE</th>\n",
       "      <th>MAE</th>\n",
       "      <th>IC</th>\n",
       "      <th>ICIR</th>\n",
       "      <th>AR</th>\n",
       "      <th>VR</th>\n",
       "      <th>SR</th>\n",
       "      <th>MDD</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Linear</th>\n",
       "      <td>0.163</td>\n",
       "      <td>0.327</td>\n",
       "      <td>0.020</td>\n",
       "      <td>0.132</td>\n",
       "      <td>-3.2%</td>\n",
       "      <td>16.8%</td>\n",
       "      <td>-0.191</td>\n",
       "      <td>32.1%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>LightGBM</th>\n",
       "      <td>0.160</td>\n",
       "      <td>0.323</td>\n",
       "      <td>0.041</td>\n",
       "      <td>0.292</td>\n",
       "      <td>7.8%</td>\n",
       "      <td>15.5%</td>\n",
       "      <td>0.503</td>\n",
       "      <td>25.7%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>MLP</th>\n",
       "      <td>0.160 (0.002)</td>\n",
       "      <td>0.323 (0.003)</td>\n",
       "      <td>0.037</td>\n",
       "      <td>0.273</td>\n",
       "      <td>3.7%</td>\n",
       "      <td>15.3%</td>\n",
       "      <td>0.264</td>\n",
       "      <td>26.2%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SFM</th>\n",
       "      <td>0.159 (0.001)</td>\n",
       "      <td>0.321 (0.001)</td>\n",
       "      <td>0.047</td>\n",
       "      <td>0.381</td>\n",
       "      <td>7.1%</td>\n",
       "      <td>14.3%</td>\n",
       "      <td>0.497</td>\n",
       "      <td>22.9%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ALSTM</th>\n",
       "      <td>0.158 (0.001)</td>\n",
       "      <td>0.320 (0.001)</td>\n",
       "      <td>0.053</td>\n",
       "      <td>0.419</td>\n",
       "      <td>12.3%</td>\n",
       "      <td>13.7%</td>\n",
       "      <td>0.897</td>\n",
       "      <td>20.2%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Trans.</th>\n",
       "      <td>0.158 (0.001)</td>\n",
       "      <td>0.322 (0.001)</td>\n",
       "      <td>0.051</td>\n",
       "      <td>0.400</td>\n",
       "      <td>14.5%</td>\n",
       "      <td>14.2%</td>\n",
       "      <td>1.028</td>\n",
       "      <td>22.5%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ALSTM+TS</th>\n",
       "      <td>0.160 (0.002)</td>\n",
       "      <td>0.321 (0.002)</td>\n",
       "      <td>0.039</td>\n",
       "      <td>0.291</td>\n",
       "      <td>6.7%</td>\n",
       "      <td>14.6%</td>\n",
       "      <td>0.480</td>\n",
       "      <td>22.3%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Trans.+TS</th>\n",
       "      <td>0.160 (0.004)</td>\n",
       "      <td>0.324 (0.005)</td>\n",
       "      <td>0.037</td>\n",
       "      <td>0.278</td>\n",
       "      <td>10.4%</td>\n",
       "      <td>14.7%</td>\n",
       "      <td>0.722</td>\n",
       "      <td>23.7%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ALSTM+TRA(Ours)</th>\n",
       "      <td>0.157 (0.000)</td>\n",
       "      <td>0.318 (0.000)</td>\n",
       "      <td>0.059</td>\n",
       "      <td>0.460</td>\n",
       "      <td>12.4%</td>\n",
       "      <td>14.0%</td>\n",
       "      <td>0.885</td>\n",
       "      <td>20.4%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Trans.+TRA(Ours)</th>\n",
       "      <td>0.157 (0.000)</td>\n",
       "      <td>0.320 (0.000)</td>\n",
       "      <td>0.056</td>\n",
       "      <td>0.442</td>\n",
       "      <td>16.1%</td>\n",
       "      <td>14.2%</td>\n",
       "      <td>1.133</td>\n",
       "      <td>23.1%</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                            MSE            MAE     IC   ICIR     AR     VR  \\\n",
       "Linear                    0.163          0.327  0.020  0.132  -3.2%  16.8%   \n",
       "LightGBM                  0.160          0.323  0.041  0.292   7.8%  15.5%   \n",
       "MLP               0.160 (0.002)  0.323 (0.003)  0.037  0.273   3.7%  15.3%   \n",
       "SFM               0.159 (0.001)  0.321 (0.001)  0.047  0.381   7.1%  14.3%   \n",
       "ALSTM             0.158 (0.001)  0.320 (0.001)  0.053  0.419  12.3%  13.7%   \n",
       "Trans.            0.158 (0.001)  0.322 (0.001)  0.051  0.400  14.5%  14.2%   \n",
       "ALSTM+TS          0.160 (0.002)  0.321 (0.002)  0.039  0.291   6.7%  14.6%   \n",
       "Trans.+TS         0.160 (0.004)  0.324 (0.005)  0.037  0.278  10.4%  14.7%   \n",
       "ALSTM+TRA(Ours)   0.157 (0.000)  0.318 (0.000)  0.059  0.460  12.4%  14.0%   \n",
       "Trans.+TRA(Ours)  0.157 (0.000)  0.320 (0.000)  0.056  0.442  16.1%  14.2%   \n",
       "\n",
       "                      SR    MDD  \n",
       "Linear            -0.191  32.1%  \n",
       "LightGBM           0.503  25.7%  \n",
       "MLP                0.264  26.2%  \n",
       "SFM                0.497  22.9%  \n",
       "ALSTM              0.897  20.2%  \n",
       "Trans.             1.028  22.5%  \n",
       "ALSTM+TS           0.480  22.3%  \n",
       "Trans.+TS          0.722  23.7%  \n",
       "ALSTM+TRA(Ours)    0.885  20.4%  \n",
       "Trans.+TRA(Ours)   1.133  23.1%  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "report\n",
    "# print(report.to_latex())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RQ1\n",
    "\n",
    "Case study"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30892 characters]",
      "text/plain": [
       "<Figure size 504x216 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df = pd.read_pickle('output/search/finetune/Transformer_tra/K3_traHs16_traSrcLR_TPE_traLamb0.0_head4_hs64_bs512_do0.1_lr0.0005_seed1000/pred.pkl')\n",
    "code = 'SH600157'\n",
    "date = '2018-09-28'\n",
    "lookbackperiod = 50\n",
    "\n",
    "prob = df.iloc[:, -3:].loc(axis=0)[:, code].reset_index(level=1, drop=True).loc[date:].iloc[:lookbackperiod]\n",
    "pred = df.loc[:,[\"score_0\",\"score_1\",\"score_2\",\"label\"]].loc(axis=0)[:, code].reset_index(level=1, drop=True).loc[date:].iloc[:lookbackperiod]\n",
    "e_all = pred.iloc[:,:-1].sub(pred.iloc[:,-1], axis=0).pow(2)\n",
    "e_all = e_all.sub(e_all.min(axis=1), axis=0)\n",
    "e_all.columns = [r'$\\theta_%d$'%d for d in range(1, 4)]\n",
    "prob = pd.Series(np.argmax(prob.values, axis=1), index=prob.index).rolling(7).mean().round()\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(7, 3))\n",
    "e_all.plot(ax=axes[0], xlabel='', rot=30)\n",
    "prob.plot(ax=axes[1], xlabel='', rot=30, color='red', linestyle='None', marker='^', markersize=5)\n",
    "plt.yticks(np.array([0, 1, 2]), e_all.columns.values)\n",
    "axes[0].set_ylabel('Predictor Loss')\n",
    "axes[1].set_ylabel('Router Selection')\n",
    "plt.tight_layout()\n",
    "# plt.savefig('select.pdf', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RQ2\n",
    "\n",
    "You could prepared the source data for this test as below:\n",
    "1. Random: Setting `src_info` = \"NONE\"\n",
    "2. LR: Setting `src_info` = \"LR\"\n",
    "3. TPE: Setting `src_info` = \"TPE\"\n",
    "4. LR+TPE: Setting `src_info` = \"LR_TPE\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "exps = {\n",
    "    'Random': glob.glob('output/search/LSTM_Attn_tra/K10_traHs16_traSrcNONE_traLamb1.0_hs256_bs1024_do0.1_lr0.0001_seed*/pred.pkl'),\n",
    "    'LR': glob.glob('output/search/LSTM_Attn_tra/K10_traHs16_traSrcLR_traLamb1.0_hs256_bs1024_do0.1_lr0.0001_seed*/pred.pkl'),\n",
    "    'TPE': glob.glob('output/search/LSTM_Attn_tra/K10_traHs16_traSrcTPE_traLamb1.0_hs256_bs1024_do0.1_lr0.0001_seed*/pred.pkl'),\n",
    "    'LR+TPE': glob.glob('output/search/finetune/LSTM_Attn_tra/K10_traHs16_traSrcLR_TPE_traLamb2.0_hs256_bs1024_do0.1_lr0.0001_seed*/pred.pkl')\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "910721fc4a7b46eea5ba6d50647320d4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "res = {\n",
    "    name: backtest_multi(exps[name])\n",
    "    for name in tqdm(exps)\n",
    "}\n",
    "report = pd.DataFrame({\n",
    "    k: v[0]\n",
    "    for k, v in res.items()\n",
    "}).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>MSE</th>\n",
       "      <th>MAE</th>\n",
       "      <th>IC</th>\n",
       "      <th>ICIR</th>\n",
       "      <th>AR</th>\n",
       "      <th>VR</th>\n",
       "      <th>SR</th>\n",
       "      <th>MDD</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Random</th>\n",
       "      <td>0.159 (0.001)</td>\n",
       "      <td>0.321 (0.002)</td>\n",
       "      <td>0.048</td>\n",
       "      <td>0.362</td>\n",
       "      <td>11.4%</td>\n",
       "      <td>14.1%</td>\n",
       "      <td>0.810</td>\n",
       "      <td>21.1%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>LR</th>\n",
       "      <td>0.158 (0.001)</td>\n",
       "      <td>0.320 (0.001)</td>\n",
       "      <td>0.053</td>\n",
       "      <td>0.409</td>\n",
       "      <td>10.3%</td>\n",
       "      <td>13.4%</td>\n",
       "      <td>0.772</td>\n",
       "      <td>20.8%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TPE</th>\n",
       "      <td>0.158 (0.001)</td>\n",
       "      <td>0.321 (0.001)</td>\n",
       "      <td>0.049</td>\n",
       "      <td>0.381</td>\n",
       "      <td>10.3%</td>\n",
       "      <td>14.0%</td>\n",
       "      <td>0.741</td>\n",
       "      <td>21.2%</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>LR+TPE</th>\n",
       "      <td>0.157 (0.000)</td>\n",
       "      <td>0.318 (0.000)</td>\n",
       "      <td>0.059</td>\n",
       "      <td>0.460</td>\n",
       "      <td>12.4%</td>\n",
       "      <td>14.0%</td>\n",
       "      <td>0.885</td>\n",
       "      <td>20.4%</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  MSE            MAE     IC   ICIR     AR     VR     SR    MDD\n",
       "Random  0.159 (0.001)  0.321 (0.002)  0.048  0.362  11.4%  14.1%  0.810  21.1%\n",
       "LR      0.158 (0.001)  0.320 (0.001)  0.053  0.409  10.3%  13.4%  0.772  20.8%\n",
       "TPE     0.158 (0.001)  0.321 (0.001)  0.049  0.381  10.3%  14.0%  0.741  21.2%\n",
       "LR+TPE  0.157 (0.000)  0.318 (0.000)  0.059  0.460  12.4%  14.0%  0.885  20.4%"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "report\n",
    "# print(report.to_latex())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RQ3\n",
    "\n",
    "Set `lamb` = 0 to obtain results without Optimal Transport(OT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7212 characters]",
      "text/plain": [
       "<Figure size 360x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "a = pd.read_pickle('output/search/finetune/Transformer_tra/K3_traHs16_traSrcLR_TPE_traLamb0.0_head4_hs64_bs512_do0.1_lr0.0005_seed3000/pred.pkl')\n",
    "b = pd.read_pickle('output/search/finetune/Transformer_tra/K3_traHs16_traSrcLR_TPE_traLamb2.0_head4_hs64_bs512_do0.1_lr0.0005_seed3000/pred.pkl')\n",
    "a = a.iloc[:, -3:]\n",
    "b = b.iloc[:, -3:]\n",
    "b = np.eye(3)[b.values.argmax(axis=1)]\n",
    "a = np.eye(3)[a.values.argmax(axis=1)]\n",
    "\n",
    "res = pd.DataFrame({\n",
    "    'with OT': b.sum(axis=0) / b.sum(),\n",
    "    'without OT': a.sum(axis=0)/ a.sum()  \n",
    "},index=[r'$\\theta_1$',r'$\\theta_2$',r'$\\theta_3$'])\n",
    "res.plot.bar(rot=30, figsize=(5, 4), color=['b', 'g'])\n",
    "del a, b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RQ4\n",
    "\n",
    "You could prepared the source data for this test as below:\n",
    "1. K=1: which is exactly the alstm model\n",
    "2. K=3: Setting `num_states` = 3\n",
    "3. K=5: Setting `num_states` = 5\n",
    "4. K=10: Setting `num_states` = 10\n",
    "5. K=20: Setting `num_states` = 20\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "exps = {\n",
    "    'K=1': glob.glob('output/search/LSTM_Attn/hs256_bs1024_do0.1_lr0.0002_seed*/info.json'),\n",
    "    'K=3': glob.glob('output/search/finetune/LSTM_Attn_tra/K3_traHs16_traSrcLR_TPE_traLamb2.0_hs256_bs1024_do0.1_lr0.0001_seed*/info.json'),\n",
    "    'K=5': glob.glob('output/search/finetune/LSTM_Attn_tra/K5_traHs16_traSrcLR_TPE_traLamb2.0_hs256_bs1024_do0.1_lr0.0001_seed*/info.json'),\n",
    "    'K=10': glob.glob('output/search/finetune/LSTM_Attn_tra/K10_traHs16_traSrcLR_TPE_traLamb2.0_hs256_bs1024_do0.1_lr0.0001_seed*/info.json'),\n",
    "    'K=20': glob.glob('output/search/finetune/LSTM_Attn_tra/K20_traHs16_traSrcLR_TPE_traLamb2.0_hs256_bs1024_do0.1_lr0.0001_seed*/info.json')\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "report = dict()\n",
    "for k, v in exps.items():\n",
    "    \n",
    "    tmp = dict()\n",
    "    for fname in v:\n",
    "        with open(fname) as f:\n",
    "            info = json.load(f)\n",
    "        tmp[fname] = (\n",
    "        {\n",
    "            \"IC\":info[\"metric\"][\"IC\"],\n",
    "            \"MSE\":info[\"metric\"][\"MSE\"]\n",
    "        })\n",
    "    tmp = pd.DataFrame(tmp).T\n",
    "    report[k] = tmp.mean()\n",
    "report = pd.DataFrame(report).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13840 characters]",
      "text/plain": [
       "<Figure size 432x216 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(6,3)); axes = axes.flatten()\n",
    "report['IC'].plot.bar(rot=30, ax=axes[0])\n",
    "axes[0].set_ylim(0.045, 0.062)\n",
    "axes[0].set_title('IC performance')\n",
    "report['MSE'].astype(float).plot.bar(rot=30, ax=axes[1], color='green')\n",
    "axes[1].set_ylim(0.155, 0.1585)\n",
    "axes[1].set_title('MSE performance')\n",
    "plt.tight_layout()\n",
    "# plt.savefig('sensitivity.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>IC</th>\n",
       "      <th>MSE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>K=1</th>\n",
       "      <td>0.053247</td>\n",
       "      <td>0.157792</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>K=3</th>\n",
       "      <td>0.055535</td>\n",
       "      <td>0.157410</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>K=5</th>\n",
       "      <td>0.059224</td>\n",
       "      <td>0.156796</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>K=10</th>\n",
       "      <td>0.059403</td>\n",
       "      <td>0.156766</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>K=20</th>\n",
       "      <td>0.059193</td>\n",
       "      <td>0.156801</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            IC       MSE\n",
       "K=1   0.053247  0.157792\n",
       "K=3   0.055535  0.157410\n",
       "K=5   0.059224  0.156796\n",
       "K=10  0.059403  0.156766\n",
       "K=20  0.059193  0.156801"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "9de784e21d4a351f53a5792b09a6ae66a23802b850ad98f62e10c0156e418c04"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('base': conda)",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}