{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "05175c21",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-11-06 20:39:57.578385: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-11-06 20:40:01.120995: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/ohpc/pub/mpi/openmpi3-gnu7/3.1.0/lib:/opt/ohpc/pub/compiler/gcc/7.3.0/lib64\n",
      "2024-11-06 20:40:01.121086: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/ohpc/pub/mpi/openmpi3-gnu7/3.1.0/lib:/opt/ohpc/pub/compiler/gcc/7.3.0/lib64\n",
      "2024-11-06 20:40:01.121097: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
     ]
    }
   ],
   "source": [
    "import deepchem as dc\n",
    "import rdkit\n",
    "from rdkit import Chem\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy.sparse as sp\n",
    "import pickle\n",
    "\n",
    "import os\n",
    "from constants import *\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1910933f",
   "metadata": {},
   "outputs": [],
   "source": [
    "global suffix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a7d93615",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_csv(path, smipos):\n",
    "    df = pd.read_pickle(path)\n",
    "    df = df.rename(columns={smipos: \"SMILES_NS\"})\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "def preprocess_graph(\n",
    "    data,\n",
    "):  # The function is to preprocessed the adjacency matrix, returning the normalized adjacency matrix in the form of numpy array for feeding into the model\n",
    "    adj_ = data + sp.eye(data.shape[0])\n",
    "    rowsum = np.array(adj_.sum(1))\n",
    "    degree_mat_inv_sqrt = np.diag(np.power(rowsum, -0.5).flatten())\n",
    "    adj_normalized = adj_.dot(degree_mat_inv_sqrt).transpose().dot(degree_mat_inv_sqrt)\n",
    "    return np.array(adj_normalized)\n",
    "\n",
    "\n",
    "def smiles_get_features(\n",
    "    a,\n",
    "):  # This function will return the smiles code into list of feature for each atoms\n",
    "    if isinstance(a, float):\n",
    "        return np.nan\n",
    "    m = rdkit.Chem.MolFromSmiles(a)\n",
    "    featurizer = dc.feat.ConvMolFeaturizer()\n",
    "    features = featurizer.featurize([m])[0]\n",
    "    if isinstance(features, np.ndarray):\n",
    "        return np.nan\n",
    "    atom_features = features.get_atom_features()  # initial atom feature vectors\n",
    "    if atom_features.shape[0] > 60:\n",
    "        return np.nan\n",
    "    return atom_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d0963037",
   "metadata": {},
   "outputs": [],
   "source": [
    "def smiles_get_adj(a):  # This function retrieve the adjacency matrix from the molecule\n",
    "    if isinstance(a, float):\n",
    "        return np.nan\n",
    "    m = rdkit.Chem.MolFromSmiles(a)\n",
    "    featurizer = dc.feat.ConvMolFeaturizer()\n",
    "    features = featurizer.featurize([m])[0]\n",
    "    if isinstance(features, np.ndarray):\n",
    "        return np.nan\n",
    "    adj_list = features.get_adjacency_list()  # adjacency list (neighbor list)\n",
    "    adj = np.zeros(\n",
    "        (len(adj_list), len(adj_list))\n",
    "    )  # convert adjacency list into adjacency matrix \"A\"\n",
    "    if len(adj_list) > 60:\n",
    "        return np.nan\n",
    "    for i in range(len(adj_list)):\n",
    "        for j in adj_list[i]:\n",
    "            adj[i][j] = 1\n",
    "    return adj\n",
    "\n",
    "\n",
    "def smiles_get_edge(a):  # This function retrieve the adjacency matrix from the molecule\n",
    "    if isinstance(a, float):\n",
    "        return np.nan\n",
    "    m = rdkit.Chem.MolFromSmiles(a)\n",
    "    featurizer = dc.feat.ConvMolFeaturizer()\n",
    "    features = featurizer.featurize([m])[0]\n",
    "    if isinstance(features, np.ndarray):\n",
    "        return np.nan\n",
    "    adj_list = features.get_adjacency_list()  # adjacency list (neighbor list)\n",
    "    node1 = []\n",
    "    node2 = []\n",
    "    for i in range(len(adj_list)):\n",
    "        for j in adj_list[i]:\n",
    "            node1.append(i)\n",
    "            node2.append(j)\n",
    "    return np.stack((np.array(node1), np.array(node2)))\n",
    "\n",
    "\n",
    "def sim_graph(smile):\n",
    "    if isinstance(smile, float):\n",
    "        return np.nan\n",
    "    mol = rdkit.Chem.MolFromSmiles(smile)\n",
    "    if mol is None:\n",
    "        return np.nan\n",
    "    Chem.Kekulize(mol)\n",
    "    atoms = [atom.GetAtomicNum() for atom in mol.GetAtoms()]\n",
    "    am = Chem.GetAdjacencyMatrix(mol, useBO=True)\n",
    "    if len(atoms) > 60:\n",
    "        return np.nan\n",
    "    for i, atom in enumerate(atoms):\n",
    "        am[i, i] = atom\n",
    "    return am"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2f3c31de",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_max_dim(\n",
    "    d,\n",
    "):  # This function is used to find the maximum dimension the set of data contain\n",
    "    maxdim = 0\n",
    "    for i in d:\n",
    "        if i.shape[0] > maxdim:\n",
    "            maxdim = i.shape[0]\n",
    "    return maxdim\n",
    "\n",
    "\n",
    "def pad_up_to(\n",
    "    t, max_in_dims, constant_values=0\n",
    "):  # This function is used to pad the data up to a given dimension\n",
    "    s = t.shape\n",
    "    size = np.subtract(max_in_dims, s)\n",
    "    return np.pad(\n",
    "        t, ((0, size[0]), (0, size[1])), \"constant\", constant_values=constant_values\n",
    "    )\n",
    "\n",
    "\n",
    "def get_np_adj_label(\n",
    "    df,\n",
    "    path,\n",
    "    smilepos,\n",
    "    labelpos,\n",
    "    fingerpos,\n",
    "    fingername,\n",
    "    idpos,\n",
    "    pad_dim=None,\n",
    "    save=True,\n",
    "    Finger=False,\n",
    "):\n",
    "    smi = df[smilepos]\n",
    "    prelabel = df[labelpos]\n",
    "    idx = df[idpos]\n",
    "    max_dim = 60\n",
    "\n",
    "    print(\"Getting numpy files for\", len(smi), \"smiles\")\n",
    "    # print(df,smilepos)\n",
    "\n",
    "    if Finger:\n",
    "        fing = df.iloc[:, fingerpos:]\n",
    "    smi_to_adj = {}\n",
    "    smi_to_edge = {}\n",
    "    padded_adj_map = {}\n",
    "    for smiles in smi.unique():\n",
    "        pre_adj_i = smiles_get_adj(smiles)\n",
    "        smi_to_adj[smiles] = pre_adj_i\n",
    "        smi_to_edge[smiles] = smiles_get_edge(smiles)\n",
    "        if pre_adj_i is not np.nan:\n",
    "            padded_adj_map[smiles] = pad_up_to(preprocess_graph(pre_adj_i), max_dim, max_dim)\n",
    "        \n",
    "    pre_adj = smi.map(smi_to_adj)\n",
    "    edge = smi.map(smi_to_edge).rename(\"edge\")\n",
    "    if Finger:\n",
    "        adj_label = pd.concat([pre_adj, prelabel, edge, fing], axis=1, sort=False)\n",
    "    else:\n",
    "        adj_label = pd.concat([idx, pre_adj, edge, prelabel], axis=1, sort=False)\n",
    "    adj_label = adj_label[adj_label[smilepos].notna()]\n",
    "    Trueedge = list(adj_label[\"edge\"].values)\n",
    "    Truelabel = adj_label[labelpos]\n",
    "    Trueidx = adj_label[idpos]\n",
    "    print(\"ID pos:\", idpos)\n",
    "    if Finger:\n",
    "        Fingerprint = adj_label.iloc[:, fingerpos:].values\n",
    "    \n",
    "    if save:\n",
    "        print(\"Saving...\")\n",
    "        if Finger:\n",
    "            np.save(\n",
    "                path + \"/\" + fingername + \"_fingerprint\", Fingerprint, fix_imports=False\n",
    "            )\n",
    "        np.save(path + f\"/label{suffix}\", Truelabel, fix_imports=False)\n",
    "        print(Trueidx)\n",
    "        np.save(path + f\"/idx{suffix}\", Trueidx, fix_imports=False)\n",
    "        with open(path + f\"/edge{suffix}\", \"wb\") as f:\n",
    "            pickle.dump(Trueedge, f)\n",
    "        print(\"Saved\")\n",
    "    \n",
    "    print(\"Preprocessing graph...\")\n",
    "    # True_adj = pre_adj.apply(preprocess_graph)\n",
    "        \n",
    "    True_adj = smi.map(padded_adj_map).dropna()\n",
    "    True_array_ADJ = np.stack(True_adj.values)\n",
    "    print(\"Done\")      \n",
    "    if save:\n",
    "        print(\"Saving...\")\n",
    "        np.save(path + f\"/adj{suffix}\", True_array_ADJ, fix_imports=False)\n",
    "        print(\"Saved\")\n",
    "    if Finger:\n",
    "        return [True_array_ADJ, max_dim, Truelabel, Fingerprint]\n",
    "    else:\n",
    "        return [True_array_ADJ, Trueedge, max_dim, Truelabel]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b87f6437",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_feature(df, max_dim, smilepos, path, save=True):\n",
    "    smi = df[smilepos]\n",
    "    print(\"Getting features for\", len(smi), \"smiles\")\n",
    "    smi_unique = pd.Series(smi.unique())\n",
    "    true_feature_map = dict(zip(smi_unique.tolist(), smi_unique\n",
    "                              .apply(smiles_get_features)\n",
    "                              .apply(lambda x: np.nan \n",
    "                                     if x is np.nan \n",
    "                                     else pad_up_to(x, (max_dim, 75)))\n",
    "                              .tolist()))\n",
    "                        \n",
    "    # pre_feature = smi.apply(smiles_get_features).dropna()\n",
    "    # True_feature = pre_feature.apply(pad_up_to, args=((max_dim, 75),))\n",
    "    True_feature = smi.map(true_feature_map).dropna()\n",
    "    True_array_feature = np.stack(True_feature.values)\n",
    "    if save:\n",
    "        np.save(path + f\"/feature{suffix}\", True_array_feature, fix_imports=False)\n",
    "    return True_array_feature\n",
    "\n",
    "\n",
    "def get_graph(df, max_dim, smilepos, path, save=True):\n",
    "    smi = df[smilepos]\n",
    "    print(\"Getting graph files for\", len(smi), \"smiles\")\n",
    "    smi_unique = pd.Series(smi.unique())\n",
    "    true_graph_map = dict(zip(smi_unique.tolist(), smi_unique\n",
    "                              .apply(sim_graph)\n",
    "                              .apply(lambda x: np.nan \n",
    "                                     if x is np.nan \n",
    "                                     else pad_up_to(x, (max_dim, max_dim)))\n",
    "                              .tolist()))\n",
    "    true_graph = smi.map(true_graph_map).dropna()\n",
    "    # true_graph = pre_graph.apply(pad_up_to, args=((max_dim, max_dim),))\n",
    "    true_array_graph = np.stack(true_graph.values)\n",
    "    if save:\n",
    "        np.save(path + f\"/graph{suffix}\", true_array_graph, fix_imports=False)\n",
    "    return true_array_graph\n",
    "\n",
    "\n",
    "def get_smiles(df, max_dim, smilepos, path, save=True):\n",
    "    smi = df[smilepos]\n",
    "    if save:\n",
    "        np.save(path + f\"/smiles{suffix}\", smi, fix_imports=False)\n",
    "    return smi\n",
    "\n",
    "\n",
    "class GraphGenerator:\n",
    "    def __init__(self, datapath, nppath):\n",
    "        os.makedirs(nppath, exist_ok=True)\n",
    "        self.df = get_csv(datapath, f\"molecule_structures{suffix}\")\n",
    "        self.save = True\n",
    "        self.smilepos = \"SMILES_NS\"\n",
    "        self.idpos = \"id\"\n",
    "        self.fingername = \"MK\"\n",
    "        self.labelpos = [\"synergy_loewe\"]\n",
    "        self.fingerpos = 2\n",
    "\n",
    "        self.adj, self.edge, self.max_dim, self.label = get_np_adj_label(\n",
    "            self.df,\n",
    "            nppath,\n",
    "            self.smilepos,\n",
    "            self.labelpos,\n",
    "            self.fingerpos,\n",
    "            self.fingername,\n",
    "            self.idpos,\n",
    "            save=self.save,\n",
    "            Finger=False,\n",
    "        )\n",
    "\n",
    "        self.feature = get_feature(\n",
    "            self.df, self.max_dim, self.smilepos, nppath, save=self.save\n",
    "        )\n",
    "        self.graph = get_graph(\n",
    "            self.df, self.max_dim, self.smilepos, nppath, save=self.save\n",
    "        )\n",
    "        self.smiles = get_smiles(\n",
    "            self.df, self.max_dim, self.smilepos, nppath, save=self.save\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7894a0dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting numpy files for 131790 smiles\n",
      "ID pos: id\n",
      "Saving...\n",
      "0              0\n",
      "1              1\n",
      "2              2\n",
      "3              3\n",
      "4              4\n",
      "           ...  \n",
      "131782    131782\n",
      "131784    131784\n",
      "131786    131786\n",
      "131787    131787\n",
      "131789    131789\n",
      "Name: id, Length: 125241, dtype: int64\n",
      "Saved\n",
      "Preprocessing graph...\n",
      "Done\n",
      "Saving...\n",
      "Saved\n",
      "Getting features for 131790 smiles\n",
      "Getting graph files for 131790 smiles\n"
     ]
    }
   ],
   "source": [
    "suffix = \"_row\"\n",
    "graph_generator = GraphGenerator(\n",
    "    f\"{EXPORT_DATA_PATH}/data_drugcomb.pkl\", f\"{DATA_PATH}/nps\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2e475827",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ True])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = get_csv(f\"{EXPORT_DATA_PATH}/data_drugcomb.pkl\", \"molecule_structures_col\")\n",
    "np.unique(x[\"SMILES_NS\"].notna())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "693df7e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "125857\n",
      "125857\n",
      "125857\n"
     ]
    }
   ],
   "source": [
    "# All three should have the same size\n",
    "np_label = np.load(f\"{DATA_PATH}/nps/idx_col.npy\")\n",
    "print(len(np_label))\n",
    "np_label = np.load(f\"{DATA_PATH}/nps/label_col.npy\")\n",
    "print(len(np_label))\n",
    "np_label = np.load(f\"{DATA_PATH}/nps/adj_col.npy\")\n",
    "print(len(np_label))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c8c6e522",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "119600"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check intersection of row and col indices\n",
    "col_idx = np.load(f\"{DATA_PATH}/nps/idx_col.npy\")\n",
    "row_idx = np.load(f\"{DATA_PATH}/nps/idx_row.npy\")\n",
    "all_idx = np.intersect1d(col_idx, row_idx)\n",
    "len(all_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d2abc165",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(f\"{DATA_PATH}/nps/idx_intersect\", all_idx, fix_imports=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "71584b86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "119600"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "col_idx = np.load(f\"{DATA_PATH}/nps/idx_col.npy\")\n",
    "col_idx = [idx for idx, x in enumerate(col_idx) if x in all_idx]\n",
    "col_label = np.load(f\"{DATA_PATH}/nps/label_col.npy\")\n",
    "len(col_label[col_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9ef4be00",
   "metadata": {},
   "outputs": [],
   "source": [
    "row_idx = np.load(f\"{DATA_PATH}/nps/idx_row.npy\")\n",
    "row_idx = [idx for idx, x in enumerate(row_idx) if x in all_idx]\n",
    "row_label = np.load(f\"{DATA_PATH}/nps/label_row.npy\")\n",
    "row_label = row_label[row_idx]\n",
    "row_adj = np.load(f\"{DATA_PATH}/nps/adj_row.npy\")\n",
    "row_adj = row_adj[row_idx]\n",
    "row_smiles = np.load(f\"{DATA_PATH}/nps/smiles_row.npy\", allow_pickle=True)\n",
    "row_smiles = row_smiles[row_idx]\n",
    "\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/idx_row\", row_idx)\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/label_row\", row_label)\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/adj_row\", row_adj)\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/smiles_row\", row_smiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3011fa60",
   "metadata": {},
   "outputs": [],
   "source": [
    "col_idx = np.load(f\"{DATA_PATH}/nps/idx_col.npy\")\n",
    "col_idx = [idx for idx, x in enumerate(col_idx) if x in all_idx]\n",
    "col_label = np.load(f\"{DATA_PATH}/nps/label_col.npy\")\n",
    "col_label = col_label[col_idx]\n",
    "col_adj = np.load(f\"{DATA_PATH}/nps/adj_col.npy\")\n",
    "col_adj = col_adj[col_idx]\n",
    "col_smiles = np.load(f\"{DATA_PATH}/nps/smiles_col.npy\", allow_pickle=True)\n",
    "col_smiles = col_smiles[col_idx]\n",
    "\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/col_row\", col_idx)\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/label_col\", col_label)\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/adj_col\", col_adj)\n",
    "np.save(f\"{DATA_PATH}/nps_intersected/smiles_col\", col_smiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bb24eef-0a7a-430f-857c-d1dbb1855436",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
