AntBO/generate_init_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from visualise_results.utils import *\n",
    "import numpy as np\n",
    "from typing import *\n",
    "import glob\n",
    "\n",
    "import sys, os, shutil\n",
    "\n",
    "ROOT_PROJECT =  os.path.join(os.path.normpath(os.path.join(os.getcwd()))) \n",
    "sys.path[0] = ROOT_PROJECT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bo.custom_init import InitialBODataset, get_initial_dataset_path\n",
    "from bo.utils import save_w_pickle\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_by_method = []\n",
    "\n",
    "for meth_file in glob.glob(\"./results_data/*_optim_res.csv\"):\n",
    "    result_by_method.append(pd.read_csv(meth_file, index_col=0))\n",
    "results = pd.concat(result_by_method)\n",
    "results = results.loc[np.sort(results.index.values)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./dataloader/core_antigens.txt\") as f:\n",
    "    core_antigens = list(map(lambda antigen: antigen.rstrip(), f.readlines()))\n",
    "\n",
    "\n",
    "with open('./utils_data/all_antigens.txt') as file:\n",
    "    antigens = file.readlines()\n",
    "    all_antigens = [antigen.rstrip() for antigen in antigens]\n",
    "\n",
    "remaining_antigens = list(set(all_antigens)-set(core_antigens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "thresholdPd = pd.read_csv(\"./utils_data/ListThresholds.txt\", sep=\" \")\n",
    "# Mascotte and above 1% "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For each of the 12 core_antigens\n",
    "# \n",
    "# 1 inital dataset Loosers (20)  # samples to get to Mascotte,Heroes,SuperHeroes\n",
    "# 1 inital dataset Loosers+Mascotte (20) # samples to get to Heroes,SuperHeroes\n",
    "# 1 inital dataset Loosers+Mascotte+Heroes (20) # samples to get to SuperHeroes\n",
    "\n",
    "AA = 'ACDEFGHIKLMNPQRSTVWY'\n",
    "AA_to_idx = {aa: i for i, aa in enumerate(AA)}\n",
    "\n",
    "categories = ['NonBinders', 'Loosers', 'Mascotte', 'Heroes', 'SuperHeroes']\n",
    "\n",
    "def get_energy_interval(antigen, category):\n",
    "    if category not in categories:\n",
    "        raise(ValueError(category))\n",
    "    ind = categories.index(category)\n",
    "    upper_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==category]['maxEnergy'].values[0]\n",
    "    if ind == (len(categories) - 1):\n",
    "        lower_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==categories[ind]]['minEnergy'].values[0]\n",
    "    else:\n",
    "        lower_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==categories[min(len(categories) - 1, ind + 1)]]['maxEnergy'].values[0] + 0.01\n",
    "    return lower_bound, upper_bound"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-26-e59a4f15015c>:16: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  upper_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==category]['maxEnergy'].values[0]\n",
      "<ipython-input-26-e59a4f15015c>:18: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  lower_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==categories[ind]]['minEnergy'].values[0]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(-125.07, -119.16)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "antigen_name_ = \"2DD8_S\"\n",
    "get_energy_interval(antigen=antigen_name_, category='SuperHeroes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_proteins_by_category(antigen, category, top_cut_ratio=0):\n",
    "    \"\"\"\n",
    "    antigen: the name of the antigen\n",
    "    category: the name of the category (Loosers,...)\n",
    "    top_cut_ratio: remove the top `top_cut_ratio` * 100 % of the protein in terms of binding energy (the lower the better)\n",
    "    \"\"\"\n",
    "    aux_results = results[results.Antigen == antigen]\n",
    "    lower_energy, upper_energy = get_energy_interval(antigen=antigen, category=category)\n",
    "    aux_results = aux_results[aux_results['Last Binding Energy'] < upper_energy][aux_results['Last Binding Energy'] > lower_energy]\n",
    "    aux_results = aux_results[['Last Protein', 'Last Binding Energy']].drop_duplicates(subset=['Last Protein']).values\n",
    "    \n",
    "    # Keep the highest energies\n",
    "    keeps = np.ceil(top_cut_ratio * len(aux_results)).astype(int)\n",
    "    inds = aux_results[:, 1].argsort()[keeps:]\n",
    "    return aux_results[inds]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-26-e59a4f15015c>:16: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  upper_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==category]['maxEnergy'].values[0]\n",
      "<ipython-input-26-e59a4f15015c>:18: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  lower_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==categories[ind]]['minEnergy'].values[0]\n",
      "<ipython-input-28-d49f2f24b56c>:9: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  aux_results = aux_results[aux_results['Last Binding Energy'] < upper_energy][aux_results['Last Binding Energy'] > lower_energy]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(1734, 2)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proteins_scores_ = get_proteins_by_category('2DD8_S', 'SuperHeroes', 0)\n",
    "proteins_scores_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getInitUpToCat(antigen: str, top_category: str, n: int, top_cat_top_cut_ratio: float=0, seed: int = 0):\n",
    "    if top_category not in categories:\n",
    "        raise(ValueError(top_category))\n",
    "    np.random.seed(seed)\n",
    "    ind = categories.index(top_category)\n",
    "    n_per_cats = [n // ind for _ in range(ind)]\n",
    "    n_per_cats[-1] = n - sum(n_per_cats[:-1])\n",
    "    top_cut_ratios = [0 for _ in range(ind)]\n",
    "    top_cut_ratios[-1] = top_cat_top_cut_ratio\n",
    "    assert sum(n_per_cats) == n\n",
    "    samples = []\n",
    "    for i in range(ind):\n",
    "        n_sample = n_per_cats[i]\n",
    "        top_cut_ratio = top_cut_ratios[i]\n",
    "        category = categories[i + 1]\n",
    "        proteins_scores = get_proteins_by_category(antigen, category=category, top_cut_ratio=top_cut_ratio)\n",
    "        sampled_inds = np.random.choice(np.arange(len(proteins_scores)), size=n_sample)\n",
    "        samples.extend([[category, np.array([AA_to_idx[el] for el in  proteins_scores[sample_ind][0]]), proteins_scores[sample_ind][0], proteins_scores[sample_ind][1]] for sample_ind in sampled_inds])\n",
    "    return pd.DataFrame(samples, columns=['type', 'AA to ind', 'Protein', 'Binding Energy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getInitDataset(antigen: str, n_per_cat: Dict[str, int], top_cut_ratio_per_cat: Dict[str, float], seed: int):\n",
    "    np.random.seed(seed)\n",
    "    samples = []\n",
    "    for category, n_sample in n_per_cat.items():\n",
    "        top_cut_ratio = top_cut_ratio_per_cat.get(category, 0)\n",
    "        proteins_scores = get_proteins_by_category(antigen, category=category, top_cut_ratio=top_cut_ratio)\n",
    "        sampled_inds = np.random.choice(np.arange(len(proteins_scores)), size=n_sample)\n",
    "        samples.extend([[category, np.array([AA_to_idx[el] for el in  proteins_scores[sample_ind][0]]), proteins_scores[sample_ind][0], proteins_scores[sample_ind][1]] for sample_ind in sampled_inds])\n",
    "    return pd.DataFrame(samples, columns=['type', 'AA to ind', 'Protein', 'Binding Energy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_n_per_cat(n_loosers: int, n_mascottes: int, n_heroes):\n",
    "    return dict(Loosers=n_loosers, Mascotte=n_mascottes, Heroes=n_heroes)\n",
    "\n",
    "def get_top_cut_ratio_per_cat(top_cut_ratio_loosers: int, top_cut_ratio_mascottes: int, top_cut_ratio_heroes):\n",
    "    return dict(Loosers=top_cut_ratio_loosers, Mascotte=top_cut_ratio_mascottes, Heroes=top_cut_ratio_heroes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                                                                                                                                             | 0/1 [00:00<?, ?it/s]/tmp/ipykernel_235928/3744245844.py:16: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  upper_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==category]['maxEnergy'].values[0]\n",
      "/tmp/ipykernel_235928/3744245844.py:20: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  lower_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==categories[min(len(categories) - 1, ind + 1)]]['maxEnergy'].values[0] + 0.01\n",
      "/tmp/ipykernel_235928/1088760158.py:9: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  aux_results = aux_results[aux_results['Last Binding Energy'] < upper_energy][aux_results['Last Binding Energy'] > lower_energy]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/0/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/1/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/2/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/3/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/4/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/5/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/6/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/7/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/8/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/9/Loosers-20_Loosers-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/0/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/1/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/2/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/3/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/4/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/5/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/6/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/7/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/8/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/9/Loosers-10_Mascotte-10_Mascotte-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/0/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/1/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/2/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/3/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/4/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/5/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/6/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/7/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/8/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n",
      "/home/rladmin/antigenbinding/bo/init_dataset/2DD8_S/9/Loosers-6_Mascotte-6_Heroes-8_Heroes-0.5/init_data\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "n_per_cat_s = [get_n_per_cat(20, 0, 0), get_n_per_cat(10, 10, 0), get_n_per_cat(6, 6, 8)] \n",
    "top_cut_ratio_per_cat_s = [get_top_cut_ratio_per_cat(0.5, 0, 0), get_top_cut_ratio_per_cat(0, 0.5, 0), get_top_cut_ratio_per_cat(0, 0, 0.5)]\n",
    "\n",
    "selected_antigens = all_antigens\n",
    "\n",
    "invalid_antigens_ = []\n",
    "for antigen_name_ in tqdm(selected_antigens):\n",
    "    for top_cut_ratio_per_cat_, n_per_cat_ in zip(top_cut_ratio_per_cat_s, n_per_cat_s):\n",
    "        if len(results[results.Antigen == antigen_name_]) == 0:\n",
    "            invalid_antigens_.append(antigen_name_)\n",
    "            continue\n",
    "        try:\n",
    "            for seed_ in range(10):\n",
    "                res = getInitDataset(antigen=antigen_name_, n_per_cat=n_per_cat_, top_cut_ratio_per_cat=top_cut_ratio_per_cat_, seed=seed_)        \n",
    "                data = InitialBODataset(res)\n",
    "                save_path = get_initial_dataset_path(\n",
    "                    antigen_name=antigen_name_, n_per_cat=n_per_cat_, top_cut_ratio_per_cat=top_cut_ratio_per_cat_, seed=seed_\n",
    "                )\n",
    "                assert data.get_index_encoded_x().shape == (sum(n_per_cat_.values()), 11)\n",
    "                save_w_pickle(data, save_path)\n",
    "                print(save_path)\n",
    "        except IndexError as e:\n",
    "            invalid_antigens_.append(antigen_name_)\n",
    "            print(antigen_name_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./dataloader/valid_antigens_init_data.txt', 'w') as f:\n",
    "    f.writelines(map(lambda a: a + '\\n', list(set(all_antigens).difference(invalid_antigens_))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "139"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set(all_antigens).difference(invalid_antigens_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_11305/3744245844.py:16: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  upper_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==category]['maxEnergy'].values[0]\n",
      "/tmp/ipykernel_11305/3744245844.py:20: UserWarning: Boolean Series key will be reindexed to match DataFrame index.\n",
      "  lower_bound = thresholdPd[thresholdPd.AGname==antigen][thresholdPd.type==categories[min(len(categories) - 1, ind + 1)]]['maxEnergy'].values[0] + 0.01\n"
     ]
    }
   ],
   "source": [
    "proteins_scores_ = get_proteins_by_category(antigen_name_, category='Loosers', top_cut_ratio=.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([], shape=(0, 2), dtype=object)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proteins_scores_"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}