--- a
+++ b/draw.ipynb
@@ -0,0 +1,990 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from sklearn.calibration import calibration_curve\n",
+    "from sklearn.metrics import roc_curve, precision_recall_curve\n",
+    "import torch\n",
+    "from sklearn.decomposition import PCA\n",
+    "from sklearn.manifold import TSNE\n",
+    "# matplotlib\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.lines as mlines\n",
+    "import matplotlib.transforms as mtransfor\n",
+    "from matplotlib.ticker import FuncFormatter\n",
+    "import seaborn as sns\n",
+    "\n",
+    "plt.style.use('default')\n",
+    "plt.rcParams['axes.facecolor']='white'\n",
+    "plt.rcParams.update({\"axes.grid\" : True, \"grid.color\": \"gainsboro\"})\n",
+    "plt.rcParams['legend.frameon']=True\n",
+    "plt.rcParams['legend.facecolor']='white'\n",
+    "plt.rcParams['legend.edgecolor']='grey'\n",
+    "plt.rcParams[\"axes.edgecolor\"] = \"black\"\n",
+    "plt.rcParams[\"axes.linewidth\"]  = 1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Read models' outcome prediction result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tj_adacare = pd.read_pickle('./saved_pkl/tongji_adacare_outcome.pkl')\n",
+    "tj_retain = pd.read_pickle('./saved_pkl/tongji_retain_outcome.pkl')\n",
+    "tj_tcn = pd.read_pickle('./saved_pkl/tongji_tcn_outcome.pkl')\n",
+    "\n",
+    "hm_concare = pd.read_pickle('./saved_pkl/hm_concare_outcome.pkl')\n",
+    "hm_tcn = pd.read_pickle('./saved_pkl/hm_tcn_outcome.pkl')\n",
+    "hm_rnn = pd.read_pickle('./saved_pkl/hm_rnn_outcome.pkl')\n",
+    "\n",
+    "tj_adacare_outcome_true, tj_adacare_outcome_pred = tj_adacare['outcome_true'], tj_adacare['outcome_pred']\n",
+    "tj_retain_outcome_true, tj_retain_outcome_pred = tj_retain['outcome_true'], tj_retain['outcome_pred']\n",
+    "tj_tcn_outcome_true, tj_tcn_outcome_pred = tj_tcn['outcome_true'], tj_tcn['outcome_pred']\n",
+    "\n",
+    "hm_concare_outcome_true, hm_concare_outcome_pred = hm_concare['outcome_true'], hm_concare['outcome_pred']\n",
+    "hm_tcn_outcome_true, hm_tcn_outcome_pred = hm_tcn['outcome_true'], hm_tcn['outcome_pred']\n",
+    "hm_rnn_outcome_true, hm_rnn_outcome_pred = hm_rnn['outcome_true'], hm_rnn['outcome_pred']"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## ROC Plot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tj_random_probs = [0 for _ in range(len(tj_adacare_outcome_true))]\n",
+    "tj_p_fpr, tj_p_tpr, _ = roc_curve(tj_adacare_outcome_true, tj_random_probs, pos_label=1)\n",
+    "\n",
+    "hm_random_probs = [0 for _ in range(len(hm_tcn_outcome_true))]\n",
+    "hm_p_fpr, hm_p_tpr, _ = roc_curve(hm_tcn_outcome_true, hm_random_probs, pos_label=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# [TJH] plot roc curves\n",
+    "\n",
+    "tj_adacare_fpr, tj_adacare_tpr, thresh1 = roc_curve(tj_adacare_outcome_true, tj_adacare_outcome_pred, pos_label=1)\n",
+    "tj_retain_fpr, tj_retain_tpr, thresh2 = roc_curve(tj_retain_outcome_true, tj_retain_outcome_pred, pos_label=1)\n",
+    "tj_tcn_fpr, tj_tcn_tpr, thresh3 = roc_curve(tj_tcn_outcome_true, tj_tcn_outcome_pred, pos_label=1)\n",
+    "\n",
+    "plt.plot(tj_p_fpr, tj_p_tpr, linestyle='-.', color='grey', label='Random')\n",
+    "plt.plot(tj_adacare_fpr, tj_adacare_tpr, linestyle='dashed',color='orange', label='AdaCare')\n",
+    "plt.plot(tj_retain_fpr, tj_retain_tpr, linestyle='solid',color='dodgerblue', label='RETAIN')\n",
+    "plt.plot(tj_tcn_fpr, tj_tcn_tpr, linestyle='dotted',color='violet', label='TCN')\n",
+    "\n",
+    "# # title\n",
+    "# plt.title('ROC curve')\n",
+    "# x label\n",
+    "plt.xlabel('False Positive Rate')\n",
+    "# y label\n",
+    "plt.ylabel('True Positive Rate')\n",
+    "\n",
+    "plt.legend(loc='lower right')\n",
+    "\n",
+    "plt.savefig('tjh_roc.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# [CDSL] plot roc curves\n",
+    "\n",
+    "hm_concare_fpr, hm_concare_tpr, thresh1 = roc_curve(hm_concare_outcome_true, hm_concare_outcome_pred, pos_label=1)\n",
+    "hm_tcn_fpr, hm_tcn_tpr, thresh2 = roc_curve(hm_tcn_outcome_true, hm_tcn_outcome_pred, pos_label=1)\n",
+    "hm_rnn_fpr, hm_rnn_tpr, thresh3 = roc_curve(hm_rnn_outcome_true, hm_rnn_outcome_pred, pos_label=1)\n",
+    "\n",
+    "plt.plot(hm_p_fpr, hm_p_tpr, linestyle='-.', color='grey', label='Random')\n",
+    "plt.plot(hm_concare_fpr, hm_concare_tpr, linestyle='solid',color='dodgerblue', label='ConCare')\n",
+    "plt.plot(hm_tcn_fpr, hm_tcn_tpr, linestyle='dotted',color='violet', label='TCN')\n",
+    "plt.plot(hm_rnn_fpr, hm_rnn_tpr, linestyle='dashed',color='orange', label='RNN')\n",
+    "\n",
+    "# # title\n",
+    "# plt.title('ROC curve')\n",
+    "# x label\n",
+    "plt.xlabel('False positive rate')\n",
+    "# y label\n",
+    "plt.ylabel('True positive rate')\n",
+    "\n",
+    "plt.legend(loc='lower right')\n",
+    "plt.savefig('cdsl_roc.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## PRC Plot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# [TJH] plot precision-recall curves\n",
+    "\n",
+    "tj_adacare_precision, tj_adacare_recall, thresh1 = precision_recall_curve(tj_adacare_outcome_true, tj_adacare_outcome_pred, pos_label=1)\n",
+    "tj_retain_precision, tj_retain_recall, thresh2 = precision_recall_curve(tj_retain_outcome_true, tj_retain_outcome_pred, pos_label=1)\n",
+    "tj_tcn_precision, tj_tcn_recall, thresh3 = precision_recall_curve(tj_tcn_outcome_true, tj_tcn_outcome_pred, pos_label=1)\n",
+    "\n",
+    "plt.plot(tj_adacare_precision, tj_adacare_recall, linestyle='dashed',color='orange', label='AdaCare')\n",
+    "plt.plot(tj_retain_precision, tj_retain_recall, linestyle='solid',color='dodgerblue', label='RETAIN')\n",
+    "plt.plot(tj_tcn_precision, tj_tcn_recall, linestyle='dotted',color='violet', label='TCN')\n",
+    "\n",
+    "# # title\n",
+    "# plt.title('PRC curve')\n",
+    "# x label\n",
+    "plt.xlabel('Recall')\n",
+    "# y label\n",
+    "plt.ylabel('Precision')\n",
+    "\n",
+    "plt.legend(loc='lower left')\n",
+    "plt.savefig('tjh_prc.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# [CDSL] plot precision-recall curves\n",
+    "\n",
+    "hm_concare_precision, hm_concare_recall, thresh1 = precision_recall_curve(hm_concare_outcome_true, hm_concare_outcome_pred, pos_label=1)\n",
+    "hm_tcn_precision, hm_tcn_recall, thresh2 = precision_recall_curve(hm_tcn_outcome_true, hm_tcn_outcome_pred, pos_label=1)\n",
+    "hm_rnn_precision, hm_rnn_recall, thresh3 = precision_recall_curve(hm_rnn_outcome_true, hm_rnn_outcome_pred, pos_label=1)\n",
+    "\n",
+    "plt.plot(hm_concare_precision, hm_concare_recall, linestyle='solid',color='dodgerblue', label='ConCare')\n",
+    "plt.plot(hm_tcn_precision, hm_tcn_recall, linestyle='dotted',color='violet', label='TCN')\n",
+    "plt.plot(hm_rnn_precision, hm_rnn_recall, linestyle='dashed',color='orange', label='RNN')\n",
+    "\n",
+    "# # title\n",
+    "# plt.title('PRC curve')\n",
+    "# x label\n",
+    "plt.xlabel('Recall')\n",
+    "# y label\n",
+    "plt.ylabel('Precision')\n",
+    "\n",
+    "plt.legend(loc='lower left')\n",
+    "plt.savefig('cdsl_prc.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Calibration Plot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tj_adacare_prob_true, tj_adacare_prob_pred = calibration_curve(tj_adacare_outcome_true, tj_adacare_outcome_pred, n_bins=10)\n",
+    "tj_retain_prob_true, tj_retain_prob_pred = calibration_curve(tj_retain_outcome_true, tj_retain_outcome_pred, n_bins=10)\n",
+    "tj_tcn_prob_true, tj_tcn_prob_pred = calibration_curve(tj_tcn_outcome_true, tj_tcn_outcome_pred, n_bins=10)\n",
+    "\n",
+    "fig, ax = plt.subplots()\n",
+    "# only these two lines are calibration curves\n",
+    "plt.plot(tj_adacare_prob_pred, tj_adacare_prob_true, marker='o', linewidth=1, label='AdaCare')\n",
+    "plt.plot(tj_retain_prob_pred, tj_retain_prob_true, marker='v', linewidth=1, label='RETAIN')\n",
+    "plt.plot(tj_tcn_prob_pred, tj_tcn_prob_true, marker='s', linewidth=1, label='TCN')\n",
+    "\n",
+    "# reference line, legends, and axis labels\n",
+    "line = mlines.Line2D([0, 1], [0, 1], linestyle='-.', color='grey')\n",
+    "transform = ax.transAxes\n",
+    "line.set_transform(transform)\n",
+    "ax.add_line(line)\n",
+    "ax.set_xlabel('Predicted probability')\n",
+    "ax.set_ylabel('True probability in each bin')\n",
+    "plt.legend(loc='lower right')\n",
+    "plt.savefig('tjh_calibration.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hm_concare_prob_true, hm_concare_prob_pred = calibration_curve(hm_concare_outcome_true, hm_concare_outcome_pred, n_bins=10)\n",
+    "hm_tcn_prob_true, hm_tcn_prob_pred = calibration_curve(hm_tcn_outcome_true, hm_tcn_outcome_pred, n_bins=10)\n",
+    "hm_rnn_prob_true, hm_rnn_prob_pred = calibration_curve(hm_rnn_outcome_true, hm_rnn_outcome_pred, n_bins=10)\n",
+    "\n",
+    "fig, ax = plt.subplots()\n",
+    "# only these two lines are calibration curves\n",
+    "plt.plot(hm_concare_prob_pred, hm_concare_prob_true, marker='o', linewidth=1, label='ConCare')\n",
+    "plt.plot(hm_tcn_prob_pred, hm_tcn_prob_true, marker='s', linewidth=1, label='TCN')\n",
+    "plt.plot(hm_rnn_prob_pred, hm_rnn_prob_true, marker='v', linewidth=1, label='RNN')\n",
+    "\n",
+    "# reference line, legends, and axis labels\n",
+    "line = mlines.Line2D([0, 1], [0, 1], linestyle='-.', color='grey')\n",
+    "transform = ax.transAxes\n",
+    "line.set_transform(transform)\n",
+    "ax.add_line(line)\n",
+    "ax.set_xlabel('Predicted probability')\n",
+    "ax.set_ylabel('True probability in each bin')\n",
+    "plt.legend(loc='lower right')\n",
+    "plt.savefig('cdsl_calibration.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Draw OSMAE/EMP scores on different threshold"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "covid_scores = pd.read_pickle('./saved_pkl/covid_evaluation_scores.pkl')\n",
+    "emp, osmae, thresholds = covid_scores[\"emp\"][1::4], covid_scores[\"osmae\"][1::4], covid_scores[\"threshold\"][1::4]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## EMP Score\n",
+    "ax = sns.regplot(x=thresholds, y=emp, marker=\"o\", color=\"g\", line_kws={\"color\": \"grey\", \"linestyle\": \"-\", \"linewidth\": \"1\"}, ci=99.9999)\n",
+    "plt.xlabel('Threshold γ')\n",
+    "plt.ylabel('ES score')\n",
+    "\n",
+    "plt.savefig('emp_trend.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## OSMAE Score\n",
+    "ax = sns.regplot(x=thresholds, y=osmae, marker=\"o\", color=\"dodgerblue\", line_kws={\"color\": \"grey\", \"linestyle\": \"-\", \"linewidth\": \"1\"}, ci=99.9999)\n",
+    "plt.xlabel('Threshold γ')\n",
+    "plt.ylabel('OSMAE score')\n",
+    "\n",
+    "plt.savefig('osmae_trend.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Draw hidden state PCA result on validation set (CDSL dataset)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_idx=[1010, 1915, 1656, 2952, 246, 2914, 3146, 2910, 914, 1335, 3046, 404, 2592, 2951, 309, 266, 471, 1112, 490, 3195, 2621, 2143, 1485, 893, 2803, 319, 3231, 2185, 771, 2811, 1950, 3615, 2537, 2546, 3750, 2284, 2122, 1817, 767, 2698, 1564, 3519, 1285, 2808, 1092, 3782, 1115, 1174, 1996, 2603, 3337, 2806, 1105, 2180, 1006, 1900, 3563, 808, 3613, 907, 2069, 1893, 1877, 2362, 2403, 693, 3425, 3501, 626, 244, 3101, 2255, 1661, 1723, 3688, 2571, 2222, 382, 1091, 1, 1884, 3559, 3450, 2648, 2246, 2757, 820, 3375, 1650, 2509, 3760, 2519, 27, 1751, 1964, 1571, 2471, 541, 815, 1094, 2749, 153, 2686, 544, 752, 3085, 1371, 2407, 3675, 1162, 1938, 1197, 3571, 2023, 2847, 1807, 1307, 793, 2610, 1469, 22, 1883, 396, 1098, 1704, 1450, 250, 1258, 3453, 2866, 1995, 2336, 1917, 1724, 3805, 41, 2461, 1241, 2376, 467, 3730, 3090, 3234, 3104, 183, 2827, 274, 1488, 1608, 2495, 3633, 3554, 2723, 3358, 2214, 2963, 3648, 3698, 1569, 3270, 1646, 2675, 2014, 2165, 3106, 2209, 2352, 3580, 3597, 659, 2349, 2074, 988, 1952, 3821, 2640, 3727, 3380, 2646, 3741, 1753, 679, 1707, 633, 1224, 1261, 1501, 1942, 935, 729, 3293, 3638, 2759, 1214, 3028, 3703, 2260, 1406, 2531, 737, 3462, 1495, 1728, 3366, 3510, 42, 3659, 2953, 2378, 1330, 3474, 1372, 89, 1153, 1825, 3218, 3068, 1888, 3287, 2071, 2082, 1460, 3761, 3480, 3424, 651, 1618, 1859, 960, 3344, 3725, 2942, 2176, 1651, 2936, 1187, 2060, 2021, 3317, 861, 1259, 241, 3528, 200, 2392, 1316, 2486, 2923, 923, 98, 3549, 1431, 534, 1840, 3208, 201, 3605, 1337, 877, 571, 3147, 2678, 460, 2970, 3161, 1409, 2304, 1955, 1338, 2364, 754, 3635, 3264, 2620, 889, 566, 744, 1848, 954, 812, 549, 1190, 745, 2371, 2590, 1759, 1710, 1203, 447, 2068, 2691, 245, 880, 122, 3182, 2997, 1934, 1139, 1491, 1166, 1368, 2030, 162, 1829, 766, 3447, 3451, 3386, 2667, 2162, 2096, 3215, 2133, 3620, 743, 1342, 3385, 446, 3557, 3305, 883, 3616, 2130, 1182, 1346, 530, 1732, 1233, 292, 3787, 2467, 3514, 230, 652, 908, 3318, 1213, 3548, 2957, 1552, 2290, 2036, 1102, 3276, 1897, 1886, 2384, 1625, 3143, 1711, 2457, 2832, 2056, 1746, 3361, 2271, 2532, 1981, 1097, 1008, 705, 3671, 2875, 2870, 1760, 582, 3014, 1524, 2682, 712, 916, 461, 474, 3245, 2337, 2077, 2715, 1765, 3409, 2826, 3734, 998, 3813, 233, 1336, 1880, 1703, 2607, 1663, 1476, 380, 3015, 1595, 132, 3737, 125, 2421, 981, 2966, 1961, 991, 3216, 723, 526, 660, 1309, 2529, 3004, 724, 2595, 2573, 354, 3352, 1436, 15, 1184, 3190, 1167, 1758, 81, 3407, 3669, 3141, 3801, 1953, 2898]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from app import models\n",
+    "\n",
+    "# model = models.RETAIN(input_dim=99, hidden_dim=128)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def extract_backbone_param(ckpt):\n",
+    "    backbone = {}\n",
+    "    for k,v in ckpt.items():\n",
+    "        if \"backbone\" in k:\n",
+    "            new_k = k.replace(\"backbone.\", \"\")\n",
+    "            backbone[new_k] = v\n",
+    "    return backbone"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = pd.read_pickle(\"datasets/hm/processed_data/x.pkl\")\n",
+    "y = pd.read_pickle(\"datasets/hm/processed_data/y.pkl\")\n",
+    "visits_length = pd.read_pickle(\"datasets/hm/processed_data/visits_length.pkl\")\n",
+    "x = x[val_idx]\n",
+    "y = y[val_idx]\n",
+    "visits_length = visits_length[val_idx]\n",
+    "device = torch.device(\"cpu\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outcome_status=[]\n",
+    "patient=[]\n",
+    "for i in range(len(visits_length)):\n",
+    "    outcome_status.append(y[i][visits_length[i]-1][0])\n",
+    "    patient.append(x[i][visits_length[i]-1].detach().numpy())\n",
+    "\n",
+    "outcome_status = torch.tensor(outcome_status)\n",
+    "patient = torch.tensor(patient)\n",
+    "# outcome_status = y[:, 0, 0]\n",
+    "outcome_status = outcome_status.unsqueeze(-1)\n",
+    "# patient = x[:, 0, :]\n",
+    "patient = torch.unsqueeze(patient, dim=1)\n",
+    "patient = patient.float()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outcome_status.shape, patient.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def remove_outliers(df,columns,n_std):\n",
+    "    for col in columns:\n",
+    "        mean = df[col].mean()\n",
+    "        sd = df[col].std()\n",
+    "        df = df[abs(df[col]-mean) <= sd*n_std]\n",
+    "    return df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_std = 3\n",
+    "approach = 'pca' # 'pca' or 'tsne'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Multitask Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model = models.RETAIN(input_dim=99, hidden_dim=128)\n",
+    "hidden_dim=128\n",
+    "model = models.ConCare(\n",
+    "    lab_dim=97,\n",
+    "    demo_dim=2,\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    d_model=hidden_dim,\n",
+    "    MHD_num_head=4,\n",
+    "    d_ff=4 * hidden_dim,\n",
+    "    drop=0.0,\n",
+    ")\n",
+    "\n",
+    "multitask_ckpt = torch.load(\"./checkpoints/hm_multitask_concare_ep100_kf10_bs64_hid128_1_seed0.pth\", map_location=torch.device('cpu'))\n",
+    "multitask_backbone = extract_backbone_param(multitask_ckpt)\n",
+    "model.load_state_dict(multitask_backbone)\n",
+    "out = model(patient, device)\n",
+    "out = torch.squeeze(out)\n",
+    "out = out.detach().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if approach == 'pca':\n",
+    "    projected = PCA(2).fit_transform(out)\n",
+    "else:\n",
+    "    projected = TSNE(n_components=2, learning_rate='auto', init='random').fit_transform(out)\n",
+    "\n",
+    "concatenated = np.concatenate([projected, outcome_status], axis=1)\n",
+    "df = pd.DataFrame(concatenated, columns = ['Component 1', 'Component 2', 'Outcome'])\n",
+    "df = remove_outliers(df, ['Component 1', 'Component 2'], n_std)\n",
+    "df['Outcome'].replace({1: 'Dead', 0: 'Alive'}, inplace=True)\n",
+    "\n",
+    "sns.scatterplot(data=df, x=\"Component 1\", y=\"Component 2\", hue=\"Outcome\", style=\"Outcome\", palette=[\"C2\", \"C3\"], alpha=0.5)\n",
+    "plt.savefig(f'multitask_{approach}.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Outcome Prediction Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hidden_dim=128\n",
+    "model = models.ConCare(\n",
+    "    lab_dim=97,\n",
+    "    demo_dim=2,\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    d_model=hidden_dim,\n",
+    "    MHD_num_head=4,\n",
+    "    d_ff=4 * hidden_dim,\n",
+    "    drop=0.0,\n",
+    ")\n",
+    "\n",
+    "outcome_ckpt = torch.load(\"./checkpoints/hm_outcome_concare_ep100_kf10_bs64_hid128_1_seed0.pth\", map_location=torch.device('cpu'))\n",
+    "outcome_backbone = extract_backbone_param(outcome_ckpt)\n",
+    "model.load_state_dict(outcome_backbone)\n",
+    "out = model(patient, device)\n",
+    "out = torch.squeeze(out)\n",
+    "out = out.detach().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if approach == 'pca':\n",
+    "    projected = PCA(2).fit_transform(out)\n",
+    "else:\n",
+    "    projected = TSNE(n_components=2, learning_rate='auto', init='random').fit_transform(out)\n",
+    "\n",
+    "concatenated = np.concatenate([projected, outcome_status], axis=1)\n",
+    "df = pd.DataFrame(concatenated, columns = ['Component 1', 'Component 2', 'Outcome'])\n",
+    "df = remove_outliers(df, ['Component 1', 'Component 2'], n_std)\n",
+    "df['Outcome'].replace({1: 'Dead', 0: 'Alive'}, inplace=True)\n",
+    "\n",
+    "sns.scatterplot(data=df, x=\"Component 1\", y=\"Component 2\", hue=\"Outcome\", style=\"Outcome\", palette=[\"C2\", \"C3\"], alpha=0.5)\n",
+    "plt.savefig(f'outcome_{approach}.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### LOS Prediction model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hidden_dim=64\n",
+    "model = models.ConCare(\n",
+    "    lab_dim=97,\n",
+    "    demo_dim=2,\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    d_model=hidden_dim,\n",
+    "    MHD_num_head=4,\n",
+    "    d_ff=4 * hidden_dim,\n",
+    "    drop=0.0,\n",
+    ")\n",
+    "\n",
+    "los_ckpt = torch.load(\"./checkpoints/hm_los_concare_ep100_kf10_bs64_hid64_1_seed0.pth\", map_location=torch.device('cpu'))\n",
+    "los_backbone = extract_backbone_param(los_ckpt)\n",
+    "model.load_state_dict(los_backbone)\n",
+    "out = model(patient, device)\n",
+    "out = torch.squeeze(out)\n",
+    "out = out.detach().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if approach == 'pca':\n",
+    "    projected = PCA(2).fit_transform(out)\n",
+    "else:\n",
+    "    projected = TSNE(n_components=2, learning_rate='auto', init='random').fit_transform(out)\n",
+    "\n",
+    "concatenated = np.concatenate([projected, outcome_status], axis=1)\n",
+    "df = pd.DataFrame(concatenated, columns = ['Component 1', 'Component 2', 'Outcome'])\n",
+    "df = remove_outliers(df, ['Component 1', 'Component 2'], n_std)\n",
+    "df['Outcome'].replace({1: 'Dead', 0: 'Alive'}, inplace=True)\n",
+    "\n",
+    "sns.scatterplot(data=df, x=\"Component 1\", y=\"Component 2\", hue=\"Outcome\", style=\"Outcome\", palette=[\"C2\", \"C3\"], alpha=0.5)\n",
+    "plt.savefig(f'los_{approach}.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Case Study"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### CDSL dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = pd.read_pickle(\"datasets/hm/processed_data/x.pkl\")\n",
+    "y = pd.read_pickle(\"datasets/hm/processed_data/y.pkl\")\n",
+    "visits_length = pd.read_pickle(\"datasets/hm/processed_data/visits_length.pkl\")\n",
+    "x = x[val_idx]\n",
+    "y = y[val_idx]\n",
+    "visits_length = visits_length[val_idx]\n",
+    "device = torch.device(\"cpu\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "long_visits_id_list = []\n",
+    "\n",
+    "for i in range(len(visits_length)):\n",
+    "    if visits_length[i] > 20:\n",
+    "        long_visits_id_list.append(i)\n",
+    "        print(f\"[{i}: {y[i][0][0].item()}]\", end=\"    \")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### ConCare Multitask"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "idx=20 # 60, 243 | 24, 114   129 20\n",
+    "outcome_status=y[idx][0][0]\n",
+    "los_status=y[idx][:visits_length[idx]][:,1]\n",
+    "patient=x[idx][:visits_length[idx]]\n",
+    "\n",
+    "outcome_status = outcome_status.unsqueeze(-1)\n",
+    "patient = patient.float()\n",
+    "\n",
+    "hidden_dim=128\n",
+    "backbone = models.ConCare(\n",
+    "    lab_dim=97,\n",
+    "    demo_dim=2,\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    d_model=hidden_dim,\n",
+    "    MHD_num_head=4,\n",
+    "    d_ff=4 * hidden_dim,\n",
+    "    drop=0.0,\n",
+    ")\n",
+    "head = models.MultitaskHead(\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    output_dim=1,\n",
+    ")\n",
+    "model = models.Model(backbone, head)\n",
+    "los_ckpt = torch.load(\"./checkpoints/hm_multitask_concare_ep100_kf10_bs64_hid128_1_seed0.pth\", map_location=torch.device('cpu'))\n",
+    "model.load_state_dict(los_ckpt)\n",
+    "\n",
+    "# ConCare model does not accept single patient input, so we need to create a batch of size 2\n",
+    "patient = torch.stack((patient, patient), dim=0)\n",
+    "risk, out = model(patient, device, None)\n",
+    "out = out[0]\n",
+    "risk = risk[0]\n",
+    "los_statistics = {'los_mean': 6.1315513, 'los_std': 5.6816683}\n",
+    "out = torch.squeeze(out)\n",
+    "risk = torch.squeeze(risk)\n",
+    "out = out * los_statistics['los_std'] + los_statistics['los_mean']\n",
+    "\n",
+    "los_status, out, outcome_status, risk"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hidden_dim=128\n",
+    "backbone = models.ConCare(\n",
+    "    lab_dim=97,\n",
+    "    demo_dim=2,\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    d_model=hidden_dim,\n",
+    "    MHD_num_head=4,\n",
+    "    d_ff=4 * hidden_dim,\n",
+    "    drop=0.0,\n",
+    ")\n",
+    "head = models.MultitaskHead(\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    output_dim=1,\n",
+    ")\n",
+    "model = models.Model(backbone, head)\n",
+    "los_ckpt = torch.load(\"./checkpoints/hm_multitask_concare_ep100_kf10_bs64_hid128_1_seed0.pth\", map_location=torch.device('cpu'))\n",
+    "model.load_state_dict(los_ckpt)\n",
+    "\n",
+    "def inference_case(idx):\n",
+    "    outcome_status=y[idx][0][0]\n",
+    "    los_status=y[idx][:visits_length[idx]][:,1]\n",
+    "    patient=x[idx][:visits_length[idx]]\n",
+    "\n",
+    "    outcome_status = outcome_status.unsqueeze(-1)\n",
+    "    patient = patient.float()\n",
+    "\n",
+    "    # ConCare model does not accept single patient input, so we need to create a batch of size 2\n",
+    "    patient = torch.stack((patient, patient), dim=0)\n",
+    "    risk, out = model(patient, device, None)\n",
+    "    out = out[0]\n",
+    "    risk = risk[0]\n",
+    "    los_statistics = {'los_mean': 6.1315513, 'los_std': 5.6816683}\n",
+    "    out = torch.squeeze(out)\n",
+    "    risk = torch.squeeze(risk)\n",
+    "    out = out * los_statistics['los_std'] + los_statistics['los_mean']\n",
+    "    # print(\"los gt:\", los_status)\n",
+    "    # print(\"los pred:\", out)\n",
+    "    # print(\"outcome:\", outcome_status)\n",
+    "    # print(\"risk pred:\", risk)\n",
+    "    # print(\"--------------------\")\n",
+    "    return los_status.cpu().detach().numpy(), out.cpu().detach().numpy(), outcome_status.cpu().detach().numpy(), risk.cpu().detach().numpy()\n",
+    "\n",
+    "\n",
+    "los_status, out, outcome_status, risk = inference_case(60)\n",
+    "los_status, out, outcome_status, risk"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_case(los_status, out, outcome_status, risk, idx):\n",
+    "    color = 'green' if outcome_status == 0 else 'red'\n",
+    "    label_info = f'Alive Case #{idx}' if outcome_status == 0 else f'Dead Case #{idx}'\n",
+    "    filename = f'case_study_los_alive_{idx}' if outcome_status == 0 else f'case_study_los_dead_{idx}'\n",
+    "    los_status = np.negative(los_status)\n",
+    "    fig = plt.figure()\n",
+    "    ax = fig.add_subplot(111)\n",
+    "    ax2 = ax.twinx()\n",
+    "    ax.plot(los_status, out, marker='o', linewidth=1, label=label_info, color=color)\n",
+    "    ax2.plot(los_status, risk, marker=',', linewidth=2, label='Risk', color='skyblue')\n",
+    "    ax.set_ylim([0, 30])\n",
+    "    ax2.set_ylim([0, 1])\n",
+    "    ax2.grid(False)\n",
+    "    ax.plot([0, -30], [0, 30], linestyle='-.', color='grey')\n",
+    "    ax.set_xlabel('True Length of Stay (days)')\n",
+    "    ax.set_ylabel('Predicted Length of Stay (days)')\n",
+    "    ax.legend(loc='upper left')\n",
+    "    ax2.legend(loc='lower left')\n",
+    "\n",
+    "    # plt.savefig(f'case_study_los_{idx}.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "    plt.savefig(f'cases/cdsl/{filename}.png', format=\"png\", bbox_inches=\"tight\")\n",
+    "    plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def abs_var(x, position):\n",
+    "    return f'{abs(int(x))}'\n",
+    "\n",
+    "# case 1\n",
+    "los_status, out, outcome_status, risk = inference_case(60)\n",
+    "los_status = np.negative(los_status)\n",
+    "fig, (ax, bx) = plt.subplots(1, 2, figsize=(12, 5))\n",
+    "ax2 = ax.twinx()\n",
+    "ax.plot(los_status, out, marker='o', linewidth=1, label=\"Alive case\", color='green')\n",
+    "ax2.plot(los_status, risk, marker=',', linewidth=2, label='Risk', color='skyblue')\n",
+    "ax.set_ylim([0, 22])\n",
+    "ax2.set_ylim([0, 1])\n",
+    "ax2.grid(False)\n",
+    "ax.plot([0, -22], [0, 22], linestyle='-.', color='grey')\n",
+    "ax.set_xlabel('True Length of Stay (days)')\n",
+    "ax.set_ylabel('Predicted Length of Stay (days)')\n",
+    "ax2.set_ylabel('Risk')\n",
+    "\n",
+    "fig.gca().xaxis.set_major_formatter(FuncFormatter(abs_var))\n",
+    "\n",
+    "# case 2\n",
+    "los_status, out, outcome_status, risk = inference_case(169)\n",
+    "los_status = np.negative(los_status)\n",
+    "\n",
+    "bx2 = bx.twinx()\n",
+    "bx.plot(los_status, out, marker='o', linewidth=1, label='Dead case', color='red')\n",
+    "bx2.plot(los_status, risk, marker=',', linewidth=2, color='skyblue')\n",
+    "bx.set_ylim([0, 22])\n",
+    "bx2.set_ylim([0, 1])\n",
+    "bx2.grid(False)\n",
+    "bx.plot([0, -22], [0, 22], linestyle='-.', color='grey')\n",
+    "bx.set_xlabel('True Length of Stay (days)')\n",
+    "bx.set_ylabel('Predicted Length of Stay (days)')\n",
+    "bx2.set_ylabel('Risk')\n",
+    "\n",
+    "fig.gca().xaxis.set_major_formatter(FuncFormatter(abs_var))\n",
+    "\n",
+    "fig.legend(bbox_to_anchor=(0.5, 1.05), loc=\"upper center\", ncol=3)\n",
+    "fig.tight_layout()\n",
+    "plt.savefig('cases/cdsl_case_study.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# for idx in long_visits_id_list:\n",
+    "#     los_status, out, outcome_status, risk = inference_case(idx)\n",
+    "#     plot_case(los_status, out, outcome_status, risk, idx)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### TJH dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_idx = [287, 116, 186, 292, 225, 290, 277, 311, 20, 71, 52, 304, 87, 74, 318, 92, 121, 236, 226, 149, 295, 103, 14, 305, 213, 165, 174, 106, 99, 102, 151, 177, 233, 130, 1, 270]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = pd.read_pickle(\"datasets/tongji/processed_data/x.pkl\")\n",
+    "y = pd.read_pickle(\"datasets/tongji/processed_data/y.pkl\")\n",
+    "visits_length = pd.read_pickle(\"datasets/tongji/processed_data/visits_length.pkl\")\n",
+    "x = x[val_idx]\n",
+    "y = y[val_idx]\n",
+    "visits_length = visits_length[val_idx]\n",
+    "device = torch.device(\"cpu\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "long_visits_id_list = []\n",
+    "\n",
+    "for i in range(len(visits_length)):\n",
+    "    if visits_length[i] > 5:\n",
+    "        long_visits_id_list.append(i)\n",
+    "        print(f\"[{i}: {y[i][0][0].item()}] len:{visits_length[i]}\", end=\"    \")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### RETAIN multitask"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hidden_dim=64\n",
+    "backbone = models.RETAIN(\n",
+    "    input_dim=75,\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    dropout=0.0,\n",
+    ")\n",
+    "head = models.MultitaskHead(\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    output_dim=1,\n",
+    ")\n",
+    "model = models.Model(backbone, head)\n",
+    "los_ckpt = torch.load(\"./checkpoints/tj_multitask_retain_ep100_kf10_bs64_hid64_1_seed42.pth\", map_location=torch.device('cpu'))\n",
+    "model.load_state_dict(los_ckpt)\n",
+    "\n",
+    "def inference_case(idx):\n",
+    "    outcome_status=y[idx][0][0]\n",
+    "    los_status=y[idx][:visits_length[idx]][:,1]\n",
+    "    patient=x[idx][:visits_length[idx]]\n",
+    "\n",
+    "    outcome_status = outcome_status.unsqueeze(-1)\n",
+    "    patient = patient.float()\n",
+    "\n",
+    "    patient = torch.stack((patient, patient), dim=0)\n",
+    "    risk, out = model(patient, device, None)\n",
+    "    out = out[0]\n",
+    "    risk = risk[0]\n",
+    "    los_statistics = {'los_mean': 7.7147756, 'los_std': 7.1851807}\n",
+    "    out = torch.squeeze(out)\n",
+    "    risk = torch.squeeze(risk)\n",
+    "    out = out * los_statistics['los_std'] + los_statistics['los_mean']\n",
+    "    # print(\"los gt:\", los_status)\n",
+    "    # print(\"los pred:\", out)\n",
+    "    # print(\"outcome:\", outcome_status)\n",
+    "    # print(\"risk pred:\", risk)\n",
+    "    # print(\"--------------------\")\n",
+    "    return los_status.cpu().detach().numpy(), out.cpu().detach().numpy(), outcome_status.cpu().detach().numpy(), risk.cpu().detach().numpy()\n",
+    "\n",
+    "\n",
+    "los_status, out, outcome_status, risk = inference_case(30)\n",
+    "los_status, out, outcome_status, risk"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_case(los_status, out, outcome_status, risk, idx):\n",
+    "    color = 'green' if outcome_status == 0 else 'red'\n",
+    "    label_info = f'Alive Case #{idx}' if outcome_status == 0 else f'Dead Case #{idx}'\n",
+    "    filename = f'case_study_los_alive_{idx}' if outcome_status == 0 else f'case_study_los_dead_{idx}'\n",
+    "    los_status = np.negative(los_status)\n",
+    "    fig = plt.figure()\n",
+    "    ax = fig.add_subplot(111)\n",
+    "    ax2 = ax.twinx()\n",
+    "    ax.plot(los_status, out, marker='o', linewidth=1, label=label_info, color=color)\n",
+    "    ax2.plot(los_status, risk, marker=',', linewidth=1, label='Risk', color='skyblue')\n",
+    "    ax.set_ylim([0, 30])\n",
+    "    ax2.set_ylim([0, 1])\n",
+    "    ax2.grid(False)\n",
+    "    ax.plot([0, -30], [0, 30], linestyle='-.', color='grey')\n",
+    "    ax.set_xlabel('True Length of Stay (days)')\n",
+    "    ax.set_ylabel('Predicted Length of Stay (days)')\n",
+    "    ax.legend(loc='upper left')\n",
+    "    ax2.legend(loc='lower left')\n",
+    "\n",
+    "    # plt.savefig(f'case_study_los_{idx}.pdf', dpi=500, format=\"pdf\", bbox_inches=\"tight\")\n",
+    "    plt.savefig(f'cases/tjh/{filename}.png', format=\"png\", bbox_inches=\"tight\")\n",
+    "    plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# for idx in long_visits_id_list:\n",
+    "#     los_status, out, outcome_status, risk = inference_case(idx)\n",
+    "#     plot_case(los_status, out, outcome_status, risk, idx)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.9.5 ('pytorch')",
+   "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.5"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "e382889b16d65b8f9d2caeea05d88db6d501b8794eac9af8ee0956d5affe33e5"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}