{

 "cells": [

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "import socket\n",

    "if socket.gethostname() == 'dlm':\n",

    "  %env CUDA_DEVICE_ORDER=PCI_BUS_ID\n",

    "  %env CUDA_VISIBLE_DEVICES=3"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 1,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "Using CPU:(\n"

     ]

    }

   ],

   "source": [

    "import os\n",

    "import sys\n",

    "import re\n",

    "import collections\n",

    "import functools\n",

    "import itertools\n",

    "import requests, zipfile, io\n",

    "import pickle\n",

    "import copy\n",

    "\n",

    "import pandas\n",

    "import numpy as np\n",

    "import matplotlib\n",

    "import matplotlib.pyplot as plt\n",

    "import sklearn\n",

    "import sklearn.decomposition\n",

    "import sklearn.metrics\n",

    "import networkx\n",

    "\n",

    "import torch\n",

    "import torch.nn as nn\n",

    "\n",

    "lib_path = 'I:/code'\n",

    "if not os.path.exists(lib_path):\n",

    "  lib_path = '/media/6T/.tianle/.lib'\n",

    "if not os.path.exists(lib_path):\n",

    "  lib_path = '/projects/academic/azhang/tianlema/lib'\n",

    "if os.path.exists(lib_path) and lib_path not in sys.path:\n",

    "  sys.path.append(lib_path)\n",

    "  \n",

    "from dl.models.basic_models import *\n",

    "from dl.utils.visualization.visualization import *\n",

    "from dl.utils.outlier import *\n",

    "from dl.utils.train import *\n",

    "from autoencoder.autoencoder import *\n",

    "from vin.vin import *\n",

    "from dl.utils.utils import get_overlap_samples, filter_clinical_dict, get_target_variable\n",

    "from dl.utils.utils import get_shuffled_data, target_to_numpy, discrete_to_id, get_mi_acc\n",

    "from dl.utils.utils import get_label_distribution, normalize_continuous_variable\n",

    "\n",

    "%load_ext autoreload\n",

    "%autoreload 2\n",

    "\n",

    "\n",

    "use_gpu = True\n",

    "if use_gpu and torch.cuda.is_available():\n",

    "  device = torch.device('cuda')\n",

    "  print('Using GPU:)')\n",

    "else:\n",

    "  device = torch.device('cpu')\n",

    "  print('Using CPU:(')"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 2,

   "metadata": {},

   "outputs": [],

   "source": [

    "# neural net models include nn (mlp), resnet, densenet; another choice is ml (machine learning)\n",

    "# model_type, dense, residual are dependent\n",

    "model_type = 'resnet'\n",

    "dense = False\n",

    "residual = True\n",

    "hidden_dim = [100, 100]\n",

    "train_portion = 0.7\n",

    "val_portion = 0.1\n",

    "test_portion = 0.2\n",

    "num_train_types = -1 # -1 means not used\n",

    "num_val_types = -1\n",

    "num_test_types = -1 # this will almost never be used \n",

    "num_sets = 10\n",

    "num_folds = 10 # no longer used anymore\n",

    "sel_set_idx = 0\n",

    "cv_type = 'instance-shuffle' # or 'group-shuffle'; cross validation shuffle method\n",

    "sel_disease_types = 'all'\n",

    "# The number of total samples and the numbers for each class in selected disease types must >=\n",

    "min_num_samples_per_type_cls = [100, 0]\n",

    "# if 'auto-search', will search for the file first; if not exist, then generate random data split\n",

    "# and write to the file;\n",

    "# if string other than 'auto-search' is provided, assume the string is a proper file name, \n",

    "# and read the file;\n",

    "# if False, will generate a random data split, but not write to file \n",

    "# if True will generate a random data split, and write to file\n",

    "predefined_sample_set_file = 'auto-search' \n",

    "target_variable = ['PFI', 'DFI', 'PFI.time'] # To do: target variable can be a list (partially handled)\n",

    "target_variable_type = ['discrete', 'discrete', 'continuous'] # or 'continuous' real numbers\n",

    "target_variable_range = [[0,1],[0,1],[0,float('Inf')]]\n",

    "data_type = ['gene', 'methy', 'rppa', 'mirna']\n",

    "normal_transform_feature = True\n",

    "additional_vars = ['age_at_initial_pathologic_diagnosis', 'gender', 'ajcc_pathologic_tumor_stage']\n",

    "additional_var_types = ['continuous', 'discrete', 'discrete']\n",

    "additional_var_ranges = [[0, 100], ['MALE', 'FEMALE'], \n",

    "                         ['I/II NOS', 'IS', 'Stage 0', 'Stage I', 'Stage IA', 'Stage IB', \n",

    "                          'Stage II', 'Stage IIA', 'Stage IIB', 'Stage IIC', 'Stage III',\n",

    "                          'Stage IIIA', 'Stage IIIB', 'Stage IIIC', 'Stage IV', 'Stage IVA',\n",

    "                          'Stage IVB', 'Stage IVC', 'Stage X']]\n",

    "randomize_labels = False\n",

    "lr = 5e-4\n",

    "weight_decay = 1e-4\n",

    "num_epochs = 100\n",

    "reduce_every = 500\n",

    "show_results_in_notebook = True"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Prepare data"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 3,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "feature_mat: rppa, max=14.141, min=-7.869, mean=0.095, 0.516\n",

      "feature_mat: mirna, max=11.813, min=0.000, mean=3.743, 1.000\n",

      "feature_mat: gene, max=16.311, min=0.000, mean=8.412, 1.000\n",

      "feature_mat: methy, max=1.000, min=0.000, mean=0.553, 1.000\n",

      "feature_interaction_mat: rppa, max=1.000, min=0.000, mean=0.198, 0.277\n",

      "feature_interaction_mat: mirna, max=1.000, min=0.010, mean=0.492, 1.000\n",

      "feature_interaction_mat: gene, max=1.000, min=0.000, mean=0.050, 0.146\n",

      "feature_interaction_mat: methy, max=1.000, min=0.000, mean=0.057, 0.127\n",

      "rppa (189,) X1433EPSILON\n",

      "mirna (662,) hsa-let-7a-2-3p\n",

      "gene (4942,) A1BG\n",

      "methy (4753,) cg00005847\n",

      "rppa (7480,) TCGA-OR-A5J2-01A-21-A39K-20\n",

      "mirna (9554,) TCGA-C4-A0F6-01A-11R-A10V-13\n",

      "gene (9702,) TCGA-OR-A5J1-01A-11R-A29S-07\n",

      "methy (10268,) TCGA-02-0001-01C-01D-0186-05\n"

     ]

    }

   ],

   "source": [

    "result_folder = 'results'\n",

    "data_split_idx_folder = f'{result_folder}/data_split_idx'\n",

    "project_folder = '../../pan-can-atlas' # on dlm or ccr\n",

    "print_stats = True\n",

    "if not os.path.exists(project_folder):\n",

    "  project_folder = 'F:/TCGA/Pan-Cancer-Atlas' # on my own desktop\n",

    "filepath = f'{project_folder}/data/processed/combined2.pkl'\n",

    "with open(filepath, 'rb') as f:\n",

    "  data = pickle.load(f)\n",

    "  patient_clinical = data['patient_clinical']\n",

    "  feature_mat_dict = data['feature_mat_dict']\n",

    "  feature_interaction_mat_dict = data['feature_interaction_mat_dict']\n",

    "  feature_id_dict = data['feature_id_dict']\n",

    "  aliquot_id_dict = data['aliquot_id_dict']\n",

    "#   sel_patient_ids = data['sample_id_sel']\n",

    "#   sample_idx_sel_dict = data['sample_idx_sel_dict']\n",

    "#   for k, v in sample_idx_sel_dict.items():\n",

    "#     assert [i[:12] for i in aliquot_id_dict[k][v]] == sel_patient_ids\n",

    "\n",

    "if print_stats:\n",

    "  for k, v in feature_mat_dict.items():\n",

    "    print(f'feature_mat: {k}, max={v.max():.3f}, min={v.min():.3f}, '\n",

    "          f'mean={v.mean():.3f}, {np.mean(v>0):.3f}')  \n",

    "  for k, v in feature_interaction_mat_dict.items():\n",

    "    print(f'feature_interaction_mat: {k}, max={v.max():.3f}, min={v.min():.3f}, '\n",

    "          f'mean={v.mean():.3f}, {np.mean(v>0):.3f}') \n",

    "  for k, v in feature_id_dict.items():\n",

    "    print(k, v.shape, v[0])\n",

    "  for k, v in aliquot_id_dict.items():\n",

    "    print(k, v.shape, v[0])"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 4,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "THCA {0.0: 240, 1.0: 24}\n",

      "BLCA {0.0: 125, 1.0: 25}\n",

      "BRCA {0.0: 666, 1.0: 62}\n",

      "KIRP {0.0: 101, 1.0: 22}\n",

      "STAD {1.0: 37, 0.0: 149}\n",

      "LIHC {0.0: 62, 1.0: 68}\n",

      "LUAD {1.0: 56, 0.0: 129}\n",

      "COAD {0.0: 108, 1.0: 17}\n",

      "LUSC {0.0: 136, 1.0: 56}\n",

      "Selected 2083 patients from 9 disease_types\n"

     ]

    }

   ],

   "source": [

    "# select samples with required clinical variables\n",

    "clinical_dict = filter_clinical_dict(target_variable, target_variable_type=target_variable_type, \n",

    "                                     target_variable_range=target_variable_range, \n",

    "                                     clinical_dict=patient_clinical)\n",

    "if len(additional_vars) > 0:\n",

    "  clinical_dict = filter_clinical_dict(additional_vars, target_variable_type=additional_var_types, \n",

    "                                       target_variable_range=additional_var_ranges, \n",

    "                                       clinical_dict=clinical_dict)\n",

    "\n",

    "# select samples with feature matrix of given type(s)\n",

    "if isinstance(data_type, str):\n",

    "  sample_list = {s[:12] for s in aliquot_id_dict[data_type]}\n",

    "  data_type_str = data_type\n",

    "elif isinstance(data_type, (list, tuple)):\n",

    "  sample_list = get_overlap_samples([aliquot_id_dict[dtype] for dtype in data_type], \n",

    "                                    common_list=None, start=0, end=12, return_common_list=True)\n",

    "  data_type_str = '-'.join(sorted(data_type))\n",

    "else:\n",

    "  raise ValueError(f'data_type must be str or list/tuple, but is {type(data_type)}')\n",

    "sample_list = set(sample_list).intersection(clinical_dict)\n",

    "\n",

    "# select samples with given disease types\n",

    "sel_disease_type_str = sel_disease_types # will be overwritten if it is a list\n",

    "if isinstance(sel_disease_types, (list, tuple)):\n",

    "  sample_list = [s for s in sample_list if clinical_dict[s]['type'] in sel_disease_types]\n",

    "  sel_disease_type_str = '-'.join(sorted(sel_disease_types))\n",

    "elif isinstance(sel_disease_types, str) and sel_disease_types!='all':\n",

    "  sample_list = [s for s in sample_list if clinical_dict[s]['type'] == sel_disease_types]\n",

    "else:\n",

    "  assert sel_disease_types == 'all'\n",

    " \n",

    "# For classification tasks with given min_num_samples_per_type_cls,\n",

    "# only keep disease types that have a minimal number of samples per type and per class\n",

    "# Reflection: it might be better to use collections.defaultdict(list) to store samples in each type\n",

    "type_cnt = collections.Counter([clinical_dict[s]['type'] for s in sample_list])\n",

    "if sum(min_num_samples_per_type_cls)>0 and (target_variable_type=='discrete' \n",

    "                                            or target_variable_type[0]=='discrete'):\n",

    "  # the number of samples in each disease type >= min_num_samples_per_type_cls[0]\n",

    "  type_cnt = {k: v for k, v in type_cnt.items() if v >= min_num_samples_per_type_cls[0]}\n",

    "  disease_type_cnt = {}\n",

    "  for k in type_cnt:\n",

    "    # collections.Counter can accept generator\n",

    "    cls_cnt = collections.Counter(clinical_dict[s][target_variable] \n",

    "                                  if isinstance(target_variable, str) \n",

    "                                  else clinical_dict[s][target_variable[0]] \n",

    "                                  for s in sample_list if clinical_dict[s]['type']==k)\n",

    "    if all([v >= min_num_samples_per_type_cls[1] for v in cls_cnt.values()]):\n",

    "      # the number of samples in each class >= min_num_samples_per_type_cls[1]\n",

    "      disease_type_cnt[k] = dict(cls_cnt)\n",

    "      print(k, disease_type_cnt[k])\n",

    "  sample_list = [s for s in sample_list if clinical_dict[s]['type'] in disease_type_cnt]\n",

    "sel_patient_ids = sorted(sample_list)\n",

    "print(f'Selected {len(sel_patient_ids)} patients from {len(disease_type_cnt)} disease_types')"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "### Split data into training, validation, and test sets"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 5,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "Read predefined_sample_set_file: results/data_split_idx/PFI-DFI-PFI.time_instance-shuffle_age_at_initial_pathologic_diagnosis-ajcc_pathologic_tumor_stage-gender_gene-methy-mirna-rppa_all_100-0_0.7-0.1-0.2_10sets.pkl\n"

     ]

    }

   ],

   "source": [

    "predefined_sample_set_filename = (target_variable if isinstance(target_variable,str) \n",

    "                                else '-'.join(target_variable))\n",

    "predefined_sample_set_filename += f'_{cv_type}'\n",

    "if len(additional_vars) > 0:\n",

    "  predefined_sample_set_filename += f\"_{'-'.join(sorted(additional_vars))}\"\n",

    "\n",

    "predefined_sample_set_filename += (f\"_{data_type_str}_{sel_disease_type_str}_\"\n",

    "                                   f\"{'-'.join(map(str, min_num_samples_per_type_cls))}\")\n",

    "predefined_sample_set_filename += f\"_{'-'.join(map(str, [train_portion, val_portion, test_portion]))}\"\n",

    "if cv_type == 'group-shuffle' and num_train_types > 0:\n",

    "  predefined_sample_set_filename += f\"_{'-'.join(map(str, [num_train_types, num_val_types, num_test_types]))}\"\n",

    "predefined_sample_set_filename += f'_{num_sets}sets'\n",

    "res_file = f\"{predefined_sample_set_filename}_{sel_set_idx}_{'-'.join(map(str, hidden_dim))}_{model_type}.pkl\"\n",

    "predefined_sample_set_filename += '.pkl'\n",

    "# This will be overwritten if predefined_sample_set_file == 'auto-search' or filepath, and the file exists\n",

    "predefined_sample_sets = [get_shuffled_data(sel_patient_ids, clinical_dict, cv_type=cv_type, \n",

    "                  instance_portions=[train_portion, val_portion, test_portion], \n",

    "                  group_sizes=[num_train_types, num_val_types, num_test_types],\n",

    "                  group_variable_name='type', seed=None, verbose=False) for i in range(num_sets)]\n",

    "if predefined_sample_set_file == 'auto-search':\n",

    "  if os.path.exists(f'{data_split_idx_folder}/{predefined_sample_set_filename}'):\n",

    "    with open(f'{data_split_idx_folder}/{predefined_sample_set_filename}', 'rb') as f:\n",

    "      print(f'Read predefined_sample_set_file: '\n",

    "            f'{data_split_idx_folder}/{predefined_sample_set_filename}')\n",

    "      tmp = pickle.load(f)\n",

    "      # overwrite calculated predefined_sample_sets\n",

    "      predefined_sample_sets = tmp['predefined_sample_sets']    \n",

    "elif isinstance(predefined_sample_set_file, str): # but not 'auto-search'; assume it's a file name\n",

    "  if os.path.exists(predefined_sample_set_file):\n",

    "    with open(f'{data_split_idx_folder}/{predefined_sample_set_file}', 'rb') as f:\n",

    "      print(f'Read predefined_sample_set_file: {data_split_idx_folder}/{predefined_sample_set_file}')\n",

    "      tmp = pickle.load(f)\n",

    "      predefined_sample_sets = tmp['predefined_sample_sets']\n",

    "  else:\n",

    "    raise ValueError(f'predefined_sample_set_file: {data_split_idx_folder}/{predefined_sample_set_file} does not exist!')\n",

    "\n",

    "if (not os.path.exists(f'{data_split_idx_folder}/{predefined_sample_set_filename}') \n",

    "    and predefined_sample_set_file == 'auto-search') or predefined_sample_set_file is True:\n",

    "  with open(f'{data_split_idx_folder}/{predefined_sample_set_filename}', 'wb') as f:\n",

    "      print(f'Write predefined_sample_set_file: {data_split_idx_folder}/{predefined_sample_set_filename}')\n",

    "      pickle.dump({'predefined_sample_sets': predefined_sample_sets}, f)\n",

    "     \n",

    "sel_patient_ids, idx_splits = predefined_sample_sets[sel_set_idx]\n",

    "train_idx, val_idx, test_idx = idx_splits"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 6,

   "metadata": {},

   "outputs": [],

   "source": [

    "if isinstance(data_type, str):\n",

    "  sample_lists = [aliquot_id_dict[data_type]]\n",

    "else:\n",

    "  assert isinstance(data_type, (list, tuple))\n",

    "  sample_lists = [aliquot_id_dict[dtype] for dtype in data_type]\n",

    "idx_lists = get_overlap_samples(sample_lists=sample_lists, common_list=sel_patient_ids, \n",

    "                    start=0, end=12, return_common_list=False)\n",

    "sample_idx_sel_dict = {}\n",

    "if isinstance(data_type, str):\n",

    "  sample_idx_sel_dict = {data_type: idx_lists[0]}\n",

    "else:\n",

    "  sample_idx_sel_dict = {dtype: idx_list for dtype, idx_list in zip(data_type, idx_lists)}"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 7,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "gene: (2083, 4942); interaction_mat: mean=0.000057, std=0.000194, 4942\n",

      "methy: (2083, 4753); interaction_mat: mean=0.000069, std=0.000199, 4753\n",

      "rppa: (2083, 189); interaction_mat: mean=0.002668, std=0.004569, 189\n",

      "mirna: (2083, 662); interaction_mat: mean=0.001408, std=0.000547, 662\n"

     ]

    }

   ],

   "source": [

    "if isinstance(data_type, str):\n",

    "  print(f'Only use one data type: {data_type}')\n",

    "  num_data_types = 1\n",

    "  mat = feature_mat_dict[data_type][sample_idx_sel_dict[data_type]]\n",

    "  # Data preprocessing: make each row have mean 0 and sd 1.\n",

    "  x = (mat - mat.mean(axis=1, keepdims=True)) / mat.std(axis=1, keepdims=True)\n",

    "  interaction_mat = feature_interaction_mat_dict[data_type]\n",

    "  interaction_mat = torch.from_numpy(interaction_mat).float().to(device)\n",

    "  # Normalize these interaction mat\n",

    "  interaction_mat = interaction_mat / interaction_mat.norm()\n",

    "else:\n",

    "  mat = []\n",

    "  interaction_mats = []\n",

    "  in_dims = []\n",

    "  num_data_types = len(data_type)\n",

    "  # do not handle the special case of [data_type] to avoid too much code complexity\n",

    "  assert num_data_types > 1 \n",

    "  for dtype in data_type: # multiple data types\n",

    "    m = feature_mat_dict[dtype][sample_idx_sel_dict[dtype]]\n",

    "    #When there are multiple data types, make sure each type is normalized to have mean 0 and std 1\n",

    "    m = (m - m.mean(axis=1, keepdims=True)) / m.std(axis=1, keepdims=True)\n",

    "    mat.append(m)\n",

    "    in_dims.append(m.shape[1])\n",

    "    # For neural network model graph laplacian regularizer\n",

    "    interaction_mat = feature_interaction_mat_dict[dtype]\n",

    "    interaction_mat = torch.from_numpy(interaction_mat).float().to(device)\n",

    "    # Normalize these interaction mat\n",

    "    interaction_mat = interaction_mat / interaction_mat.norm()\n",

    "    interaction_mats.append(interaction_mat)\n",

    "    print(f'{dtype}: {m.shape}; '\n",

    "          f'interaction_mat: mean={interaction_mat.mean().item():2f}, '\n",

    "          f'std={interaction_mat.std().item():2f}, {interaction_mat.shape[0]}')\n",

    "  # Later interaction_mat will be passed to Loss_feature_interaction\n",

    "  interaction_mat = interaction_mats\n",

    "  mat = np.concatenate(mat, axis=1)\n",

    "  # For machine learing methods that use concatenated features without knowing underlying views,\n",

    "  # it might be good to make each row have mean 0 and sd 1.\n",

    "  x = (mat - mat.mean(axis=1, keepdims=True)) / mat.std(axis=1, keepdims=True)\n",

    "\n",

    "if normal_transform_feature:\n",

    "  X = x\n",

    "else:\n",

    "  X = mat"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 8,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "torch.Size([1458, 10546]) torch.Size([208, 10546]) torch.Size([417, 10546])\n"

     ]

    }

   ],

   "source": [

    "# sklearn classifiers also accept torch.Tensor\n",

    "X = torch.tensor(X).float().to(device)\n",

    "x_train = X[train_idx]\n",

    "x_val = X[val_idx]\n",

    "x_test = X[test_idx]\n",

    "print(x_train.shape, x_val.shape, x_test.shape)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 12,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "Changed class labels for the model: {0.0: 0, 1.0: 1}\n",

      "Changed class labels for the model: {0.0: 0, 1.0: 1}\n",

      "PFI:\n",

      "train:torch.Size([1458]), val:torch.Size([208]), test:torch.Size([417])\n",

      "label distribution:\n",

      " [[0.82304525 0.84615386 0.81534773]\n",

      " [0.17695473 0.15384616 0.18465228]]\n",

      "DFI:\n",

      "train:torch.Size([1458]), val:torch.Size([208]), test:torch.Size([417])\n",

      "label distribution:\n",

      " [[0.8360768  0.84615386 0.82254195]\n",

      " [0.16392319 0.15384616 0.17745803]]\n",

      "PFI.time:\n",

      "train:torch.Size([1458, 1]), val:torch.Size([208, 1]), test:torch.Size([417, 1])\n"

     ]

    }

   ],

   "source": [

    "y_targets = get_target_variable(target_variable, clinical_dict, sel_patient_ids)\n",

    "y_targets = normalize_continuous_variable(y_targets, target_variable_type, transform=True, \n",

    "                                          forced=False, threshold=10, rm_outlier=True, whis=1.5, \n",

    "                                          only_positive=True, max_val=1)\n",

    "y_true = target_to_numpy(y_targets, target_variable_type, target_variable_range)\n",

    "if len(additional_vars) > 0:\n",

    "  additional_variables = get_target_variable(additional_vars, clinical_dict, sel_patient_ids)\n",

    "  # to do handle additional variables such as age and gender\n",

    "\n",

    "# should have written a recursive function instead\n",

    "if isinstance(target_variable_type, list):\n",

    "  y_targets = []\n",

    "  num_cls = []\n",

    "  y_train = []\n",

    "  y_val = []\n",

    "  y_test = []\n",

    "  for i, var_type in enumerate(target_variable_type):\n",

    "    y = torch.tensor(y_true[i]).to(device)\n",

    "    if var_type == 'discrete':\n",

    "      y = y.long()\n",

    "    elif var_type == 'continuous':\n",

    "      y = y.float()\n",

    "      if y.dim()==1:\n",

    "        y = y.unsqueeze(-1)\n",

    "    else:\n",

    "      raise ValueError(f'target type should be either discrete or continuous but is {var_type}')\n",

    "    y_targets.append(y)\n",

    "    num_cls.append(len(torch.unique(y))) # include continous target variables\n",

    "    y_train.append(y[train_idx])\n",

    "    y_val.append(y[val_idx])\n",

    "    y_test.append(y[test_idx])\n",

    "    print(f'{target_variable[i]}:\\ntrain:{y_train[-1].shape}, val:{y_val[-1].shape}, '\n",

    "         f'test:{y_test[-1].shape}')\n",

    "    if var_type == 'discrete':\n",

    "      label_probs = get_label_distribution([y_train[-1], y_val[-1], y_test[-1]])\n",

    "      if randomize_labels: # Optionally randomize true class labels\n",

    "        print('Randomize class labels!')\n",

    "        y_train[-1] = torch.multinomial(label_probs[0], len(y_train[-1]), replacement=True)\n",

    "        if len(y_val) > 0:\n",

    "          y_val[-1] = torch.multinomial(label_probs[1], len(y_val[-1]), replacement=True)\n",

    "        if len(y_test) > 0:\n",

    "          y_test[-1] = torch.multinomial(label_probs[2], len(y_test[-1]), replacement=True)\n",

    "        get_label_distribution([y_train[-1], y_val[-1], y_test[-1]])\n",

    "  y_true = y_targets\n",

    "elif isinstance(target_variable_type, str):\n",

    "  y = torch.tensor(y_true).to(device)\n",

    "  if var_type == 'discrete':\n",

    "    y = y.long()\n",

    "  elif var_type == 'continuous':\n",

    "    y = y.float()\n",

    "    if y.dim()==1:\n",

    "      y = y.unsqueeze(-1)\n",

    "  else:\n",

    "    raise ValueError(f'target type should be either discrete or continuous but is {var_type}')\n",

    "  y_true = y\n",

    "  num_cls = len(torch.unique(y_true))\n",

    "  y_train = y_true[train_idx]\n",

    "  y_val = y_true[val_idx]\n",

    "  y_test = y_true[test_idx]\n",

    "  print(f'{target_variable}:\\ntrain:{y_train.shape}, val:{y_val.shape}, '\n",

    "         f'test:{y_test.shape}')\n",

    "  label_probs = get_label_distribution([y_train, y_val, y_test])\n",

    "  if randomize_labels: # Optionally randomize true class labels\n",

    "    print('Randomize class labels!')\n",

    "    y_train = torch.multinomial(label_probs[0], len(y_train), replacement=True)\n",

    "    if len(y_val) > 0:\n",

    "      y_val = torch.multinomial(label_probs[1], len(y_val), replacement=True)\n",

    "    if len(y_test) > 0:\n",

    "      y_test = torch.multinomial(label_probs[2], len(y_test), replacement=True)\n",

    "    get_label_distribution([y_train, y_val, y_test])\n",

    "else:\n",

    "  raise ValueError(f'target_variable_type should be str or list, but is {type(target_variable_type)}')"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Use additional variables for prediction"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 13,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "age_at_initial_pathologic_diagnosis\n",

      "Counter({5: 577, 4: 468, 6: 429, 3: 285, 7: 137, 2: 130, 1: 48, 0: 9})\n",

      "gender\n",

      "Counter({0: 1322, 1: 761})\n",

      "ajcc_pathologic_tumor_stage\n",

      "Counter({0: 398, 4: 379, 5: 276, 8: 220, 3: 185, 1: 164, 7: 162, 2: 144, 10: 77, 9: 76, 11: 1, 6: 1})\n"

     ]

    }

   ],

   "source": [

    "embedding_dim = 50\n",

    "input_list = []\n",

    "xs = []\n",

    "for v, n, t in zip(additional_variables, additional_vars, additional_var_types):\n",

    "  if n.startswith('age'):    \n",

    "    bins = [0, 20, 30, 40, 50, 60, 70, 80, 100]\n",

    "    v = np.digitize(v, bins)\n",

    "    t = 'discrete'\n",

    "  if t=='discrete':\n",

    "    target_ids, cls_id_dict = discrete_to_id(v, start=0, sort=True)\n",

    "    # some target_ids may have very few instances\n",

    "    print(n)\n",

    "    print(collections.Counter(target_ids))\n",

    "    xs.append(torch.tensor(target_ids, device=device).long())\n",

    "    # did not handle missing value yet\n",

    "    input_list.append({'in_dim': len(cls_id_dict), 'in_type': 'discrete', 'padding_idx':None, \n",

    "                       'embedding_dim':embedding_dim, 'hidden_dim': hidden_dim})\n",

    "  else: # t=='continuous'\n",

    "    xs.append(torch.tensor(v, device=device).float())\n",

    "    input_list.append({'in_dim': len(v[0]), 'in_type': 'continuous', 'hidden_dim': hidden_dim})"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 15,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "target: PFI\n",

      "             Variable              \t MI  \tAdj_MI\tBayes_ACC\n",

      "age_at_initial_pathologic_diagnosis\t0.003\t0.001 \t  0.824  \n",

      "              gender               \t0.009\t0.013 \t  0.824  \n",

      "    ajcc_pathologic_tumor_stage    \t0.012\t0.004 \t  0.824  \n",

      "                0-1                \t0.011\t0.003 \t  0.824  \n",

      "                0-2                \t0.029\t0.003 \t  0.825  \n",

      "                1-2                \t0.033\t0.010 \t  0.831  \n",

      "               0-1-2               \t0.061\t0.006 \t  0.836  \n",

      "target: DFI\n",

      "             Variable              \t MI  \tAdj_MI\tBayes_ACC\n",

      "age_at_initial_pathologic_diagnosis\t0.002\t0.000 \t  0.834  \n",

      "              gender               \t0.009\t0.013 \t  0.834  \n",

      "    ajcc_pathologic_tumor_stage    \t0.011\t0.004 \t  0.835  \n",

      "                0-1                \t0.011\t0.003 \t  0.834  \n",

      "                0-2                \t0.028\t0.003 \t  0.836  \n",

      "                1-2                \t0.029\t0.009 \t  0.838  \n",

      "               0-1-2               \t0.056\t0.005 \t  0.843  \n"

     ]

    }

   ],

   "source": [

    "if isinstance(target_variable, list):\n",

    "  for i, var_name in enumerate(target_variable):\n",

    "    if target_variable_type[i]=='discrete':\n",

    "      print(f'target: {var_name}')\n",

    "      get_mi_acc(xs, y_true=y_true[i], var_names=additional_vars, var_name_length=35)\n",

    "elif isinstance(target_variable, str):\n",

    "  if target_variable_type=='discrete':\n",

    "    print(f'target: {target_variable}')\n",

    "    get_mi_acc(xs, y_true, var_names=additional_vars, var_name_length=35)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 16,

   "metadata": {},

   "outputs": [],

   "source": [

    "xs_train = [x[train_idx] for x in xs]\n",

    "xs_val = [x[val_idx] for x in xs]\n",

    "xs_test = [x[test_idx] for x in xs]"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 17,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "torch.Size([1458])\n",

      "torch.Size([1458])\n",

      "torch.Size([1458])\n",

      "torch.Size([1458, 10546])\n"

     ]

    }

   ],

   "source": [

    "last_nonlinearity = True\n",

    "input_list.append({'in_dim': x_train.size(1), 'in_type': 'continuous', 'hidden_dim': hidden_dim,\n",

    "                  'last_nonlinearity':last_nonlinearity, })\n",

    "xs_train.append(x_train)\n",

    "xs_val.append(x_val)\n",

    "xs_test.append(x_test)\n",

    "\n",

    "for i in xs_train:\n",

    "  print(i.shape)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 19,

   "metadata": {},

   "outputs": [],

   "source": [

    "hidden_dim = [100]\n",

    "output_info = [{'hidden_dim': hidden_dim+[2]}, {'hidden_dim': hidden_dim+[2]},\n",

    "              {'hidden_dim': hidden_dim+[1]}]\n",

    "            \n",

    "fusion_lists = [[{'fusion_type': 'repr-weighted-avg_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info}, \n",

    "                 {'fusion_type': 'repr-loss-avg_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info},\n",

    "                 {'fusion_type': 'repr-avg_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info},\n",

    "                 {'fusion_type': 'repr-cat_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info},\n",

    "                 {'fusion_type': 'out-weighted-avg', 'output_info': output_info},\n",

    "                 {'fusion_type': 'out-loss-avg', 'output_info': output_info},\n",

    "                 {'fusion_type': 'out-avg', 'output_info': output_info}],\n",

    "                [{'fusion_type': 'repr-weighted-avg_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info}, \n",

    "                 {'fusion_type': 'repr-loss-avg_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info},\n",

    "                 {'fusion_type': 'repr-avg_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info},\n",

    "                 {'fusion_type': 'repr-cat_repr', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info},\n",

    "                 {'fusion_type': 'out-weighted-avg', 'output_info': output_info},\n",

    "                 {'fusion_type': 'out-loss-avg', 'output_info': output_info},\n",

    "                 {'fusion_type': 'out-avg', 'output_info': output_info}],\n",

    "                [{'fusion_type': 'repr0', 'hidden_dim': hidden_dim, \n",

    "                  'last_nonlinearity':last_nonlinearity, 'output_info': output_info}]\n",

    "               ]"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 20,

   "metadata": {},

   "outputs": [],

   "source": [

    "model = VIN(input_list, output_info, fusion_lists, nonlinearity=nn.ReLU())\n",

    "if target_variable_type=='discrete':\n",

    "  loss_fn = nn.CrossEntropyLoss()\n",

    "elif target_variable_type=='dontinuous':\n",

    "  loss_fn = nn.MSELoss()\n",

    "else:\n",

    "  loss_fn = []\n",

    "  for var_type in target_variable_type:\n",

    "    if var_type == 'discrete':\n",

    "      loss_fn.append(nn.CrossEntropyLoss())\n",

    "    elif var_type == 'continuous':\n",

    "      loss_fn.append(nn.MSELoss())\n",

    "    else:\n",

    "      raise ValueError(f'target type should be either discrete or continous, but is {var_type}')"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 21,

   "metadata": {},

   "outputs": [],

   "source": [

    "optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), \n",

    "                             lr=1e-2, weight_decay=weight_decay, amsgrad=True)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 22,

   "metadata": {

    "scrolled": true

   },

   "outputs": [],

   "source": [

    "for n, p in model.named_parameters():\n",

    "#   print(n, p.size())\n",

    "  if p.grad is not None and p.grad.norm()==0:\n",

    "    print(n, p.grad if p.grad is None else p.grad.norm())"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 23,

   "metadata": {},

   "outputs": [],

   "source": [

    "target_loss_weight = [1., 1., 1.]"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 24,

   "metadata": {

    "scrolled": true

   },

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "0 6.476935386657715\n",

      "99 1.0573850870132446\n"

     ]

    }

   ],

   "source": [

    "loss_his = []\n",

    "num_iters = 100\n",

    "print_every = 100\n",

    "for i in range(num_iters):\n",

    "  pred = model(xs_test)\n",

    "  losses = get_vin_loss(pred, y_test, loss_fn, model, valid_loc=None, target_id=None, \n",

    "                        level_weight=None)\n",

    "  loss = sum(losses[j][0]*target_loss_weight[j] for j in range(len(losses)))\n",

    "  losses = losses[0]\n",

    "  optimizer.zero_grad()\n",

    "  loss.backward()\n",

    "  optimizer.step()\n",

    "  loss_his.append([losses[0].item()] + \n",

    "                [[[v.item() for v in losses[1][i][0]], losses[1][i][1].item()] for i in range(2)])\n",

    "  if i%print_every == 0 or i==num_iters-1:\n",

    "    print(i, loss.item())"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 25,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "acc=1.000, precision=1.000, recall=1.000, fl=1.000, adj_MI=1.000, auc=1.000, ap=1.000, confusion_mat=\n",

      "[[340   0]\n",

      " [  0  77]]\n",

      "report              precision    recall  f1-score   support\n",

      "\n",

      "          0       1.00      1.00      1.00       340\n",

      "          1       1.00      1.00      1.00        77\n",

      "\n",

      "avg / total       1.00      1.00      1.00       417\n",

      "\n"

     ]

    },

    {

     "data": {

      "text/plain": [

       "[(tensor(0.3750, grad_fn=<ThAddBackward>),\n",

       "  [[[tensor(0.4638, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4475, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4651, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0441, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.3546, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0055, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0074, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0078, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0036, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0263, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0265, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0267, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.0128, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0036, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0036, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0037, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0031, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0083, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0080, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0091, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.0055, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0021, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.0021, grad_fn=<AddBackward>)]]),\n",

       " (tensor(0.3739, grad_fn=<ThAddBackward>),\n",

       "  [[[tensor(0.4564, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4383, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4547, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0552, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.3481, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0104, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0119, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0157, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0020, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0307, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0302, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0317, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.0178, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0038, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0047, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0053, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0008, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0114, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0114, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.0120, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.0068, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0013, grad_fn=<NllLossBackward>)],\n",

       "    tensor(0.0013, grad_fn=<AddBackward>)]]),\n",

       " (tensor(0.3023, grad_fn=<ThAddBackward>),\n",

       "  [[[tensor(0.0833, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0855, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0839, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0756, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0823, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0728, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0724, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0717, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0723, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0770, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0770, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0767, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0741, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0739, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0723, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0734, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0714, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0717, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0717, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0718, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0724, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0735, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0735, grad_fn=<AddBackward>)]])]"

      ]

     },

     "execution_count": 25,

     "metadata": {},

     "output_type": "execute_result"

    }

   ],

   "source": [

    "pred = model(xs_test)\n",

    "eval_classification(y_test[0], pred[0][-1][0])\n",

    "get_vin_loss(pred, y_test, loss_fn, model, valid_loc=None, target_id=None, \n",

    "                        level_weight=None)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 26,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "acc=0.765, precision=0.736, recall=0.765, fl=0.749, adj_MI=0.007, auc=0.627, ap=0.253, confusion_mat=\n",

      "[[1067  133]\n",

      " [ 209   49]]\n",

      "report              precision    recall  f1-score   support\n",

      "\n",

      "          0       0.84      0.89      0.86      1200\n",

      "          1       0.27      0.19      0.22       258\n",

      "\n",

      "avg / total       0.74      0.77      0.75      1458\n",

      "\n"

     ]

    },

    {

     "data": {

      "text/plain": [

       "[(tensor(6.2221, grad_fn=<ThAddBackward>),\n",

       "  [[[tensor(0.4945, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4714, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4632, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.8643, grad_fn=<NllLossBackward>)],\n",

       "    tensor(1.0757, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(1.6178, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.8301, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.5669, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.8302, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.8695, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.8655, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.8631, grad_fn=<NllLossBackward>)],\n",

       "    tensor(1.4329, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(1.7378, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.7598, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.7488, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.0182, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.3899, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.4158, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.3304, grad_fn=<NllLossBackward>)],\n",

       "    tensor(1.6415, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(2.0721, grad_fn=<NllLossBackward>)],\n",

       "    tensor(2.0721, grad_fn=<AddBackward>)]]),\n",

       " (tensor(9.2298, grad_fn=<ThAddBackward>),\n",

       "  [[[tensor(0.4708, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4488, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.4466, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.9583, grad_fn=<NllLossBackward>)],\n",

       "    tensor(1.1003, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(2.5054, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.8771, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.2349, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.5306, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.9175, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.9234, grad_fn=<NllLossBackward>),\n",

       "     tensor(0.9039, grad_fn=<NllLossBackward>)],\n",

       "    tensor(1.9370, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(2.1666, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.5449, grad_fn=<NllLossBackward>),\n",

       "     tensor(2.4155, grad_fn=<NllLossBackward>),\n",

       "     tensor(4.3624, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.9025, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.9214, grad_fn=<NllLossBackward>),\n",

       "     tensor(1.8252, grad_fn=<NllLossBackward>)],\n",

       "    tensor(2.4909, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(3.7016, grad_fn=<NllLossBackward>)],\n",

       "    tensor(3.7016, grad_fn=<AddBackward>)]]),\n",

       " (tensor(0.3275, grad_fn=<ThAddBackward>),\n",

       "  [[[tensor(0.0827, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0782, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0821, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0795, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0807, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0871, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0904, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0796, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0926, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0761, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0761, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0760, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0827, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0821, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0815, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0803, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0861, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0802, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0802, grad_fn=<MseLossBackward>),\n",

       "     tensor(0.0800, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0815, grad_fn=<ThAddBackward>)],\n",

       "   [[tensor(0.0826, grad_fn=<MseLossBackward>)],\n",

       "    tensor(0.0826, grad_fn=<AddBackward>)]])]"

      ]

     },

     "execution_count": 26,

     "metadata": {},

     "output_type": "execute_result"

    }

   ],

   "source": [

    "pred = model(xs_train)\n",

    "eval_classification(y_train[0], pred[0][-1][0])\n",

    "get_vin_loss(pred, y_train, loss_fn, model, valid_loc=None, target_id=None, \n",

    "                   level_weight=None)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 27,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "Train\n",

      "acc=0.765, precision=0.736, recall=0.765, fl=0.749, adj_MI=0.007, auc=0.627, ap=0.253, confusion_mat=\n",

      "[[1067  133]\n",

      " [ 209   49]]\n",

      "report              precision    recall  f1-score   support\n",

      "\n",

      "          0       0.84      0.89      0.86      1200\n",

      "          1       0.27      0.19      0.22       258\n",

      "\n",

      "avg / total       0.74      0.77      0.75      1458\n",

      "\n",

      "Validataion\n",

      "acc=0.808, precision=0.798, recall=0.808, fl=0.802, adj_MI=0.041, auc=0.675, ap=0.315, confusion_mat=\n",

      "[[158  18]\n",

      " [ 22  10]]\n",

      "report              precision    recall  f1-score   support\n",

      "\n",

      "          0       0.88      0.90      0.89       176\n",

      "          1       0.36      0.31      0.33        32\n",

      "\n",

      "avg / total       0.80      0.81      0.80       208\n",

      "\n",

      "Test\n",

      "acc=1.000, precision=1.000, recall=1.000, fl=1.000, adj_MI=1.000, auc=1.000, ap=1.000, confusion_mat=\n",

      "[[340   0]\n",

      " [  0  77]]\n",

      "report              precision    recall  f1-score   support\n",

      "\n",

      "          0       1.00      1.00      1.00       340\n",

      "          1       1.00      1.00      1.00        77\n",

      "\n",

      "avg / total       1.00      1.00      1.00       417\n",

      "\n"

     ]

    },

    {

     "data": {

      "text/plain": [

       "[(array([0.7654321 , 0.73587779, 0.7654321 , 0.74877395, 0.00738225,\n",

       "         0.62744186, 0.25277694]), array([[1067,  133],\n",

       "         [ 209,   49]], dtype=int64)),\n",

       " (array([0.80769231, 0.7976801 , 0.80769231, 0.80236243, 0.04149792,\n",

       "         0.67542614, 0.31500486]), array([[158,  18],\n",

       "         [ 22,  10]], dtype=int64)),\n",

       " (array([1., 1., 1., 1., 1., 1., 1.]), array([[340,   0],\n",

       "         [  0,  77]], dtype=int64))]"

      ]

     },

     "execution_count": 27,

     "metadata": {},

     "output_type": "execute_result"

    }

   ],

   "source": [

    "eval_classification_multi_splits(model, [xs_train, xs_val, xs_test], [y_train[0], y_val[0], y_test[0]], \n",

    "  batch_size=None, multi_heads=False, cls_head=0, average='weighted', return_result=True, \n",

    "  split_names=['Train', 'Validataion', 'Test'], verbose=True, \n",

    "  predict_func=predict_func, pred_kwargs={'target_idx':0, 'level':-1, 'loc':0, 'train':False})"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Neural network models"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "# loss_fn_cls = torch.nn.CrossEntropyLoss(weight=torch.tensor([0.3, 0.6], device=device))\n",

    "loss_fn_cls = torch.nn.CrossEntropyLoss()\n",

    "loss_fn_reg = torch.nn.MSELoss()\n",

    "loss_fns = [loss_fn_cls, loss_fn_reg]\n",

    "# For multiple data types, there are multiple interaction mats\n",

    "feat_interact_loss_type = 'graph_laplacian'\n",

    "if num_data_types > 1:\n",

    "  weight_path = ['decoders', range(num_data_types), 'weight']  \n",

    "else:\n",

    "  weight_path = ['decoder', 'weight']\n",

    "loss_feat_interact = Loss_feature_interaction(interaction_mat=interaction_mat, \n",

    "                                              loss_type=feat_interact_loss_type, \n",

    "                                              weight_path=weight_path, \n",

    "                                              normalize=True)\n",

    "other_loss_fns = [loss_feat_interact]\n",

    "if num_data_types > 1:\n",

    "  view_sim_loss_type = 'hub'\n",

    "  explicit_target = True\n",

    "  cal_target='mean-feature'\n",

    "  # In this set of experiments, the encoders for all views will have the same hidden_dim\n",

    "  loss_view_sim = Loss_view_similarity(sections=hidden_dim[-1], loss_type=view_sim_loss_type, \n",

    "    explicit_target=explicit_target, cal_target=cal_target, target=None)\n",

    "  loss_fns.append(loss_view_sim)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "model_names = []\n",