{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"/home/a03-sgoel/mESMerize/benchmarks/DeepLoc\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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>Unnamed: 0</th>\n",
       "      <th>ACC</th>\n",
       "      <th>Kingdom</th>\n",
       "      <th>Partition</th>\n",
       "      <th>Peripheral</th>\n",
       "      <th>Transmembrane</th>\n",
       "      <th>LipidAnchor</th>\n",
       "      <th>Soluble</th>\n",
       "      <th>Sequence</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>I3R9M8</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MSTDSDAETVDLADGVDHQVAMVMDLNKCIGCQTCTVACKSLWTEG...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>I3R9M9</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MSRNDASQLDDGETTAESPPDDQANDAPEVGDPPGDPVDADSGVSR...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>Q7ZAG8</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MTKVLVLGGRFGALTAAYTLKRLVGSKADVKVINKSRFSYFRPALP...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>Q8PZ67</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MPPKIAEVIQHDVCAACGACEAVCPIGAVTVKKAAEIRDPNDLSLY...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>Q9YGA6</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MAGVRLVDVWKVFGEVTAVREMSLEVKDGEFMILLGPSGCGKTTTL...</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",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28021</th>\n",
       "      <td>28021</td>\n",
       "      <td>P86949</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MLRFIAIVALIATVNAKGGTYGIGVLPSVTYVSGGGGGYPGIYGTY...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28022</th>\n",
       "      <td>28022</td>\n",
       "      <td>P86950</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MKPFISLASLIVLIASASAGGDDDYGKYGYGSYGPGIGGIGGGGGG...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28023</th>\n",
       "      <td>28023</td>\n",
       "      <td>P86951</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MLKLVCAVVLIATVNAKGSSPGFGIGQLPGITVVSGGVSGGSLSGG...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28024</th>\n",
       "      <td>28024</td>\n",
       "      <td>P86983</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MHQSSLGVLVLFSLIYLCISVHVPFDLNGWKALRLDNNRVQDSTNL...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28025</th>\n",
       "      <td>28025</td>\n",
       "      <td>P86984</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MLMLLCIIATVIPFSLVEGRKGCWADPTPPGKECLYGKEIHGGRNL...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>28026 rows × 9 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       Unnamed: 0     ACC    Kingdom  Partition  Peripheral  Transmembrane  \\\n",
       "0               0  I3R9M8    Archaea          0           1              0   \n",
       "1               1  I3R9M9    Archaea          1           1              0   \n",
       "2               2  Q7ZAG8    Archaea          2           1              0   \n",
       "3               3  Q8PZ67    Archaea          0           1              0   \n",
       "4               4  Q9YGA6    Archaea          0           1              0   \n",
       "...           ...     ...        ...        ...         ...            ...   \n",
       "28021       28021  P86949  Eukaryota          0           0              0   \n",
       "28022       28022  P86950  Eukaryota          0           0              0   \n",
       "28023       28023  P86951  Eukaryota          0           0              0   \n",
       "28024       28024  P86983  Eukaryota          3           0              0   \n",
       "28025       28025  P86984  Eukaryota          4           0              0   \n",
       "\n",
       "       LipidAnchor  Soluble                                           Sequence  \n",
       "0                0        0  MSTDSDAETVDLADGVDHQVAMVMDLNKCIGCQTCTVACKSLWTEG...  \n",
       "1                0        0  MSRNDASQLDDGETTAESPPDDQANDAPEVGDPPGDPVDADSGVSR...  \n",
       "2                0        0  MTKVLVLGGRFGALTAAYTLKRLVGSKADVKVINKSRFSYFRPALP...  \n",
       "3                0        1  MPPKIAEVIQHDVCAACGACEAVCPIGAVTVKKAAEIRDPNDLSLY...  \n",
       "4                0        0  MAGVRLVDVWKVFGEVTAVREMSLEVKDGEFMILLGPSGCGKTTTL...  \n",
       "...            ...      ...                                                ...  \n",
       "28021            0        1  MLRFIAIVALIATVNAKGGTYGIGVLPSVTYVSGGGGGYPGIYGTY...  \n",
       "28022            0        1  MKPFISLASLIVLIASASAGGDDDYGKYGYGSYGPGIGGIGGGGGG...  \n",
       "28023            0        1  MLKLVCAVVLIATVNAKGSSPGFGIGQLPGITVVSGGVSGGSLSGG...  \n",
       "28024            0        1  MHQSSLGVLVLFSLIYLCISVHVPFDLNGWKALRLDNNRVQDSTNL...  \n",
       "28025            0        1  MLMLLCIIATVIPFSLVEGRKGCWADPTPPGKECLYGKEIHGGRNL...  \n",
       "\n",
       "[28026 rows x 9 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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>ACC</th>\n",
       "      <th>Kingdom</th>\n",
       "      <th>Partition</th>\n",
       "      <th>Peripheral</th>\n",
       "      <th>Transmembrane</th>\n",
       "      <th>LipidAnchor</th>\n",
       "      <th>Soluble</th>\n",
       "      <th>Sequence</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>I3R9M8</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MSTDSDAETVDLADGVDHQVAMVMDLNKCIGCQTCTVACKSLWTEG...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>I3R9M9</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MSRNDASQLDDGETTAESPPDDQANDAPEVGDPPGDPVDADSGVSR...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Q7ZAG8</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MTKVLVLGGRFGALTAAYTLKRLVGSKADVKVINKSRFSYFRPALP...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Q8PZ67</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MPPKIAEVIQHDVCAACGACEAVCPIGAVTVKKAAEIRDPNDLSLY...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Q9YGA6</td>\n",
       "      <td>Archaea</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>MAGVRLVDVWKVFGEVTAVREMSLEVKDGEFMILLGPSGCGKTTTL...</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",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28021</th>\n",
       "      <td>P86949</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MLRFIAIVALIATVNAKGGTYGIGVLPSVTYVSGGGGGYPGIYGTY...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28022</th>\n",
       "      <td>P86950</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MKPFISLASLIVLIASASAGGDDDYGKYGYGSYGPGIGGIGGGGGG...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28023</th>\n",
       "      <td>P86951</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MLKLVCAVVLIATVNAKGSSPGFGIGQLPGITVVSGGVSGGSLSGG...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28024</th>\n",
       "      <td>P86983</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MHQSSLGVLVLFSLIYLCISVHVPFDLNGWKALRLDNNRVQDSTNL...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28025</th>\n",
       "      <td>P86984</td>\n",
       "      <td>Eukaryota</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>MLMLLCIIATVIPFSLVEGRKGCWADPTPPGKECLYGKEIHGGRNL...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>28026 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          ACC    Kingdom  Partition  Peripheral  Transmembrane  LipidAnchor  \\\n",
       "0      I3R9M8    Archaea          0           1              0            0   \n",
       "1      I3R9M9    Archaea          1           1              0            0   \n",
       "2      Q7ZAG8    Archaea          2           1              0            0   \n",
       "3      Q8PZ67    Archaea          0           1              0            0   \n",
       "4      Q9YGA6    Archaea          0           1              0            0   \n",
       "...       ...        ...        ...         ...            ...          ...   \n",
       "28021  P86949  Eukaryota          0           0              0            0   \n",
       "28022  P86950  Eukaryota          0           0              0            0   \n",
       "28023  P86951  Eukaryota          0           0              0            0   \n",
       "28024  P86983  Eukaryota          3           0              0            0   \n",
       "28025  P86984  Eukaryota          4           0              0            0   \n",
       "\n",
       "       Soluble                                           Sequence  \n",
       "0            0  MSTDSDAETVDLADGVDHQVAMVMDLNKCIGCQTCTVACKSLWTEG...  \n",
       "1            0  MSRNDASQLDDGETTAESPPDDQANDAPEVGDPPGDPVDADSGVSR...  \n",
       "2            0  MTKVLVLGGRFGALTAAYTLKRLVGSKADVKVINKSRFSYFRPALP...  \n",
       "3            1  MPPKIAEVIQHDVCAACGACEAVCPIGAVTVKKAAEIRDPNDLSLY...  \n",
       "4            0  MAGVRLVDVWKVFGEVTAVREMSLEVKDGEFMILLGPSGCGKTTTL...  \n",
       "...        ...                                                ...  \n",
       "28021        1  MLRFIAIVALIATVNAKGGTYGIGVLPSVTYVSGGGGGYPGIYGTY...  \n",
       "28022        1  MKPFISLASLIVLIASASAGGDDDYGKYGYGSYGPGIGGIGGGGGG...  \n",
       "28023        1  MLKLVCAVVLIATVNAKGSSPGFGIGQLPGITVVSGGVSGGSLSGG...  \n",
       "28024        1  MHQSSLGVLVLFSLIYLCISVHVPFDLNGWKALRLDNNRVQDSTNL...  \n",
       "28025        1  MLMLLCIIATVIPFSLVEGRKGCWADPTPPGKECLYGKEIHGGRNL...  \n",
       "\n",
       "[28026 rows x 8 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(path + \"/OG_membrane_type_all.csv\")\n",
    "df = df.drop(columns=['Unnamed: 0'])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = df[df['Partition'] != 4]\n",
    "test = df[df['Partition'] == 4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "train.to_csv(path + \"/membrane_type_train.csv\", index=False)\n",
    "test.to_csv(path + \"/membrane_type_test.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}