bench_neural_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm\n",
    "from argparse import Namespace\n",
    "\n",
    "\n",
    "config = Namespace(\n",
    "    data_folder='./wm_bench_data', \n",
    "    max_seq_len=20, \n",
    "    rs_img_size=32, \n",
    "    batch_size=1, \n",
    "    num_workers=4, \n",
    "    use_cnn=1, \n",
    "    model_path='./model.pt', \n",
    ")\n",
    "\n",
    "device = torch.device('cuda:3' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.utils.data_utils import get_test_multitask_dataloader\n",
    "\n",
    "test_loader = get_test_multitask_dataloader(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.model import WM_Model\n",
    "\n",
    "model_data = torch.load(config.model_path)\n",
    "model = WM_Model(Namespace(**model_data['config']), device).to(device)\n",
    "model.load_state_dict(model_data['model_state_dict'])\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dict = {\n",
    "    'model_1': model, \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "rnn_out_dict = {'SC_Task': [], \n",
    "                'SFR_Task': [], \n",
    "                'SI_Task': [], \n",
    "                'SMU_Task': [], \n",
    "                'STSC_Task': [], \n",
    "                'VIRec_2C_Task': [], \n",
    "                'VSR_Task': [], \n",
    "                'VSRec_Task': [], \n",
    "                'CD_Color_Task': [], \n",
    "                'CD_Orientation_Task': [], \n",
    "                'CD_Size_Task': [], \n",
    "                'CD_Gap_Task': [], \n",
    "                'CD_Conj_Task': [], \n",
    "                'Complex_WM_Task': []}\n",
    "\n",
    "dataloader = zip(*test_loader.values())\n",
    "\n",
    "model = model_dict['model_1']\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch_index, multi_task_batch in tqdm(enumerate(dataloader)):\n",
    "\n",
    "        # SC Task\n",
    "        stim_batch, resp_batch, seq_len, _ = multi_task_batch[0]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SC_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SC_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # SFR Task\n",
    "        stim_batch, resp_batch, seq_len, _ = multi_task_batch[1]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SFR_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SFR_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # SI Task\n",
    "        stim_batch, resp_batch, seq_len, _ = multi_task_batch[2]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SI_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SI_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # SMU Task\n",
    "        stim_batch, resp_batch, seq_len, _ = multi_task_batch[3]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SMU_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SMU_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # STSC Task\n",
    "        stim_batch, resp_batch, seq_len = multi_task_batch[4]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'STSC_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['STSC_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # VIRec 2C Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[5]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'VIRec_2C_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['VIRec_2C_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # VSR Task\n",
    "        stim_batch, resp_batch, seq_len, list_length = multi_task_batch[6]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'VSR_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['VSR_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # VSRec Task\n",
    "        stim_batch, resp_batch, seq_len, list_length, _ = multi_task_batch[7]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'VSRec_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['VSRec_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # CD Color Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[8]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Color_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Color_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # CD Orientation Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[9]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Orientation_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Orientation_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # CD Size Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[10]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Size_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Size_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # CD Gap Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[11]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Gap_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Gap_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # CD Conj Task\n",
    "        stim_batch, resp_batch, seq_len, _, _, _ = multi_task_batch[12]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Conj_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Conj_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # Complex_WM Task\n",
    "        stim_batch, resp_batch, seq_len, _, variation = multi_task_batch[13]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'Complex_WM_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['Complex_WM_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "rnn_out_dict = {'SC_Task': [], \n",
    "                'SFR_Task': [], \n",
    "                'SI_Task': [], \n",
    "                'SMU_Task': [], \n",
    "                'STSC_Task': [], \n",
    "                'VIRec_2C_Task': [], \n",
    "                'VSR_Task': [], \n",
    "                'VSRec_Task': [], \n",
    "                'CD_Color_Task': [], \n",
    "                'CD_Orientation_Task': [], \n",
    "                'CD_Size_Task': [], \n",
    "                'CD_Gap_Task': [], \n",
    "                'CD_Conj_Task': [], \n",
    "                'Complex_WM_Task': []}\n",
    "\n",
    "dataloader = zip(*test_loader.values())\n",
    "\n",
    "model = model_dict['lstm_256']\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch_index, multi_task_batch in tqdm(enumerate(dataloader)):\n",
    "\n",
    "        # SC Task\n",
    "        stim_batch, resp_batch, seq_len, _ = multi_task_batch[0]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SC_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SC_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # SFR Task\n",
    "        stim_batch, resp_batch, seq_len, _ = multi_task_batch[1]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SFR_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SFR_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # SI Task\n",
    "        stim_batch, resp_batch, seq_len, _ = multi_task_batch[2]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SI_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SI_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # SMU Task\n",
    "        stim_batch, resp_batch, seq_len, set_size = multi_task_batch[3]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'SMU_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['SMU_Task'].append(rnn_out[:, :seq_len[0]-set_size[0], :][0].cpu().numpy())\n",
    "\n",
    "        # STSC Task\n",
    "        stim_batch, resp_batch, seq_len = multi_task_batch[4]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'STSC_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['STSC_Task'].append(rnn_out[:, :seq_len[0], :][0].cpu().numpy())\n",
    "\n",
    "        # VIRec 2C Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[5]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'VIRec_2C_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['VIRec_2C_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # VSR Task\n",
    "        stim_batch, resp_batch, seq_len, list_length = multi_task_batch[6]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'VSR_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['VSR_Task'].append(rnn_out[:, :seq_len[0]-list_length[0], :][0].cpu().numpy())\n",
    "\n",
    "        # VSRec Task\n",
    "        stim_batch, resp_batch, seq_len, list_length, _ = multi_task_batch[7]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'VSRec_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['VSRec_Task'].append(rnn_out[:, :seq_len[0]-list_length[0], :][0].cpu().numpy())\n",
    "\n",
    "        # CD Color Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[8]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Color_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Color_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # CD Orientation Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[9]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Orientation_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Orientation_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # CD Size Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[10]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Size_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Size_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # CD Gap Task\n",
    "        stim_batch, resp_batch, seq_len, _, _ = multi_task_batch[11]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Gap_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Gap_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # CD Conj Task\n",
    "        stim_batch, resp_batch, seq_len, _, _, _ = multi_task_batch[12]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'CD_Conj_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['CD_Conj_Task'].append(rnn_out[:, :seq_len[0]-1, :][0].cpu().numpy())\n",
    "\n",
    "        # Complex_WM Task\n",
    "        stim_batch, resp_batch, seq_len, _, variation = multi_task_batch[13]\n",
    "\n",
    "        stim_batch = stim_batch.to(device)\n",
    "        resp_batch = resp_batch.to(device)\n",
    "\n",
    "        _, rnn_out, _, _, _ = model(stim_batch, 'Complex_WM_Task', seq_len)\n",
    "\n",
    "        rnn_out_dict['Complex_WM_Task'].append(rnn_out[:, :seq_len[0]-2, :][0].cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rnn_out_ts_dict = {'SC_Task': {},\n",
    "                   'SFR_Task': {}, \n",
    "                   'SI_Task': {}, \n",
    "                   'SMU_Task': {}, \n",
    "                   'STSC_Task': {}, \n",
    "                   'VIRec_2C_Task': {}, \n",
    "                   'VSR_Task': {}, \n",
    "                   'VSRec_Task': {}, \n",
    "                   'CD_Task': {}, \n",
    "                   'Complex_WM_Task': {}\n",
    "                   }\n",
    "\n",
    "for trial in range(len(rnn_out_dict['SC_Task'])):\n",
    "    for task in list(rnn_out_dict.keys()):\n",
    "        if task in ['CD_Color_Task', 'CD_Orientation_Task', 'CD_Size_Task', 'CD_Gap_Task', 'CD_Conj_Task']:\n",
    "            tasksss = 'CD_Task'\n",
    "        else:\n",
    "            tasksss = task\n",
    "        for ts in range(len(rnn_out_dict[task][trial])):\n",
    "            if ts not in rnn_out_ts_dict[tasksss].keys():\n",
    "                rnn_out_ts_dict[tasksss][ts] = [rnn_out_dict[task][trial][ts]]\n",
    "            else:\n",
    "                rnn_out_ts_dict[tasksss][ts].append(rnn_out_dict[task][trial][ts])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "rnn_out_ts_dict_var = {'SC_Task': [], \n",
    "                       'SFR_Task': [], \n",
    "                       'SI_Task': [], \n",
    "                       'SMU_Task': [], \n",
    "                       'STSC_Task': [], \n",
    "                       'VIRec_2C_Task': [], \n",
    "                       'VSR_Task': [], \n",
    "                       'VSRec_Task': [], \n",
    "                       'CD_Task': [], \n",
    "                       'Complex_WM_Task': []\n",
    "                       }\n",
    "\n",
    "for task in rnn_out_ts_dict.keys():\n",
    "    for ts in range(len(rnn_out_ts_dict[task])):\n",
    "        rnn_out_ts_dict_var[task].append(np.var(rnn_out_ts_dict[task][ts], axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc_task_var = np.array(rnn_out_ts_dict_var['SC_Task'])\n",
    "sfr_task_var = np.array(rnn_out_ts_dict_var['SFR_Task'])\n",
    "si_task_var = np.array(rnn_out_ts_dict_var['SI_Task'])\n",
    "smu_task_var = np.array(rnn_out_ts_dict_var['SMU_Task'])\n",
    "stsc_task_var = np.array(rnn_out_ts_dict_var['STSC_Task'])\n",
    "virec_2c_task_var = np.array(rnn_out_ts_dict_var['VIRec_2C_Task'])\n",
    "vsr_task_var = np.array(rnn_out_ts_dict_var['VSR_Task'])\n",
    "vsrec_task_var = np.array(rnn_out_ts_dict_var['VSRec_Task'])\n",
    "cd_task_var = np.array(rnn_out_ts_dict_var['CD_Task'])\n",
    "complex_wm_task_var = np.array(rnn_out_ts_dict_var['Complex_WM_Task'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc_task_var = np.mean(sc_task_var, axis=0)\n",
    "sfr_task_var = np.mean(sfr_task_var, axis=0)\n",
    "si_task_var = np.mean(si_task_var, axis=0)\n",
    "smu_task_var = np.mean(smu_task_var, axis=0)\n",
    "stsc_task_var = np.mean(stsc_task_var, axis=0)\n",
    "virec_2c_task_var = np.mean(virec_2c_task_var, axis=0)\n",
    "vsr_task_var = np.mean(vsr_task_var, axis=0)\n",
    "vsrec_task_var = np.mean(vsrec_task_var, axis=0)\n",
    "cd_task_var = np.mean(cd_task_var, axis=0)\n",
    "complex_wm_task_var = np.mean(complex_wm_task_var, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_var_array = np.array([sc_task_var, sfr_task_var, si_task_var, smu_task_var, \n",
    "                           stsc_task_var, virec_2c_task_var, vsr_task_var, vsrec_task_var, \n",
    "                           cd_task_var, complex_wm_task_var])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_var_array_normalized = task_var_array / np.max(task_var_array, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_var_array_normalized = task_var_array_normalized.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "\n",
    "num_clust = []\n",
    "silhouette_scores = []\n",
    "\n",
    "for num_clusters in range(2, 100):\n",
    "    kmeans = KMeans(n_clusters=num_clusters, random_state=45).fit(task_var_array_normalized)\n",
    "\n",
    "    # Calculate the silhouette score\n",
    "    from sklearn.metrics import silhouette_score\n",
    "\n",
    "    num_clust.append(num_clusters)\n",
    "    silhouette_scores.append(silhouette_score(task_var_array_normalized, kmeans.labels_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(num_clust, silhouette_scores)\n",
    "plt.xlabel('Number of Clusters')\n",
    "plt.ylabel('Silhouette Score')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_silhoette_score = max(silhouette_scores)\n",
    "apt_num_clusters = num_clust[silhouette_scores.index(max_silhoette_score)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "\n",
    "kmeans = KMeans(n_clusters=apt_num_clusters, random_state=45).fit(task_var_array_normalized)\n",
    "\n",
    "from sklearn.metrics import silhouette_score\n",
    "\n",
    "silhouette_score(task_var_array_normalized, kmeans.labels_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clustered_task_var_array_normalized = [[] for _ in range(apt_num_clusters)]\n",
    "\n",
    "for i in range(len(kmeans.labels_)):\n",
    "    clustered_task_var_array_normalized[kmeans.labels_[i]].append(task_var_array_normalized[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "concatenated_clusters = np.concatenate(clustered_task_var_array_normalized, axis=0)\n",
    "concatenated_clusters = concatenated_clusters.T\n",
    "concatenated_clusters.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"Serif\"\n",
    "sns.set_context(\"paper\", font_scale=1.5, rc={\"lines.linewidth\": 1})\n",
    "\n",
    "plt.figure(figsize=(15, 5))\n",
    "plt.imshow(concatenated_clusters, cmap='hot', aspect='auto')\n",
    "\n",
    "plt.ylabel('Task', fontsize=28)\n",
    "\n",
    "plt.xticks([])\n",
    "plt.yticks(np.arange(0, 10, 1), ['SC', 'SFR', 'SI', 'SMU', 'STSC', 'VIR', 'VSR', 'VSRec', \n",
    "                                'CD', 'CS'], fontsize=20)\n",
    "\n",
    "cb = plt.colorbar()\n",
    "cb.ax.set_yticks([0, 1])\n",
    "cb.ax.set_ylabel('Normalized Task Variance', fontsize=18)\n",
    "cb.ax.tick_params(labelsize=20)\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "memnet",
   "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.8.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}