00.generate_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2cf37bb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import gzip\n",
    "\n",
    "from pypdb.clients.pdb import pdb_client\n",
    "import pandas as pd\n",
    "from Bio.PDB import *\n",
    "from Bio.PDB.PDBExceptions import PDBConstructionWarning\n",
    "from tqdm import tqdm\n",
    "\n",
    "import warnings\n",
    "warnings.simplefilter('ignore', PDBConstructionWarning)\n",
    "warnings.simplefilter('ignore', UserWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "af83f889",
   "metadata": {},
   "outputs": [],
   "source": [
    "## This code for pdb file manipulation is taken is taken from https://github.com/compbiomed-unito/acdc-nn/blob/master/acdc_nn/util.py\n",
    "\n",
    "def magic_open(path):\n",
    "    return (gzip.open if path.endswith('.gz') else open)(path, 'rt')\n",
    "\n",
    "def pdb2seq(pp):\n",
    "    ''' pdb2seq(pp) takes a pdb_structure_chain \n",
    "    and return its sequence '''\n",
    "    seq = [] # pp.get_sequence()\n",
    "    reslist = []\n",
    "    for ppc  in pp:\n",
    "        reslist += [res for res in ppc]\n",
    "        seq += [str(ppc.get_sequence())]\n",
    "    return \"\".join(seq)\n",
    "\n",
    "def map_pdb_pos(pp):\n",
    "    ''' map_pdb_pos\n",
    "    Returns two dicts seq2pdb[seq_pos], pdb2seq[pdb_pos]'''\n",
    "    reslist = []\n",
    "    for ppc  in pp:\n",
    "        reslist += [res for res in ppc]\n",
    "    seq2pdb = dict(zip( map(str,range(1,len(reslist)+1)), [str(r.get_id()[1])+r.get_id()[2].strip() for r in reslist]))\n",
    "    pdb2seq = dict(zip( [str(r.get_id()[1])+r.get_id()[2].strip() for r in reslist], map(str,range(1,len(reslist)+1)) ))\n",
    "    return seq2pdb, pdb2seq\n",
    "\n",
    "def pdb2info(pdb_file, chain):\n",
    "    ''' pdb2info(pdb_file) \n",
    "    Returns structure, polypeptide '''\n",
    "    parser=PDBParser(QUIET=True)\n",
    "    assert pdb_file.endswith(\".pdb\")\n",
    "    pdb_file = pdb_file[:-4].upper() + \".pdb\"\n",
    "    with magic_open(pdb_file) as f:\n",
    "        structure = parser.get_structure('X', f)\n",
    "    pchain=structure[0][chain]\n",
    "    ppb=PPBuilder()\n",
    "    pp = ppb.build_peptides(pchain, aa_only=False) #[0]\n",
    "    return (structure, pchain, pdb2seq(pp), *map_pdb_pos(pp)) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86e36f0a",
   "metadata": {},
   "source": [
    "# S2648"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e12146e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_S2648 = pd.read_csv('DATA/S2648.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "86a16815",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 2648\n",
      "Total number of different chains in dataset 132\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_S2648))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_S2648.PDB_CHAIN.to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0a257eb9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4f512fe3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing s2648\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/2648 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2648/2648 [01:47<00:00, 24.63it/s] \n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "\n",
    "verbatim_pdb_ids = {'1LVEA'}\n",
    "\n",
    "\n",
    "print('Processing s2648')\n",
    "\n",
    "for idx in tqdm(range(len(df_S2648))):\n",
    "    pdb_id = df_S2648.iloc[idx]['PDB_CHAIN'].upper()\n",
    "    wild_aa = df_S2648.iloc[idx]['WILD_RES']\n",
    "    pos = str(df_S2648.iloc[idx]['POSITION'])\n",
    "    mutant_aa = df_S2648.iloc[idx]['MUTANT_RES']\n",
    "    exp_ddg = df_S2648.iloc[idx]['EXP_DDG']\n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    if pdb_id in verbatim_pdb_ids:\n",
    "        seq2pdb_pos = {str(i):str(i) for i in range(len(sequence))}\n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "            print(f'Sequence is {sequence}')\n",
    "            print(f'Mapping is {seq2pdb_pos}')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1a36cc2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': ddg, \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss}).to_csv('DATASETS/S2648.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49a7a64a",
   "metadata": {},
   "source": [
    "# S3488"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6c3f1359",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_1744 = pd.read_csv('DATA/Q1744.txt', sep = ' ', names = ['PDB_CHAIN', 'POSITION', 'WILD_RES', 'MUTANT_RES', 'EXP_DDG'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3917321c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 1744\n",
      "Total number of different chains in dataset 127\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_1744))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_1744.PDB_CHAIN.to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "546ab418",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3c716243",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing S3488\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1744 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1744/1744 [01:27<00:00, 19.82it/s]\n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "\n",
    "verbatim_pdb_ids = {'1LVEA'}\n",
    "\n",
    "\n",
    "print('Processing S3488')\n",
    "\n",
    "for idx in tqdm(range(len(df_1744))):\n",
    "    pdb_id = df_1744.iloc[idx]['PDB_CHAIN']\n",
    "    wild_aa = df_1744.iloc[idx]['WILD_RES']\n",
    "    pos = str(df_1744.iloc[idx]['POSITION'])\n",
    "    mutant_aa = df_1744.iloc[idx]['MUTANT_RES']\n",
    "    exp_ddg = df_1744.iloc[idx]['EXP_DDG']\n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    if pdb_id in verbatim_pdb_ids:\n",
    "        seq2pdb_pos = {str(i):str(i) for i in range(len(sequence))}\n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "            print(f'Sequence is {sequence}')\n",
    "            print(f'Mapping is {seq2pdb_pos}')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(-1*exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))\n",
    "            \n",
    "            mut.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            wt.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f7c63e26",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': ddg, \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss}).to_csv('DATASETS/S3488.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e0501e5",
   "metadata": {},
   "source": [
    "# S3421"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a1176540",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_17354/2764013025.py:1: ParserWarning: Length of header or names does not match length of data. This leads to a loss of data with index_col=False.\n",
      "  df_3421 = pd.read_csv('DATA/Q3421.txt', sep = '\\t', skiprows = 2, index_col=False,\n"
     ]
    }
   ],
   "source": [
    "df_3421 = pd.read_csv('DATA/Q3421.txt', sep = '\\t', skiprows = 2, index_col=False,\n",
    "                      names = ['PDB_ID', 'PDB_CHAIN', 'POSITION', 'WILD_RES', 'MUTANT_RES', 'EXP_DDG', 'T', 'PH', 'POS2'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f27fa5bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 3421\n",
      "Total number of different chains in dataset 148\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_3421))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_3421.PDB_ID.to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d80a2c00",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "de585b1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing S3421\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 3/3421 [00:00<02:05, 27.25it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███       | 1048/3421 [00:53<02:53, 13.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1LVEA expected Q at position 89 \n",
      "Error for 1LVEA expected Q at position 89 \n",
      "Error for 1LVEA expected Q at position 38 \n",
      "Error for 1LVEA expected S at position 97 \n",
      "Error for 1LVEA expected I at position 106 \n",
      "Error for 1LVEA expected N at position 28 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███       | 1060/3421 [00:54<01:46, 22.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1LVEA expected K at position 39 \n",
      "Error for 1LVEA expected K at position 30 \n",
      "Error for 1LVEA expected K at position 30 \n",
      "Error for 1LVEA expected S at position 29 \n",
      "Error for 1LVEA expected T at position 94 \n",
      "Error for 1LVEA expected Y at position 96 \n",
      "Error for 1LVEA expected Y at position 96 \n",
      "Error for 1LVEA expected P at position 40 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3421/3421 [02:08<00:00, 26.71it/s] \n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "\n",
    "verbatim_pdb_ids = {'1LVEA'}\n",
    "\n",
    "\n",
    "print('Processing S3421')\n",
    "\n",
    "for idx in tqdm(range(len(df_3421))):\n",
    "    pdb_id = df_3421.iloc[idx]['PDB_ID'].upper() + df_3421.iloc[idx]['PDB_CHAIN'].upper()\n",
    "    wild_aa = df_3421.iloc[idx]['WILD_RES']\n",
    "    pos = str(df_3421.iloc[idx]['POSITION'])\n",
    "    mutant_aa = df_3421.iloc[idx]['MUTANT_RES']\n",
    "    exp_ddg = df_3421.iloc[idx]['EXP_DDG']\n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    if pdb_id in verbatim_pdb_ids:\n",
    "        seq2pdb_pos = {str(i):str(i) for i in range(len(sequence))}\n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4b848c20",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': ddg, \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss}).to_csv('DATASETS/S3421.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f285dc0c",
   "metadata": {},
   "source": [
    "# ACDC-varibench"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "74b70548",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_acdc_varibench = pd.concat([pd.read_csv(os.path.join('DATA/varibench/', f), sep = ' ',\n",
    "            names = ['PDB_CHAIN', 'MUTATION', 'EXP_DDG']) for f in sorted(os.listdir('DATA/varibench/'))]).drop_duplicates()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "92d4871a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 1387\n",
      "Total number of different chains in dataset 78\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_acdc_varibench))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_acdc_varibench.PDB_CHAIN.to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e54dd0f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "99bc4cbd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing ACDC-varibench\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 30/1387 [00:00<00:39, 34.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1AM7A expected H at position 30 \n",
      "Error for 1AM7A expected H at position 47 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  9%|▉         | 126/1387 [00:03<00:35, 35.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexing error for 1BNIA position 108 not present in mapping {'0': '0', '1': '1', '2': '2', '3': '3', '4': '4', '5': '5', '6': '6', '7': '7', '8': '8', '9': '9', '10': '10', '11': '11', '12': '12', '13': '13', '14': '14', '15': '15', '16': '16', '17': '17', '18': '18', '19': '19', '20': '20', '21': '21', '22': '22', '23': '23', '24': '24', '25': '25', '26': '26', '27': '27', '28': '28', '29': '29', '30': '30', '31': '31', '32': '32', '33': '33', '34': '34', '35': '35', '36': '36', '37': '37', '38': '38', '39': '39', '40': '40', '41': '41', '42': '42', '43': '43', '44': '44', '45': '45', '46': '46', '47': '47', '48': '48', '49': '49', '50': '50', '51': '51', '52': '52', '53': '53', '54': '54', '55': '55', '56': '56', '57': '57', '58': '58', '59': '59', '60': '60', '61': '61', '62': '62', '63': '63', '64': '64', '65': '65', '66': '66', '67': '67', '68': '68', '69': '69', '70': '70', '71': '71', '72': '72', '73': '73', '74': '74', '75': '75', '76': '76', '77': '77', '78': '78', '79': '79', '80': '80', '81': '81', '82': '82', '83': '83', '84': '84', '85': '85', '86': '86', '87': '87', '88': '88', '89': '89', '90': '90', '91': '91', '92': '92', '93': '93', '94': '94', '95': '95', '96': '96', '97': '97', '98': '98', '99': '99', '100': '100', '101': '101', '102': '102', '103': '103', '104': '104', '105': '105', '106': '106', '107': '107'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 377/1387 [00:09<00:18, 54.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1ONCA expected M at position 22 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 596/1387 [00:13<00:10, 74.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexing error for 1STNA position 136 not present in mapping {'0': '0', '1': '1', '2': '2', '3': '3', '4': '4', '5': '5', '6': '6', '7': '7', '8': '8', '9': '9', '10': '10', '11': '11', '12': '12', '13': '13', '14': '14', '15': '15', '16': '16', '17': '17', '18': '18', '19': '19', '20': '20', '21': '21', '22': '22', '23': '23', '24': '24', '25': '25', '26': '26', '27': '27', '28': '28', '29': '29', '30': '30', '31': '31', '32': '32', '33': '33', '34': '34', '35': '35', '36': '36', '37': '37', '38': '38', '39': '39', '40': '40', '41': '41', '42': '42', '43': '43', '44': '44', '45': '45', '46': '46', '47': '47', '48': '48', '49': '49', '50': '50', '51': '51', '52': '52', '53': '53', '54': '54', '55': '55', '56': '56', '57': '57', '58': '58', '59': '59', '60': '60', '61': '61', '62': '62', '63': '63', '64': '64', '65': '65', '66': '66', '67': '67', '68': '68', '69': '69', '70': '70', '71': '71', '72': '72', '73': '73', '74': '74', '75': '75', '76': '76', '77': '77', '78': '78', '79': '79', '80': '80', '81': '81', '82': '82', '83': '83', '84': '84', '85': '85', '86': '86', '87': '87', '88': '88', '89': '89', '90': '90', '91': '91', '92': '92', '93': '93', '94': '94', '95': '95', '96': '96', '97': '97', '98': '98', '99': '99', '100': '100', '101': '101', '102': '102', '103': '103', '104': '104', '105': '105', '106': '106', '107': '107', '108': '108', '109': '109', '110': '110', '111': '111', '112': '112', '113': '113', '114': '114', '115': '115', '116': '116', '117': '117', '118': '118', '119': '119', '120': '120', '121': '121', '122': '122', '123': '123', '124': '124', '125': '125', '126': '126', '127': '127', '128': '128', '129': '129', '130': '130', '131': '131', '132': '132', '133': '133', '134': '134', '135': '135'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 645/1387 [00:16<00:36, 20.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1YCCA expected C at position 106 \n",
      "Error for 1YCCA expected C at position 106 \n",
      "Error for 1YCCA expected C at position 106 \n",
      "Error for 1YCCA expected F at position 86 \n",
      "Error for 1YCCA expected L at position 89 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 656/1387 [00:16<00:21, 33.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1YCCA expected P at position 80 \n",
      "Error for 1YCCA expected P at position 80 \n",
      "Error for 1YCCA expected P at position 80 \n",
      "Error for 1YCCA expected P at position 80 \n",
      "Error for 1YCCA expected P at position 80 \n",
      "Error for 1YCCA expected P at position 80 \n",
      "Error for 1YCCA expected P at position 80 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1387/1387 [00:40<00:00, 34.02it/s] \n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "\n",
    "no_verbatim_pdb_ids = {'1C9OA', '1VQBA'}\n",
    "\n",
    "\n",
    "print('Processing ACDC-varibench')\n",
    "\n",
    "for idx in tqdm(range(len(df_acdc_varibench))):\n",
    "    pdb_id = df_acdc_varibench.iloc[idx]['PDB_CHAIN'].upper()\n",
    "    wild_aa = df_acdc_varibench.iloc[idx]['MUTATION'][0]\n",
    "    pos = df_acdc_varibench.iloc[idx]['MUTATION'][1:-1]\n",
    "    mutant_aa = df_acdc_varibench.iloc[idx]['MUTATION'][-1]\n",
    "    exp_ddg = df_acdc_varibench.iloc[idx]['EXP_DDG']\n",
    "    \n",
    "    #if pdb_id!= '1CLWA':\n",
    "    #    continue\n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    if pdb_id not in no_verbatim_pdb_ids:\n",
    "        seq2pdb_pos = {str(i):str(i) for i in range(len(sequence))}\n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2fef3221",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': ddg, \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss}).to_csv('DATASETS/ACDC_varibench.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be020955",
   "metadata": {},
   "source": [
    "# Ssym"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d2441109",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_ssym = pd.read_csv('DATA/s_sym.txt', sep= ' ', names = ['PDB_ID', '_', 'POSITION', 'WILD_RES', 'MUTANT_RES', 'EXP_DDG'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b081f3cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 342\n",
      "Total number of different chains in dataset 15\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_ssym))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_ssym['PDB_ID'].to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3055a12b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "ba064e61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing Ssym\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 342/342 [00:08<00:00, 41.72it/s]\n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "\n",
    "print('Processing Ssym')\n",
    "\n",
    "for idx in tqdm(range(len(df_ssym))):\n",
    "    pdb_id = df_ssym.iloc[idx]['PDB_ID'].upper()\n",
    "    wild_aa = df_ssym.iloc[idx]['WILD_RES']\n",
    "    pos = str(df_ssym.iloc[idx]['POSITION'])\n",
    "    mutant_aa = df_ssym.iloc[idx]['MUTANT_RES']\n",
    "    exp_ddg = df_ssym.iloc[idx]['EXP_DDG']\n",
    "    \n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    #if pdb_id not in no_verbatim_pdb_ids:\n",
    "    #  seq2pdb_pos = {str(i):str(i) for i in range(len(sequence))}\n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "            print(f'Sequence is {sequence}')\n",
    "            print(f'Mapping is {seq2pdb_pos}')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "7fc79923",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': [-t for t in ddg], \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss}).to_csv('DATASETS/ssym.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e547246c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': mut, \n",
    "              'mut_seq': wt ,\n",
    "              'ddg': ddg, \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': [t[-1] + t[1:-1] + t[0] for t in mut_infos],\n",
    "              'pos': poss}).to_csv('DATASETS/ssym_r.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "868c5000",
   "metadata": {},
   "source": [
    "# Myoglobin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "e1c768fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_myoglobin = pd.read_csv('DATA/myoglobin.txt', sep= ' ', names = ['PDB_ID', 'POSITION', 'WILD_RES', 'MUTANT_RES', 'EXP_DDG'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "c5cf2f3a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 134\n",
      "Total number of different chains in dataset 1\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_myoglobin))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_myoglobin['PDB_ID'].to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "9c95e116",
   "metadata": {},
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "b7d5f7cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing myoglobin\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 134/134 [00:07<00:00, 18.76it/s]\n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "\n",
    "#no_verbatim_pdb_ids = {'1C9OA', '1VQBA'}\n",
    "\n",
    "\n",
    "print('Processing myoglobin')\n",
    "\n",
    "for idx in tqdm(range(len(df_myoglobin))):\n",
    "    pdb_id = df_myoglobin.iloc[idx]['PDB_ID'].upper()\n",
    "    wild_aa = df_myoglobin.iloc[idx]['WILD_RES']\n",
    "    pos = str(df_myoglobin.iloc[idx]['POSITION'])\n",
    "    mutant_aa = df_myoglobin.iloc[idx]['MUTANT_RES']\n",
    "    exp_ddg = df_myoglobin.iloc[idx]['EXP_DDG']\n",
    "    \n",
    "    #if pdb_id!= '1CLWA':\n",
    "    #    continue\n",
    "        \n",
    "    \n",
    "    pdb = PDBParser().get_structure(pdb_id[:4], f'PDB/{pdb_id[:4]}.pdb')\n",
    "    chain = next(pdb.get_chains()).get_id()\n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    #if pdb_id not in no_verbatim_pdb_ids:\n",
    "    #  seq2pdb_pos = {str(i):str(i) for i in range(len(sequence))}\n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "            print(f'Sequence is {sequence}')\n",
    "            print(f'Mapping is {seq2pdb_pos}')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "ff237702",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': [-t for t in ddg], \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss}).to_csv('DATASETS/myoglobin.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "68a5b215",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': mut, \n",
    "              'mut_seq': wt ,\n",
    "              'ddg': ddg, \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': [t[-1] + t[1:-1] + t[0] for t in mut_infos],\n",
    "              'pos': poss}).to_csv('DATASETS/myoglobin_r.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "376691a0",
   "metadata": {},
   "source": [
    "# P53"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "25d79625",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_p53 = pd.read_csv('DATA/p53.txt', sep= ' ', names = ['PDB_ID', 'POSITION', 'WILD_RES', 'MUTANT_RES', 'EXP_DDG'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "4563a086",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 42\n",
      "Total number of different chains in dataset 1\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_p53))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_p53['PDB_ID'].to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "15756c9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8807d652",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing p53\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/42 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 42/42 [00:05<00:00,  8.01it/s]\n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "\n",
    "print('Processing p53')\n",
    "\n",
    "for idx in tqdm(range(len(df_p53))):\n",
    "    pdb_id = df_p53.iloc[idx]['PDB_ID'].upper()\n",
    "    wild_aa = df_p53.iloc[idx]['WILD_RES']\n",
    "    pos = str(df_p53.iloc[idx]['POSITION'])\n",
    "    mutant_aa = df_p53.iloc[idx]['MUTANT_RES']\n",
    "    exp_ddg = df_p53.iloc[idx]['EXP_DDG']\n",
    "    \n",
    "        \n",
    "    \n",
    "    pdb = PDBParser().get_structure(pdb_id[:4], f'PDB/{pdb_id[:4]}.pdb')\n",
    "    chain = next(pdb.get_chains()).get_id()\n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "            print(f'Sequence is {sequence}')\n",
    "            print(f'Mapping is {seq2pdb_pos}')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "1d85e1a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': [-t for t in ddg], \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss}).to_csv('DATASETS/p53.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ea033b9",
   "metadata": {},
   "source": [
    "# PROSTATA dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "50e83df0",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_article_dataset = pd.read_pickle('DATA/dataset_our_w_clusters_v2.5.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fa68c29b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dataset length 6547\n",
      "Total number of different chains in dataset 663\n"
     ]
    }
   ],
   "source": [
    "print('Total dataset length', len(df_article_dataset))\n",
    "pdb_ids = list(set([t.split()[0].upper() for t in df_article_dataset['pdb_id'].to_list()]))\n",
    "print('Total number of different chains in dataset', len(pdb_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "418a70be",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for pdb_id in pdb_ids:\n",
    "    if not os.path.isfile(f\"PDB/{pdb_id[:4]}.pdb\"):\n",
    "        with open(f\"PDB/{pdb_id[:4]}.pdb\", \"w\") as fh:\n",
    "            fh.write(pdb_client.get_pdb_file(f\"{pdb_id[:4]}\", compression=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "c2940d20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing article dataset\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/6547 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▍         | 313/6547 [00:55<03:09, 32.82it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1AYEA expected A at position 31 \n",
      "Error for 1AYEA expected E at position 14 \n",
      "Error for 1AYEA expected E at position 20 \n",
      "Error for 1AYEA expected I at position 15 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▍         | 318/6547 [00:55<03:56, 26.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for 1AYEA expected I at position 23 \n",
      "Error for 1AYEA expected V at position 12 \n",
      "Error for 1AYEA expected V at position 64 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 37%|███▋      | 2443/6547 [02:57<03:34, 19.10it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexing error for 1LRPA position 15 not present in mapping {}\n",
      "Indexing error for 1LRPA position 20 not present in mapping {}\n",
      "Indexing error for 1LRPA position 37 not present in mapping {}\n",
      "Indexing error for 1LRPA position 49 not present in mapping {}\n",
      "Indexing error for 1LRPA position 49 not present in mapping {}\n",
      "Indexing error for 1LRPA position 63 not present in mapping {}\n",
      "Indexing error for 1LRPA position 66 not present in mapping {}\n",
      "Indexing error for 1LRPA position 66 not present in mapping {}\n",
      "Indexing error for 1LRPA position 81 not present in mapping {}\n",
      "Indexing error for 1LRPA position 46 not present in mapping {}\n",
      "Indexing error for 1LRPA position 48 not present in mapping {}\n",
      "Indexing error for 1LRPA position 48 not present in mapping {}\n",
      "Indexing error for 1LRPA position 48 not present in mapping {}\n",
      "Indexing error for 1LRPA position 84 not present in mapping {}\n",
      "Indexing error for 1LRPA position 4 not present in mapping {}\n",
      "Indexing error for 1LRPA position 40 not present in mapping {}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 2456/6547 [02:57<01:31, 44.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexing error for 1LRPA position 78 not present in mapping {}\n",
      "Indexing error for 1LRPA position 33 not present in mapping {}\n",
      "Indexing error for 1LRPA position 44 not present in mapping {}\n",
      "Indexing error for 1LRPA position 36 not present in mapping {}\n",
      "Indexing error for 1LRPA position 22 not present in mapping {}\n",
      "Indexing error for 1LRPA position 88 not present in mapping {}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6547/6547 [09:00<00:00, 12.10it/s]\n"
     ]
    }
   ],
   "source": [
    "wt = []\n",
    "mut = []\n",
    "ddg = []\n",
    "pdb_ids = []\n",
    "mut_infos = []\n",
    "poss = []\n",
    "ssym_splits = []\n",
    "s669_splits = []\n",
    "clusters = []\n",
    "\n",
    "print('Processing article dataset')\n",
    "\n",
    "for idx in tqdm(range(len(df_article_dataset))):\n",
    "    pdb_id = df_article_dataset.iloc[idx]['pdb_id'].upper() + df_article_dataset.iloc[idx]['pdb_chain'].upper()\n",
    "    wild_aa = df_article_dataset.iloc[idx]['Mut_code'][0]\n",
    "    pos = df_article_dataset.iloc[idx]['Mut_code'][1:-1]\n",
    "    mutant_aa = df_article_dataset.iloc[idx]['Mut_code'][-1]\n",
    "    exp_ddg = df_article_dataset.iloc[idx]['mean_ddG']\n",
    "    \n",
    "    ssym_split = df_article_dataset.iloc[idx]['ssym_split']\n",
    "    s669_split = df_article_dataset.iloc[idx]['s669_split']\n",
    "    cluster = df_article_dataset.iloc[idx]['cluster_id'] \n",
    "    \n",
    "        \n",
    "    \n",
    "    pdb = PDBParser().get_structure(pdb_id[:4], f'PDB/{pdb_id[:4]}.pdb')\n",
    "    chain = next(pdb.get_chains()).get_id()\n",
    "        \n",
    "    _, _, sequence, pdb2seq_pos, seq2pdb_pos = pdb2info(f'PDB/{pdb_id[:4]}.pdb', pdb_id[-1])\n",
    "    \n",
    "    \n",
    "    if pos not in seq2pdb_pos:\n",
    "        print(f'Indexing error for {pdb_id} position {pos} not present in mapping {seq2pdb_pos}')\n",
    "        \n",
    "    else:\n",
    "        if sequence[int(seq2pdb_pos[pos])-1]!=wild_aa:\n",
    "            print(f'Error for {pdb_id} expected {wild_aa} at position {pos} ')\n",
    "        \n",
    "        else:\n",
    "            wt.append(sequence)\n",
    "            tt = list(sequence)\n",
    "            tt[int(seq2pdb_pos[pos])-1] = mutant_aa\n",
    "            poss.append(int(seq2pdb_pos[pos])-1)\n",
    "            mut.append(''.join(tt))\n",
    "            ddg.append(exp_ddg)\n",
    "            pdb_ids.append(pdb_id)\n",
    "            mut_infos.append(str(wild_aa) + pos + str(mutant_aa))\n",
    "            ssym_splits.append(ssym_split)\n",
    "            s669_splits.append(s669_split)\n",
    "            clusters.append(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "e63c5760",
   "metadata": {},
   "outputs": [],
   "source": [
    "article_dataset = pd.DataFrame({'wt_seq': wt, \n",
    "              'mut_seq': mut ,\n",
    "              'ddg': [t for t in ddg], \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': mut_infos,\n",
    "              'pos': poss,\n",
    "              'ssym_split': ssym_splits,\n",
    "              's669_split': s669_splits,\n",
    "              'cluster': clusters})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "99283846",
   "metadata": {},
   "outputs": [],
   "source": [
    "article_dataset_rev = pd.DataFrame({'wt_seq': mut, \n",
    "              'mut_seq': wt ,\n",
    "              'ddg': [-t for t in ddg], \n",
    "              'pdb_id': pdb_ids, \n",
    "              'mut_info': [t[-1] + t[1:-1] + t[0] for t in mut_infos],\n",
    "              'pos': poss,\n",
    "              'ssym_split': ssym_splits,\n",
    "              's669_split': s669_splits,\n",
    "              'cluster': clusters})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "47fca711",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10542\n"
     ]
    }
   ],
   "source": [
    "article_dataset_with_rev = pd.concat([article_dataset, article_dataset_rev]).drop_duplicates(['wt_seq', 'mut_seq'])\n",
    "print(len(article_dataset_with_rev))\n",
    "#10542"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "1ad7ad7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "article_dataset_with_rev[article_dataset_with_rev.ssym_split == 'train'][['wt_seq', 'mut_seq', 'ddg', 'pdb_id', 'mut_info', 'pos']].to_csv('DATASETS/new_ds_for_ssym.csv')\n",
    "article_dataset_with_rev[article_dataset_with_rev.s669_split == 'train'][['wt_seq', 'mut_seq', 'ddg', 'pdb_id', 'mut_info', 'pos']].to_csv('DATASETS/new_ds_for_s669.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "4ac58052",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import GroupKFold\n",
    "group_kfold = GroupKFold(n_splits=5)\n",
    "article_dataset_with_rev = article_dataset_with_rev.reset_index(drop = True)\n",
    "article_dataset_with_rev['fold'] = None\n",
    "for fold, (train_index, test_index) in enumerate(group_kfold.split(article_dataset_with_rev.cluster, \n",
    "                                                                   article_dataset_with_rev.cluster, \n",
    "                                                                   article_dataset_with_rev.cluster)):\n",
    "    article_dataset_with_rev.loc[test_index, 'fold'] = fold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "8f0ef7a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "article_dataset_with_rev[['wt_seq', 'mut_seq', 'ddg', 'pdb_id', 'mut_info', 'pos', 'fold']].to_csv('DATASETS/new_ds_with_folds.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56e80bea",
   "metadata": {},
   "source": [
    "# Case studies datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "f14fc010",
   "metadata": {},
   "outputs": [],
   "source": [
    "dimer_clusters = {'cluster_162_ref_1AV1_A', 'cluster_4_ref_1BFM_A', 'cluster_161_ref_1ARR_A', 'cluster_34_ref_1R6R_A', 'cluster_139_ref_1SAK_A', 'cluster_222_ref_1ZNJ_B', 'cluster_223_ref_1ZNJ_A', 'cluster_210_ref_1UWO_A', 'cluster_69_ref_3MON_B', 'cluster_47_ref_2KJ3_A', 'cluster_183_ref_1CDC_B', 'cluster_140_ref_1SCE_A'}\n",
    "gem_clusters = {'cluster_2_ref_1A7V_A', 'cluster_189_ref_1CYC_A', 'cluster_217_ref_1YCC_A', 'cluster_92_ref_1I5T_A', 'cluster_257_ref_451C_A', 'cluster_168_ref_1B5M_A', 'cluster_190_ref_1CYO_A', 'cluster_179_ref_1C52_A', 'cluster_178_ref_1C2R_A', 'cluster_218_ref_1YEA_A', 'cluster_177_ref_1BVC_A'}\n",
    "transmembrane_clusters = {'cluster_145_ref_1THQ_A', 'cluster_200_ref_1FEP_A', 'cluster_230_ref_2BRD_A', 'cluster_129_ref_1QJP_A'}\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "7be8fc2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "article_dataset_full = article_dataset_with_rev[['wt_seq', 'mut_seq', 'ddg', 'pdb_id', 'mut_info', 'pos', 'cluster']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "bc8004ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "article_dataset_full[article_dataset_full.cluster.isin(dimer_clusters)].to_csv('DATASETS/case_study_dimer_test.csv')\n",
    "article_dataset_full[article_dataset_full.cluster.isin(gem_clusters)].to_csv('DATASETS/case_study_gem_test.csv')\n",
    "article_dataset_full[article_dataset_full.cluster.isin(transmembrane_clusters)].to_csv('DATASETS/case_study_transmembrane_test.csv')"
   ]
  }
 ],
 "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
}