02.test_models_by_folds.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e10590fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: TORCH_HOME=/ayb/vol2/home/dumerenkov/torch_hub\n"
     ]
    }
   ],
   "source": [
    "%env TORCH_HOME=/ayb/vol2/home/dumerenkov/torch_hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "14e770d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import Bio\n",
    "from Bio import SeqIO\n",
    "import os\n",
    "import torch\n",
    "import math\n",
    "import esm\n",
    "\n",
    "from torch import nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from sklearn.model_selection import train_test_split\n",
    "import scipy\n",
    "from scipy import stats\n",
    "\n",
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6de1bfa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "81437c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "HIDDEN_UNITS_POS_CONTACT = 5\n",
    "class ESMForSingleMutationPosConcat(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm2, _ = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.fc1 = nn.Linear(1280 * 2, HIDDEN_UNITS_POS_CONTACT)\n",
    "        self.fc2 = nn.Linear(HIDDEN_UNITS_POS_CONTACT, 1)\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm2.forward(token_ids1, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs2 = self.esm2.forward(token_ids2, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs1_pos = outputs1[:, pos + 1]\n",
    "        outputs2_pos = outputs2[:, pos + 1]\n",
    "        outputs_pos_concat = torch.cat((outputs1_pos, outputs2_pos), 2)\n",
    "        fc1_outputs = F.relu(self.fc1(outputs_pos_concat))\n",
    "        logits = self.fc2(fc1_outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8f39f27b",
   "metadata": {},
   "outputs": [],
   "source": [
    "HIDDEN_UNITS_POS_OUTER = 5\n",
    "class ESMForSingleMutationPosOuter(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm2, _ = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self._freeze_esm2_layers()\n",
    "        self.fc1 = nn.Linear(1280 * 1280, HIDDEN_UNITS_POS_OUTER)\n",
    "        self.fc2 = nn.Linear(HIDDEN_UNITS_POS_OUTER, 1)\n",
    "\n",
    "    def _freeze_esm2_layers(self):\n",
    "        total_blocks = 33\n",
    "        initial_layers = 2\n",
    "        layers_per_block = 16\n",
    "        num_freeze_blocks = total_blocks - 3\n",
    "        for _, param in list(self.esm2.named_parameters())[\n",
    "            :initial_layers + layers_per_block * num_freeze_blocks]:\n",
    "            param.requires_grad = False\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm2.forward(token_ids1, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs2 = self.esm2.forward(token_ids2, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs1_pos = outputs1[:, pos + 1]\n",
    "        outputs2_pos = outputs2[:, pos + 1]\n",
    "        outer_prod = outputs1_pos.unsqueeze(3) @ outputs2_pos.unsqueeze(2)\n",
    "        outer_prod_view = outer_prod.view(outer_prod.shape[0], outer_prod.shape[1], -1)\n",
    "        fc1_outputs = F.relu(self.fc1(outer_prod_view))\n",
    "        logits = self.fc2(fc1_outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9aede66e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ESMForSingleMutation_pos(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__() \n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()        \n",
    "        self.classifier = nn.Linear(1280, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1,1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1,1280)))\n",
    "        \n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):                \n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])['representations'][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])['representations'][33]\n",
    "        outputs = self.const1 * outputs1[:,pos + 1,:] + self.const2 * outputs2[:,pos + 1,:]        \n",
    "        logits = self.classifier(outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "838b0855",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ESMForSingleMutation_cls(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__() \n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()        \n",
    "        self.classifier = nn.Linear(1280, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1,1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1,1280)))\n",
    "        \n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):                \n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])['representations'][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])['representations'][33]\n",
    "        outputs = self.const1 * outputs1[:,0,:] + self.const2 * outputs2[:,0,:]        \n",
    "        logits = self.classifier(outputs.unsqueeze(0))\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1df241db",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ESMForSingleMutation_pos_cat_cls(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__() \n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()        \n",
    "        self.classifier = nn.Linear(1280*2, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1,1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1,1280)))\n",
    "        \n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):                \n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])['representations'][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])['representations'][33]\n",
    "        cls_out = self.const1 * outputs1[:,0,:] + self.const2 * outputs2[:,0,:]\n",
    "        pos_out = self.const1 * outputs1[:,pos+1,:] + self.const2 * outputs2[:,pos+1,:]\n",
    "        outputs = torch.cat([cls_out.unsqueeze(0), pos_out], axis = -1)        \n",
    "        logits = self.classifier(outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1a1cb550",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProteinDataset(Dataset):\n",
    "    def __init__(self, df):\n",
    "        self.df = df\n",
    "        _, esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.esm1v_batch_converter = esm1v_alphabet.get_batch_converter()\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        _, _, esm1b_batch_tokens1 = self.esm1v_batch_converter([('' , ''.join(self.df.iloc[idx]['wt_seq'])[:1022])])\n",
    "        _, _, esm1b_batch_tokens2 = self.esm1v_batch_converter([('' , ''.join(self.df.iloc[idx]['mut_seq'])[:1022])])\n",
    "        pos = self.df.iloc[idx]['pos']\n",
    "        return esm1b_batch_tokens1, esm1b_batch_tokens2, pos, torch.unsqueeze(torch.FloatTensor([self.df.iloc[idx]['ddg']]), 0)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "642d008e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    tr_loss, tr_accuracy = 0, 0\n",
    "    nb_tr_examples, nb_tr_steps = 0, 0\n",
    "    tr_preds, tr_labels = [], []\n",
    "    model.train()\n",
    "    \n",
    "    for idx, batch in enumerate(training_loader):\n",
    "        input_ids1, input_ids2, pos, labels = batch            \n",
    "        input_ids1 = input_ids1[0].to(device)\n",
    "        input_ids2 = input_ids2[0].to(device)\n",
    "        logits = model(token_ids1 = input_ids1, token_ids2 = input_ids2, pos = pos).to('cpu')        \n",
    "        loss = torch.nn.functional.mse_loss(logits, labels)\n",
    "        tr_loss += loss.item()\n",
    "\n",
    "        nb_tr_steps += 1\n",
    "        nb_tr_examples += labels.size(0)\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(\n",
    "            parameters=model.parameters(), max_norm=0.1\n",
    "        )\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "    epoch_loss = tr_loss / nb_tr_steps\n",
    "    print(f\"Training loss epoch: {epoch_loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "929405ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def valid(model, testing_loader):\n",
    "    model.eval()\n",
    "    \n",
    "    eval_loss, eval_accuracy = 0, 0\n",
    "    nb_eval_examples, nb_eval_steps = 0, 0\n",
    "    eval_preds, eval_labels, eval_scores = [], [], []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for idx, batch in enumerate(testing_loader):\n",
    "            \n",
    "            input_ids1, input_ids2, pos, labels = batch            \n",
    "            input_ids1 = input_ids1[0].to(device)\n",
    "            input_ids2 = input_ids2[0].to(device)\n",
    "            labels = labels.to(device)\n",
    "            logits = model(token_ids1 = input_ids1, token_ids2 = input_ids2, pos = pos)        \n",
    "            loss = torch.nn.functional.mse_loss(logits, labels)\n",
    "            eval_loss += loss.item()\n",
    "\n",
    "            nb_eval_steps += 1\n",
    "            nb_eval_examples += labels.size(0)\n",
    "                     \n",
    "            eval_labels.extend(labels.cpu().detach())\n",
    "            eval_preds.extend(logits.cpu().detach())\n",
    "            \n",
    "  \n",
    "    labels = [id.item() for id in eval_labels]\n",
    "    predictions = [id.item() for id in eval_preds]\n",
    "    \n",
    "    eval_loss = eval_loss / nb_eval_steps\n",
    "    print(f\"Validation Loss: {eval_loss}\")\n",
    "\n",
    "    return labels, predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f9dfa131",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training model ESMForSingleMutationPosOuter on fold 0\n",
      "Training loss epoch: 2.635423413551298\n",
      "Training loss epoch: 1.1322193411048722\n",
      "Training loss epoch: 0.6197464244398997\n",
      "Validation Loss: 3.172407191069145\n",
      "MAE 1.1664041202748194 Correlation SpearmanrResult(correlation=0.6642448778966968, pvalue=1.8245666909273552e-268)\n",
      "Training model ESMForSingleMutationPosConcat on fold 0\n",
      "Training loss epoch: 3.6268955301182415\n",
      "Training loss epoch: 1.5482527004603415\n",
      "Training loss epoch: 0.6599749537010731\n",
      "Validation Loss: 3.353341689513026\n",
      "MAE 1.2339659470389288 Correlation SpearmanrResult(correlation=0.613092901597194, pvalue=7.370066835129871e-218)\n",
      "Training model ESMForSingleMutation_pos_cat_cls on fold 0\n",
      "Training loss epoch: 2.5174619745669298\n",
      "Training loss epoch: 0.7621416057688709\n",
      "Training loss epoch: 0.19051986832222012\n",
      "Validation Loss: 3.228121882597974\n",
      "MAE 1.1842531303281363 Correlation SpearmanrResult(correlation=0.6476683091905255, pvalue=5.402042441777921e-251)\n",
      "Training model ESMForSingleMutation_pos on fold 0\n",
      "Training loss epoch: 2.4769797154396995\n",
      "Training loss epoch: 0.7541997423927302\n",
      "Training loss epoch: 0.24432987727469754\n",
      "Validation Loss: 3.283616643486261\n",
      "MAE 1.1852712030255235 Correlation SpearmanrResult(correlation=0.6450460732501254, pvalue=2.493925679166963e-248)\n",
      "Training model ESMForSingleMutation_cls on fold 0\n",
      "Training loss epoch: 3.440278009140623\n",
      "Training loss epoch: 1.9350419008513802\n",
      "Training loss epoch: 0.891156534245437\n",
      "Validation Loss: 4.298389999941283\n",
      "MAE 1.3850962093525852 Correlation SpearmanrResult(correlation=0.5870341963435821, pvalue=1.761977104505598e-195)\n",
      "\n",
      "Training model ESMForSingleMutationPosOuter on fold 1\n",
      "Training loss epoch: 3.0403295606570127\n",
      "Training loss epoch: 1.3113204625467225\n",
      "Training loss epoch: 0.7027697870940723\n",
      "Validation Loss: 2.11417320095372\n",
      "MAE 0.9713469774884648 Correlation SpearmanrResult(correlation=0.7433860968014595, pvalue=0.0)\n",
      "Training model ESMForSingleMutationPosConcat on fold 1\n",
      "Training loss epoch: 3.918982719163384\n",
      "Training loss epoch: 1.5452872551874821\n",
      "Training loss epoch: 0.6941021446107163\n",
      "Validation Loss: 2.5525579193364023\n",
      "MAE 1.0670988453813675 Correlation SpearmanrResult(correlation=0.6356055083485904, pvalue=5.96081453213227e-239)\n",
      "Training model ESMForSingleMutation_pos_cat_cls on fold 1\n",
      "Training loss epoch: 2.823819295836094\n",
      "Training loss epoch: 0.8391675335688937\n",
      "Training loss epoch: 0.3157066082720597\n",
      "Validation Loss: 2.3566652883957215\n",
      "MAE 0.9944679572784202 Correlation SpearmanrResult(correlation=0.6854149933648048, pvalue=1.7555969678793751e-292)\n",
      "Training model ESMForSingleMutation_pos on fold 1\n",
      "Training loss epoch: 2.740738289692721\n",
      "Training loss epoch: 0.8742920969787655\n",
      "Training loss epoch: 0.2930991777171697\n",
      "Validation Loss: 2.313170445259252\n",
      "MAE 0.9863319716158286 Correlation SpearmanrResult(correlation=0.697960501121269, pvalue=1.039863469907121e-307)\n",
      "Training model ESMForSingleMutation_cls on fold 1\n",
      "Training loss epoch: 4.0321554556106225\n",
      "Training loss epoch: 2.4944623875761467\n",
      "Training loss epoch: 1.1323999142675756\n",
      "Validation Loss: 2.4943226574156925\n",
      "MAE 1.0671829977010923 Correlation SpearmanrResult(correlation=0.6565721324319559, pvalue=3.061411473796465e-260)\n",
      "\n",
      "Training model ESMForSingleMutationPosOuter on fold 2\n",
      "Training loss epoch: 2.872127168906501\n",
      "Training loss epoch: 1.184510158625728\n",
      "Training loss epoch: 0.6920238167326511\n",
      "Validation Loss: 2.088054274346603\n",
      "MAE 0.9740557395532946 Correlation SpearmanrResult(correlation=0.6571908715065021, pvalue=5.113725120945134e-261)\n",
      "Training model ESMForSingleMutationPosConcat on fold 2\n",
      "Training loss epoch: 4.084319781591863\n",
      "Training loss epoch: 1.7189699855529021\n",
      "Training loss epoch: 0.8124657553161575\n",
      "Validation Loss: 2.2387282528784125\n",
      "MAE 0.9937048408067257 Correlation SpearmanrResult(correlation=0.6226926161651507, pvalue=9.503006019328431e-227)\n",
      "Training model ESMForSingleMutation_pos_cat_cls on fold 2\n",
      "Training loss epoch: 2.944175185958653\n",
      "Training loss epoch: 0.8916680296439314\n",
      "Training loss epoch: 0.36897549374894684\n",
      "Validation Loss: 2.2467366992300946\n",
      "MAE 0.9906155266964078 Correlation SpearmanrResult(correlation=0.6403251979253304, pvalue=1.0310219467050842e-243)\n",
      "Training model ESMForSingleMutation_pos on fold 2\n",
      "Training loss epoch: 2.8528178783323423\n",
      "Training loss epoch: 0.8354474142909082\n",
      "Training loss epoch: 0.26442301521546197\n",
      "Validation Loss: 2.226006064438612\n",
      "MAE 1.0002656335416509 Correlation SpearmanrResult(correlation=0.6286339811399014, pvalue=2.4208632503455274e-232)\n",
      "Training model ESMForSingleMutation_cls on fold 2\n",
      "Training loss epoch: 4.06304127029205\n",
      "Training loss epoch: 2.721917298385394\n",
      "Training loss epoch: 1.2433942237871731\n",
      "Validation Loss: 2.220547264601859\n",
      "MAE 1.01369134363581 Correlation SpearmanrResult(correlation=0.645755667110443, pvalue=3.641024277188849e-249)\n",
      "\n",
      "Training model ESMForSingleMutationPosOuter on fold 3\n",
      "Training loss epoch: 3.5099051278692417\n",
      "Training loss epoch: 2.9106325820471937\n",
      "Training loss epoch: 2.5860694610270927\n",
      "Validation Loss: 4.673355655075913\n",
      "MAE 1.5515301051744357 Correlation SpearmanrResult(correlation=0.6725347117894619, pvalue=9.213302941959879e-278)\n",
      "Training model ESMForSingleMutationPosConcat on fold 3\n",
      "Training loss epoch: 3.4384211400345297\n",
      "Training loss epoch: 1.4407134821565624\n",
      "Training loss epoch: 0.7402251170556432\n",
      "Validation Loss: 3.4840959660857074\n",
      "MAE 1.3021785023857584 Correlation SpearmanrResult(correlation=0.6726285994576253, pvalue=7.225041875683031e-278)\n",
      "Training model ESMForSingleMutation_pos_cat_cls on fold 3\n",
      "Training loss epoch: 2.506272209797038\n",
      "Training loss epoch: 0.8531412728216953\n",
      "Training loss epoch: 0.28465209813304987\n",
      "Validation Loss: 3.188406802149826\n",
      "MAE 1.2366966885619877 Correlation SpearmanrResult(correlation=0.7064917892315333, pvalue=1.12321e-318)\n",
      "Training model ESMForSingleMutation_pos on fold 3\n",
      "Training loss epoch: 2.3858776319090094\n",
      "Training loss epoch: 0.718571719438181\n",
      "Training loss epoch: 0.25222080060520174\n",
      "Validation Loss: 3.093592798683901\n",
      "MAE 1.2258662197980879 Correlation SpearmanrResult(correlation=0.7157019340561402, pvalue=0.0)\n",
      "Training model ESMForSingleMutation_cls on fold 3\n",
      "Training loss epoch: 3.4261173913625327\n",
      "Training loss epoch: 2.2451298931844437\n",
      "Training loss epoch: 1.0340366083043382\n",
      "Validation Loss: 3.6145928786320014\n",
      "MAE 1.318943848199312 Correlation SpearmanrResult(correlation=0.7187052096972885, pvalue=0.0)\n",
      "\n",
      "Training model ESMForSingleMutationPosOuter on fold 4\n",
      "Training loss epoch: 3.1649036781705977\n",
      "Training loss epoch: 1.311469840057423\n",
      "Training loss epoch: 0.707558824100786\n",
      "Validation Loss: 1.6251482764449607\n",
      "MAE 0.9188618701864988 Correlation SpearmanrResult(correlation=0.7101240490747646, pvalue=3e-323)\n",
      "Training model ESMForSingleMutationPosConcat on fold 4\n",
      "Training loss epoch: 4.283814802257693\n",
      "Training loss epoch: 2.0105970728784635\n",
      "Training loss epoch: 0.9698082053238632\n",
      "Validation Loss: 1.762680560357898\n",
      "MAE 0.9208942983516861 Correlation SpearmanrResult(correlation=0.6557006290630941, pvalue=2.5402722953841667e-259)\n",
      "Training model ESMForSingleMutation_pos_cat_cls on fold 4\n",
      "Training loss epoch: 2.9842594084610945\n",
      "Training loss epoch: 0.9320307261245315\n",
      "Training loss epoch: 0.28574557136997364\n",
      "Validation Loss: 1.7787370953639992\n",
      "MAE 0.9017537063830647 Correlation SpearmanrResult(correlation=0.6820483077528149, pvalue=1.5859588797307158e-288)\n",
      "Training model ESMForSingleMutation_pos on fold 4\n",
      "Training loss epoch: 2.890180034494858\n",
      "Training loss epoch: 0.8979511637920148\n",
      "Training loss epoch: 0.28300524551732115\n",
      "Validation Loss: 1.8130640315207103\n",
      "MAE 0.91204060766929 Correlation SpearmanrResult(correlation=0.6653619783615994, pvalue=1.1046375561335693e-269)\n",
      "Training model ESMForSingleMutation_cls on fold 4\n",
      "Training loss epoch: 4.028261523084231\n",
      "Training loss epoch: 2.5983444396653854\n",
      "Training loss epoch: 1.2163689472167507\n",
      "Validation Loss: 1.7519382091621183\n",
      "MAE 0.9333104577776579 Correlation SpearmanrResult(correlation=0.6955293398674175, pvalue=1.070218936895318e-304)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "lr = 1e-5\n",
    "EPOCHS = 3\n",
    "device = 'cuda:3'\n",
    "\n",
    "models = ['ESMForSingleMutationPosOuter',\n",
    "          'ESMForSingleMutationPosConcat',\n",
    "          'ESMForSingleMutation_pos_cat_cls',  \n",
    "              'ESMForSingleMutation_pos', \n",
    "              'ESMForSingleMutation_cls']\n",
    "\n",
    "full_df = pd.read_csv('DATASETS/new_ds_with_folds.csv')\n",
    "\n",
    "preds = {n:[] for n in models} \n",
    "true = [None]*5\n",
    "\n",
    "for fold_no in range(5):\n",
    "    for model_name in models:\n",
    "        model_class = globals()[model_name]\n",
    "        print(f'Training model {model_name} on fold {fold_no}')\n",
    "        train_df, test_df = full_df[full_df.fold!=fold_no], full_df[full_df.fold==fold_no]\n",
    "        train_ds, test_ds = ProteinDataset(train_df), ProteinDataset(test_df)\n",
    "        \n",
    "        model = model_class()                        \n",
    "        model.to(device) \n",
    "        optimizer = torch.optim.Adam(params=model.parameters(), lr=lr)\n",
    "        training_loader = DataLoader(train_ds, batch_size=1, num_workers = 2, shuffle = True)\n",
    "        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=lr, steps_per_epoch=len(training_loader), epochs=EPOCHS)\n",
    "        \n",
    "        for epoch in range(EPOCHS):\n",
    "            train(epoch)\n",
    "            \n",
    "        testing_loader = DataLoader(test_ds, batch_size=1, num_workers = 2)\n",
    "        labels, predictions = valid(model, testing_loader)\n",
    "        print(f'MAE {np.mean(np.abs(np.array(labels) - np.array(predictions)))} Correlation {stats.spearmanr(labels, predictions)}')     \n",
    "        preds[model_name].append(predictions)\n",
    "        if true[fold_no] is None:\n",
    "            true[fold_no] = labels\n",
    "\n",
    "         \n",
    "        model.to('cpu')\n",
    "        del model\n",
    "    print()\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "62e82205",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ESMForSingleMutationPosOuter RMSE 1.6537080398020116 MAE 1.1164675282649332 Correlation SpearmanrResult(correlation=0.6460236725156133, pvalue=0.0)\n",
      "ESMForSingleMutationPosConcat RMSE 1.636556025307776 MAE 1.103576905154647 Correlation SpearmanrResult(correlation=0.637497712577878, pvalue=0.0)\n",
      "ESMForSingleMutation_pos_cat_cls RMSE 1.5999260928945713 MAE 1.0615672858765015 Correlation SpearmanrResult(correlation=0.6696735701582707, pvalue=0.0)\n",
      "ESMForSingleMutation_pos RMSE 1.595591303549888 MAE 1.0619648237340489 Correlation SpearmanrResult(correlation=0.6673067138634392, pvalue=0.0)\n",
      "ESMForSingleMutation_cls RMSE 1.6958673546219474 MAE 1.143649272722892 Correlation SpearmanrResult(correlation=0.6615472579470065, pvalue=0.0)\n",
      "\n",
      "Ensemble RMSE 1.566587274006694 MAE 1.0343388491591312 Correlation SpearmanrResult(correlation=0.6945659312361261, pvalue=0.0)\n"
     ]
    }
   ],
   "source": [
    "all_true = np.concatenate(true)\n",
    "for model in models:\n",
    "    all_pred = np.concatenate(preds[model])\n",
    "    print(f'{model} RMSE {np.sqrt(np.mean((all_true-all_pred)**2))} MAE {np.mean(np.abs(all_true - all_pred))} Correlation {stats.spearmanr(all_true, all_pred)}')     \n",
    "\n",
    "print()\n",
    "ens_pred = np.mean(np.stack([np.concatenate(preds[model]) for model in models], axis = 0), axis = 0)\n",
    "print(f'Ensemble RMSE {np.sqrt(np.mean((all_true-ens_pred)**2))} MAE {np.mean(np.abs(all_true - ens_pred))} Correlation {stats.spearmanr(all_true, ens_pred)}')     \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a5f9c8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}