generate_dataset_confirm.ipynb

{
 "metadata": {
  "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.5"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python385jvsc74a57bd023bc902ea7de6fb6674954bc35d9f64533aca6659ce812252d037bb8d15d2114",
   "display_name": "Python 3.8.5 64-bit ('base': conda)"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import pandas as pd\r\n",
    "import seaborn as sns\r\n",
    "import numpy as np\r\n",
    "import tensorflow as tf\r\n",
    "import scipy\r\n",
    "import sys\r\n",
    "\r\n",
    "from rdkit import Chem\r\n",
    "from rdkit.Chem.Lipinski import NumAromaticRings\r\n",
    "from rdkit.Chem import PandasTools\r\n",
    "\r\n",
    "import matplotlib.pyplot as plt\r\n",
    "\r\n",
    "%matplotlib inline"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "# 读入各个数据\r\n",
    "latent = pd.read_table('./data/2105_confirm/latent_features.txt', sep='.\\s', header=None, engine='python')\r\n",
    "smiles = pd.read_table('./data/2105_confirm/z5.txt',sep='\\n', header=None)\r\n",
    "data_all = pd.concat([smiles, latent], axis=1)\r\n",
    "col = ['smiles']\r\n",
    "name = 'jtvae_ld_'\r\n",
    "for i in range(56):\r\n",
    "    col.append(name + str(i))\r\n",
    "data_all.columns = col\r\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "data_all"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "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>smiles</th>\n",
       "      <th>jtvae_ld_0</th>\n",
       "      <th>jtvae_ld_1</th>\n",
       "      <th>jtvae_ld_2</th>\n",
       "      <th>jtvae_ld_3</th>\n",
       "      <th>jtvae_ld_4</th>\n",
       "      <th>jtvae_ld_5</th>\n",
       "      <th>jtvae_ld_6</th>\n",
       "      <th>jtvae_ld_7</th>\n",
       "      <th>jtvae_ld_8</th>\n",
       "      <th>...</th>\n",
       "      <th>jtvae_ld_46</th>\n",
       "      <th>jtvae_ld_47</th>\n",
       "      <th>jtvae_ld_48</th>\n",
       "      <th>jtvae_ld_49</th>\n",
       "      <th>jtvae_ld_50</th>\n",
       "      <th>jtvae_ld_51</th>\n",
       "      <th>jtvae_ld_52</th>\n",
       "      <th>jtvae_ld_53</th>\n",
       "      <th>jtvae_ld_54</th>\n",
       "      <th>jtvae_ld_55</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cc1c(C(=O)N[C@H](C[NH+](C)C)c2ccccc2)cnn1-c1cc...</td>\n",
       "      <td>-1.261103</td>\n",
       "      <td>1.888711</td>\n",
       "      <td>-6.242961</td>\n",
       "      <td>-6.217049</td>\n",
       "      <td>-1.963670</td>\n",
       "      <td>2.881199</td>\n",
       "      <td>-1.815423</td>\n",
       "      <td>2.818454</td>\n",
       "      <td>-5.236242</td>\n",
       "      <td>...</td>\n",
       "      <td>-3.111375</td>\n",
       "      <td>1.120334</td>\n",
       "      <td>-2.613182</td>\n",
       "      <td>2.387751</td>\n",
       "      <td>5.268334</td>\n",
       "      <td>2.352951</td>\n",
       "      <td>-3.541683</td>\n",
       "      <td>-2.524042</td>\n",
       "      <td>7.874137</td>\n",
       "      <td>0.189427</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>COc1ccc(C[NH+](C)[C@@H](C)C(=O)Nc2ccc(F)c(F)c2...</td>\n",
       "      <td>2.526336</td>\n",
       "      <td>-2.564474</td>\n",
       "      <td>-3.490833</td>\n",
       "      <td>4.425713</td>\n",
       "      <td>-2.487651</td>\n",
       "      <td>1.911674</td>\n",
       "      <td>1.005830</td>\n",
       "      <td>-2.580270</td>\n",
       "      <td>-3.847111</td>\n",
       "      <td>...</td>\n",
       "      <td>-4.105083</td>\n",
       "      <td>2.813607</td>\n",
       "      <td>4.779819</td>\n",
       "      <td>1.010575</td>\n",
       "      <td>6.866753</td>\n",
       "      <td>-1.441126</td>\n",
       "      <td>-4.302085</td>\n",
       "      <td>-5.931015</td>\n",
       "      <td>-1.912545</td>\n",
       "      <td>-0.369789</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CC[NH+](CCS(C)(=O)=O)Cc1ccc(-c2ccc(Cl)cc2)o1</td>\n",
       "      <td>1.788244</td>\n",
       "      <td>-9.902813</td>\n",
       "      <td>-2.343237</td>\n",
       "      <td>2.489923</td>\n",
       "      <td>2.706184</td>\n",
       "      <td>3.073043</td>\n",
       "      <td>2.884255</td>\n",
       "      <td>-2.205066</td>\n",
       "      <td>7.632130</td>\n",
       "      <td>...</td>\n",
       "      <td>3.991690</td>\n",
       "      <td>2.795439</td>\n",
       "      <td>3.964744</td>\n",
       "      <td>5.913845</td>\n",
       "      <td>6.573098</td>\n",
       "      <td>-2.178797</td>\n",
       "      <td>7.590621</td>\n",
       "      <td>-8.152710</td>\n",
       "      <td>7.147285</td>\n",
       "      <td>-0.221404</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CN1C(=O)/C(=C(\\Sc2nnc3c(n2)[nH]c2ccccc23)C(=O)...</td>\n",
       "      <td>-1.428424</td>\n",
       "      <td>2.735042</td>\n",
       "      <td>-4.686931</td>\n",
       "      <td>4.330184</td>\n",
       "      <td>2.785543</td>\n",
       "      <td>-6.069830</td>\n",
       "      <td>6.248152</td>\n",
       "      <td>1.152485</td>\n",
       "      <td>-1.735201</td>\n",
       "      <td>...</td>\n",
       "      <td>1.654738</td>\n",
       "      <td>-2.257859</td>\n",
       "      <td>-3.779827</td>\n",
       "      <td>5.830238</td>\n",
       "      <td>3.428591</td>\n",
       "      <td>2.706142</td>\n",
       "      <td>3.741402</td>\n",
       "      <td>3.237207</td>\n",
       "      <td>1.102386</td>\n",
       "      <td>-0.089020</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>COc1cccc(C[NH2+]C[C@H]2CCN(CC(F)(F)F)C2)c1</td>\n",
       "      <td>4.793871</td>\n",
       "      <td>-2.767154</td>\n",
       "      <td>-2.875158</td>\n",
       "      <td>2.710938</td>\n",
       "      <td>-1.471237</td>\n",
       "      <td>9.194351</td>\n",
       "      <td>-1.084160</td>\n",
       "      <td>-1.066807</td>\n",
       "      <td>2.845645</td>\n",
       "      <td>...</td>\n",
       "      <td>-5.026968</td>\n",
       "      <td>5.103010</td>\n",
       "      <td>8.291696</td>\n",
       "      <td>4.514052</td>\n",
       "      <td>2.214745</td>\n",
       "      <td>-3.881510</td>\n",
       "      <td>6.646420</td>\n",
       "      <td>-5.560030</td>\n",
       "      <td>-1.020045</td>\n",
       "      <td>0.095746</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>49995</th>\n",
       "      <td>CSc1ccccc1NC(=O)[C@H]1CC[C@@H]2CCCC[C@@H]2[NH2+]1</td>\n",
       "      <td>-1.071066</td>\n",
       "      <td>1.715748</td>\n",
       "      <td>1.629091</td>\n",
       "      <td>3.196286</td>\n",
       "      <td>-2.792036</td>\n",
       "      <td>-5.045714</td>\n",
       "      <td>-9.558347</td>\n",
       "      <td>-3.738821</td>\n",
       "      <td>-1.076215</td>\n",
       "      <td>...</td>\n",
       "      <td>1.651217</td>\n",
       "      <td>2.902797</td>\n",
       "      <td>-4.159595</td>\n",
       "      <td>9.998109</td>\n",
       "      <td>-1.511694</td>\n",
       "      <td>2.492025</td>\n",
       "      <td>4.853007</td>\n",
       "      <td>2.112207</td>\n",
       "      <td>3.294874</td>\n",
       "      <td>0.295081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49996</th>\n",
       "      <td>Cc1cc([C@]2(O)CS[C@H](C)C2)ccc1F</td>\n",
       "      <td>5.500004</td>\n",
       "      <td>-4.974154</td>\n",
       "      <td>-1.010477</td>\n",
       "      <td>1.012926</td>\n",
       "      <td>-2.121687</td>\n",
       "      <td>-2.464846</td>\n",
       "      <td>8.930106</td>\n",
       "      <td>-4.507516</td>\n",
       "      <td>-6.026894</td>\n",
       "      <td>...</td>\n",
       "      <td>3.552136</td>\n",
       "      <td>4.714711</td>\n",
       "      <td>-3.645862</td>\n",
       "      <td>-3.396438</td>\n",
       "      <td>6.527954</td>\n",
       "      <td>7.050245</td>\n",
       "      <td>9.533513</td>\n",
       "      <td>-7.430381</td>\n",
       "      <td>1.574888</td>\n",
       "      <td>0.413131</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49997</th>\n",
       "      <td>CN(CC1CC1)C(=O)NCC(C)(C)c1ccc(Cl)cc1</td>\n",
       "      <td>3.877814</td>\n",
       "      <td>-1.258259</td>\n",
       "      <td>-2.414311</td>\n",
       "      <td>-9.483868</td>\n",
       "      <td>-9.757860</td>\n",
       "      <td>4.956900</td>\n",
       "      <td>1.118094</td>\n",
       "      <td>1.629370</td>\n",
       "      <td>-9.229904</td>\n",
       "      <td>...</td>\n",
       "      <td>-2.426497</td>\n",
       "      <td>-3.229744</td>\n",
       "      <td>1.192198</td>\n",
       "      <td>5.310479</td>\n",
       "      <td>6.923873</td>\n",
       "      <td>-4.689966</td>\n",
       "      <td>1.704111</td>\n",
       "      <td>-4.871596</td>\n",
       "      <td>-3.132816</td>\n",
       "      <td>0.115074</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49998</th>\n",
       "      <td>O=C([O-])CCc1ccc(S(=O)(=O)Nc2cccc(I)c2)cc1</td>\n",
       "      <td>1.091118</td>\n",
       "      <td>-4.024869</td>\n",
       "      <td>-1.806850</td>\n",
       "      <td>5.785818</td>\n",
       "      <td>-3.052545</td>\n",
       "      <td>4.511624</td>\n",
       "      <td>-1.724618</td>\n",
       "      <td>-1.059059</td>\n",
       "      <td>4.428439</td>\n",
       "      <td>...</td>\n",
       "      <td>-2.388021</td>\n",
       "      <td>-1.569831</td>\n",
       "      <td>-1.008806</td>\n",
       "      <td>2.462220</td>\n",
       "      <td>2.819120</td>\n",
       "      <td>-4.666101</td>\n",
       "      <td>5.199871</td>\n",
       "      <td>-2.651287</td>\n",
       "      <td>-4.078031</td>\n",
       "      <td>-0.273533</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49999</th>\n",
       "      <td>C[C@H]1CCN(C(=O)c2ccc(C(=O)OC(C)(C)C)cc2)C[C@H...</td>\n",
       "      <td>1.491479</td>\n",
       "      <td>3.813991</td>\n",
       "      <td>-4.410887</td>\n",
       "      <td>-2.458951</td>\n",
       "      <td>-7.702771</td>\n",
       "      <td>-3.231463</td>\n",
       "      <td>9.329567</td>\n",
       "      <td>4.035351</td>\n",
       "      <td>-1.499851</td>\n",
       "      <td>...</td>\n",
       "      <td>-3.168039</td>\n",
       "      <td>3.012530</td>\n",
       "      <td>4.459123</td>\n",
       "      <td>9.976082</td>\n",
       "      <td>4.769394</td>\n",
       "      <td>-3.395801</td>\n",
       "      <td>-4.322295</td>\n",
       "      <td>-2.730070</td>\n",
       "      <td>-2.213073</td>\n",
       "      <td>0.352081</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>50000 rows × 57 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  smiles  jtvae_ld_0  \\\n",
       "0      Cc1c(C(=O)N[C@H](C[NH+](C)C)c2ccccc2)cnn1-c1cc...   -1.261103   \n",
       "1      COc1ccc(C[NH+](C)[C@@H](C)C(=O)Nc2ccc(F)c(F)c2...    2.526336   \n",
       "2           CC[NH+](CCS(C)(=O)=O)Cc1ccc(-c2ccc(Cl)cc2)o1    1.788244   \n",
       "3      CN1C(=O)/C(=C(\\Sc2nnc3c(n2)[nH]c2ccccc23)C(=O)...   -1.428424   \n",
       "4             COc1cccc(C[NH2+]C[C@H]2CCN(CC(F)(F)F)C2)c1    4.793871   \n",
       "...                                                  ...         ...   \n",
       "49995  CSc1ccccc1NC(=O)[C@H]1CC[C@@H]2CCCC[C@@H]2[NH2+]1   -1.071066   \n",
       "49996                   Cc1cc([C@]2(O)CS[C@H](C)C2)ccc1F    5.500004   \n",
       "49997               CN(CC1CC1)C(=O)NCC(C)(C)c1ccc(Cl)cc1    3.877814   \n",
       "49998         O=C([O-])CCc1ccc(S(=O)(=O)Nc2cccc(I)c2)cc1    1.091118   \n",
       "49999  C[C@H]1CCN(C(=O)c2ccc(C(=O)OC(C)(C)C)cc2)C[C@H...    1.491479   \n",
       "\n",
       "       jtvae_ld_1  jtvae_ld_2  jtvae_ld_3  jtvae_ld_4  jtvae_ld_5  jtvae_ld_6  \\\n",
       "0        1.888711   -6.242961   -6.217049   -1.963670    2.881199   -1.815423   \n",
       "1       -2.564474   -3.490833    4.425713   -2.487651    1.911674    1.005830   \n",
       "2       -9.902813   -2.343237    2.489923    2.706184    3.073043    2.884255   \n",
       "3        2.735042   -4.686931    4.330184    2.785543   -6.069830    6.248152   \n",
       "4       -2.767154   -2.875158    2.710938   -1.471237    9.194351   -1.084160   \n",
       "...           ...         ...         ...         ...         ...         ...   \n",
       "49995    1.715748    1.629091    3.196286   -2.792036   -5.045714   -9.558347   \n",
       "49996   -4.974154   -1.010477    1.012926   -2.121687   -2.464846    8.930106   \n",
       "49997   -1.258259   -2.414311   -9.483868   -9.757860    4.956900    1.118094   \n",
       "49998   -4.024869   -1.806850    5.785818   -3.052545    4.511624   -1.724618   \n",
       "49999    3.813991   -4.410887   -2.458951   -7.702771   -3.231463    9.329567   \n",
       "\n",
       "       jtvae_ld_7  jtvae_ld_8  ...  jtvae_ld_46  jtvae_ld_47  jtvae_ld_48  \\\n",
       "0        2.818454   -5.236242  ...    -3.111375     1.120334    -2.613182   \n",
       "1       -2.580270   -3.847111  ...    -4.105083     2.813607     4.779819   \n",
       "2       -2.205066    7.632130  ...     3.991690     2.795439     3.964744   \n",
       "3        1.152485   -1.735201  ...     1.654738    -2.257859    -3.779827   \n",
       "4       -1.066807    2.845645  ...    -5.026968     5.103010     8.291696   \n",
       "...           ...         ...  ...          ...          ...          ...   \n",
       "49995   -3.738821   -1.076215  ...     1.651217     2.902797    -4.159595   \n",
       "49996   -4.507516   -6.026894  ...     3.552136     4.714711    -3.645862   \n",
       "49997    1.629370   -9.229904  ...    -2.426497    -3.229744     1.192198   \n",
       "49998   -1.059059    4.428439  ...    -2.388021    -1.569831    -1.008806   \n",
       "49999    4.035351   -1.499851  ...    -3.168039     3.012530     4.459123   \n",
       "\n",
       "       jtvae_ld_49  jtvae_ld_50  jtvae_ld_51  jtvae_ld_52  jtvae_ld_53  \\\n",
       "0         2.387751     5.268334     2.352951    -3.541683    -2.524042   \n",
       "1         1.010575     6.866753    -1.441126    -4.302085    -5.931015   \n",
       "2         5.913845     6.573098    -2.178797     7.590621    -8.152710   \n",
       "3         5.830238     3.428591     2.706142     3.741402     3.237207   \n",
       "4         4.514052     2.214745    -3.881510     6.646420    -5.560030   \n",
       "...            ...          ...          ...          ...          ...   \n",
       "49995     9.998109    -1.511694     2.492025     4.853007     2.112207   \n",
       "49996    -3.396438     6.527954     7.050245     9.533513    -7.430381   \n",
       "49997     5.310479     6.923873    -4.689966     1.704111    -4.871596   \n",
       "49998     2.462220     2.819120    -4.666101     5.199871    -2.651287   \n",
       "49999     9.976082     4.769394    -3.395801    -4.322295    -2.730070   \n",
       "\n",
       "       jtvae_ld_54  jtvae_ld_55  \n",
       "0         7.874137     0.189427  \n",
       "1        -1.912545    -0.369789  \n",
       "2         7.147285    -0.221404  \n",
       "3         1.102386    -0.089020  \n",
       "4        -1.020045     0.095746  \n",
       "...            ...          ...  \n",
       "49995     3.294874     0.295081  \n",
       "49996     1.574888     0.413131  \n",
       "49997    -3.132816     0.115074  \n",
       "49998    -4.078031    -0.273533  \n",
       "49999    -2.213073     0.352081  \n",
       "\n",
       "[50000 rows x 57 columns]"
      ]
     },
     "metadata": {},
     "execution_count": 3
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "def get_aroma(smiles):\r\n",
    "    mol = Chem.MolFromSmiles(smiles)\r\n",
    "    return NumAromaticRings(mol)\r\n",
    "\r\n",
    "data_all.set_index(['smiles'], inplace=True)\r\n",
    "data_index = pd.DataFrame(smiles)\r\n",
    "data_index.columns = ['smiles']\r\n",
    "data_index['aroma'] = data_index['smiles'].apply(lambda x: get_aroma(x))\r\n",
    "data_index.set_index(['smiles'], inplace=True)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "data_index"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "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>aroma</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>smiles</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Cc1c(C(=O)N[C@H](C[NH+](C)C)c2ccccc2)cnn1-c1ccccc1</th>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>COc1ccc(C[NH+](C)[C@@H](C)C(=O)Nc2ccc(F)c(F)c2F)cc1</th>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CC[NH+](CCS(C)(=O)=O)Cc1ccc(-c2ccc(Cl)cc2)o1</th>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CN1C(=O)/C(=C(\\Sc2nnc3c(n2)[nH]c2ccccc23)C(=O)[O-])c2ccccc21</th>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>COc1cccc(C[NH2+]C[C@H]2CCN(CC(F)(F)F)C2)c1</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CSc1ccccc1NC(=O)[C@H]1CC[C@@H]2CCCC[C@@H]2[NH2+]1</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Cc1cc([C@]2(O)CS[C@H](C)C2)ccc1F</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CN(CC1CC1)C(=O)NCC(C)(C)c1ccc(Cl)cc1</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>O=C([O-])CCc1ccc(S(=O)(=O)Nc2cccc(I)c2)cc1</th>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>C[C@H]1CCN(C(=O)c2ccc(C(=O)OC(C)(C)C)cc2)C[C@H]1n1cc[nH+]c1</th>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>50000 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    aroma\n",
       "smiles                                                   \n",
       "Cc1c(C(=O)N[C@H](C[NH+](C)C)c2ccccc2)cnn1-c1ccccc1      3\n",
       "COc1ccc(C[NH+](C)[C@@H](C)C(=O)Nc2ccc(F)c(F)c2F...      2\n",
       "CC[NH+](CCS(C)(=O)=O)Cc1ccc(-c2ccc(Cl)cc2)o1            2\n",
       "CN1C(=O)/C(=C(\\Sc2nnc3c(n2)[nH]c2ccccc23)C(=O)[...      4\n",
       "COc1cccc(C[NH2+]C[C@H]2CCN(CC(F)(F)F)C2)c1              1\n",
       "...                                                   ...\n",
       "CSc1ccccc1NC(=O)[C@H]1CC[C@@H]2CCCC[C@@H]2[NH2+]1       1\n",
       "Cc1cc([C@]2(O)CS[C@H](C)C2)ccc1F                        1\n",
       "CN(CC1CC1)C(=O)NCC(C)(C)c1ccc(Cl)cc1                    1\n",
       "O=C([O-])CCc1ccc(S(=O)(=O)Nc2cccc(I)c2)cc1              2\n",
       "C[C@H]1CCN(C(=O)c2ccc(C(=O)OC(C)(C)C)cc2)C[C@H]...      2\n",
       "\n",
       "[50000 rows x 1 columns]"
      ]
     },
     "metadata": {},
     "execution_count": 5
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "# set A = 2 rings, set B = 1 or 3 rings\r\n",
    "mols_to_stay = np.logical_or.reduce((data_index['aroma'].values == 1, data_index['aroma'].values == 2, data_index['aroma'].values == 3))\r\n",
    "df_experiment = data_index.iloc[mols_to_stay]\r\n",
    "print(\"Number of molecules with 1/2/3 aromatic rings: %d\" % df_experiment.shape[0])"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Number of molecules with 1/2/3 aromatic rings: 44411\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "# Split molecules into dataset A and B\r\n",
    "A_indices = df_experiment['aroma'].values == 2\r\n",
    "\r\n",
    "df_A = df_experiment.iloc[A_indices]\r\n",
    "df_B = df_experiment.iloc[~A_indices]\r\n",
    "\r\n",
    "print(\"Number of molecules in dataset A: %d\" % df_A.shape[0])\r\n",
    "print(\"Number of molecules in dataset B: %d\" % df_B.shape[0])"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Number of molecules in dataset A: 19762\n",
      "Number of molecules in dataset B: 24649\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "# Sanity check of datasets\r\n",
    "assert set(df_A['aroma'].values) == set([2])\r\n",
    "assert set(df_B['aroma'].values) == set([1,3])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "# Split datasets into train and test datasets.\r\n",
    "train_size = 12000\r\n",
    "\r\n",
    "# split set A to train/test\r\n",
    "train_sample_A = np.zeros(df_A.shape[0], dtype=bool)\r\n",
    "train_sample_A[np.random.choice(df_A.shape[0], train_size, replace=False)] = True\r\n",
    "\r\n",
    "df_train_A = df_A.iloc[train_sample_A]\r\n",
    "df_test_A = df_A.iloc[~train_sample_A]\r\n",
    "print(\"Number of molecules in dataset train_A: %d\" % df_train_A.shape[0])\r\n",
    "print(\"Number of molecules in dataset test_A: %d\" % df_test_A.shape[0])\r\n",
    "\r\n",
    "# split set B to train/test\r\n",
    "train_sample_B = np.zeros(df_B.shape[0], dtype=bool)\r\n",
    "train_sample_B[np.random.choice(df_B.shape[0], train_size, replace=False)] = True\r\n",
    "\r\n",
    "df_train_B = df_B.iloc[train_sample_B]\r\n",
    "df_test_B = df_B.iloc[~train_sample_B]\r\n",
    "print(\"Number of molecules in dataset train_B: %d\" % df_train_B.shape[0])\r\n",
    "print(\"Number of molecules in dataset test_B: %d\" % df_test_B.shape[0])"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Number of molecules in dataset train_A: 12000\n",
      "Number of molecules in dataset test_A: 7762\n",
      "Number of molecules in dataset train_B: 12000\n",
      "Number of molecules in dataset test_B: 12649\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "source": [
    "X_JTVAE_zinc_train_A = data_all.loc[df_train_A.index]\r\n",
    "X_JTVAE_zinc_test_A = data_all.loc[df_test_A.index]\r\n",
    "\r\n",
    "X_JTVAE_zinc_train_B = data_all.loc[df_train_B.index]\r\n",
    "X_JTVAE_zinc_test_B = data_all.loc[df_test_B.index]\r\n",
    "\r\n",
    "print(X_JTVAE_zinc_train_A.shape, X_JTVAE_zinc_test_A.shape, \r\n",
    "      X_JTVAE_zinc_train_B.shape, X_JTVAE_zinc_test_B.shape)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(12000, 56) (7762, 56) (12000, 56) (12649, 56)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "X_JTVAE_zinc_train_A.to_csv('./data/2105_confirm/X_JTVAE_zinc_train_A.csv')\r\n",
    "X_JTVAE_zinc_test_A.to_csv('./data/2105_confirm/X_JTVAE_zinc_test_A.csv')\r\n",
    "X_JTVAE_zinc_train_B.to_csv('./data/2105_confirm/X_JTVAE_zinc_train_B.csv')\r\n",
    "X_JTVAE_zinc_test_B.to_csv('./data/2105_confirm/X_JTVAE_zinc_test_B.csv')"
   ],
   "outputs": [],
   "metadata": {}
  }
 ]
}