Multimodal_Integration_for_Indentifying_Stratification_of_ASD.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/felixkoerber/Multimodal-Integration-ABIDE/blob/main/Multimodal_Integration_for_Indentifying_Stratification_of_ASD.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U21C3w94iMbs"
      },
      "source": [
        "# Multimodal Integration for Indentifying Stratification of ASD\n",
        "\n",
        "Code for Bachelor Thesis of Felix Körber"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qDcvVH8rjZtZ"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-hKGC86q6N0b"
      },
      "outputs": [],
      "source": [
        "## Torch\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "import torch.nn as nn\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data import DataLoader\n",
        "import torch.utils.data\n",
        "from torch import nn,optim\n",
        "from torch.autograd import Variable\n",
        "\n",
        "## Visualization\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import cm\n",
        "import seaborn as sns\n",
        "\n",
        "# Data Handling\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "#scipy\n",
        "from scipy import spatial\n",
        "from scipy import stats\n",
        "from scipy.stats import kendalltau, pearsonr, permutation_test,bootstrap, norm\n",
        "\n",
        "# Misc\n",
        "import cv2\n",
        "import os\n",
        "import sys\n",
        "\n",
        "import argparse\n",
        "# Link Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f2DQZ-fcjxJ2"
      },
      "source": [
        "Set Up GPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YGQTYUvSjw0L"
      },
      "outputs": [],
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SUBrTu2Vk3gO"
      },
      "source": [
        "# Model Architecture"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oukdi9J9k64r"
      },
      "outputs": [],
      "source": [
        "# adapted from pytorch/examples/vae and ethanluoyc/pytorch-vae\n",
        "class FC_VAE(nn.Module):\n",
        "    \"\"\"Fully connected variational Autoencoder\"\"\"\n",
        "    def __init__(self, n_input, nz, n_hidden=1024):\n",
        "        super(FC_VAE, self).__init__()\n",
        "        self.nz = nz\n",
        "        self.n_input = n_input\n",
        "        self.n_hidden = n_hidden\n",
        "\n",
        "        self.encoder = nn.Sequential(nn.Linear(n_input, n_hidden*5),\n",
        "                                nn.ReLU(inplace=True),\n",
        "                                nn.BatchNorm1d(n_hidden*5),\n",
        "\n",
        "                                nn.Linear(n_hidden*5, n_hidden*2),\n",
        "                                nn.ReLU(inplace=True),\n",
        "                                nn.BatchNorm1d(n_hidden*2),\n",
        "\n",
        "                                nn.Linear(n_hidden*2, n_hidden),\n",
        "                                nn.ReLU(inplace=True),\n",
        "                                nn.BatchNorm1d(n_hidden),\n",
        "\n",
        "                                nn.Linear(n_hidden, n_hidden),\n",
        "                                nn.ReLU(inplace=True),\n",
        "                                nn.BatchNorm1d(n_hidden),\n",
        "\n",
        "                                nn.Linear(n_hidden, n_hidden),\n",
        "                                nn.ReLU(inplace=True),\n",
        "                                nn.BatchNorm1d(n_hidden),\n",
        "                                nn.Linear(n_hidden, n_hidden),\n",
        "                                )\n",
        "\n",
        "        self.fc1 = nn.Linear(n_hidden, nz)\n",
        "        self.fc2 = nn.Linear(n_hidden, nz)\n",
        "\n",
        "        self.decoder = nn.Sequential(nn.Linear(nz, n_hidden),\n",
        "                                     nn.ReLU(inplace=True),\n",
        "                                     nn.BatchNorm1d(n_hidden),\n",
        "\n",
        "                                     nn.Linear(n_hidden, n_hidden),\n",
        "                                     nn.ReLU(inplace=True),\n",
        "                                     nn.BatchNorm1d(n_hidden),\n",
        "\n",
        "                                     nn.Linear(n_hidden, n_hidden),\n",
        "                                     nn.BatchNorm1d(n_hidden),\n",
        "                                     nn.ReLU(inplace=True),\n",
        "\n",
        "\n",
        "                                     nn.Linear(n_hidden, n_hidden*2),\n",
        "                                     nn.ReLU(inplace=True),\n",
        "                                     nn.BatchNorm1d(n_hidden*2),\n",
        "\n",
        "                                     nn.Linear(n_hidden*2, n_hidden*5),\n",
        "                                     nn.ReLU(inplace=True),\n",
        "                                     nn.BatchNorm1d(n_hidden*5),\n",
        "                                     nn.Linear(n_hidden*5, n_input),\n",
        "                                    )\n",
        "    def forward(self, x):\n",
        "        mu, logvar = self.encode(x)\n",
        "        z = self.reparametrize(mu, logvar)\n",
        "        res = self.decode(z)\n",
        "        return res, z, mu, logvar\n",
        "\n",
        "    def encode(self, x):\n",
        "        h = self.encoder(x)\n",
        "        return self.fc1(h), self.fc2(h)\n",
        "\n",
        "    def reparametrize(self, mu, logvar):\n",
        "        std = logvar.mul(0.5).exp_()\n",
        "        if torch.cuda.is_available():\n",
        "            eps = torch.cuda.FloatTensor(std.size()).normal_()\n",
        "        else:\n",
        "            eps = torch.FloatTensor(std.size()).normal_()\n",
        "        eps = Variable(eps)\n",
        "        return eps.mul(std).add_(mu)\n",
        "\n",
        "    def decode(self, z):\n",
        "        return self.decoder(z)\n",
        "\n",
        "    def get_latent_var(self, x):\n",
        "        mu, logvar = self.encode(x)\n",
        "        z = self.reparametrize(mu, logvar)\n",
        "        return z\n",
        "\n",
        "    def generate(self, z):\n",
        "        res = self.decode(z)\n",
        "        return res\n",
        "\n",
        "\n",
        "\n",
        "class FC_Classifier(nn.Module):\n",
        "    \"\"\"Latent space discriminator\"\"\"\n",
        "    def __init__(self, nz, n_hidden=512, n_out=3):\n",
        "        super(FC_Classifier, self).__init__()\n",
        "        self.nz = nz\n",
        "        self.n_hidden = n_hidden\n",
        "        self.n_out = n_out\n",
        "\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(nz, n_hidden),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Linear(n_hidden, n_hidden),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Linear(n_hidden, n_hidden),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Linear(n_hidden, n_hidden),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Linear(n_hidden, n_hidden),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Linear(n_hidden,n_out)\n",
        "        )\n",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n",
        "class Simple_Classifier(nn.Module):\n",
        "    \"\"\"Latent space discriminator\"\"\"\n",
        "    def __init__(self, nz, n_out=3):\n",
        "        super(Simple_Classifier, self).__init__()\n",
        "        self.nz = nz\n",
        "        self.n_out = n_out\n",
        "\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(nz, n_out),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kf-WLMrloGpN"
      },
      "source": [
        "# Arguments"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_KYM02M8oJUf"
      },
      "outputs": [],
      "source": [
        "# parse arguments\n",
        "class setup_args():\n",
        "\n",
        "    batch_size=32\n",
        "    max_epochs=300\n",
        "    nz=512\n",
        "    lamb=0.0000001 #beta weight\n",
        "    alpha =.1\n",
        "    dist_factor =1\n",
        "    learning_rate_D = 1e-4\n",
        "    learning_rate_AE =1e-4\n",
        "    weight_decay=0\n",
        "    n_hidden=512\n",
        "    save_freq=10\n",
        "\n",
        "args = setup_args()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKIehI7IdGua"
      },
      "source": [
        "# Set-Up Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VL9h98WZh5PU"
      },
      "outputs": [],
      "source": [
        "### Dataloader\n",
        "class Combined_Dataset_Train():\n",
        "    def __init__(self):\n",
        "        isfunc=np.load(\"/content/drive/My Drive/BA/isfunc.npy\", allow_pickle=True)\n",
        "        #Pheno Data\n",
        "        pheno_data=pd.read_csv('/content/drive/My Drive/BA/pheno_data.csv', index_col=0)\n",
        "        pheno_data.iloc[isfunc==1]\n",
        "\n",
        "        #Functional Data\n",
        "        func_data= np.load(\"/content/drive/My Drive/BA/func_flat_Tal_new.npy\", allow_pickle=True)\n",
        "        func_data[np.where(np.isnan(func_data))]=0\n",
        "        func_data= (func_data - np.mean(func_data)) / np.std(func_data)\n",
        "\n",
        "        #Area Data\n",
        "        area_data=np.load('/content/drive/My Drive/BA/area_data_red.npy')[:,:,0].T\n",
        "        area_data=(area_data - np.mean(area_data)) / np.std(area_data)\n",
        "        area_data=area_data[isfunc==1]\n",
        "        #Cortical Thickness Data\n",
        "        thick_data=np.load('/content/drive/My Drive/BA/thick_data_red.npy')[:,:,0].T\n",
        "        thick_data=(thick_data - np.mean(thick_data)) / np.std(thick_data)\n",
        "        thick_data=thick_data[isfunc==1]\n",
        "\n",
        "        self.label = pheno_data.iloc             [isfunc==1]\n",
        "        self.func_data = func_data              [:]\n",
        "        self.area_data =area_data               [:]\n",
        "        self.thick_data =thick_data             [:]\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.label)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        return torch.tensor(self.func_data[idx,:]).float(),torch.tensor(self.area_data[idx,:]).float(),torch.tensor(self.thick_data[idx,:]).float(),torch.tensor(self.label.iloc[idx,6]-1)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o-mpYm7wdJ9u"
      },
      "source": [
        "# Initiate Dataloader and Models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oOIqRL_75wSw"
      },
      "outputs": [],
      "source": [
        "# retrieve dataloader\n",
        "\n",
        "# Define the train and test dataloaders\n",
        "dataset_Train = (Combined_Dataset_Train())\n",
        "train_dataloader = DataLoader(dataset_Train, batch_size=args.batch_size, drop_last=False, shuffle=True)\n",
        "dataset_Test = (Combined_Dataset_Train())\n",
        "test_dataloader = DataLoader(dataset_Test, batch_size=args.batch_size,drop_last=False, shuffle=False)\n",
        "\n",
        "print('Data loaded')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nXfeBhaRjlw-"
      },
      "outputs": [],
      "source": [
        "#============= TRAINING INITIALIZATION ==============\n",
        "\n",
        "# initialize autoencoder\n",
        "model_func = FC_VAE  (n_input=4656,    nz=args.nz,n_hidden=args.n_hidden).to(device)\n",
        "model_thick = FC_VAE (n_input=163842,  nz=args.nz,n_hidden=args.n_hidden).to(device)\n",
        "model_area = FC_VAE  (n_input=163842,  nz=args.nz,n_hidden=args.n_hidden).to(device)\n",
        "netClf = FC_Classifier(nz=args.nz).to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "63faGTdUvnzn"
      },
      "outputs": [],
      "source": [
        "# setup optimizer\n",
        "opt_netthick  = torch.optim.Adam(params=model_thick.parameters (), lr=args.learning_rate_AE)\n",
        "opt_netarea   = torch.optim.Adam(params=model_area.parameters  (), lr=args.learning_rate_AE)\n",
        "opt_netfunc   = torch.optim.Adam(params=model_func.parameters  (), lr=args.learning_rate_AE)\n",
        "opt_netClf    = torch.optim.Adam(params=netClf.parameters      (), lr=args.learning_rate_D)\n",
        "\n",
        "# loss criteria\n",
        "criterion_reconstruct = nn.MSELoss()\n",
        "criterion_classify    = nn.CrossEntropyLoss()\n",
        "criterion_latent_dis  = nn.L1Loss()\n",
        "\n",
        "criterion_reconstruct = criterion_reconstruct.to(device)\n",
        "criterion_classify    = criterion_classify.to(device)\n",
        "criterion_latent_dis    = criterion_latent_dis.to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BT1xQUfVdDCx"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I35amxs56Qv3"
      },
      "outputs": [],
      "source": [
        "def compute_KL_loss(mu, logvar):\n",
        "    if args.lamb>0:\n",
        "        KLloss = -0.5*torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n",
        "        return args.lamb * KLloss\n",
        "    return 0\n",
        "\n",
        "def train_autoencoders(func_inputs, area_inputs, thick_inputs,labels):\n",
        "    model_thick.train()\n",
        "    model_area.train()\n",
        "    model_func.train()\n",
        "    netClf.eval()\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "      thick_inputs, area_inputs,func_inputs = thick_inputs.to(device), area_inputs.to(device),func_inputs.to(device)\n",
        "      labels=labels.to(device)\n",
        "\n",
        "    # forward pass\n",
        "    thick_recon, thick_latents,thick_mu, thick_logvar = model_thick(thick_inputs)\n",
        "    area_recon, area_latents, area_mu, area_logvar = model_area(area_inputs)\n",
        "    func_recon, func_latents,func_mu, func_logvar = model_func(func_inputs)\n",
        "\n",
        "    thick_to_area=model_area.decode(thick_latents)\n",
        "    thick_to_func=model_func.decode(thick_latents)\n",
        "\n",
        "    area_to_thick=model_thick.decode(area_latents)\n",
        "    area_to_func=model_func.decode(area_latents)\n",
        "\n",
        "    func_to_thick=model_thick.decode(func_latents)\n",
        "    func_to_area=model_area.decode(area_latents)\n",
        "\n",
        "    thick_scores = netClf(thick_latents)\n",
        "    area_scores = netClf(area_latents)\n",
        "    func_scores = netClf(func_latents)\n",
        "\n",
        "\n",
        "    thick_labels = torch.zeros  (thick_scores.size(0),).long()\n",
        "    area_labels  = torch.zeros(area_scores.size (0),).long()\n",
        "    func_labels  = torch.zeros (func_scores.size(0),).long()\n",
        "\n",
        "    thick_labels[:]= 0\n",
        "    area_labels [:]= 1\n",
        "    func_labels [:]= 2\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "        thick_labels, area_labels,func_labels = thick_labels.to(device), area_labels.to(device),func_labels.to(device)\n",
        "\n",
        "\n",
        "    # compute recon losses\n",
        "    thick_recon_loss = criterion_reconstruct(thick_inputs, thick_recon)\n",
        "    area_recon_loss = criterion_reconstruct(area_inputs, area_recon)\n",
        "    func_recon_loss = criterion_reconstruct(func_inputs, func_recon)\n",
        "\n",
        "\n",
        "    distances_metric =  (1/3)*criterion_latent_dis(func_latents,thick_latents) +   (1/3)*criterion_latent_dis(func_latents,area_latents)+(1/3)*criterion_latent_dis(thick_latents,area_latents)\n",
        "    distances_metric = args.dist_factor*distances_metric\n",
        "\n",
        "    # compute cross modal latent recon losses\n",
        "    cmrl_ttoa= criterion_reconstruct(area_inputs, thick_to_area)\n",
        "    cmrl_ttof= criterion_reconstruct(func_inputs, thick_to_func)\n",
        "    cmrl_atot= criterion_reconstruct(thick_inputs, area_to_thick)\n",
        "    cmrl_atof= criterion_reconstruct(func_inputs, area_to_func)\n",
        "    cmrl_ftot= criterion_reconstruct(thick_inputs, func_to_thick)\n",
        "    cmrl_ftoa= criterion_reconstruct(area_inputs, func_to_area)\n",
        "\n",
        "    cmrl = (cmrl_ttoa+cmrl_ttof+cmrl_atot+cmrl_atof+cmrl_ftot+cmrl_ftoa)/6\n",
        "\n",
        "    kl_loss = compute_KL_loss(thick_mu, thick_logvar) + compute_KL_loss(area_mu, area_logvar)+compute_KL_loss(func_mu,func_logvar)\n",
        "    clf_loss =  (1/6) * criterion_classify(thick_scores, area_labels) +  \\\n",
        "      (1/6) * criterion_classify(thick_scores, func_labels) +(1/6) * criterion_classify(area_scores, func_labels) + \\\n",
        "      (1/6)* criterion_classify(area_scores, thick_labels) +  (1/6) * criterion_classify(func_scores, area_labels) +  (1/6) * criterion_classify(func_scores, thick_labels)\n",
        "    loss = thick_recon_loss + area_recon_loss + func_recon_loss+ kl_loss   +  clf_loss +distances_metric + cmrl\n",
        "\n",
        "    # reset parameter gradients\n",
        "    opt_netthick.zero_grad()\n",
        "    opt_netarea.zero_grad()\n",
        "    opt_netfunc.zero_grad()\n",
        "\n",
        "    # backpropagate and update model\n",
        "    loss.backward()\n",
        "\n",
        "    opt_netthick.step()\n",
        "    opt_netarea.step()\n",
        "    opt_netfunc.step()\n",
        "\n",
        "    #Validation Loss\n",
        "    summary_stats = {'loss': loss,'thick_recon_loss': thick_recon_loss*thick_scores.size(0), 'area_recon_loss': area_recon_loss*area_scores.size(0),\n",
        "                     'func_recon_loss': func_recon_loss*func_scores.size(0),'clf_loss': clf_loss*(thick_scores.size(0)+area_scores.size(0)+func_scores.size(0)),\n",
        "                      'cmr_loss':cmrl*(thick_scores.size(0)+area_scores.size(0)+func_scores.size(0)),\n",
        "                      'KL_loss':kl_loss,'distances_metric':distances_metric*(thick_scores.size(0)+area_scores.size(0)+func_scores.size(0))\n",
        "                      }\n",
        "\n",
        "    return summary_stats\n",
        "\n",
        "def train_classifier(func_inputs,area_inputs,thick_inputs):\n",
        "\n",
        "    model_thick.eval()\n",
        "    model_area.eval()\n",
        "    model_func.eval()\n",
        "    netClf.train()\n",
        "\n",
        "    # process input data\n",
        "    if torch.cuda.is_available():\n",
        "      thick_inputs, area_inputs,func_inputs = thick_inputs.to(device), area_inputs.to(device),func_inputs.to(device)\n",
        "\n",
        "\n",
        "    # forward pass\n",
        "    _, thick_latents, _, _  = model_thick(thick_inputs)\n",
        "    _, area_latents, _, _   = model_area(area_inputs)\n",
        "    _, func_latents, _, _   = model_func(func_inputs)\n",
        "\n",
        "    thick_scores  = netClf(thick_latents)\n",
        "    area_scores   = netClf(area_latents)\n",
        "    func_scores   = netClf(func_latents)\n",
        "\n",
        "    thick_labels = torch.zeros  (thick_scores.size(0),).long()\n",
        "    area_labels  = torch.zeros(area_scores.size (0),).long()\n",
        "    func_labels  = torch.zeros (func_scores.size(0),).long()\n",
        "\n",
        "    thick_labels[:]= 0\n",
        "    area_labels [:]= 1\n",
        "    func_labels [:]= 2\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "        thick_labels, area_labels,func_labels = thick_labels.to(device), area_labels.to(device),func_labels.to(device)\n",
        "\n",
        "    clf_loss = (1/3)*criterion_classify(thick_scores, thick_labels) + (1/3)* criterion_classify(area_scores, area_labels) + (1/3)* criterion_classify(func_scores, func_labels)\n",
        "    loss = clf_loss\n",
        "\n",
        "\n",
        "    # backpropagate and update model\n",
        "    opt_netClf.zero_grad()\n",
        "    loss.backward()\n",
        "    opt_netClf.step()\n",
        "    summary_stats = {'clf_loss': clf_loss*(thick_scores.size(0)+area_scores.size(0)+area_scores.size(0)),\n",
        "                     'thick_n_samples': thick_scores.size(0),'area_n_samples': area_scores.size(0),\n",
        "                     'func_n_samples':func_scores.size(0)}\n",
        "\n",
        "    return summary_stats\n",
        "\n",
        "def mean_loss(input):\n",
        "  mean=criterion_reconstruct(input,torch.mean(input))\n",
        "  return mean"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fxACEt-DzJFC"
      },
      "outputs": [],
      "source": [
        "# Mean Loss\n",
        "mean_func        = 0\n",
        "mean_thick       = 0\n",
        "mean_area        = 0\n",
        "for idx, (sample) in enumerate(train_dataloader):\n",
        "\n",
        "  func_inputs   = sample[0]\n",
        "  area_inputs   = sample[1]\n",
        "  thick_inputs  = sample[2]\n",
        "  labels        = sample[3]\n",
        "  #print(func_inputs.shape,area_inputs.shape,thick_inputs.shape,labels.shape)\n",
        "  mean_func     +=mean_loss(func_inputs)\n",
        "  mean_area     +=mean_loss(area_inputs)\n",
        "  mean_thick    +=mean_loss(thick_inputs)\n",
        "mean_func /= 883\n",
        "mean_area /= 883\n",
        "mean_thick/= 883\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training Loop"
      ],
      "metadata": {
        "id": "ZJfoD-VGlxYZ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MaRdo6VCH9e2"
      },
      "outputs": [],
      "source": [
        "\n",
        "print('Training Model')\n",
        "combined_loss=torch.empty(size=(0,11))\n",
        "\n",
        "for epoch in range(args.max_epochs):\n",
        "    #print(epoch)\n",
        "\n",
        "    recon_thick_loss = 0\n",
        "    recon_area_loss  = 0\n",
        "    recon_func_loss  = 0\n",
        "    KL_loss          = 0\n",
        "    clf_loss         = 0\n",
        "    AE_clf_loss      = 0\n",
        "    cmr_loss         = 0\n",
        "    distances_metric = 0\n",
        "    n_thick_total    = 0\n",
        "    n_area_total     = 0\n",
        "    n_func_total     = 0\n",
        "    total_loss       = 0\n",
        "\n",
        "    for idx, (sample) in enumerate(train_dataloader):\n",
        "        func_inputs   = sample[0]\n",
        "        area_inputs   = sample[1]\n",
        "        thick_inputs  = sample[2]\n",
        "        labels        = sample[3]\n",
        "\n",
        "        out = train_autoencoders(func_inputs, area_inputs, thick_inputs,labels)\n",
        "\n",
        "        recon_thick_loss += out['thick_recon_loss']\n",
        "        recon_area_loss += out['area_recon_loss']\n",
        "        recon_func_loss += out['func_recon_loss']\n",
        "        AE_clf_loss += out['clf_loss']\n",
        "        cmr_loss    += out['cmr_loss']\n",
        "        KL_loss     += out['KL_loss']\n",
        "        distances_metric  += out['distances_metric']\n",
        "        total_loss  += out['loss']\n",
        "        out = train_classifier(func_inputs, area_inputs,thick_inputs)\n",
        "\n",
        "        clf_loss      += out['clf_loss']\n",
        "        n_thick_total += out['thick_n_samples']\n",
        "        n_area_total  += out['area_n_samples']\n",
        "        n_func_total += out['func_n_samples']\n",
        "\n",
        "\n",
        "\n",
        "    recon_thick_loss  /= n_thick_total\n",
        "    recon_area_loss   /= n_area_total\n",
        "    recon_func_loss   /= n_func_total\n",
        "    cmr_loss          /= (n_thick_total+n_area_total+n_func_total)\n",
        "    distances_metric  /= ((n_thick_total+n_area_total+n_func_total)*args.dist_factor)\n",
        "    KL_loss           /= (n_thick_total+n_area_total+n_func_total)\n",
        "    clf_loss /= (n_thick_total+n_area_total+n_func_total)\n",
        "    AE_clf_loss /= (n_thick_total+n_area_total+n_func_total)\n",
        "\n",
        "    if epoch==0:\n",
        "      combined_loss=torch.tensor([recon_thick_loss,mean_thick,recon_area_loss,mean_area,recon_func_loss,mean_func,KL_loss,AE_clf_loss,cmr_loss,AE_clf_loss,clf_loss,distances_metric,total_loss])[None,:]\n",
        "      print(combined_loss.size())\n",
        "    else:\n",
        "      combined_loss=torch.cat((combined_loss,torch.tensor([recon_thick_loss,mean_thick,recon_area_loss,mean_area,recon_func_loss,mean_func,KL_loss,AE_clf_loss,cmr_loss,AE_clf_loss,clf_loss,distances_metric,total_loss])[None,:]))\n",
        "      if epoch==1:\n",
        "        print(combined_loss.size())\n",
        "    print('Epoch: ', epoch ,'total loss: ',total_loss,  'distance metric: %.8f' % float(distances_metric), ', thick recon loss: %.8f' % float(recon_thick_loss), ', area recon loss: %.8f' % float(recon_area_loss), ', func recon loss: %.8f' % float(recon_func_loss),\n",
        "                ', KL Div. loss: %.8f' %float (KL_loss), ', Cross-Modal loss: %.8f' % float(cmr_loss), ', AE clf loss: %.8f' % float(AE_clf_loss), ', clf loss: %.8f' % float(clf_loss))\n",
        "    print('Meanthick: ',mean_thick,'Mean area: ',mean_area,'Mean func: ',mean_func)\n",
        "for idx in range(len(combined_loss)):\n",
        "  if idx<7:\n",
        "    plt.plot(combined_loss[:,idx])\n",
        "    #print(combined_loss[:,idx])\n",
        "\n",
        "'''torch.save(model_thick.cpu().state_dict(),os.path.join(\"/content/drive/My Drive/BA/weights/\",  \"big_model_thick_complicated_loss_%s.pth\" % epoch))\n",
        "torch.save  (model_area.cpu().state_dict(),os.path.join(\"/content/drive/My Drive/BA/weights/\", \"big_model_area_complicated_loss_%s.pth\" % epoch))\n",
        "torch.save  (model_func.cpu().state_dict(), os.path.join(\"/content/drive/My Drive/BA/weights/\",\"big_model_func_complicated_loss_%s.pth\" % epoch))\n",
        "torch.save  (netClf.cpu().state_dict(),os.path.join(\"/content/drive/My Drive/BA/weights/\",     \"big_netClf_complicated_loss_%s.pth\" % epoch))'''\n",
        "for idx in range(len(combined_loss)):\n",
        "  if idx<12:\n",
        "    plt.plot(combined_loss[:,idx])\n",
        "    #print(combined_loss[:,idx])\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R5I162UuN_mu"
      },
      "outputs": [],
      "source": [
        "model_thick.to(device)\n",
        "model_area.to(device)\n",
        "model_func.to(device)\n",
        "netClf.to(device)\n",
        "for idx, (sample) in enumerate(test_dataloader):\n",
        "    func_inputs   = sample[0]\n",
        "    area_inputs   = sample[1]\n",
        "    thick_inputs  = sample[2]\n",
        "    labels        = sample[3]\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "        thick_inputs, area_inputs,func_inputs = thick_inputs.cuda(), area_inputs.cuda(), func_inputs.cuda()\n",
        "\n",
        "    # forward pass\n",
        "    _, thick_latents, _, _ = model_thick(thick_inputs)\n",
        "    _, area_latents, _, _ = model_area(area_inputs)\n",
        "    _, func_latents, _, _ = model_func(func_inputs)\n",
        "\n",
        "    thick_scores = netClf(thick_latents)\n",
        "    area_scores = netClf(area_latents)\n",
        "\n",
        "    func_scores = netClf(func_latents)\n",
        "    if idx ==0:\n",
        "      thick_latent_ful,area_latent_ful,func_latent_ful,labels_ful=thick_latents,area_latents,func_latents,labels\n",
        "\n",
        "    else:\n",
        "      thick_latent_ful=torch.cat((thick_latent_ful,thick_latents))\n",
        "      area_latent_ful=torch.cat((area_latent_ful,area_latents))\n",
        "      func_latent_ful=torch.cat((func_latent_ful,func_latents))\n",
        "      labels_ful=torch.cat((labels_ful,labels))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AMitHYA7PVe4"
      },
      "outputs": [],
      "source": [
        "def nearest_neighbour(points_a, points_b):\n",
        "    tree = spatial.KDTree(points_b)\n",
        "    return tree.query(points_a,k=50)[1]\n",
        "n=0\n",
        "neighbors = np.zeros((6,func_latent_ful.size(0),50))\n",
        "for i, lat_vec1 in enumerate([area_latent_ful.detach().numpy(), thick_latent_ful.detach().numpy(), func_latent_ful.detach().numpy()]):\n",
        "  for j, lat_vec2 in enumerate([area_latent_ful.detach().numpy(), thick_latent_ful.detach().numpy(), func_latent_ful.detach().numpy()]):\n",
        "    if i!=j:\n",
        "      neighbors[n,:]=nearest_neighbour(lat_vec1,lat_vec2)\n",
        "      n+=1\n",
        "\n",
        "neighbors_acc = np.zeros((7,50))\n",
        "\n",
        "for m in range(6):\n",
        "  for k in range(50):\n",
        "    acc=0\n",
        "    for i in range(func_latent_ful.size(0)):\n",
        "      if i in neighbors[m,i,:k]:\n",
        "        acc+=1\n",
        "    if k==0:\n",
        "      neighbors_acc[m,k]=0\n",
        "    else:\n",
        "      neighbors_acc[m,k]=acc/func_latent_ful.size(0)\n",
        "for k in range(50):\n",
        "  if k==0:\n",
        "    neighbors_acc[6,k]=0\n",
        "  else:\n",
        "    neighbors_acc[6,k]=1/(func_latent_ful.size(0)/k)\n",
        "\n",
        "neighbors_acc = pd.DataFrame(neighbors_acc.T, columns = ['Goal: FC, Neighbors: CSA',\n",
        "                                              'Goal: FC, Neighbors: CT',\n",
        "                                              'Goal: CSA, Neighbors: FC',\n",
        "                                              'Goal: CSA, Neighbors: CT',\n",
        "                                              'Goal: CT, Neighbors: FC',\n",
        "                                              'Goal: CT, Neighbors: CSA',\n",
        "                                              'Random Accuracy'])"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Plotting K-NN"
      ],
      "metadata": {
        "id": "1f-pheAtmJsa"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jH2TbeGguqXf"
      },
      "outputs": [],
      "source": [
        "\n",
        "fig, ax = plt.subplots(1,1,figsize=(10,4), dpi=300, constrained_layout=True)\n",
        "\n",
        "ax1 = sns.lineplot(data=neighbors_acc,ax=ax[0],legend=None)\n",
        "ax1.set(ylim=(0, 1))\n",
        "box = ax1.get_position()\n",
        "ax1.set_position([box.x0, box.y0 + box.height * 0.1,\n",
        "                 box.width, box.height * 0.9])\n",
        "ax1.legend()\n",
        "\n",
        "ax1.set(xlabel=\"Number of k-nearest neighbors\", ylabel=\"k-nearest neighbors accuracy\")\n",
        "ax1.set_title(\"Our Method\")\n",
        "#plt.suptitle(\"k-nearest Neighbor Accuracy of different Modality Pairs\")\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "SUBrTu2Vk3gO",
        "kf-WLMrloGpN",
        "o-mpYm7wdJ9u"
      ],
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}