03.test_models_on_other_datasets_ensemble.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\n",
    "import pickle\n",
    "import logging\n",
    "\n",
    "logging.basicConfig(filename='PROSTATA_experiments_pearson.log', encoding='utf-8', level=logging.DEBUG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bab2dd4f",
   "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": "784eef4a",
   "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": null,
   "id": "82532b3a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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",
    "    logging.info(f\"Training loss epoch: {epoch_loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fa0a19ef",
   "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",
    "    logging.info(f\"Validation Loss: {eval_loss}\")\n",
    "\n",
    "    return labels, predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bd04a35",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5c7d22fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "experiments = [{\"train\":[\"Q3488\"],             \"test\":[\"ssym\",           \"ssym_r\"]},\n",
    "               {\"train\":[\"Q3488\"],             \"test\":[\"p53\",             \"p53_r\"]},\n",
    "               {\"train\":[\"Q3488\"],             \"test\":[\"myoglobin\", \"myoglobin_r\"]},\n",
    "               {\"train\":[\"S2648\",\"Vb\",\"Vb_r\"], \"test\":[\"ssym\",           \"ssym_r\"]},\n",
    "               {\"train\":[\"S2648\",\"Vb\",\"Vb_r\"], \"test\":[\"p53\",             \"p53_r\"]},\n",
    "               {\"train\":[\"S2648\",\"Vb\",\"Vb_r\"], \"test\":[\"L20\",             \"L20_r\"]},\n",
    "               {\"train\":[\"S2648\",\"Vb\",\"Vb_r\"], \"test\":[\"myoglobin\", \"myoglobin_r\"]},\n",
    "               {\"train\":[\"Q3421\"],             \"test\":[\"ssym\",           \"ssym_r\"]},\n",
    "               {\"train\":[\"prostata\"],          \"test\":[\"s669\",\"s669_r\",\"s669_dssp\"]},\n",
    "               {\"train\":[\"prostata\"],          \"test\":[\"hem_test\"]},\n",
    "               {\"train\":[\"prostata\"],          \"test\":[\"oligo_test\"]},\n",
    "               {\"train\":[\"Q3421\"],             \"test\":[\"ssym\", \"ssym_r\"]}]\n",
    "              "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a00d92",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f9dfa131",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "lr = 1e-5\n",
    "EPOCHS = 3\n",
    "device = 'cuda:1'\n",
    "\n",
    "\n",
    "models = ['ESMForSingleMutationPosOuter',\n",
    "          'ESMForSingleMutationPosConcat',\n",
    "          'ESMForSingleMutation_pos_cat_cls',  \n",
    "              'ESMForSingleMutation_pos', \n",
    "              'ESMForSingleMutation_cls']\n",
    "                        \n",
    "for experiment in experiments:\n",
    "    train_datasets, \n",
    "    test_datasets\n",
    "    \n",
    "    logging.info('Train on ' + str(train_datasets))\n",
    "    train_df = pd.concat([pd.read_csv(os.path.join('DATASETS', t)) for t in train_datasets])\n",
    "    train_ds = ProteinDataset(train_df)\n",
    "    \n",
    "    \n",
    "    all_preds = {}\n",
    "    all_true = {}\n",
    "    for model_no, model_name in enumerate(models):\n",
    "        model_class = globals()[model_name]\n",
    "        logging.info(f'Training model {model_name}')\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",
    "        \n",
    "        for test_no, test_dataset in enumerate(test_datasets):\n",
    "            if test_no not in all_preds:\n",
    "                all_preds[test_no] = {}\n",
    "            logging.info('Test on ' + str(test_dataset))\n",
    "            test_df = pd.read_csv(os.path.join('DATASETS', test_dataset))\n",
    "            test_ds = ProteinDataset(test_df)\n",
    "            testing_loader = DataLoader(test_ds, batch_size=1, num_workers = 2)\n",
    "            labels, predictions = valid(model, testing_loader)\n",
    "            logging.info(f'MAE {np.mean(np.abs(np.array(labels) - np.array(predictions)))} Correlation {stats.pearsonr(labels, predictions)}')     \n",
    "            all_preds[test_no][model_no] = predictions\n",
    "            all_true[test_no] = labels\n",
    "            \n",
    "        del model, optimizer, scheduler\n",
    "        #torch.cuda.empty_cache()    \n",
    "        logging.info('')\n",
    "    \n",
    "    for test_idx in all_true.keys():\n",
    "        logging.info(f'Ensemble result for dataset {test_datasets[test_idx]}')\n",
    "        ens_preds = np.mean(np.stack([all_preds[test_idx][t] for t in all_preds[test_idx].keys()], axis = -1), axis = -1)\n",
    "\n",
    "        logging.info(f'Correlation {stats.pearsonr(all_true[test_idx], ens_preds)}')\n",
    "        logging.info(f'RMSE {np.sqrt(np.mean((np.array(ens_preds)-np.array(all_true[test_idx]))**2))}')        \n",
    "        logging.info(f'MAE {np.mean(np.abs(np.array(all_true[test_idx]) - np.array(ens_preds)))}')\n",
    "        \n",
    "        with open(os.path.join('PREDS', 'train_on_' + '_and_'.join(train_datasets) + 'test_on_' + test_datasets[test_idx] + '.preds.pkl'),\"wb\") as fh:\n",
    "              pickle.dump(ens_preds, fh)\n",
    "    \n",
    "    logging.info('*****************')\n",
    "    logging.info('')\n"
   ]
  }
 ],
 "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
}