UMM_discovery_BBBC021.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "Copyright 2020 Novartis Institutes for BioMedical Research Inc.\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "This file incorporates work covered by the following copyright and  \n",
    "permission notice:\n",
    "\n",
    "   Copyright (c) 2017-present, Facebook, Inc.\n",
    "   All rights reserved.\n",
    "\n",
    "   This source code is licensed as Creative Commons Attribution-Noncommercial \n",
    "   and can be found under https://creativecommons.org/licenses/by-nc/4.0/."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# UMM Discovery"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check for GPU device\n",
    "import os\n",
    "print(os.getenv('CUDA_VISIBLE_DEVICES'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba.errors import NumbaDeprecationWarning, NumbaPendingDeprecationWarning, NumbaPerformanceWarning\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(action='once')\n",
    "warnings.simplefilter('ignore', category=NumbaDeprecationWarning)\n",
    "warnings.simplefilter('ignore', category=NumbaPendingDeprecationWarning)\n",
    "warnings.simplefilter('ignore', category=NumbaPerformanceWarning)\n",
    "warnings.simplefilter(\"ignore\", category=PendingDeprecationWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import time\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim\n",
    "import torch.utils.data\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "from torchsummary import summary\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import clustering\n",
    "from model import MultiScaleNet\n",
    "from util import Logger, UnifLabelSampler\n",
    "from training import compute_features, train\n",
    "from correction import do_batch_correction\n",
    "from evaluation import evaluate_epoch, evaluate_training\n",
    "from my_dataset import bbbc021_dataset\n",
    "from plot import plot_training, plot_training_summary, plot_inter_part_validation, plot_inter_hier_validation, plot_external_validation, plot_NSC_NSB, plot_total_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class args:\n",
    "    # Experiment ID\n",
    "    run_ID = \"EXP01\"\n",
    "    \n",
    "    # path to dataset\n",
    "    data = \"/path/to/dataset\"\n",
    "    \n",
    "    # file name of the metadata csv file\n",
    "    csv = \"dataset.csv\"\n",
    "    \n",
    "    # path to the output folder\n",
    "    save_path = \"/path/to/output/folder/\"\n",
    "    output_path = \"\" #dont assign\n",
    "\n",
    "    # CNN architecture \n",
    "    arch = \"msnet\" \n",
    "\n",
    "    # input channels\n",
    "    n_input_dim = 3\n",
    "    \n",
    "    # number of features\n",
    "    n_features = 64\n",
    "    \n",
    "    # Batch correction [\"TVN\", \"COMBAT\", MNN]\n",
    "    batch_corrections = [\"COMBAT\", \"TVN\"] \n",
    "    \n",
    "    # dimension reduction method (default: PCA)\n",
    "    # 'PCA' or 'UMAP' or 'TSNE' or 'AdaptiveTSNE'\n",
    "    dim_method = 'PCA' \n",
    "    \n",
    "    # dimensions after reduction\n",
    "    n_components = 16\n",
    "    \n",
    "    # clustering algorithm (default: Kmeans)\n",
    "    # 'Kmeans' or 'AdaptiveKmeans' or 'PIC' or 'HDBscan' \n",
    "    clustering = 'Kmeans'\n",
    "\n",
    "    # number of cluster for k-means (default: 10000)\n",
    "    nmb_cluster = 104\n",
    "    \n",
    "    # parameter for adaptive k-means\n",
    "    end_k = 39\n",
    "    end_epochs = 150\n",
    "    \n",
    "    # parameter for PIC\n",
    "    sigma = 0.2 \n",
    "    nnn = 5\n",
    "    \n",
    "    # paramter for HDBscan\n",
    "    min_cluster_size = 12\n",
    "    min_samples = 3\n",
    "    nnn = 25\n",
    "    \n",
    "    # mini-batch size (default: 256)\n",
    "    batch = 16\n",
    "    \n",
    "    # learning rate (default: 0.05)\n",
    "    lr = 0.05\n",
    "    \n",
    "    # weight decay pow (default: -5)\n",
    "    wd=-5\n",
    "    \n",
    "    # number of total epochs to run (default: 200)\n",
    "    epochs = 50\n",
    "\n",
    "    # momentum (default: 0.9)\n",
    "    momentum = 0.9\n",
    "\n",
    "    # how many epochs of training between two consecutive reassignments of clusters (default: 1)\n",
    "    reassign = 1\n",
    "\n",
    "    # number of data loading workers (default: 4)\n",
    "    workers=4\n",
    "    \n",
    "    # cuda version\n",
    "    cuda = 10\n",
    "    \n",
    "    # Epoch to continue from (default: None)\n",
    "    resume = None\n",
    "    \n",
    "    # manual epoch number (useful on restarts) (default: 0)\n",
    "    start_epoch = 0\n",
    "    \n",
    "    # random seed (default: 42)\n",
    "    seed = 12\n",
    "    \n",
    "    # chatty\n",
    "    verbose = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_name = \"\"\n",
    "exp_name +='{}_'.format(args.run_ID)\n",
    "\n",
    "## Add clustering algorithm\n",
    "if args.clustering == \"Kmeans\":\n",
    "    clustering_parameter = 'Kmeans_cl{}'.format(args.nmb_cluster)\n",
    "elif args.clustering == \"AdaptiveKmeans\":\n",
    "    clustering_parameter = 'AdaptiveKmeans_cl{}-{}_ep{}'.format(args.nmb_cluster, args.end_k, args.end_epochs)\n",
    "elif args.clustering == \"PIC\":\n",
    "    clustering_parameter = 'PIC_sig{}_nn{}'.format(args.sigma, args.nnn)\n",
    "elif args.clustering == \"HDBscan\":\n",
    "    clustering_parameter = 'HDBscan_mincl{}-minsam{}_nn{}'.format(args.min_cluster_size, args.min_samples, args.nnn)\n",
    "\n",
    "exp_name +='{}'.format(clustering_parameter)\n",
    "    \n",
    "## Add dimensionality reduction method\n",
    "exp_name +='_{}-{}d'.format(args.dim_method, args.n_components)\n",
    "\n",
    "## Add batch correction\n",
    "batch_corrections_parameter = None\n",
    "if args.batch_corrections == None:\n",
    "    batch_corrections_parameter = \"NoCorrection\"\n",
    "else:\n",
    "    for corr in args.batch_corrections:\n",
    "        if batch_corrections_parameter is None:\n",
    "            batch_corrections_parameter = corr\n",
    "        else:\n",
    "            batch_corrections_parameter +='_{}'.format(corr)\n",
    "exp_name +='_{}'.format(batch_corrections_parameter)\n",
    "    \n",
    "print('Experiment: {}'.format(exp_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args.output_path = os.path.join(args.save_path, exp_name)\n",
    "print('Output path: {}'.format(args.output_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_path = os.path.join(args.output_path, 'y-embed_{}.txt'.format(exp_name))\n",
    "metadata_path = os.path.join(args.output_path, 'metadata_{}.txt'.format(exp_name))\n",
    "assign_epoch_path = os.path.join(args.output_path, 'epoch_assignment_{}.txt'.format(exp_name))\n",
    "metrices_path = os.path.join(args.output_path, 'metrices_{}.txt'.format(exp_name))\n",
    "loss_path = os.path.join(args.output_path, 'loss_{}.txt'.format(exp_name))\n",
    "\n",
    "best_labeled_sil_embed_path = os.path.join(args.output_path, 'best_labeled_sil_y-embed_{}.txt'.format(exp_name))\n",
    "best_labeled_nsc_embed_path = os.path.join(args.output_path, 'best_labeled_nsc_y-embed_{}.txt'.format(exp_name))\n",
    "best_unlabeled_embed_path = os.path.join(args.output_path, 'best_unlabeled_y-embed_{}.txt'.format(exp_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeds_cols = ['Z{:03d}'.format(i) for i in range(args.n_features)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Neural Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fix random seeds\n",
    "torch.manual_seed(args.seed)\n",
    "torch.cuda.manual_seed_all(args.seed)\n",
    "np.random.seed(args.seed)\n",
    "\n",
    "# CNN\n",
    "if args.verbose:\n",
    "    print('Create Multi-Scale Neural Network')\n",
    "    \n",
    "model = MultiScaleNet(input_dim=3, num_features=args.n_features, num_classes=args.nmb_cluster)\n",
    "fd = int(model.top_layer.weight.size()[1])\n",
    "model.top_layer = None\n",
    "model.features = torch.nn.DataParallel(model.features)\n",
    "model.cuda()\n",
    "cudnn.benchmark = True\n",
    "\n",
    "# create optimizer\n",
    "optimizer = torch.optim.SGD(\n",
    "    filter(lambda x: x.requires_grad, model.parameters()),\n",
    "    lr=args.lr,\n",
    "    momentum=args.momentum,\n",
    "    weight_decay=10**args.wd,\n",
    ")\n",
    "\n",
    "# define loss function\n",
    "criterion = nn.CrossEntropyLoss().cuda()\n",
    "\n",
    "# optionally resume from a checkpoint\n",
    "if args.resume:\n",
    "    checkpoint_path = os.path.join(args.output_path, 'checkpoints', 'checkpoint_epoch{}.pth.tar'.format(args.resume))\n",
    "    if os.path.isfile(checkpoint_path):\n",
    "        print(\"=> loading checkpoint '{}'\".format(checkpoint_path))\n",
    "        checkpoint = torch.load(checkpoint_path)\n",
    "        args.start_epoch = checkpoint['epoch']\n",
    "        # remove top_layer parameters from checkpoint\n",
    "        for key in list(checkpoint['state_dict']):\n",
    "            if 'top_layer' in key:\n",
    "                del checkpoint['state_dict'][key]\n",
    "        model.load_state_dict(checkpoint['state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "        print(\"=> loaded checkpoint '{}' (epoch {})\"\n",
    "              .format(args.resume, checkpoint['epoch']))\n",
    "    else:\n",
    "        print(\"=> no checkpoint found at '{}'\".format(args.resume))\n",
    "\n",
    "# creating checkpoint repo\n",
    "exp_check = os.path.join(args.output_path, 'checkpoints')\n",
    "if not os.path.isdir(exp_check):\n",
    "    os.makedirs(exp_check)\n",
    "\n",
    "# creating cluster assignments log\n",
    "cluster_log = Logger(os.path.join(args.output_path, 'train_data.pkl'))\n",
    "if args.resume:\n",
    "    cluster_log.load_data(args.resume)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if args.verbose:\n",
    "    summary(model, input_size=(3, 512, 640))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.cuda==10:\n",
    "    # create tensorboard writer\n",
    "    tensorboard_path = os.path.join(args.save_path, 'runs', exp_name)\n",
    "\n",
    "    if os.path.isdir(tensorboard_path) and not args.resume:\n",
    "        shutil.rmtree(tensorboard_path)\n",
    "\n",
    "    if not os.path.isdir(tensorboard_path):\n",
    "        os.makedirs(tensorboard_path)\n",
    "    \n",
    "    from torch.utils.tensorboard import SummaryWriter\n",
    "    tensorboard_writer = SummaryWriter(log_dir=tensorboard_path)\n",
    "else:\n",
    "    tensorboard_writer = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tra = [transforms.ToTensor()]\n",
    "    \n",
    "# load the data\n",
    "end = time.time()\n",
    "\n",
    "dataset = bbbc021_dataset(args.data, args.csv, 'pseudoclass', transform=transforms.Compose(tra))\n",
    "\n",
    "if args.verbose: \n",
    "    print('Load dataset: {0:.2f} s'.format(time.time() - end))\n",
    "    \n",
    "dataloader = torch.utils.data.DataLoader(dataset,\n",
    "                                         batch_size=args.batch,\n",
    "                                         num_workers=args.workers,\n",
    "                                         pin_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.get_df().head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Number of compounds: {}'.format(len(dataset.get_df()['compound'].unique())))\n",
    "print('Number of treatments: {}'.format(len(dataset.get_df()['pseudoclass'].unique())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.get_df().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(3):\n",
    "    img = dataset[i][0]\n",
    "\n",
    "    plt.figure(figsize=(20,10))\n",
    "    ax = plt.subplot(1, 3, 1)\n",
    "    plt.imshow(img[0,:,:])\n",
    "    \n",
    "    ax = plt.subplot(1, 3, 2)\n",
    "    plt.imshow(img[1,:,:])\n",
    "    \n",
    "    ax = plt.subplot(1, 3, 3)\n",
    "    plt.imshow(img[2,:,:])\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup clustering algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# clustering algorithm to use\n",
    "if args.clustering == \"Kmeans\":\n",
    "    deepcluster = clustering.Kmeans(k=args.nmb_cluster, \n",
    "                                    dim_method=args.dim_method, \n",
    "                                    n_components=args.n_components,\n",
    "                                    n_jobs=args.workers+1)\n",
    "elif args.clustering == \"AdaptiveKmeans\":\n",
    "    deepcluster = clustering.AdaptiveKmeans(initial_k=args.nmb_cluster, \n",
    "                                            end_k=args.end_k, \n",
    "                                            end_epochs=args.end_epochs, \n",
    "                                            initial_epoch=args.start_epoch, \n",
    "                                            dim_method=args.dim_method, \n",
    "                                            n_components=args.n_components,\n",
    "                                            n_jobs=args.workers+1)\n",
    "elif args.clustering == \"PIC\":\n",
    "    deepcluster = clustering.PIC(sigma=args.sigma, \n",
    "                                 nnn=args.nnn, \n",
    "                                 dim_method=args.dim_method, \n",
    "                                 n_components=args.n_components,\n",
    "                                 n_jobs=args.workers+1)\n",
    "elif args.clustering == \"HDBscan\":\n",
    "    deepcluster = clustering.HDBscan(min_cluster_size=args.min_cluster_size, \n",
    "                                    min_samples=args.min_samples, \n",
    "                                    nnn=args.nnn, \n",
    "                                    dim_method=args.dim_method, \n",
    "                                    n_components=args.n_components,\n",
    "                                    n_jobs=args.workers+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_list = []\n",
    "nmi_list = [0]\n",
    "best_NSC_acc = 0\n",
    "df_meta = dataset.get_df()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# training convnet with DeepCluster\n",
    "start_train_time = time.time()\n",
    "for epoch in range(args.start_epoch, args.epochs):\n",
    "    end = time.time()\n",
    "\n",
    "    # remove head\n",
    "    model.top_layer = None\n",
    "    model.classifier = nn.Sequential(*list(model.classifier.children())[:-1])\n",
    "\n",
    "    # compute features for the whole dataset\n",
    "    features = compute_features(dataloader, model, len(dataset), args)\n",
    "    #print(features.shape)\n",
    "    embeds = pd.DataFrame(features, columns=embeds_cols)\n",
    "    embeds = pd.concat([df_meta, embeds], axis=1)\n",
    "    \n",
    "    # batch correction \n",
    "    embeds = do_batch_correction(embeds, embeds_cols, args.batch_corrections, verbose=args.verbose)\n",
    "        \n",
    "    # evaluate features\n",
    "    df_eval = evaluate_epoch(embeds.copy(), embeds_cols, verbose=args.verbose)\n",
    "    \n",
    "    # best epoch\n",
    "    NSC_acc = df_eval[\"NSC_1-NN_treatment\"][0]\n",
    "    if NSC_acc > best_NSC_acc:\n",
    "        best_NSC_acc = NSC_acc\n",
    "    \n",
    "    # cluster the features\n",
    "    clustering_loss = deepcluster.cluster(embeds[embeds_cols].to_numpy(copy=True), verbose=args.verbose)\n",
    "\n",
    "    # save cluster assignment\n",
    "    df_assign_epoch = pd.DataFrame(clustering.arrange_clustering(deepcluster.images_lists), columns=['Assignment_{}'.format(epoch)])\n",
    "\n",
    "    # assign pseudo-labels\n",
    "    train_dataset = clustering.cluster_assign(deepcluster.images_lists, dataset.imgs)\n",
    "\n",
    "    # uniformely sample per target\n",
    "    sampler = UnifLabelSampler(int(args.reassign * len(train_dataset)),\n",
    "                               deepcluster.images_lists)\n",
    "\n",
    "    train_dataloader = torch.utils.data.DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=args.batch,\n",
    "        num_workers=args.workers,\n",
    "        sampler=sampler,\n",
    "        pin_memory=True,\n",
    "    )\n",
    "\n",
    "    # set last fully connected layer\n",
    "    mlp = list(model.classifier.children())\n",
    "    mlp.append(nn.ReLU(inplace=True).cuda())\n",
    "    model.classifier = nn.Sequential(*mlp)\n",
    "    model.top_layer = nn.Linear(args.n_features, len(deepcluster.images_lists))\n",
    "    model.top_layer.weight.data.normal_(0, 0.01)\n",
    "    model.top_layer.bias.data.zero_()\n",
    "    model.top_layer.cuda()\n",
    "\n",
    "    # train network with clusters as pseudo-labels\n",
    "    start_epoch_train_time = time.time()\n",
    "    loss = train(train_dataloader, model, criterion, optimizer, epoch, args)\n",
    "    loss_list.append(float(loss))\n",
    "\n",
    "    # print log\n",
    "    if args.verbose:\n",
    "        print('###### Epoch [{0}] ###### \\n'\n",
    "              'Total Time: {1:.3f} s\\n'\n",
    "              'Training time: {2:.3f} s\\n'\n",
    "              'ConvNet loss: {3:.3f} \\n'\n",
    "              'NSC acc: {4:.3f} \\n'\n",
    "              'Silhoutte score: {5:.3f} \\n'\n",
    "              'Best NSC acc: {6:.3f}'\n",
    "              .format(epoch, \n",
    "                      time.time() - end, \n",
    "                      time.time() - start_epoch_train_time, \n",
    "                      loss, \n",
    "                      df_eval[\"NSC_1-NN_treatment\"][0]*100, \n",
    "                      df_eval[\"silhou_moa\"][0],\n",
    "                      best_NSC_acc*100\n",
    "                     ))\n",
    "        try:\n",
    "            if epoch > 0:\n",
    "                nmi = normalized_mutual_info_score(\n",
    "                    clustering.arrange_clustering(deepcluster.images_lists),\n",
    "                    clustering.arrange_clustering(cluster_log.data['clustering'][-1])\n",
    "                )\n",
    "                print('NMI against previous assignment: {0:.3f}'.format(nmi))\n",
    "            else:\n",
    "                nmi = 0\n",
    "            nmi_list.append(nmi)\n",
    "        except IndexError:\n",
    "            pass\n",
    "        print('####################### \\n')\n",
    "    # save running checkpoint\n",
    "    torch.save({'epoch': epoch + 1,\n",
    "                'arch': args.arch,\n",
    "                'state_dict': model.state_dict(),\n",
    "                'optimizer' : optimizer.state_dict()},\n",
    "               os.path.join(args.output_path, 'checkpoint.pth.tar'))\n",
    "\n",
    "    # save cluster assignments, evaluation data and features\n",
    "    cluster_log.log(epoch, embeds, df_eval, df_assign_epoch, loss, nmi, deepcluster.images_lists)\n",
    "    \n",
    "    # write to tensorboard\n",
    "    if tensorboard_writer is not None:\n",
    "        tensorboard_writer.add_scalar('loss', loss, epoch)\n",
    "        tensorboard_writer.add_scalar('nmi', nmi, epoch)\n",
    "        for eval_metric, data in df_eval.iteritems():\n",
    "            tensorboard_writer.add_scalar(eval_metric, data[0], epoch)\n",
    "        \n",
    "print('Total training time {}'.format(time.time() - start_train_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate training"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# load training data (for loading a completed training)\n",
    "cluster_log.load_data(args.resume)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_log.data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_log.data['features'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_metric = pd.concat(cluster_log.data['eval_metrices'], axis=0)\n",
    "df_metric[\"epoch\"] = cluster_log.data['epoch']\n",
    "df_metric[\"loss\"] = cluster_log.data['loss']\n",
    "df_metric[\"loss\"] = df_metric[\"loss\"].map(float)\n",
    "df_metric[\"nmi\"] = cluster_log.data['nmi']\n",
    "df_metric[\"nmi\"] = df_metric[\"nmi\"].map(float)\n",
    "df_metric = df_metric.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_metric = df_metric.rename(columns={'n_clusters_well':'n_clusters'})\n",
    "df_metric['total_score'] = (df_metric['adj-mut_comp_pred'] + df_metric['comple_treat-pred'] + df_metric['silhou_pred_tsne'])/3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "df_metric.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_training(df_metric, args.output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overall summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_training_summary(df_metric, args.output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Internal partitional validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_inter_part_validation(df_metric, args.output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Internal hierarchical validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_inter_hier_validation(df_metric, args.output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### External Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_external_validation(df_metric, args.output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_NSC_NSB(df_metric, args.output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Total score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_total_score(df_metric, args.output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find best epoch for labeled and unlabeled situation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Best epoch labeled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_epoch_labeled_silh = df_metric['silhou_moa_tsne'].idxmax()\n",
    "print('Best epoch (Silhouette): ' , best_epoch_labeled_silh)\n",
    "print('NSC: ', df_metric[\"NSC_1-NN_treatment\"][best_epoch_labeled_silh])\n",
    "print('Silhouette moa tsne: ', df_metric['silhou_moa_tsne'][best_epoch_labeled_silh])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_epoch_labeled_nsc = df_metric['NSC_1-NN_treatment'].idxmax()\n",
    "print('Best epoch (NSC): ' , best_epoch_labeled_nsc)\n",
    "print('NSC: ', df_metric[\"NSC_1-NN_treatment\"][best_epoch_labeled_nsc])\n",
    "print('Silhouette moa tsne: ', df_metric['silhou_moa_tsne'][best_epoch_labeled_nsc])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Best epoch unlabeled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_epoch_unlabeled = df_metric.loc[:,'total_score'].idxmax()\n",
    "print('Best epoch (unlabeled): ' , best_epoch_unlabeled)\n",
    "print('NSC: ', df_metric['NSC_1-NN_treatment'][best_epoch_unlabeled])\n",
    "print('Silhouette pred tsne: ', df_metric['silhou_pred_tsne'][best_epoch_unlabeled])\n",
    "print('Adjusted Rand index (treat-pred): ', df_metric['adj-rand_treat-pred'][best_epoch_unlabeled])\n",
    "print('Adjusted Mutual information (treat-pred): ', df_metric['adj-mut_treat-pred'][best_epoch_unlabeled])\n",
    "print('Adjusted Mutual information (moa-pred): ', df_metric['adj-mut_moa-pred'][best_epoch_unlabeled])\n",
    "print('Jaccard (treat-pred): ', df_metric['jaccard_treat-pred'][best_epoch_unlabeled])\n",
    "print('Completeness (treat-pred): ', df_metric['comple_treat-pred'][best_epoch_unlabeled])\n",
    "print('Number of clusters: ', df_metric['n_clusters'][best_epoch_unlabeled])\n",
    "print('Total score: ', df_metric['total_score'][best_epoch_unlabeled])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save hyperparameter with best values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if tensorboard_writer is not None:\n",
    "    tensorboard_writer.add_hparams(hparam_dict={\n",
    "                                    'cuda_version': args.cuda,\n",
    "                                    'lr': args.lr, \n",
    "                                    'batch_size': args.batch, \n",
    "                                    'n_features': args.n_features,\n",
    "                                    'batch_correction': batch_corrections_parameter,\n",
    "                                    'dim_method': args.dim_method, \n",
    "                                    'n_features_reduced': args.n_components,\n",
    "                                    'clustering_algorithmn': args.clustering,\n",
    "                                    'clustering_parameter': clustering_parameter,\n",
    "        \n",
    "                                   }, \n",
    "                                   metric_dict={\n",
    "                                       # labeled best silhoutte score moa\n",
    "                                       'labeled_best_silh/epoch': best_epoch_labeled_silh,\n",
    "                                       'labeled_best_silh/silhou_moa_tsne': df_metric['silhou_moa_tsne'][best_epoch_labeled_silh],\n",
    "                                       'labeled_best_silh/NSC_1-NN_treatment': df_metric['NSC_1-NN_treatment'][best_epoch_labeled_silh],\n",
    "                                       'labeled_best_silh/NSC_1-NN_avg': df_metric['NSC_1-NN_avg'][best_epoch_labeled_silh],\n",
    "                                       # labeled NSC\n",
    "                                       'labeled_best_nsc/epoch': best_epoch_labeled_nsc,\n",
    "                                       'labeled_best_nsc/silhou_moa_tsne': df_metric['silhou_moa_tsne'][best_epoch_labeled_nsc],\n",
    "                                       'labeled_best_nsc/NSC_1-NN_treatment': df_metric['NSC_1-NN_treatment'][best_epoch_labeled_nsc],\n",
    "                                       'labeled_best_nsc/NSC_1-NN_avg': df_metric['NSC_1-NN_avg'][best_epoch_labeled_nsc],\n",
    "                                       # unlabeled\n",
    "                                       'unlabeled/epoch': best_epoch_unlabeled,\n",
    "                                       'unlabeled/NSC_1-NN_treatment': df_metric['NSC_1-NN_treatment'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/silhou_moa_tsne': df_metric['silhou_moa_tsne'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/silhou_pred_tsne': df_metric['silhou_pred_tsne'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/adj-rand_treat-pred': df_metric['adj-rand_treat-pred'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/adj-mut_treat-pred': df_metric['adj-mut_treat-pred'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/jaccard_treat-pred': df_metric['jaccard_treat-pred'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/comple_treat-pred': df_metric['comple_treat-pred'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/n_clusters': df_metric['n_clusters'][best_epoch_unlabeled],\n",
    "                                       'unlabeled/total_score': df_metric['total_score'][best_epoch_unlabeled]\n",
    "                                   })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save metrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_metric.to_csv(metrices_path, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_meta = dataset.get_df()\n",
    "df_meta.to_csv(metadata_path, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save cluster assignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assign_list = cluster_log.data['cluster_assignments']\n",
    "df_assignments = pd.concat(assign_list, axis=1)\n",
    "df_assignments = pd.concat([df_meta, df_assignments], axis=1)\n",
    "df_assignments.to_csv(assign_epoch_path, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_df = pd.DataFrame({'model_loss':cluster_log.data['loss'], \n",
    "                        'NMI':cluster_log.data['nmi']})\n",
    "loss_df.to_csv(loss_path, index=False, sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save the best embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_best_labeled_sil_embeds = cluster_log.data['features'][best_epoch_labeled_silh]\n",
    "df_best_labeled_nsc_embeds = cluster_log.data['features'][best_epoch_labeled_nsc]\n",
    "df_best_unlabeled_embeds = cluster_log.data['features'][best_epoch_unlabeled]\n",
    "\n",
    "df_best_labeled_sil_embeds.to_csv(best_labeled_sil_embed_path, sep='\\t', index=False)\n",
    "df_best_labeled_nsc_embeds.to_csv(best_labeled_nsc_embed_path, sep='\\t', index=False)\n",
    "df_best_unlabeled_embeds.to_csv(best_unlabeled_embed_path, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Best epoch evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "df_best_labeled_sil_save_well = evaluate_training(df_best_labeled_sil_embeds, embeds_cols, os.path.join(args.output_path,'best_labeled_sil'), verbose=args.verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "df_best_labeled_ncs_save_well = evaluate_training(df_best_labeled_nsc_embeds, embeds_cols, os.path.join(args.output_path,'best_labeled_nsc'), verbose=args.verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "df_best_unlabeled_save_well = evaluate_training(df_best_unlabeled_embeds, embeds_cols, os.path.join(args.output_path,'best_unlabeled'), verbose=args.verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"finished\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:deepcluster_gpu]",
   "language": "python",
   "name": "conda-env-deepcluster_gpu-py"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}