{

 "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

}