{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "069a0e55-8080-43d0-b206-d53656e0c009",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import glob\n",
    "from tqdm import tqdm, trange\n",
    "import joblib\n",
    "\n",
    "# math imports\n",
    "import numpy as np\n",
    "import scipy\n",
    "import sklearn\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import warnings\n",
    "import matplotlib\n",
    "from matplotlib import pyplot\n",
    "\n",
    "import copy\n",
    "from functools import reduce\n",
    "\n",
    "from dataclasses import dataclass\n",
    "\n",
    "import sys\n",
    "sys.path.append('../../')\n",
    "from src.data_class import matrix_class\n",
    "from utils.util import reduce_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec7618d1-d9f6-4c47-ae8a-4ccfcdeab0a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_questionnaires(pickle_list, identity_col, threshold):\n",
    "    response_list = []\n",
    "    dataname_list = []\n",
    "    description_list = []\n",
    "    df_list = []\n",
    "    for file in pickle_list:\n",
    "        questionnaire = pickle.load(open(file, 'rb'))\n",
    "        df_list.append(questionnaire['data'])\n",
    "        response_list.append([r[0] for r in questionnaire['response_info']])\n",
    "        dataname_list.append(file.replace('.pickle', '').split('/')[-1])\n",
    "        description_list.append(questionnaire['data_info'])\n",
    "    dfmerge = reduce(lambda left,right: pd.merge(left, right, how='outer', left_on=identity_col, right_on=identity_col), df_list)\n",
    "    \n",
    "    include_list = []\n",
    "    for (idx, q) in enumerate(dataname_list):\n",
    "        Mvalue = dfmerge[ response_list[idx] ].values\n",
    "        include = np.sum(pd.isnull(Mvalue))/(Mvalue.shape[0]*Mvalue.shape[1]) < 1.0 - threshold\n",
    "        include_list.append([q, Mvalue.shape[1], include])\n",
    "            \n",
    "    filter_list = pd.DataFrame(include_list, columns=['Questionnaire', 'dimension', '>20% available'])\n",
    "    filter_list = filter_list.sort_values(by=['>20% available', 'Questionnaire']).reset_index(drop=True)\n",
    "    display(filter_list)\n",
    "    return [q[0] for q in include_list if q[2] == True]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "39398704-da4a-4051-814b-6a1878a0b142",
   "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>Questionnaire</th>\n",
       "      <th>dimension</th>\n",
       "      <th>&gt;20% available</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ASR</td>\n",
       "      <td>126</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>CAARS</td>\n",
       "      <td>26</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CDI2_P</td>\n",
       "      <td>17</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CDI2_SR</td>\n",
       "      <td>10</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>ARI_P</td>\n",
       "      <td>7</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>ARI_S</td>\n",
       "      <td>7</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>ASSQ</td>\n",
       "      <td>27</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>C3SR</td>\n",
       "      <td>39</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>CBCL</td>\n",
       "      <td>119</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>ESWAN</td>\n",
       "      <td>65</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>ICU_P</td>\n",
       "      <td>24</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>ICU_SR</td>\n",
       "      <td>24</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>MFQ_P</td>\n",
       "      <td>34</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>MFQ_SR</td>\n",
       "      <td>33</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>PANAS</td>\n",
       "      <td>20</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>RBS</td>\n",
       "      <td>43</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>SCARED_P</td>\n",
       "      <td>41</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>SCARED_SR</td>\n",
       "      <td>41</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>SCQ</td>\n",
       "      <td>40</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>SDQ</td>\n",
       "      <td>33</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>SRS</td>\n",
       "      <td>65</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>SWAN</td>\n",
       "      <td>18</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>SympChck</td>\n",
       "      <td>63</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>TRF</td>\n",
       "      <td>116</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>YSR</td>\n",
       "      <td>119</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Questionnaire  dimension  >20% available\n",
       "0            ASR        126           False\n",
       "1          CAARS         26           False\n",
       "2         CDI2_P         17           False\n",
       "3        CDI2_SR         10           False\n",
       "4          ARI_P          7            True\n",
       "5          ARI_S          7            True\n",
       "6           ASSQ         27            True\n",
       "7           C3SR         39            True\n",
       "8           CBCL        119            True\n",
       "9          ESWAN         65            True\n",
       "10         ICU_P         24            True\n",
       "11        ICU_SR         24            True\n",
       "12         MFQ_P         34            True\n",
       "13        MFQ_SR         33            True\n",
       "14         PANAS         20            True\n",
       "15           RBS         43            True\n",
       "16      SCARED_P         41            True\n",
       "17     SCARED_SR         41            True\n",
       "18           SCQ         40            True\n",
       "19           SDQ         33            True\n",
       "20           SRS         65            True\n",
       "21          SWAN         18            True\n",
       "22      SympChck         63            True\n",
       "23           TRF        116            True\n",
       "24           YSR        119            True"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pickle_list = glob.glob('./processed/pickle/*.pickle')\n",
    "pickle_list = [ p for p in pickle_list if 'Barratt' not in p]\n",
    "pickle_list = [ p for p in pickle_list if '_pos.' not in p]\n",
    "pickle_list.sort()\n",
    "identity_col = ['EID', 'Age', 'Sex']\n",
    "threshold = 0.2\n",
    "\n",
    "dataname_list  = filter_questionnaires(pickle_list, identity_col, threshold)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "586fe242-581d-4b46-a3fa-fa05f464c952",
   "metadata": {},
   "source": [
    "### Detect duplicate subject"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "00127929-553f-4618-b195-26b9c3703d33",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 21/21 [00:00<00:00, 56.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The following subjects have duplicated responses: \n",
      "['NDARPL406KX4']\n"
     ]
    }
   ],
   "source": [
    "df_list = []\n",
    "sdf_list = [] # scaled df list\n",
    "norm_df_list = []\n",
    "response_list = []\n",
    "minmax_list = []\n",
    "\n",
    "max_age = 0\n",
    "min_age = 100\n",
    "\n",
    "for dataname in tqdm(dataname_list):\n",
    "    data = pickle.load(open('./processed/pickle/{}.pickle'.format(dataname), 'rb'))\n",
    "\n",
    "    df = data['data']\n",
    "    response_info = data['response_info']\n",
    "    itemlist = [r[0] for r in response_info]\n",
    "    for i in itemlist: minmax_list.append([data['min_response'], data['max_response']])\n",
    "    if data['max_age'] > max_age: max_age = data['max_age']\n",
    "    if data['min_age'] < min_age: min_age = data['min_age']\n",
    "\n",
    "    # Drop duplicate rows\n",
    "    df = df.drop_duplicates(subset=['Age', 'EID'] + itemlist)\n",
    "    df_list.append(df)\n",
    "    response_list.append(response_info)\n",
    "    \n",
    "minmax_list = np.array(minmax_list)\n",
    "df_merge = reduce(lambda left,right: pd.merge(left,right, how='outer'), df_list)\n",
    "\n",
    "unique_EID, unique_idx = np.unique(df_merge['EID'].values, return_index=True)\n",
    "duplicates = df_merge.iloc[list(set(np.arange(df_merge.shape[0])) - set(unique_idx))]['EID'].values\n",
    "print('The following subjects have duplicated responses: ')\n",
    "print(duplicates)\n",
    "# display(df_merge.loc[df_merge['EID'] == duplicates[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceba5657-40c5-4659-84b3-ab16dce8ddec",
   "metadata": {},
   "source": [
    "### drop duplicated subjects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f43b1ed5-593a-4221-bad1-1dca70b1a448",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unique subjects : 3572\n",
      "Age range : 5.00 -- 21.90\n"
     ]
    }
   ],
   "source": [
    "df_merge = df_merge.drop([3568])\n",
    "df_merge.reset_index(inplace=True)\n",
    "unique_EID, unique_idx = np.unique(df_merge['EID'].values, return_index=True)\n",
    "duplicates = df_merge.iloc[list(set(np.arange(df_merge.shape[0])) - set(unique_idx))]['EID'].values\n",
    "assert len(duplicates) == 0\n",
    "\n",
    "print('Unique subjects : {}'.format(df_merge.shape[0]))\n",
    "print('Age range : {:.2f} -- {:.2f}'.format(min_age, max_age))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19e792af-cbd5-4350-9c86-676cb4049482",
   "metadata": {},
   "source": [
    "### merge questionnaires"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0586189c-d5a6-4de8-abd5-8e613b4828fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "response_merge_list = [r for sublist in response_list for r in sublist]\n",
    "itemlist = [r[0] for r in response_merge_list]\n",
    "\n",
    "M_raw = df_merge[itemlist].values\n",
    "nan_mask = 1.0 - np.isnan(M_raw)\n",
    "\n",
    "M = copy.deepcopy(M_raw)\n",
    "for i in range(M.shape[1]):\n",
    "    lb, ub = minmax_list[i][0], minmax_list[i][1]\n",
    "    M[:,i][np.isnan(M[:,i])] = lb\n",
    "    M[:,i] = (M[:,i] - lb)/(ub - lb)\n",
    "\n",
    "assert np.sum(np.isnan(M)) == 0\n",
    "assert np.max(M) <= 1.0\n",
    "assert np.min(M) >= 0.0\n",
    "\n",
    "confound_raw = df_merge[['Age', 'Sex']].values\n",
    "confound = copy.deepcopy(confound_raw)\n",
    "confound[:,0] = (confound[:,0] - min_age)/(max_age - min_age)\n",
    "assert np.sum(np.isnan(confound)) == 0\n",
    "\n",
    "subjlist = list(df_merge['EID'].values)\n",
    "\n",
    "matrix_merge = matrix_class(M, M_raw,\n",
    "                            confound, confound_raw,\n",
    "                            nan_mask,\n",
    "                            None, None, None, # row_idx, col_idx, mask\n",
    "                            'Merged questionnaires',\n",
    "                            subjlist,\n",
    "                            response_merge_list,\n",
    "                            None, None, # W, Q\n",
    "                            None, None, # C, Qc\n",
    "                            None, None) # Z, aZ\n",
    "\n",
    "os.makedirs('../../output/full-HBN_split', exist_ok=True)\n",
    "\n",
    "with open('../../output/full-HBN_split/merge.pickle', 'wb') as handle:\n",
    "    pickle.dump(matrix_merge, handle, protocol=4)\n",
    "    \n",
    "np.savez('../../output/full-HBN_split/merge.npz', M=M, nan_mask=nan_mask, confound=confound)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12ffbe92-1a11-42ba-8fd2-52c7c61c4316",
   "metadata": {},
   "source": [
    "### update matrix class for each questionnaire on row_idx, col_idx and mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "abb6c92b-1a1a-4294-894c-7e8317583dd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 21/21 [00:05<00:00,  3.57it/s]\n"
     ]
    }
   ],
   "source": [
    "matrix_merge = pickle.load(open('../../output/full-HBN_split/merge.pickle', 'rb'))\n",
    "itemlist_merge = matrix_merge.itemlist\n",
    "subjlist_merge = matrix_merge.subjlist\n",
    "\n",
    "matrix_list = []\n",
    "for dataname in tqdm(dataname_list):\n",
    "    matrixpath = './processed/dataclass/{}.pickle'.format(dataname)\n",
    "    matrix = pickle.load(open(matrixpath, 'rb'))\n",
    "    \n",
    "    dataname = matrix.dataname\n",
    "    itemlist = matrix.itemlist\n",
    "    subjlist = matrix.subjlist\n",
    "    \n",
    "    row_idx = [subjlist_merge.index(subj) for subj in subjlist]\n",
    "    col_idx = np.array([itemlist_merge.index(item) for item in itemlist])\n",
    "    \n",
    "    matrix.row_idx = row_idx\n",
    "    matrix.col_idx = col_idx\n",
    "    \n",
    "    assert matrix.M.shape[0] == len(row_idx)\n",
    "    assert matrix.M.shape[1] == len(col_idx)\n",
    "    \n",
    "    mask = np.zeros_like(matrix_merge.M)\n",
    "    index = np.ix_(row_idx, col_idx)\n",
    "    \n",
    "    mask[index] = np.ones((len(row_idx), len(col_idx)))\n",
    "    \n",
    "    matrix.mask = mask\n",
    "    matrix_list.append(matrix)\n",
    "    \n",
    "    os.makedirs('../../output/full-HBN_split/dataclass', exist_ok=True)\n",
    "    with open(os.path.join('../../output/full-HBN_split/dataclass', '{}.pickle'.format(dataname)), 'wb') as handle:\n",
    "        pickle.dump(matrix, handle, protocol=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ed3e8a2-a4aa-4746-a2fb-444e74e8669a",
   "metadata": {},
   "source": [
    "### visualize full-HBN questionnaires response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6df4e3f7-af7e-4b65-9105-a4b42f1f94e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 440508 characters]\n",
      "text/plain": [
       "<Figure size 1600x800 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "fig = pyplot.figure(figsize=(16, 8), constrained_layout=True)\n",
    "spec = fig.add_gridspec(ncols=3, nrows=1, width_ratios=[30, 1, 1])\n",
    "\n",
    "ax = fig.add_subplot(spec[0,0])\n",
    "\n",
    "age_sort = np.argsort(matrix_merge.confound[:,0])\n",
    "\n",
    "nquestion_list = [m.M.shape[1] for m in matrix_list]\n",
    "xtickslocation = np.cumsum(np.asarray([0]+nquestion_list[:-1])) + np.asarray(nquestion_list)//2\n",
    "sns.heatmap(matrix_merge.M[age_sort], ax=ax, cmap=\"Blues\", cbar=False)\n",
    "ax.set_xticks(xtickslocation)\n",
    "ax.set_xticklabels(dataname_list, rotation=90, fontsize=8)\n",
    "\n",
    "range_list = []\n",
    "for dataname in dataname_list:\n",
    "    datapath = './processed/pickle/{}.pickle'.format(dataname)\n",
    "    data = pickle.load(open(datapath, 'rb'))\n",
    "    range_list.append('{}-{}'.format(int(data['min_response']), int(data['max_response'])))\n",
    "    \n",
    "for y in np.cumsum(nquestion_list[:-1]):\n",
    "    ax.axvline(x=y, c='r')\n",
    "ax.set_xlim([0, matrix_merge.M.shape[1]])\n",
    "ax.set_title('Normalized responses')\n",
    "ax.set_xlabel('Questionnaires')\n",
    "ax.set_ylabel('Unique subjects')\n",
    "ax_a = ax.twiny()\n",
    "ax_a.set_xlim(ax.get_xlim())\n",
    "ax_a.set_xticks(xtickslocation)\n",
    "ax_a.set_xticklabels(range_list, rotation=90, fontsize=8)\n",
    "ax_a.set_xlabel('Original Value ranges')\n",
    "\n",
    "ax2 = fig.add_subplot(spec[0,1])\n",
    "\n",
    "sns.heatmap(matrix_merge.confound[:,0][age_sort][:,np.newaxis], ax=ax2, cmap=\"coolwarm\", cbar=False)\n",
    "ax2.set_xticks([0.5])\n",
    "ax2.set_xticklabels(['Age'])\n",
    "ax2.set_yticks([0, matrix_merge.M.shape[0]])\n",
    "ax2.set_yticklabels([int(data['min_age']), int(np.round(data['max_age']))])\n",
    "\n",
    "ax3 = fig.add_subplot(spec[0,2])\n",
    "sns.heatmap(matrix_merge.confound[:,1][age_sort][:,np.newaxis], ax=ax3, cmap=\"vlag\", cbar=False)\n",
    "ax3.set_xticks([0.5])\n",
    "ax3.set_xticklabels(['Sex\\n Blue=M\\n Red=F'])\n",
    "ax3.set_yticks([])\n",
    "\n",
    "pyplot.savefig('../../output/figure/full-HBN.png', dpi=300)  \n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b2c277-deec-482d-b47d-0e2ce47eaf13",
   "metadata": {},
   "source": [
    "### split dataset conditoned on population distribution and data availability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2fd6015a-5387-4f3b-8e02-fe587792c35f",
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix_merge = pickle.load(open('../../output/full-HBN_split/merge.pickle', 'rb'))\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "gender_split = matrix_merge.confound[:,1]\n",
    "\n",
    "\n",
    "age_split = np.zeros(matrix_merge.M.shape[0])\n",
    "ages = matrix_merge.confound[:,0]\n",
    "n_age_split = 10\n",
    "for i in range(n_age_split):\n",
    "    age_group = np.where(np.logical_and(ages >= i/n_age_split, ages < (i+1)/n_age_split))[0]\n",
    "    age_split[age_group] += i\n",
    "\n",
    "NaN_split = np.zeros(matrix_merge.nan_mask.shape[0])\n",
    "for res in response_list:\n",
    "    sub_col = [r[0] for r in res]\n",
    "    NaN_group = np.where(df_merge[sub_col].isnull().all(1) == 1)[0]\n",
    "    NaN_split[NaN_group] += 1\n",
    "    \n",
    "classes = age_split\n",
    "classes[gender_split == 1] *= 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa80a63f-901d-4220-b2e4-8998cb3e6825",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b2ce679a-9d63-4f06-9060-9e4af19ec850",
   "metadata": {},
   "outputs": [],
   "source": [
    "nsubj = matrix_merge.M.shape[0]\n",
    "split_idx = 0\n",
    "while split_idx < 30:\n",
    "    try: \n",
    "        _Xtrain, _Xvalid, _ytrain, _yvalid = train_test_split(np.arange(nsubj), classes, stratify=classes, test_size=0.3)\n",
    "        _Xvalid, _Xtest, _yvalid, _ytest = train_test_split(_Xvalid, _yvalid, stratify=_yvalid, test_size=0.5)\n",
    "\n",
    "        subject_train = []\n",
    "        subject_valid = []\n",
    "        subject_test = []\n",
    "\n",
    "        for group in [_Xtrain, _Xvalid, _Xtest]:\n",
    "            Xtrain, Xvalid, ytrain, yvalid = train_test_split(group, NaN_split[group], stratify=NaN_split[group], test_size=0.3)\n",
    "            try:\n",
    "                Xvalid, Xtest, yvalid, ytest = train_test_split(Xvalid, yvalid, stratify=yvalid, test_size=0.5)\n",
    "            except:\n",
    "                new_split = reduce_class(yvalid, 2)\n",
    "                Xvalid, Xtest, yvalid, ytest = train_test_split(Xvalid, new_split, stratify=new_split, test_size=0.5)\n",
    "\n",
    "            subject_train.append(Xtrain)\n",
    "            subject_valid.append(Xvalid)\n",
    "            subject_test.append(Xtest)\n",
    "        \n",
    "        subject_train = np.concatenate(subject_train)\n",
    "        subject_valid = np.concatenate(subject_valid)\n",
    "        subject_test = np.concatenate(subject_test)\n",
    "        \n",
    "        ttt_size = len(list(set( list(subject_train) + list(subject_valid) + list(subject_test) )))\n",
    "        assert nsubj == ttt_size\n",
    "        \n",
    "        os.makedirs('../../output/full-HBN_split/repeat_{}'.format(split_idx+1), exist_ok=True)\n",
    "        np.savez('../../output/full-HBN_split/repeat_{}/split.npz'.format(split_idx+1),\n",
    "                 train=subject_train, valid=subject_valid, test=subject_test)\n",
    "        \n",
    "        split_idx +=1\n",
    "        \n",
    "    except:\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5f69dab-1407-4590-8fcf-c1dcdda4f881",
   "metadata": {},
   "source": [
    "### visualize split data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "a363dfcf-868b-4160-b8f1-061406ab6252",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 83932 characters]\n",
      "text/plain": [
       "<Figure size 1600x450 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fulldata = np.load('../../output/full-HBN_split/repeat_{}/split.npz'.format(1))\n",
    "subject_train = fulldata['train']\n",
    "\n",
    "fig = pyplot.figure(figsize=(16, 4.5), constrained_layout=True)\n",
    "spec = fig.add_gridspec(ncols=3, nrows=1, width_ratios=[30, 1, 1])\n",
    "\n",
    "_split_type = subject_train\n",
    "_age_sort = np.argsort(matrix_merge.confound[:,0][_split_type])\n",
    "\n",
    "ax = fig.add_subplot(spec[0,0])\n",
    "\n",
    "nquestion_list = [m.M.shape[1] for m in matrix_list]\n",
    "xtickslocation = np.cumsum(np.asarray([0]+nquestion_list[:-1])) + np.asarray(nquestion_list)//2\n",
    "sns.heatmap(matrix_merge.nan_mask[_split_type][_age_sort], ax=ax, cmap=\"Blues\", cbar=False)\n",
    "ax.set_xticks(xtickslocation)\n",
    "ax.set_xticklabels(dataname_list, rotation=90, fontsize=12)\n",
    "\n",
    "range_list = []\n",
    "for dataname in dataname_list:\n",
    "    datapath = './processed/pickle/{}.pickle'.format(dataname)\n",
    "    data = pickle.load(open(datapath, 'rb'))\n",
    "    range_list.append('{}-{}'.format(int(data['min_response']), int(data['max_response'])))\n",
    "\n",
    "for y in np.cumsum(nquestion_list[:-1]):\n",
    "    ax.axvline(x=y, c='r')\n",
    "ax.set_xlim([0, matrix_merge.M[_split_type].shape[1]])\n",
    "ax.set_title('Data avaliability (train)', fontsize=26)\n",
    "ax.set_ylabel('Subjects', fontsize=22)\n",
    "ax.set_xticks([])\n",
    "ax_a = ax.twiny()\n",
    "ax_a.set_xlim(ax.get_xlim())\n",
    "ax_a.set_xticks(xtickslocation)\n",
    "ax_a.set_xticklabels(range_list, rotation=90, fontsize=12)\n",
    "ax_a.set_xlabel('Original Value ranges', fontsize=22)\n",
    "\n",
    "split_min_age = np.min(matrix_merge.confound[:,0][_split_type][_age_sort])\n",
    "split_max_age = np.max(matrix_merge.confound[:,0][_split_type][_age_sort])\n",
    "split_size = len(_age_sort)\n",
    "\n",
    "ax2 = fig.add_subplot(spec[0,1])\n",
    "sns.heatmap(matrix_merge.confound[:,0][_split_type][_age_sort][:,np.newaxis], ax=ax2, cmap=\"coolwarm\", cbar=False)\n",
    "ax2.set_xticks([0.5])\n",
    "ax2.set_xticklabels(['Age'], fontsize=18)\n",
    "ax2.set_yticks([0, split_size])\n",
    "ax2.set_yticklabels([5, int(np.round(split_max_age*data['max_age']))], fontsize=18)\n",
    "\n",
    "ax3 = fig.add_subplot(spec[0,2])\n",
    "sns.heatmap(matrix_merge.confound[:,1][_split_type][_age_sort][:,np.newaxis], ax=ax3, cmap=\"vlag\", cbar=False)\n",
    "ax3.set_xticks([0.5])\n",
    "ax3.set_xticklabels(['Sex'], fontsize=18)\n",
    "ax3.set_yticks([])\n",
    "\n",
    "pyplot.savefig('../../output/figure/merge_train.png', dpi=300)  \n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8cc8d8b3-f75c-45a1-a2ac-428cfd5d99b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53984 characters]\n",
      "text/plain": [
       "<Figure size 1600x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "subject_valid = fulldata['valid']\n",
    "\n",
    "fig = pyplot.figure(figsize=(16, 3), constrained_layout=True)\n",
    "spec = fig.add_gridspec(ncols=3, nrows=1, width_ratios=[30, 1, 1])\n",
    "\n",
    "_split_type = subject_valid\n",
    "_age_sort = np.argsort(matrix_merge.confound[:,0][_split_type])\n",
    "\n",
    "ax = fig.add_subplot(spec[0,0])\n",
    "\n",
    "nquestion_list = [m.M.shape[1] for m in matrix_list]\n",
    "xtickslocation = np.cumsum(np.asarray([0]+nquestion_list[:-1])) + np.asarray(nquestion_list)//2\n",
    "sns.heatmap(matrix_merge.nan_mask[_split_type][_age_sort], ax=ax, cmap=\"Blues\", cbar=False)\n",
    "ax.set_xticks(xtickslocation)\n",
    "ax.set_xticklabels(dataname_list, rotation=90, fontsize=12)\n",
    "\n",
    "range_list = []\n",
    "for dataname in dataname_list:\n",
    "    datapath = './processed/pickle/{}.pickle'.format(dataname)\n",
    "    data = pickle.load(open(datapath, 'rb'))\n",
    "    range_list.append('{}-{}'.format(int(data['min_response']), int(data['max_response'])))\n",
    "\n",
    "for y in np.cumsum(nquestion_list[:-1]):\n",
    "    ax.axvline(x=y, c='r')\n",
    "ax.set_xlim([0, matrix_merge.M[_split_type].shape[1]])\n",
    "ax.set_title('Data avaliability (validation)', fontsize=26)\n",
    "ax.set_xticks([])\n",
    "ax.set_ylabel('subjects', fontsize=22)\n",
    "\n",
    "split_min_age = np.min(matrix_merge.confound[:,0][_split_type][_age_sort])\n",
    "split_max_age = np.max(matrix_merge.confound[:,0][_split_type][_age_sort])\n",
    "split_size = len(_age_sort)\n",
    "\n",
    "ax2 = fig.add_subplot(spec[0,1])\n",
    "sns.heatmap(matrix_merge.confound[:,0][_split_type][_age_sort][:,np.newaxis], ax=ax2, cmap=\"coolwarm\", cbar=False)\n",
    "ax2.set_xticks([0.5])\n",
    "ax2.set_xticklabels(['Age'], fontsize=18)\n",
    "ax2.set_yticks([0, split_size])\n",
    "ax2.set_yticklabels([5, int(np.round(split_max_age*data['max_age']))], fontsize=18)\n",
    "\n",
    "ax3 = fig.add_subplot(spec[0,2])\n",
    "sns.heatmap(matrix_merge.confound[:,1][_split_type][_age_sort][:,np.newaxis], ax=ax3, cmap=\"vlag\", cbar=False)\n",
    "ax3.set_xticks([0.5])\n",
    "ax3.set_xticklabels(['Sex'], fontsize=18)\n",
    "ax3.set_yticks([])\n",
    "\n",
    "pyplot.savefig('../../output/figure/merge_valid.png', dpi=300)   \n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "545b84b8-73ea-4a6a-94cd-ed088020ec76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 55040 characters]\n",
      "text/plain": [
       "<Figure size 1600x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "subject_valid = fulldata['test']\n",
    "\n",
    "fig = pyplot.figure(figsize=(16, 3), constrained_layout=True)\n",
    "spec = fig.add_gridspec(ncols=3, nrows=1, width_ratios=[30, 1, 1])\n",
    "\n",
    "_split_type = subject_valid\n",
    "_age_sort = np.argsort(matrix_merge.confound[:,0][_split_type])\n",
    "\n",
    "ax = fig.add_subplot(spec[0,0])\n",
    "\n",
    "nquestion_list = [m.M.shape[1] for m in matrix_list]\n",
    "xtickslocation = np.cumsum(np.asarray([0]+nquestion_list[:-1])) + np.asarray(nquestion_list)//2\n",
    "sns.heatmap(matrix_merge.nan_mask[_split_type][_age_sort], ax=ax, cmap=\"Blues\", cbar=False)\n",
    "ax.set_xticks(xtickslocation)\n",
    "ax.set_xticklabels(dataname_list, rotation=90, fontsize=12)\n",
    "\n",
    "range_list = []\n",
    "for dataname in dataname_list:\n",
    "    datapath = './processed/pickle/{}.pickle'.format(dataname)\n",
    "    data = pickle.load(open(datapath, 'rb'))\n",
    "    range_list.append('{}-{}'.format(int(data['min_response']), int(data['max_response'])))\n",
    "\n",
    "for y in np.cumsum(nquestion_list[:-1]):\n",
    "    ax.axvline(x=y, c='r')\n",
    "ax.set_xlim([0, matrix_merge.M[_split_type].shape[1]])\n",
    "ax.set_title('Data avaliability (validation)', fontsize=26)\n",
    "ax.set_xticks([])\n",
    "ax.set_ylabel('subjects', fontsize=22)\n",
    "\n",
    "split_min_age = np.min(matrix_merge.confound[:,0][_split_type][_age_sort])\n",
    "split_max_age = np.max(matrix_merge.confound[:,0][_split_type][_age_sort])\n",
    "split_size = len(_age_sort)\n",
    "\n",
    "ax2 = fig.add_subplot(spec[0,1])\n",
    "sns.heatmap(matrix_merge.confound[:,0][_split_type][_age_sort][:,np.newaxis], ax=ax2, cmap=\"coolwarm\", cbar=False)\n",
    "ax2.set_xticks([0.5])\n",
    "ax2.set_xticklabels(['Age'], fontsize=18)\n",
    "ax2.set_yticks([0, split_size])\n",
    "ax2.set_yticklabels([5, int(np.round(split_max_age*data['max_age']))], fontsize=18)\n",
    "\n",
    "ax3 = fig.add_subplot(spec[0,2])\n",
    "sns.heatmap(matrix_merge.confound[:,1][_split_type][_age_sort][:,np.newaxis], ax=ax3, cmap=\"vlag\", cbar=False)\n",
    "ax3.set_xticks([0.5])\n",
    "ax3.set_xticklabels(['Sex'], fontsize=18)\n",
    "ax3.set_yticks([])\n",
    "\n",
    "pyplot.savefig('../../output/figure/merge_test.png', dpi=300)   \n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ddc8713-71a4-4081-8f5c-81bdecb6865f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "“deepMF_test”",
   "language": "python",
   "name": "deepmf_test"
  },
  "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
