{

 "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": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "import os\n",

    "import sys\n",

    "import re\n",

    "import collections\n",

    "import functools\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 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\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": null,

   "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' # To do: target variable can be a list (partially handled)\n",

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

    "target_variable_range = [0, 1]\n",

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

    "normal_transform_feature = True\n",

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

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

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

    "randomize_labels = False\n",

    "lr = 5e-4\n",

    "weight_decay = 1e-4\n",

    "num_epochs = 1000\n",

    "reduce_every = 500\n",

    "show_results_in_notebook = True"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Prepare data"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "result_folder = 'results'\n",

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

    "project_folder = '../../pan-can-atlas'\n",

    "print_stats = True\n",

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

    "  project_folder = 'F:/TCGA/Pan-Cancer-Atlas'\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": null,

   "metadata": {},

   "outputs": [],

   "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 = 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": null,

   "metadata": {},

   "outputs": [],

   "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": null,

   "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": null,

   "metadata": {},

   "outputs": [],

   "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": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "y_targets = get_target_variable(target_variable, clinical_dict, sel_patient_ids)\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"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "### To do: handle multiple inputs, multiple targets"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

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

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

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

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

    "\n",

    "x_train, y_train = X[train_idx], y_true[train_idx]\n",

    "x_val, y_val = X[val_idx], y_true[val_idx]\n",

    "x_test, y_test = X[test_idx], y_true[test_idx]\n",

    "print(x_train.shape, x_val.shape, x_test.shape, y_train.shape, y_val.shape, y_test.shape)\n",

    "\n",

    "label_prob_train = get_label_prob(y_train, verbose=False)\n",

    "label_probs = [label_prob_train]\n",

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

    "  label_prob_val = get_label_prob(y_val, verbose=False)\n",

    "  assert len(label_prob_train) == len(label_prob_val)\n",

    "  label_probs.append(label_prob_val)\n",

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

    "  label_prob_test = get_label_prob(y_test, verbose=False)\n",

    "  assert len(label_prob_train) == len(label_prob_test)\n",

    "  label_probs.append(label_prob_test)\n",

    "if isinstance(label_probs, torch.Tensor):\n",

    "  print('label distribution:\\n', torch.stack(label_probs, dim=1))\n",

    "else:\n",

    "  print('label distribution:\\n', np.stack(label_probs, axis=1))"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "### Optionally randomize true class labels"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {

    "scrolled": true

   },

   "outputs": [],

   "source": [

    "if randomize_labels:\n",

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

    "  y_train = torch.multinomial(label_prob_train, len(y_train), replacement=True)\n",

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

    "    y_val = torch.multinomial(label_prob_val, len(y_val), replacement=True)\n",

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

    "    y_test = torch.multinomial(label_prob_test, len(y_test), replacement=True)"

   ]

  },

  {

   "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",

    "split_names = ['train', 'val', 'test']\n",

    "metric_names = ['acc', 'precision', 'recall', 'f1_score', 'adjusted_mutual_info', 'auc', \n",

    "                'average_precision']\n",

    "metric_all = []\n",

    "confusion_mat_all = []\n",

    "loss_his_all = []\n",

    "acc_his_all = []"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "def get_result(model, best_model, best_val_acc, best_epoch, x_train, y_train, x_val, y_val, \n",

    "               x_test, y_test, batch_size, multi_heads, show_results_in_notebook=True, \n",

    "               loss_idx=0, acc_idx=0):\n",

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

    "    print(f'Best model on validation set: best_val_acc={best_val_acc:.2f}, epoch={best_epoch}')\n",

    "    metric = eval_classification_multi_splits(best_model, xs=[x_train, x_val, x_test], \n",

    "      ys=[y_train, y_val, y_test], batch_size=batch_size, multi_heads=multi_heads)\n",

    "\n",

    "  if show_results_in_notebook:\n",

    "    print('\\nModel after the last training epoch:')\n",

    "    eval_classification_multi_splits(model, xs=[x_train, x_val, x_test], \n",

    "                                     ys=[y_train, y_val, y_test], batch_size=batch_size, \n",

    "                                     multi_heads=multi_heads, return_result=False)\n",

    "\n",

    "    plot_history_multi_splits([loss_train_his, loss_val_his, loss_test_his], title='Loss', \n",

    "                              idx=loss_idx)\n",

    "    plot_history_multi_splits([acc_train_his, acc_val_his, acc_test_his], title='Acc', idx=acc_idx)\n",

    "    # scatter plot\n",

    "    plot_data_multi_splits(best_model, [x_train, x_val, x_test], [y_train, y_val, y_test], \n",

    "                           num_heads=2 if multi_heads else 1, \n",

    "                           titles=['Training', 'Validation', 'Test'], batch_size=batch_size)\n",

    "    return metric"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Plain deep learning model"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "batch_size = 1000\n",

    "print_every = 100\n",

    "eval_every = 1"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "in_dim = x_train.shape[1]\n",

    "print('Plain deep learning model')\n",

    "model_names.append('NN')\n",

    "model = DenseLinear(in_dim, hidden_dim+[num_cls], dense=dense, residual=residual).to(device)\n",

    "multi_heads = False\n",

    "\n",

    "loss_train_his = []\n",

    "loss_val_his = []\n",

    "loss_test_his = []\n",

    "acc_train_his = []\n",

    "acc_val_his = []\n",

    "acc_test_his = []\n",

    "best_model = model\n",

    "best_val_acc = 0\n",

    "best_epoch = 0"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "best_model, best_val_acc, best_epoch = train_single_loss(model, x_train, y_train, \n",

    "    x_val, y_val, x_test, y_test, loss_fn=loss_fn_cls, lr=lr, weight_decay=weight_decay, \n",

    "    amsgrad=True, batch_size=batch_size, num_epochs=num_epochs, \n",

    "    reduce_every=reduce_every, eval_every=eval_every, print_every=print_every, verbose=False, \n",

    "    loss_train_his=loss_train_his, loss_val_his=loss_val_his, loss_test_his=loss_test_his, \n",

    "    acc_train_his=acc_train_his, acc_val_his=acc_val_his, acc_test_his=acc_test_his, \n",

    "    return_best_val=True)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "metric = get_result(model, best_model, best_val_acc, best_epoch, x_train, y_train, x_val, y_val, \n",

    "                    x_test, y_test, batch_size, multi_heads, show_results_in_notebook, \n",

    "                    loss_idx=0, acc_idx=0)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "loss_his_all.append([loss_train_his, loss_val_his, loss_test_his])\n",

    "acc_his_all.append([acc_train_his, acc_val_his, acc_test_his])\n",

    "metric_all.append([v[0] for v in metric])\n",

    "confusion_mat_all.append([v[1] for v in metric])"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Factorization AutoEncoder"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "def run_one_model(model, loss_weights, other_loss_weights, \n",

    "                  loss_his_all=[], acc_his_all=[], metric_all=[], confusion_mat_all=[],\n",

    "                  heads=[0,1], multi_heads=True, return_results=False, \n",

    "                  loss_fns=loss_fns, other_loss_fns=other_loss_fns, \n",

    "                  lr=lr, weight_decay=weight_decay, batch_size=batch_size, \n",

    "                  num_epochs=num_epochs, reduce_every=reduce_every, eval_every=eval_every, \n",

    "                  print_every=print_every, x_train=x_train, y_train=y_train,\n",

    "                  x_val=x_val, y_val=y_val, x_test=x_test, y_test=y_test,\n",

    "                  show_results_in_notebook=show_results_in_notebook):\n",

    "  \"\"\"Train a model and get results  \n",

    "    Most of the parameters are from the context; handle it properly\n",

    "  \"\"\"\n",

    "  loss_train_his = []\n",

    "  loss_val_his = []\n",

    "  loss_test_his = []\n",

    "  acc_train_his = []\n",

    "  acc_val_his = []\n",

    "  acc_test_his = []\n",

    "  best_model = model\n",

    "  best_val_acc = 0\n",

    "  best_epoch = 0\n",

    "\n",

    "  best_model, best_val_acc, best_epoch = train_multiloss(model, x_train, [y_train, x_train], \n",

    "    x_val, [y_val, x_val], x_test, [y_test, x_test], heads=heads, loss_fns=loss_fns, \n",

    "    loss_weights=loss_weights, other_loss_fns=other_loss_fns, \n",

    "    other_loss_weights=other_loss_weights, \n",

    "    lr=lr, weight_decay=weight_decay, batch_size=batch_size, num_epochs=num_epochs, \n",

    "    reduce_every=reduce_every, eval_every=eval_every, print_every=print_every,\n",

    "    loss_train_his=loss_train_his, loss_val_his=loss_val_his, loss_test_his=loss_test_his, \n",

    "    acc_train_his=acc_train_his, acc_val_his=acc_val_his, acc_test_his=acc_test_his, \n",

    "    return_best_val=True, amsgrad=True, verbose=False)\n",

    "\n",

    "  metric = get_result(model, best_model, best_val_acc, best_epoch, x_train, y_train, x_val, y_val, \n",

    "                      x_test, y_test, batch_size, multi_heads, show_results_in_notebook, \n",

    "                      loss_idx=0, acc_idx=0)\n",

    "\n",

    "  loss_his_all.append([loss_train_his, loss_val_his, loss_test_his])\n",

    "  acc_his_all.append([acc_train_his, acc_val_his, acc_test_his])\n",

    "  metric_all.append([v[0] for v in metric])\n",

    "  confusion_mat_all.append([v[1] for v in metric])\n",

    "  \n",

    "  if return_results:\n",

    "    return loss_his_all, acc_his_all, metric_all, confusion_mat_all"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "decoder_norm = False\n",

    "uniform_decoder_norm = False\n",

    "print('Plain AutoEncoder model')\n",

    "model_names.append('AE')\n",

    "model = AutoEncoder(in_dim, hidden_dim, num_cls, dense=dense, residual=residual,\n",

    "          decoder_norm=decoder_norm, uniform_decoder_norm=uniform_decoder_norm).to(device)\n",

    "loss_weights = [1,1]\n",

    "other_loss_weights = [0]\n",

    "# heads = None should work for all the following; keep this for clarity\n",

    "heads = [0,1] \n",

    "run_one_model(model, loss_weights, other_loss_weights,\n",

    "              loss_his_all, acc_his_all, metric_all, confusion_mat_all,\n",

    "              heads=heads, multi_heads=True, return_results=False, \n",

    "              loss_fns=loss_fns, other_loss_fns=other_loss_fns, \n",

    "              lr=lr, weight_decay=weight_decay, batch_size=batch_size, \n",

    "              num_epochs=num_epochs, reduce_every=reduce_every, eval_every=eval_every, \n",

    "              print_every=print_every, x_train=x_train, y_train=y_train,\n",

    "              x_val=x_val, y_val=y_val, x_test=x_test, y_test=y_test,\n",

    "              show_results_in_notebook=show_results_in_notebook)"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Add feature interaction network regularizer"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "if num_data_types > 1:\n",

    "  fuse_type = 'sum'\n",

    "  print('MultiviewAE with feature interaction network regularizer')\n",

    "  model_names.append('MultiviewAE + feat_int')\n",

    "  model = MultiviewAE(in_dims=in_dims, hidden_dims=hidden_dim, out_dim=num_cls, \n",

    "                      fuse_type=fuse_type, dense=dense, residual=residual, \n",

    "                      residual_layers='all', decoder_norm=decoder_norm, \n",

    "                      decoder_norm_dim=0, uniform_decoder_norm=uniform_decoder_norm, \n",

    "                      nonlinearity=nn.ReLU(), last_nonlinearity=True, bias=True).to(device)\n",

    "else:\n",

    "  print('AutoEncoder with feature interaction network regularizer')\n",

    "  model_names.append('AE + feat_int')\n",

    "  model = AutoEncoder(in_dim, hidden_dim, num_cls, dense=dense, residual=residual, \n",

    "          decoder_norm=decoder_norm, uniform_decoder_norm=uniform_decoder_norm).to(device)\n",

    "\n",

    "loss_weights = [1,1]\n",

    "other_loss_weights = [1]\n",

    "heads = [0,1]\n",

    "run_one_model(model, loss_weights, other_loss_weights, \n",

    "              loss_his_all, acc_his_all, metric_all, confusion_mat_all,\n",

    "              heads=heads, multi_heads=True, return_results=False, \n",

    "              loss_fns=loss_fns, other_loss_fns=other_loss_fns, \n",

    "              lr=lr, weight_decay=weight_decay, batch_size=batch_size, \n",

    "              num_epochs=num_epochs, reduce_every=reduce_every, eval_every=eval_every, \n",

    "              print_every=print_every, x_train=x_train, y_train=y_train,\n",

    "              x_val=x_val, y_val=y_val, x_test=x_test, y_test=y_test,\n",

    "              show_results_in_notebook=show_results_in_notebook)"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## For multi-view data, add view similarity network regularizer"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "if num_data_types > 1:\n",

    "  # plain multiviewAE; compare it with plain AutoEncoder to see \n",

    "  # if separating views in lower layers in MultiviewAE is better than combining them all the way\n",

    "  print('Run plain MultiviewAE model')\n",

    "  model_names.append('MultiviewAE')\n",

    "  model = MultiviewAE(in_dims=in_dims, hidden_dims=hidden_dim, out_dim=num_cls, \n",

    "                    fuse_type=fuse_type, dense=dense, residual=residual, \n",

    "                    residual_layers='all', decoder_norm=decoder_norm, \n",

    "                    decoder_norm_dim=0, uniform_decoder_norm=uniform_decoder_norm, \n",

    "                    nonlinearity=nn.ReLU(), last_nonlinearity=True, bias=True).to(device)\n",

    "\n",

    "  loss_weights = [1,1]\n",

    "  other_loss_weights = [0]\n",

    "  heads = [0,1]\n",

    "  run_one_model(model, loss_weights, other_loss_weights, \n",

    "                loss_his_all, acc_his_all, metric_all, confusion_mat_all,\n",

    "                heads=heads, multi_heads=True, return_results=False, \n",

    "                loss_fns=loss_fns, other_loss_fns=other_loss_fns, \n",

    "                lr=lr, weight_decay=weight_decay, batch_size=batch_size, \n",

    "                num_epochs=num_epochs, reduce_every=reduce_every, eval_every=eval_every, \n",

    "                print_every=print_every, x_train=x_train, y_train=y_train,\n",

    "                x_val=x_val, y_val=y_val, x_test=x_test, y_test=y_test,\n",

    "                show_results_in_notebook=show_results_in_notebook)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "if num_data_types > 1:\n",

    "  print('MultiviewAE with view similarity regularizers')\n",

    "  model_names.append('MultiviewAE + view_sim')\n",

    "  model = MultiviewAE(in_dims=in_dims, hidden_dims=hidden_dim, out_dim=num_cls, \n",

    "                      fuse_type=fuse_type, dense=dense, residual=residual, \n",

    "                      residual_layers='all', decoder_norm=decoder_norm, \n",

    "                      decoder_norm_dim=0, uniform_decoder_norm=uniform_decoder_norm, \n",

    "                      nonlinearity=nn.ReLU(), last_nonlinearity=True, bias=True).to(device)\n",

    "  loss_weights = [1,1,1]\n",

    "  other_loss_weights = [0]\n",

    "  heads = [0,1,2]\n",

    "  run_one_model(model, loss_weights, other_loss_weights, \n",

    "                loss_his_all, acc_his_all, metric_all, confusion_mat_all,\n",

    "                heads=heads, multi_heads=True, return_results=False, \n",

    "                loss_fns=loss_fns, other_loss_fns=other_loss_fns, \n",

    "                lr=lr, weight_decay=weight_decay, batch_size=batch_size, \n",

    "                num_epochs=num_epochs, reduce_every=reduce_every, eval_every=eval_every, \n",

    "                print_every=print_every, x_train=x_train, y_train=y_train,\n",

    "                x_val=x_val, y_val=y_val, x_test=x_test, y_test=y_test,\n",

    "                show_results_in_notebook=show_results_in_notebook)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

    "if num_data_types > 1:\n",

    "  print('MultiviewAE with both feature interaction and view similarity regularizers')\n",

    "  model_names.append('MultiviewAE + feat_int + view_sim')\n",

    "  model = MultiviewAE(in_dims=in_dims, hidden_dims=hidden_dim, out_dim=num_cls, \n",

    "                      fuse_type=fuse_type, dense=dense, residual=residual, \n",

    "                      residual_layers='all', decoder_norm=decoder_norm, \n",

    "                      decoder_norm_dim=0, uniform_decoder_norm=uniform_decoder_norm, \n",

    "                      nonlinearity=nn.ReLU(), last_nonlinearity=True, bias=True).to(device)\n",

    "  loss_weights = [1,1,1]\n",

    "  other_loss_weights = [1]\n",

    "  heads = [0,1,2]\n",

    "  run_one_model(model, loss_weights, other_loss_weights,\n",

    "                loss_his_all, acc_his_all, metric_all, confusion_mat_all,\n",

    "                heads=heads, multi_heads=True, return_results=False, \n",

    "                loss_fns=loss_fns, other_loss_fns=other_loss_fns, \n",

    "                lr=lr, weight_decay=weight_decay, batch_size=batch_size, \n",

    "                num_epochs=num_epochs, reduce_every=reduce_every, eval_every=eval_every, \n",

    "                print_every=print_every, x_train=x_train, y_train=y_train,\n",

    "                x_val=x_val, y_val=y_val, x_test=x_test, y_test=y_test,\n",

    "                show_results_in_notebook=show_results_in_notebook)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": [

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

    "  print(f'Write result to file {result_folder}/{res_file}')\n",

    "  pickle.dump({'loss_his_all': loss_his_all,\n",

    "               'acc_his_all': acc_his_all,\n",

    "               'metric_all': metric_all,\n",

    "               'confusion_mat_all': confusion_mat_all,\n",

    "               'model_names': model_names,\n",

    "               'split_names': split_names,\n",

    "               'metric_names': metric_names\n",

    "              }, f)"

   ]

  }

 ],

 "metadata": {

  "kernelspec": {

   "display_name": "Python 3",

   "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.6.5"

  }

 },

 "nbformat": 4,

 "nbformat_minor": 2

}