notebooks/0.2-test_ensembles/0.2-test_ensembles.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate predictions for the ensembles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-31T10:31:54.039063854Z",
     "start_time": "2024-01-31T10:31:54.021466263Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# define the directory where the data is\n",
    "data_path = \"/home/jcapela/ec_numbers_prediction/required_data_ec_number_paper\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Load models predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-31T10:31:56.826990828Z",
     "start_time": "2024-01-31T10:31:54.817072339Z"
    }
   },
   "outputs": [],
   "source": [
    "# read predictions from pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/predictions_prot_bert.pkl\", \"rb\") as f:\n",
    "    predictions_prot_bert = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read predictions from pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/predictions_esm2_3b.pkl\", \"rb\") as f:\n",
    "    predictions_esm2_3b = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read predictions from pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/predictions_esm1b.pkl\", \"rb\") as f:\n",
    "    predictions_esm1b = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jcapela/miniconda3/envs/ec_number_prediction/lib/python3.9/site-packages/plants_sm/io/csv.py:56: DtypeWarning: Columns (4,5) have mixed types. Specify dtype option on import or set low_memory=False.\n",
      "  return pd.read_csv(self.path, **self.kwargs)\n"
     ]
    }
   ],
   "source": [
    "from plants_sm.data_structures.dataset.single_input_dataset import SingleInputDataset\n",
    "\n",
    "test_dataset = SingleInputDataset.from_csv(f'{data_path}/data/test.csv',\n",
    "                                            instances_ids_field=\"accession\", representation_field=\"sequence\",\n",
    "                                            labels_field=slice(8, 2779))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create ensemble of predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a voting classifier for the 3 models\n",
    "import numpy as np\n",
    "\n",
    "def determine_ensemble_predictions(threshold=3, *model_predictions):\n",
    "    model_predictions = list(model_predictions)\n",
    "\n",
    "    for i, model_prediction in enumerate(model_predictions):\n",
    "        model_predictions[i] = np.array(model_prediction)\n",
    "\n",
    "    predictions_voting = np.zeros_like(model_predictions[0])\n",
    "\n",
    "    for i in range(model_predictions[0].shape[0]):\n",
    "        # Combine conditions into a single array and sum along the second axis\n",
    "        combined_conditions = np.sum(np.array([model_predictions[j][i] for j in range(len(model_predictions))]), axis=0)\n",
    "\n",
    "        # Apply the threshold condition\n",
    "        predictions_voting[i] = (combined_conditions >= threshold).astype(int)\n",
    "\n",
    "    # If you want to ensure the resulting array is of integer type\n",
    "    predictions_voting = predictions_voting.astype(int)\n",
    "    return predictions_voting\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_voting = determine_ensemble_predictions(2, predictions_esm1b, predictions_prot_bert, predictions_esm2_3b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save predictions in pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/predictions_models_voting.pkl\", \"wb\") as f:\n",
    "    pickle.dump(predictions_voting, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9011672051082994\n",
      "0.885065468521875\n",
      "0.9274756923443451\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import f1_score, recall_score, precision_score\n",
    "print(f1_score(test_dataset.y, predictions_voting, average='macro'))\n",
    "print(recall_score(test_dataset.y, predictions_voting, average='macro'))\n",
    "print(precision_score(test_dataset.y, predictions_voting, average='macro'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get BLASTp predictions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3010274/2952861442.py:3: DtypeWarning: Columns (2,3) have mixed types. Specify dtype option on import or set low_memory=False.\n",
      "  blast_results = pd.read_csv(f\"{data_path}/test_blast_predictions_right_format.csv\")\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "blast_results = pd.read_csv(f\"{data_path}/test_blast_predictions_right_format.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qseqid</th>\n",
       "      <th>EC</th>\n",
       "      <th>EC1</th>\n",
       "      <th>EC2</th>\n",
       "      <th>EC3</th>\n",
       "      <th>EC4</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>...</th>\n",
       "      <th>7.4.2.5</th>\n",
       "      <th>7.4.2.8</th>\n",
       "      <th>7.5.2.11</th>\n",
       "      <th>7.6.2.1</th>\n",
       "      <th>7.6.2.11</th>\n",
       "      <th>7.6.2.13</th>\n",
       "      <th>7.6.2.2</th>\n",
       "      <th>7.6.2.3</th>\n",
       "      <th>7.6.2.5</th>\n",
       "      <th>7.6.2.8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A0A009IHW8</td>\n",
       "      <td>3.2.2.6</td>\n",
       "      <td>3</td>\n",
       "      <td>3.2</td>\n",
       "      <td>3.2.2</td>\n",
       "      <td>3.2.2.6</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>A0A009YU83</td>\n",
       "      <td>2.3.3.5</td>\n",
       "      <td>2</td>\n",
       "      <td>2.3</td>\n",
       "      <td>2.3.3</td>\n",
       "      <td>2.3.3.5</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A0A010ZGY3</td>\n",
       "      <td>3.7.1.12</td>\n",
       "      <td>3</td>\n",
       "      <td>3.7</td>\n",
       "      <td>3.7.1</td>\n",
       "      <td>3.7.1.12</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>A0A011MZP2</td>\n",
       "      <td>2.6.1.50</td>\n",
       "      <td>2</td>\n",
       "      <td>2.6</td>\n",
       "      <td>2.6.1</td>\n",
       "      <td>2.6.1.50</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>A0A011N9Z1</td>\n",
       "      <td>7.4.2.5</td>\n",
       "      <td>7</td>\n",
       "      <td>7.4</td>\n",
       "      <td>7.4.2</td>\n",
       "      <td>7.4.2.5</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 2777 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       qseqid        EC EC1  EC2    EC3       EC4    1    2    3    4  ...  \\\n",
       "0  A0A009IHW8   3.2.2.6   3  3.2  3.2.2   3.2.2.6  0.0  0.0  1.0  0.0  ...   \n",
       "1  A0A009YU83   2.3.3.5   2  2.3  2.3.3   2.3.3.5  0.0  0.0  0.0  0.0  ...   \n",
       "2  A0A010ZGY3  3.7.1.12   3  3.7  3.7.1  3.7.1.12  0.0  0.0  1.0  0.0  ...   \n",
       "3  A0A011MZP2  2.6.1.50   2  2.6  2.6.1  2.6.1.50  0.0  1.0  0.0  0.0  ...   \n",
       "4  A0A011N9Z1   7.4.2.5   7  7.4  7.4.2   7.4.2.5  0.0  0.0  0.0  0.0  ...   \n",
       "\n",
       "   7.4.2.5  7.4.2.8  7.5.2.11  7.6.2.1  7.6.2.11  7.6.2.13  7.6.2.2  7.6.2.3  \\\n",
       "0      0.0      0.0       0.0      0.0       0.0       0.0      0.0      0.0   \n",
       "1      0.0      0.0       0.0      0.0       0.0       0.0      0.0      0.0   \n",
       "2      0.0      0.0       0.0      0.0       0.0       0.0      0.0      0.0   \n",
       "3      0.0      0.0       0.0      0.0       0.0       0.0      0.0      0.0   \n",
       "4      1.0      0.0       0.0      0.0       0.0       0.0      0.0      0.0   \n",
       "\n",
       "   7.6.2.5  7.6.2.8  \n",
       "0      0.0      0.0  \n",
       "1      0.0      0.0  \n",
       "2      0.0      0.0  \n",
       "3      0.0      0.0  \n",
       "4      0.0      0.0  \n",
       "\n",
       "[5 rows x 2777 columns]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "blast_results.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "accessions = test_dataset.dataframe[\"accession\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qseqid</th>\n",
       "      <th>EC</th>\n",
       "      <th>EC1</th>\n",
       "      <th>EC2</th>\n",
       "      <th>EC3</th>\n",
       "      <th>EC4</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>...</th>\n",
       "      <th>7.4.2.5</th>\n",
       "      <th>7.4.2.8</th>\n",
       "      <th>7.5.2.11</th>\n",
       "      <th>7.6.2.1</th>\n",
       "      <th>7.6.2.11</th>\n",
       "      <th>7.6.2.13</th>\n",
       "      <th>7.6.2.2</th>\n",
       "      <th>7.6.2.3</th>\n",
       "      <th>7.6.2.5</th>\n",
       "      <th>7.6.2.8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>64342</th>\n",
       "      <td>Q7XQ85</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>4</td>\n",
       "      <td>4.4</td>\n",
       "      <td>4.4.1</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1935</th>\n",
       "      <td>A0A0P0WIY3</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>4</td>\n",
       "      <td>4.4</td>\n",
       "      <td>4.4.1</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>54941</th>\n",
       "      <td>Q43309</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>4</td>\n",
       "      <td>4.4</td>\n",
       "      <td>4.4.1</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42238</th>\n",
       "      <td>P29535</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>4</td>\n",
       "      <td>4.4</td>\n",
       "      <td>4.4.1</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1932</th>\n",
       "      <td>A0A0P0UZP7</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>4</td>\n",
       "      <td>4.4</td>\n",
       "      <td>4.4.1</td>\n",
       "      <td>4.4.1.14</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>968</th>\n",
       "      <td>A0A0D0TJI9</td>\n",
       "      <td>1.14.15.13</td>\n",
       "      <td>1</td>\n",
       "      <td>1.14</td>\n",
       "      <td>1.14.15</td>\n",
       "      <td>1.14.15.13</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12834</th>\n",
       "      <td>A0A5C6EG74</td>\n",
       "      <td>1.3.99.26</td>\n",
       "      <td>1</td>\n",
       "      <td>1.3</td>\n",
       "      <td>1.3.99</td>\n",
       "      <td>1.3.99.26</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15759</th>\n",
       "      <td>A0A7C9CV01</td>\n",
       "      <td>2.4.1.-</td>\n",
       "      <td>2</td>\n",
       "      <td>2.4</td>\n",
       "      <td>2.4.1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13578</th>\n",
       "      <td>A0A644YY06</td>\n",
       "      <td>3.1.4.57</td>\n",
       "      <td>3</td>\n",
       "      <td>3.1</td>\n",
       "      <td>3.1.4</td>\n",
       "      <td>3.1.4.57</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17103</th>\n",
       "      <td>A0A7W7ABF6</td>\n",
       "      <td>1.3.99.5</td>\n",
       "      <td>1</td>\n",
       "      <td>1.3</td>\n",
       "      <td>1.3.99</td>\n",
       "      <td>1.3.99.5</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>75566 rows × 2777 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "           qseqid          EC EC1   EC2      EC3         EC4    1    2    3  \\\n",
       "64342      Q7XQ85    4.4.1.14   4   4.4    4.4.1    4.4.1.14  0.0  0.0  0.0   \n",
       "1935   A0A0P0WIY3    4.4.1.14   4   4.4    4.4.1    4.4.1.14  0.0  0.0  0.0   \n",
       "54941      Q43309    4.4.1.14   4   4.4    4.4.1    4.4.1.14  0.0  0.0  0.0   \n",
       "42238      P29535    4.4.1.14   4   4.4    4.4.1    4.4.1.14  0.0  0.0  0.0   \n",
       "1932   A0A0P0UZP7    4.4.1.14   4   4.4    4.4.1    4.4.1.14  0.0  0.0  0.0   \n",
       "...           ...         ...  ..   ...      ...         ...  ...  ...  ...   \n",
       "968    A0A0D0TJI9  1.14.15.13   1  1.14  1.14.15  1.14.15.13  1.0  0.0  0.0   \n",
       "12834  A0A5C6EG74   1.3.99.26   1   1.3   1.3.99   1.3.99.26  1.0  0.0  0.0   \n",
       "15759  A0A7C9CV01     2.4.1.-   2   2.4    2.4.1           0  0.0  1.0  0.0   \n",
       "13578  A0A644YY06    3.1.4.57   3   3.1    3.1.4    3.1.4.57  0.0  0.0  1.0   \n",
       "17103  A0A7W7ABF6    1.3.99.5   1   1.3   1.3.99    1.3.99.5  1.0  0.0  0.0   \n",
       "\n",
       "         4  ...  7.4.2.5  7.4.2.8  7.5.2.11  7.6.2.1  7.6.2.11  7.6.2.13  \\\n",
       "64342  1.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "1935   1.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "54941  1.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "42238  1.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "1932   1.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "...    ...  ...      ...      ...       ...      ...       ...       ...   \n",
       "968    0.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "12834  0.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "15759  0.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "13578  0.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "17103  0.0  ...      0.0      0.0       0.0      0.0       0.0       0.0   \n",
       "\n",
       "       7.6.2.2  7.6.2.3  7.6.2.5  7.6.2.8  \n",
       "64342      0.0      0.0      0.0      0.0  \n",
       "1935       0.0      0.0      0.0      0.0  \n",
       "54941      0.0      0.0      0.0      0.0  \n",
       "42238      0.0      0.0      0.0      0.0  \n",
       "1932       0.0      0.0      0.0      0.0  \n",
       "...        ...      ...      ...      ...  \n",
       "968        0.0      0.0      0.0      0.0  \n",
       "12834      0.0      0.0      0.0      0.0  \n",
       "15759      0.0      0.0      0.0      0.0  \n",
       "13578      0.0      0.0      0.0      0.0  \n",
       "17103      0.0      0.0      0.0      0.0  \n",
       "\n",
       "[75566 rows x 2777 columns]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a new column with the custom order as a categorical type\n",
    "blast_results['CustomOrder'] = pd.Categorical(blast_results['qseqid'], categories=accessions, ordered=True)\n",
    "blast_results.sort_values('CustomOrder', inplace=True)\n",
    "blast_results.drop(columns=[\"CustomOrder\"], inplace=True)\n",
    "blast_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write to pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/blast_predictions.pkl\", \"wb\") as f:\n",
    "    pickle.dump(blast_results.iloc[:, 6:], f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create ensemble with BLASTp and models "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/blast_predictions.pkl\", \"rb\") as f:\n",
    "    blast_predictions = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a voting classifier for the 3 models and blast\n",
    "\n",
    "voting_predictions_blast = determine_ensemble_predictions(2, predictions_esm1b, predictions_prot_bert, predictions_esm2_3b, blast_predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save predictions in pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/predictions_models_voting_blast.pkl\", \"wb\") as f:\n",
    "    pickle.dump(voting_predictions_blast, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read predictions from pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/predictions_models_voting_blast.pkl\", \"rb\") as f:\n",
    "    voting_predictions_blast = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9090857401638838\n",
      "0.9049003181814923\n",
      "0.9204416806398337\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import f1_score, recall_score, precision_score\n",
    "print(f1_score(test_dataset.y, voting_predictions_blast, average='macro'))\n",
    "print(recall_score(test_dataset.y, voting_predictions_blast, average='macro'))\n",
    "print(precision_score(test_dataset.y, voting_predictions_blast, average='macro'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Generate metrics for all the ensembles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from sklearn.metrics import f1_score, recall_score, precision_score\n",
    "\n",
    "\n",
    "def get_ec_levels(labels):\n",
    "    level_1 = []\n",
    "    level_2 = []\n",
    "    level_3 = []\n",
    "    level_4 = []\n",
    "    for i, label in enumerate(labels):\n",
    "        if re.match(r\"^\\d+.\\d+.\\d+.n*\\d+$\", label):\n",
    "            level_4.append(i)\n",
    "        elif re.match(r\"^\\d+.\\d+.\\d+$\", label):\n",
    "            level_3.append(i)\n",
    "        elif re.match(r\"^\\d+.\\d+$\", label):\n",
    "            level_2.append(i)\n",
    "        elif re.match(r\"^\\d+$\", label):\n",
    "            level_1.append(i)\n",
    "    return level_1, level_2, level_3, level_4\n",
    "\n",
    "def get_metrics(y_true_, predictions_, model_name, labels_):\n",
    "    level_1, level_2, level_3, level_4 = get_ec_levels(labels_)  \n",
    "    metrics = {}\n",
    "    metrics[\"mF1 overall\"] = f1_score(y_true_, predictions_, average=\"macro\")\n",
    "    metrics[\"mF1 level 1\"] = f1_score(y_true_[:, level_1], predictions_[:, level_1], average=\"macro\")\n",
    "    metrics[\"mF1 level 2\"] = f1_score(y_true_[:, level_2], predictions_[:, level_2], average=\"macro\")\n",
    "    metrics[\"mF1 level 3\"] = f1_score(y_true_[:, level_3], predictions_[:, level_3], average=\"macro\")\n",
    "    metrics[\"mF1 level 4\"] = f1_score(y_true_[:, level_4], predictions_[:, level_4], average=\"macro\")\n",
    "\n",
    "    metrics[\"mPrecision overall\"] = precision_score(y_true_, predictions_, average=\"macro\")\n",
    "    metrics[\"mPrecision level 1\"] = precision_score(y_true_[:, level_1], predictions_[:, level_1], average=\"macro\")\n",
    "    metrics[\"mPrecision level 2\"] = precision_score(y_true_[:, level_2], predictions_[:, level_2], average=\"macro\")\n",
    "    metrics[\"mPrecision level 3\"] = precision_score(y_true_[:, level_3], predictions_[:, level_3], average=\"macro\")\n",
    "    metrics[\"mPrecision level 4\"] = precision_score(y_true_[:, level_4], predictions_[:, level_4], average=\"macro\")\n",
    "\n",
    "    metrics[\"mRecall overall\"] = recall_score(y_true_, predictions_, average=\"macro\")\n",
    "    metrics[\"mRecall level 1\"] = recall_score(y_true_[:, level_1], predictions_[:, level_1], average=\"macro\")\n",
    "    metrics[\"mRecall level 2\"] = recall_score(y_true_[:, level_2], predictions_[:, level_2], average=\"macro\")\n",
    "    metrics[\"mRecall level 3\"] = recall_score(y_true_[:, level_3], predictions_[:, level_3], average=\"macro\")\n",
    "    metrics[\"mRecall level 4\"] = recall_score(y_true_[:, level_4], predictions_[:, level_4], average=\"macro\")\n",
    "    return pd.DataFrame(metrics, index=[model_name])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels_ = test_dataset._labels_names\n",
    "true_values = test_dataset.y\n",
    "del test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = get_metrics(true_values, voting_predictions_blast, \"ensemble_blast\", labels_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.index.name = \"ensemble\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>mF1 overall</th>\n",
       "      <th>mF1 level 1</th>\n",
       "      <th>mF1 level 2</th>\n",
       "      <th>mF1 level 3</th>\n",
       "      <th>mF1 level 4</th>\n",
       "      <th>mPrecision overall</th>\n",
       "      <th>mPrecision level 1</th>\n",
       "      <th>mPrecision level 2</th>\n",
       "      <th>mPrecision level 3</th>\n",
       "      <th>mPrecision level 4</th>\n",
       "      <th>mRecall overall</th>\n",
       "      <th>mRecall level 1</th>\n",
       "      <th>mRecall level 2</th>\n",
       "      <th>mRecall level 3</th>\n",
       "      <th>mRecall level 4</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ensemble</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ensemble_blast</th>\n",
       "      <td>0.909086</td>\n",
       "      <td>0.984783</td>\n",
       "      <td>0.963836</td>\n",
       "      <td>0.943704</td>\n",
       "      <td>0.904108</td>\n",
       "      <td>0.920442</td>\n",
       "      <td>0.985747</td>\n",
       "      <td>0.969475</td>\n",
       "      <td>0.953077</td>\n",
       "      <td>0.915839</td>\n",
       "      <td>0.9049</td>\n",
       "      <td>0.983822</td>\n",
       "      <td>0.958839</td>\n",
       "      <td>0.936733</td>\n",
       "      <td>0.900196</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                mF1 overall  mF1 level 1  mF1 level 2  mF1 level 3  \\\n",
       "ensemble                                                             \n",
       "ensemble_blast     0.909086     0.984783     0.963836     0.943704   \n",
       "\n",
       "                mF1 level 4  mPrecision overall  mPrecision level 1  \\\n",
       "ensemble                                                              \n",
       "ensemble_blast     0.904108            0.920442            0.985747   \n",
       "\n",
       "                mPrecision level 2  mPrecision level 3  mPrecision level 4  \\\n",
       "ensemble                                                                     \n",
       "ensemble_blast            0.969475            0.953077            0.915839   \n",
       "\n",
       "                mRecall overall  mRecall level 1  mRecall level 2  \\\n",
       "ensemble                                                            \n",
       "ensemble_blast           0.9049         0.983822         0.958839   \n",
       "\n",
       "                mRecall level 3  mRecall level 4  \n",
       "ensemble                                          \n",
       "ensemble_blast         0.936733         0.900196  "
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.to_csv(\"results_ensemble_blast.csv\", index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save predictions in pickle file\n",
    "import pickle\n",
    "\n",
    "with open(f\"{data_path}/predictions/predictions_models_voting.pkl\", \"rb\") as f:\n",
    "    voting_predictions = pickle.load(f)\n",
    "\n",
    "results = get_metrics(true_values, voting_predictions, \"ensemble\", labels_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.index.name = \"ensemble\"\n",
    "results.to_csv(\"results_models_ensemble.csv\", index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform columns into rows \n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "results = pd.read_csv(\"results_models_ensemble.csv\")\n",
    "results = results.set_index(\"ensemble\")\n",
    "results = results.transpose()\n",
    "new_results = pd.DataFrame(columns=[\"model\", \"metric\", \"train\", \"test\"])\n",
    "new_results[\"model\"] = [\"models ensemble\"] * 15\n",
    "new_results[\"metric\"] = results.index\n",
    "new_results[\"train\"] = np.NaN \n",
    "new_results[\"test\"] = results[\"ensemble\"].values\n",
    "new_results.to_csv(\"results_models_ensemble2.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = pd.read_csv(\"results_ensemble_blast.csv\")\n",
    "results = results.set_index(\"ensemble\")\n",
    "results = results.transpose()\n",
    "new_results = pd.DataFrame(columns=[\"model\", \"metric\", \"train\", \"test\"])\n",
    "new_results[\"model\"] = [\"models + BLASTp ensemble\"] * 15\n",
    "new_results[\"metric\"] = results.index\n",
    "new_results[\"train\"] = np.NaN \n",
    "new_results[\"test\"] = results[\"ensemble_blast\"].values\n",
    "new_results.to_csv(\"results_ensemble_blast2.csv\", index=False)"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}