test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0 环境配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "import pickle\n",
    "import timeit\n",
    "import random\n",
    "import argparse\n",
    "import cairosvg\n",
    "import statistics\n",
    "from statistics import mean, median\n",
    "from importlib import reload\n",
    "from joblib import Parallel, delayed\n",
    "from collections import Counter\n",
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_path = os.path.abspath('')\n",
    "sys.path.append(base_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e1bf55cc83d04f6a9290fde6c5cf693a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import oddt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import nglview as nv\n",
    "import matplotlib.pyplot as plt\n",
    "from openbabel import openbabel\n",
    "from icecream import ic\n",
    "from scipy import stats\n",
    "from sklearn import metrics\n",
    "from sklearn.manifold import TSNE\n",
    "from sklearn.decomposition import PCA\n",
    "from rdkit import Chem\n",
    "from rdkit import RDLogger   \n",
    "from rdkit.Chem import AllChem, Draw, rdMolAlign\n",
    "from rdkit.Chem.MolStandardize import rdMolStandardize\n",
    "from rdkit.Chem.Draw import DrawingOptions\n",
    "from espsim.helpers import mlCharges\n",
    "from pandarallel import pandarallel\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import utils.functions as fn\n",
    "import utils.alignment as al\n",
    "import utils.similarity as sm\n",
    "import utils.metrics as mt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.style.use('ggplot')\n",
    "# plt.rcParams[\"font.sans-serif\"]=[\"SimHei\"] #设置字体\n",
    "plt.rcParams[\"axes.unicode_minus\"]=False #该语句解决图像中的“-”负号的乱码问题\n",
    "\n",
    "sns.set_theme()\n",
    "\n",
    "RDLogger.EnableLog('rdApp.*') \n",
    "\n",
    "random.seed = 2022"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reload(fn)\n",
    "reload(al)\n",
    "reload(sm)\n",
    "reload(mt)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1 测试"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Mol2MolSupplier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1 Mol2MolSupplier\n",
    "database=fn.Mol2MolSupplier(os.path.join(base_path, 'demo-data/pde5a/actives_final.mol2'), sanitize=True)\n",
    "len(database)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 DUD-E Correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 auc_dude\n",
    "auc_dude = 'ROCSComboscore 0.598 ± 0.152 0.681 ± 0.166 0.674 ± 0.115 0.727 ± 0.128 ROCSColorscore 0.620 ± 0.139 0.712 ± 0.159 0.677 ± 0.117 0.752 ± 0.136 ROCSShapeTanimoto 0.547 ± 0.138 0.611 ± 0.163 0.618 ± 0.105 0.667 ± 0.127 Phase Shape_Mmod 0.677 ± 0.143 0.686 ± 0.145 0.772 ± 0.105 0.769 ± 0.108 Phase Shape_Ele 0.674 ± 0.153 0.688 ± 0.158 0.753 ± 0.105 0.750 ± 0.111 Phase Shape_Pharm 0.692 ± 0.160 0.694 ± 0.168 0.761 ± 0.145 0.767 ± 0.143 Shape-it 0.541 ± 0.133 0.590 ± 0.141 0.612 ± 0.110 0.639 ± 0.115 Align-it 0.659 ± 0.137 0.680 ± 0.157 0.729 ± 0.132 0.746 ± 0.139 ShaEPbest 0.658 ± 0.122 0.660 ± 0.139 0.709 ± 0.099 0.699 ± 0.120 ShaEPshape 0.625 ± 0.139 0.632 ± 0.143 0.681 ± 0.105 0.676 ± 0.116 ShaEPESP 0.606 ± 0.109 0.591 ± 0.117 0.627 ± 0.105 0.585 ± 0.125 SHAFTS 0.733 ± 0.144 0.731 ± 0.157 0.792 ± 0.129 0.782 ± 0.135 WEGA 0.645 ± 0.143 0.659 ± 0.154 0.716 ± 0.107 0.716 ± 0.125 LIGSIFT 0.718 ± 0.133 0.755 ± 0.143 0.758 ± 0.117 0.784 ± 0.120 LS-align 0.699 ± 0.126 0.759 ± 0.119 0.773 ± 0.098 0.786 ± 0.096'\n",
    "auc_dude_list = auc_dude.split(' ')\n",
    "while 'Phase' in auc_dude_list:\n",
    "    auc_dude_list.remove('Phase')\n",
    "\n",
    "auc_dude_dict = []\n",
    "for i in range(len(auc_dude_list)//13):\n",
    "    auc_dude_dict.append([auc_dude_list[i*13]]+[float(ele) for ele in auc_dude_list[i*13+1:i*13+13:3]])\n",
    "\n",
    "df = pd.DataFrame(auc_dude_dict)\n",
    "df = df.sort_values(by=[1])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in [2, 3, 4]:\n",
    "    pe, pv = stats.pearsonr(df[1].to_numpy(),df[i].to_numpy())\n",
    "    ic(f'pe: {pe:>10}, pv: {pv:>10}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pccs = np.corrcoef([df[i].to_numpy() for i in [1,2,3,4]])\n",
    "ic(pccs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.heatmap(pccs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5, 5), dpi=200)\n",
    "for i in [1,2,3,4]:\n",
    "    plt.scatter(list(range(df.shape[0])), df[i].tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 去重"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# table = Chem.GetPeriodicTable()\n",
    "\n",
    "# # for ele in [14, 16, 19]:\n",
    "# #     ligand.GetAtomWithIdx(ele).SetNumExplicitHs(1)\n",
    "# ligand = Chem.RemoveAllHs(ligand, sanitize=True)\n",
    "# ligand.RemoveAllConformers()\n",
    "# for atom in ligand.GetAtoms():\n",
    "#     if atom.GetTotalValence() != table.GetDefaultValence(atom.GetAtomicNum()):\n",
    "#         atom.SetFormalCharge(atom.GetTotalValence()-table.GetDefaultValence(atom.GetAtomMapNum()))\n",
    "# Chem.SanitizeMol(ligand)\n",
    "# ligand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = Chem.MolFromMol2File(os.path.join(base_path, f'demo-data/pde5a/crystal_ligand.mol2'))\n",
    "m2 = Chem.MolFromSmiles('CCCc1nn(c2c1nc([nH]c2=O)c1cc(ccc1OCC)S(=O)(=O)n1ccn(cc1)C)C')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1.RemoveAllConformers()\n",
    "m1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ic(Chem.MolToInchiKey(m1))\n",
    "ic(Chem.MolToInchiKey(m2))\n",
    "ic(Chem.MolToSmiles(m1))\n",
    "ic(Chem.MolToSmiles(m2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.4 绘图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parser = argparse.ArgumentParser( 'smiles to png inmage' )\n",
    "# parser.add_argument( 'smiles' )\n",
    "# parser.add_argument( '--filename', default=\"mol.\" )\n",
    " \n",
    "# parser.add_argument( 'smiles' )\n",
    " \n",
    "# param = parser.parse_args()\n",
    "# smiles = param.smiles\n",
    "# fname = param.filename\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_mol(mol, save_dir, fname):\n",
    "    DrawingOptions.atomLabelFontSize = 55\n",
    "    DrawingOptions.dotsPerAngstrom = 100\n",
    "    DrawingOptions.bondLineWidth = 3.0\n",
    "\n",
    "    Draw.MolToFile(mol, os.path.join(save_dir, fname+\".png\"))\n",
    "    Draw.MolToFile(mol, os.path.join(save_dir, fname+'.svg'))\n",
    "    cairosvg.svg2png(url=os.path.join(save_dir, fname+'.svg'), write_to= os.path.join(save_dir, \"svg_\"+fname+\".png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_mol(m1, '/home/jovyan/work-home/DUD-E/src/img', 'pde5a_cry_lig')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 DUD-E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "dude_dir = 'dude/all'\n",
    "n_jobs = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_list = ['pde5a', 'akt1', 'ada', 'andr', 'def', 'gria2', 'egfr', 'gcr', 'igf1r']\n",
    "target_flag = 6\n",
    "target_list[target_flag]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gz_a_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'actives_final.sdf.gz')\n",
    "gz_d_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'decoys_final.sdf.gz')\n",
    "for gz_path in [gz_a_path, gz_d_path]:\n",
    "    if os.path.exists(gz_path):\n",
    "        os.system(f\"gzip -d {gz_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# decoys_path \n",
    "d_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'decoys_final.sdf')\n",
    "# actives_path \n",
    "a_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'actives_final.sdf')\n",
    "# target_rec = os.path.join(base_path, f'demo-data/{target_list[target_flag]}/receptor_FH.pdb')\n",
    "# ligand_path \n",
    "l_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'crystal_ligand.mol2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.0 数据去重"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# act_405 = [mol for mol in act if mol.title=='CHEMBL405920']\n",
    "act_405_path = os.path.join(base_path, dude_dir, 'ada', 'actives_CHE405.mol2')\n",
    "# w = oddt.toolkits.rdk.Outputfile('mol2', act_405_path, overwrite=True)\n",
    "# for m in act_405:\n",
    "#     m.Mol = ''\n",
    "#     w.write(m)\n",
    "# w.close()\n",
    "act_405 = fn.Mol2MolSupplier(act_405_path, sanitize=True)\n",
    "Draw.MolsToGridImage(act_405, molsPerRow=8, subImgSize=(150,150), legends=[mol.GetProp('_Name') for mol in act_405],maxMols=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# act = fn.Mol2MolSupplier(a_path.split('.')[0]+'.mol2', sanitize=True)\n",
    "# dec = fn.Mol2MolSupplier(d_path.split('.')[0]+'.mol2', sanitize=True)\n",
    "act = [m for m in Chem.SDMolSupplier(a_path, removeHs=False) if m]\n",
    "dec = [m for m in Chem.SDMolSupplier(d_path, removeHs=False) if m]\n",
    "\n",
    "ic(len(act))\n",
    "ic(len(dec))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def de_duplicate(m_list):\n",
    "    name_l = []\n",
    "    inck_l = []\n",
    "    de_m_l = []\n",
    "\n",
    "    for m in m_list:\n",
    "        name = m.GetProp('_Name')\n",
    "        inck = Chem.MolToInchiKey(m)\n",
    "        if name in name_l:\n",
    "            # ic('name duplicate')\n",
    "            pass\n",
    "        elif inck in inck_l:\n",
    "            # ic('icnk duplicate')\n",
    "            pass\n",
    "        else:\n",
    "            name_l.append(name)\n",
    "            inck_l.append(inck)\n",
    "            de_m_l.append(m)\n",
    "    \n",
    "    return de_m_l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "act_dd = de_duplicate(act)\n",
    "ic(len(act_dd))\n",
    "with Chem.SDWriter(os.path.join(base_path, dude_dir, target_list[target_flag], 'actives_dd.sdf')) as w:\n",
    "    for m in act_dd:\n",
    "        w.write(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_dd = de_duplicate(dec)\n",
    "ic(len(dec_dd))\n",
    "with Chem.SDWriter(os.path.join(base_path, dude_dir, target_list[target_flag], 'decoys_dd.sdf')) as w:\n",
    "    for m in dec_dd:\n",
    "        w.write(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'actives_dd.sdf')\n",
    "d_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'decoys_dd.sdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sdf2mol2(sdf_path, mol_path):\n",
    "    sdf_l = [mol for mol in list(oddt.toolkit.readfile('sdf', sdf_path)) if mol]\n",
    "    w = oddt.toolkits.rdk.Outputfile('mol2', mol_path, overwrite=True)\n",
    "    for m in sdf_l:\n",
    "        m.Mol = ''\n",
    "        w.write(m)\n",
    "    w.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sdf2mol2(a_path, a_path.split('.')[0]+'.mol2')\n",
    "sdf2mol2(d_path, d_path.split('.')[0]+'.mol2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_l = [mol for mol in list(oddt.toolkit.readfile('mol2',a_path.split('.')[0]+'.mol2')) if mol]\n",
    "len(t_l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# m = AllChem.MolFromSmiles('c1ccc(cc1)COC(=O)c1cc2c(c(c1)CCN1C=NC3=C1NC=[NH]C[C@@H]3O)CCCC2')\n",
    "# type(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 数据加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'actives_dd.sdf')\n",
    "d_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'decoys_dd.sdf')\n",
    "l_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'crystal_ligand.mol2')\n",
    "ic(a_path)\n",
    "\n",
    "# decoys = fn.Mol2MolSupplier(d_path, sanitize=True)\n",
    "decoys = [m for m in Chem.SDMolSupplier(d_path, removeHs=False) if m]\n",
    "ic(len(decoys))\n",
    "Draw.MolsToGridImage(random.sample(decoys, 7), molsPerRow=7, subImgSize=(150,150), legends=[mol.GetProp('_Name') for mol in decoys[:7]],maxMols=100)\n",
    "\n",
    "actives = [m for m in Chem.SDMolSupplier(a_path, removeHs=False) if m]\n",
    "ic(len(actives))\n",
    "Draw.MolsToGridImage(random.sample(actives, 7), molsPerRow=7, subImgSize=(150,150), legends=[mol.GetProp('_Name') for mol in actives[:7]],maxMols=100)\n",
    "\n",
    "ligand = AllChem.MolFromMol2File(l_path, removeHs=False)\n",
    "ligand = Chem.AddHs(ligand, addCoords=True)\n",
    "ligand"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 RDKIT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1.1 测试分子对齐"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prbs, actives_o3a_score = al.rdkit_o3a(actives[10:20], ligand, n_jobs=32)\n",
    "statistics.mean(actives_o3a_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = fn.show_ligands([ligand]+prbs)\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives_crippeno3a_score = al.rdkit_crippeno3a(actives[:10], ligand)\n",
    "statistics.mean(actives_crippeno3a_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = fn.show_ligands([ligand]+actives[:10])\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rmsd = rdkit_alignmol(actives[:10], ligand)\n",
    "# statistics.mean(rmsd)\n",
    "# # RuntimeError: No sub-structure match found between the probe and query mol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoys_o3a_score = al.rdkit_o3a(decoys[:10], ligand)\n",
    "statistics.mean(decoys_o3a_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = fn.show_ligands([ligand]+decoys[:10])\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoys_crippeno3a_score = al.rdkit_crippeno3a(decoys[:10], ligand)\n",
    "statistics.mean(decoys_crippeno3a_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = fn.show_ligands([ligand]+decoys[:10])\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prb_crippen = [Chem.rdMolDescriptors._CalcCrippenContribs(mol) for mol in actives[:10]]\n",
    "ref_crippen = Chem.rdMolDescriptors._CalcCrippenContribs(ligand)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives_crippeno3a_score = al.rdkit_crippeno3a(actives[:10], ligand, prb_crippen=prb_crippen, ref_crippen=ref_crippen)\n",
    "statistics.mean(actives_crippeno3a_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prb_mmff = [AllChem.MMFFGetMoleculeProperties(mol) for mol in actives[0:10]]\n",
    "# ref_mmff = AllChem.MMFFGetMoleculeProperties(ligand)\n",
    "# decoys_o3a_score = rdkit_o3a(decoys[:10], ligand, prb_mmff=prb_mmff, ref_mmff=ref_mmff)\n",
    "# statistics.mean(decoys_o3a_score)\n",
    "# # RuntimeError: Invariant Violation\n",
    "# # \tNegative weight specified for a point\n",
    "# # \tViolation occurred on line 57 in file Code/Numerics/Alignment/AlignPoints.cpp\n",
    "# # \tFailed Expression: wData[i] > 0.0\n",
    "# # \tRDKIT: 2022.03.4\n",
    "# # \tBOOST: 1_74"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1.2 分子三维结构相似度计算"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# actives_sampled = random.sample(actives, 100)\n",
    "# decoys_sampled = random.sample(decoys, 500)\n",
    "# RDLogger.DisableLog('rdApp.*') \n",
    "RDLogger.EnableLog('rdApp.*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# actives_sampled_path = os.path.join(base_path, 'demo-data/pde5a/actives_sample_100.sdf')\n",
    "# decoys_sampled_path = os.path.join(base_path, 'demo-data/pde5a/decoys_sample_500.sdf')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# writer = Chem.SDWriter(actives_sampled_path)\n",
    "# for cid in range(len(actives_sampled)):\n",
    "#     writer.write(actives_sampled[cid])\n",
    "    \n",
    "# writer = Chem.SDWriter(decoys_sampled_path)\n",
    "# for cid in range(len(decoys_sampled)):\n",
    "#     writer.write(decoys_sampled[cid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import dcg_score\n",
    "\n",
    "\n",
    "# actives = [m for m in Chem.SDMolSupplier(actives_sampled_path) if m]\n",
    "# decoys = [m for m in Chem.SDMolSupplier(decoys_sampled_path) if m]\n",
    "# ligand = AllChem.MolFromMol2File(l_path)\n",
    "\n",
    "len_a = len(actives)\n",
    "len_d = len(decoys)\n",
    "ic(len_a)\n",
    "ic(len_d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_o_a_time = 0\n",
    "start = timeit.default_timer()\n",
    "actives_o, align_score =al.rdkit_o3a(actives, ligand, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "a_o_a_time = end - start\n",
    "ic(a_o_a_time)\n",
    "a_o_d_time = 0\n",
    "start = timeit.default_timer()\n",
    "decoys_o, align_score =al.rdkit_o3a(decoys, ligand, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "a_o_d_time = end - start\n",
    "ic(a_o_d_time)\n",
    "\n",
    "a_c_a_time = 0\n",
    "start = timeit.default_timer()\n",
    "actives_c, align_score =al.rdkit_crippeno3a(actives, ligand, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "a_c_a_time = end - start\n",
    "ic(a_c_a_time)\n",
    "a_c_d_time = 0\n",
    "start = timeit.default_timer()\n",
    "decoys_c, align_score =al.rdkit_crippeno3a(decoys, ligand, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "a_c_d_time = end - start\n",
    "ic(a_c_d_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# writer = Chem.SDWriter(actives_sampled_path)\n",
    "# for cid in range(len(actives_sampled)):\n",
    "#     actives_sampled[cid].SetProp('_Name', f'activeInChIKey{Chem.MolToInchiKey(actives_sampled[cid])}')\n",
    "#     writer.write(actives_sampled[cid])\n",
    "    \n",
    "# writer = Chem.SDWriter(decoys_sampled_path)\n",
    "# for cid in range(len(decoys_sampled)):\n",
    "#     decoys_sampled[cid].SetProp('_Name', f'InChIKey{Chem.MolToInchiKey(decoys_sampled[cid])}')\n",
    "#     writer.write(decoys_sampled[cid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sampled_path = os.path.join(base_path, 'demo-data/pde5a/sample_600.sdf')\n",
    "# writer = Chem.SDWriter(sampled_path)\n",
    "# for cid in range(len(actives)):\n",
    "#     writer.write(actives[cid])\n",
    "# for cid in range(len(decoys)):\n",
    "#     writer.write(decoys[cid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# actives = actives[:10]\n",
    "# decoys = decoys[:50]\n",
    "# lan_a = len(actives)\n",
    "# len_d = len(decoys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def toy_test_similarity(s_func, prb_mols, ref_mol, align_time=0):\n",
    "    start = timeit.default_timer()\n",
    "    similarity_list = s_func(prb_mols, ref_mol, n_jobs=n_jobs, verbose=1)\n",
    "    end = timeit.default_timer()\n",
    "\n",
    "    print(f'mean   : {mean(similarity_list)}')\n",
    "    print(f'median : {median(similarity_list)}')\n",
    "    time = len(prb_mols) / (end - start + align_time)\n",
    "    return similarity_list, time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_dict = dict()\n",
    "tim_dict = dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_methods = {\n",
    "    'rdkit_protrude': sm.rdkit_shape_protrude_dist, \n",
    "    'rdkit_tanimoto': sm.rdkit_shape_tanimoto_dist,\n",
    "    'rdkit_tversky': sm.rdkit_shape_tversky_index,\n",
    "    'rdkit_sc': sm.rdkit_sc_score,\n",
    "    'rdkit_sc_tanimoto': sm.rdkit_sc_score_tanimoto,\n",
    "    'rdkit_sc_tversky': sm.rdkit_sc_score_tversky,\n",
    "}\n",
    "\n",
    "for name, func in similarity_methods.items():\n",
    "    _name = name\n",
    "    name = _name + '-o3a'    \n",
    "    print(f'🌟 {name} 🌟')\n",
    "    result_tmp = sim_dict.get(name, [])\n",
    "    t_tmp = tim_dict.get(name, [])\n",
    "    print('----👇 actives 👇----')\n",
    "    similarity_score, time_a = toy_test_similarity(func, actives_o, ligand, a_o_a_time)\n",
    "    result_tmp += [(ele, 1) for ele in similarity_score]\n",
    "    t_tmp.append(time_a)\n",
    "    print('----👇 decoys 👇----')\n",
    "    similarity_score, time_d = toy_test_similarity(func, decoys_o, ligand, a_o_d_time)\n",
    "    result_tmp += [(ele, 0) for ele in similarity_score]\n",
    "    t_tmp.append(time_d)\n",
    "    # assert len(result_tmp) == 600, f\"len(result_tmp) = {len(result_tmp)}\"\n",
    "    sim_dict[name] = result_tmp\n",
    "    tim_dict[name] = t_tmp\n",
    "    print(f'speed: {mean([time_a, time_d])}')\n",
    "\n",
    "\n",
    "    name = _name + '-crippeno3a'\n",
    "    print(f'🌟 {name} 🌟')\n",
    "    result_tmp = sim_dict.get(name, [])\n",
    "    t_tmp = tim_dict.get(name, [])\n",
    "    print('----👇 actives 👇----')\n",
    "    similarity_score, time_a = toy_test_similarity(func, actives_c, ligand, a_c_a_time)\n",
    "    result_tmp += [(ele, 1) for ele in similarity_score]\n",
    "    t_tmp.append(time_a)\n",
    "    print('----👇 decoys 👇----')\n",
    "    similarity_score, time_d = toy_test_similarity(func, decoys_c, ligand, a_c_d_time)\n",
    "    result_tmp += [(ele, 0) for ele in similarity_score]\n",
    "    t_tmp.append(time_d)\n",
    "    # assert len(result_tmp) == 600, f\"len(result_tmp) = {len(result_tmp)}\"\n",
    "    sim_dict[name] = result_tmp\n",
    "    tim_dict[name] = t_tmp\n",
    "    print(f'speed: {mean([time_a, time_d])}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_methods = {\n",
    "    'rdkit_pharm_tanimoto': sm.rdkit_pharm_tanimoto,\n",
    "    'rdkit_pharm_tversky': sm.rdkit_pharm_tversky,\n",
    "    'rdkit_fp_maccs': sm.rdkit_fp_maccs,\n",
    "    'rdkit_fp_maccs_tanimoto': sm.rdkit_fp_maccs_tainimoto,\n",
    "    'rdkit_fp_maccs_tversky': sm.rdkit_fp_maccs_tversky,\n",
    "    'rdkit_fp_margan': sm.rdkit_fp_morgan,\n",
    "    'rdkit_fp_margan_tanimoto': sm.rdkit_fp_morgan_tanimoto,\n",
    "    'rdkit_fp_margan_tversky': sm.rdkit_fp_morgan_tversky,\n",
    "}\n",
    "\n",
    "for name, func in similarity_methods.items():  \n",
    "    print(f'🌟 {name} 🌟')\n",
    "    result_tmp = sim_dict.get(name, [])\n",
    "    t_tmp = tim_dict.get(name, [])\n",
    "    print('----👇 actives 👇----')\n",
    "    similarity_score, time_a = toy_test_similarity(func, actives, ligand)\n",
    "    result_tmp += [(ele, 1) for ele in similarity_score]\n",
    "    t_tmp.append(time_a)\n",
    "    print('----👇 decoys 👇----')\n",
    "    similarity_score, time_d = toy_test_similarity(func, decoys, ligand)\n",
    "    result_tmp += [(ele, 0) for ele in similarity_score]\n",
    "    t_tmp.append(time_d)\n",
    "    # assert len(result_tmp) == 600, f\"len(result_tmp) = {len(result_tmp)}\"\n",
    "    sim_dict[name] = result_tmp\n",
    "    tim_dict[name] = t_tmp\n",
    "    print(f'speed: {mean([time_a, time_d])}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# similarity_methods = {\n",
    "#     'rdkit_protrude': sm.rdkit_shape_protrude_dist, \n",
    "#     'rdkit_tanimoto': sm.rdkit_shape_tanimoto_dist,\n",
    "#     'rdkit_tversky': sm.rdkit_shape_tversky_index,\n",
    "#     'rdkit_sc': sm.rdkit_sc_score,\n",
    "#     'rdkit_sc_tanimoto': sm.rdkit_sc_score_tanimoto,\n",
    "#     'rdkit_sc_tversky': sm.rdkit_sc_score_tversky,\n",
    "#     'rdkit_pharm_tanimoto': sm.rdkit_pharm_tanimoto,\n",
    "#     'rdkit_pharm_tversky': sm.rdkit_pharm_tversky,\n",
    "#     'rdkit_fp_maccs': sm.rdkit_fp_maccs,\n",
    "#     'rdkit_fp_maccs_tanimoto': sm.rdkit_fp_maccs_tainimoto,\n",
    "#     'rdkit_fp_maccs_tversky': sm.rdkit_fp_maccs_tversky,\n",
    "#     'rdkit_fp_margan': sm.rdkit_fp_morgan,\n",
    "#     'rdkit_fp_margan_tanimoto': sm.rdkit_fp_morgan_tanimoto,\n",
    "#     'rdkit_fp_margan_tversky': sm.rdkit_fp_morgan_tversky,\n",
    "# }\n",
    "\n",
    "# do_not_align = ['rdkit_pharm_tanimoto', 'rdkit_pharm_tversky','rdkit_fp_maccs', 'rdkit_fp_maccs_tanimoto', 'rdkit_fp_maccs_tversky', 'rdkit_fp_margan', 'rdkit_fp_margan_tanimoto', 'rdkit_fp_margan_tversky']\n",
    "\n",
    "# align_funcs_map_a = {\n",
    "#     'o3a': a_o_a_time, \n",
    "#     'crippeno3a': a_c_a_time,\n",
    "# }\n",
    "# align_funcs_map_d = {\n",
    "#     'o3a': a_o_d_time, \n",
    "#     'crippeno3a': a_c_d_time\n",
    "# }\n",
    "# align_funcs_list = ['o3a', 'crippeno3a']\n",
    "\n",
    "# sim_dict = dict()\n",
    "# tim_dict = dict()\n",
    "\n",
    "# for name, func in similarity_methods.items():\n",
    "#     ic(f'🌟 {name} 🌟')\n",
    "#     actives_sim_list = []\n",
    "#     decoys_sim_list = []\n",
    "#     ic('----👇 actives 👇----')\n",
    "#     align_funcs = align_funcs_list if name not in do_not_align else [None]\n",
    "#     tmp_time_list = []\n",
    "#     for align_name in align_funcs:\n",
    "#         if align_name:\n",
    "#             result_name = f'{name}-{align_name}'\n",
    "#             align_func = align_funcs_map_a[align_name]\n",
    "#             actives = actives_o if align_name=='o3a' else actives_c\n",
    "#         else:\n",
    "#             result_name = name\n",
    "#         tmp_sim, time = toy_test_similarity(\n",
    "#                             func, \n",
    "#                             actives, \n",
    "#                             ligand, \n",
    "#                             align_func=align_func\n",
    "#                             )\n",
    "#         result_tmp = sim_dict.get(result_name, [])\n",
    "#         result_tmp += [(ele, 1) for ele in tmp_sim]\n",
    "#         sim_dict[result_name] = result_tmp\n",
    "#         tmp_time_list.append(time)\n",
    "        \n",
    "#         t_tmp = tim_dict.get(result_name, [])\n",
    "#         t_tmp.append(time)\n",
    "#         tim_dict[result_name] = t_tmp\n",
    "\n",
    "#     ic('----👇 decoys 👇----')\n",
    "#     for align_name in align_funcs:\n",
    "#         if align_name:\n",
    "#             result_name = f'{name}-{align_name}'\n",
    "#             align_func = align_funcs_map_d[align_name]\n",
    "#             decoys = decoys_o \n",
    "#         else:\n",
    "#             result_name = name\n",
    "#         tmp_sim, time = toy_test_similarity(\n",
    "#                         func, \n",
    "#                         decoys, \n",
    "#                         ligand, \n",
    "#                         align_func=align_func\n",
    "#                         )\n",
    "#         result_tmp = sim_dict.get(result_name, [])\n",
    "#         result_tmp += [(ele, 0) for ele in tmp_sim]\n",
    "#         assert len(result_tmp) == (len_a + len_d), f\"len(result_tmp) = {len(result_tmp)}\"\n",
    "#         sim_dict[result_name] = result_tmp\n",
    "#         tmp_time_list.append(time)\n",
    "\n",
    "#         t_tmp = tim_dict.get(result_name, [])\n",
    "#         t_tmp.append(time)\n",
    "#         tim_dict[result_name] = t_tmp\n",
    "    \n",
    "#     for i in range(len(tmp_time_list)//2):\n",
    "#         ic(mean(tmp_time_list[i::len(tmp_time_list)//2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, score_list in sim_dict.items():\n",
    "    mean_actives_score = mean([ele[0] for ele in score_list if ele[1]==1])\n",
    "    mean_decoys_score = mean([ele[0] for ele in score_list if ele[1]==0])\n",
    "    ic(name)\n",
    "    ic(mean_actives_score)\n",
    "    ic(mean_decoys_score)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 ODDT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2.1 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "next(oddt.toolkit.readfile('mol2', l_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = [mol for mol in list(oddt.toolkit.readfile('sdf', d_path)) if mol]\n",
    "len(tmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2.2 相似度计算"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ligand = next(oddt.toolkit.readfile('mol2', l_path))\n",
    "actives = [mol for mol in oddt.toolkit.readfile('sdf', a_path) if mol]\n",
    "decoys = [mol for mol in oddt.toolkit.readfile('sdf', d_path) if mol]\n",
    "# actives = [mol for mol in oddt.toolkit.readfile('sdf', actives_sampled_path) if mol]\n",
    "# decoys = [mol for mol in oddt.toolkit.readfile('sdf', decoys_sampled_path) if mol]\n",
    "\n",
    "len_a = len(actives)\n",
    "len_d = len(decoys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# actives = actives[:10]\n",
    "# decoys = decoys[:50]\n",
    "# len_a = len(actives)\n",
    "# len_d = len(decoys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ic(len_a)\n",
    "ic(len_d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def toy_test_similarity(s_func, prb_mols, ref_mol):\n",
    "    start = timeit.default_timer()\n",
    "    score_list = s_func(prb_mols, ref_mol, n_jobs=n_jobs, verbose=1)\n",
    "    end = timeit.default_timer()\n",
    "    \n",
    "    print(f'mean   : {mean(score_list)}')\n",
    "    print(f'median : {median(score_list)}')\n",
    "\n",
    "    return score_list, len(prb_mols)/(end - start)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_methods = {\n",
    "    'oddt_usr': sm.oddt_usr,\n",
    "    'oddt_usr_cat': sm.oddt_usr_cat,\n",
    "    'oddt_electroshape': sm.oddt_electroshape,\n",
    "}\n",
    "for name, func in similarity_methods.items():\n",
    "    print(f'🌟 {name} 🌟')\n",
    "    result_tmp = sim_dict.get(name, [])\n",
    "    t_tmp = tim_dict.get(name, [])\n",
    "    print('----👇 actives 👇----')\n",
    "    similarity_score, time_a = toy_test_similarity(func, actives, ligand)\n",
    "    result_tmp += [(ele, 1) for ele in similarity_score]\n",
    "    t_tmp.append(time_a)\n",
    "    print('----👇 decoys 👇----')\n",
    "    similarity_score, time_d = toy_test_similarity(func, decoys, ligand)\n",
    "    result_tmp += [(ele, 0) for ele in similarity_score]\n",
    "    t_tmp.append(time_d)\n",
    "    # assert len(result_tmp) == 600, f\"len(result_tmp) = {len(result_tmp)}\"\n",
    "    sim_dict[name] = result_tmp\n",
    "    tim_dict[name] = t_tmp\n",
    "    print(f'speed: {mean([time_a, time_d])}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_dict.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = os.path.join(base_path, f'demo-data/result/{target_list[target_flag]}/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(save_dir):\n",
    "    os.system(f\"mkdir {save_dir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = os.path.join(base_path, f'demo-data/result/{target_list[target_flag]}/')\n",
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-sim_dict.pickle'), 'wb') as f:\n",
    "    pickle.dump(sim_dict, f)\n",
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-tim_dict.pickle'), 'wb') as f:\n",
    "    pickle.dump(tim_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = os.path.join(base_path, f'demo-data/result/{target_list[target_flag]}/')\n",
    "ic(save_dir)\n",
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-sim_dict.pickle'), 'rb') as f:\n",
    "    sim_dict = pickle.load(f)\n",
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-tim_dict.pickle'), 'rb') as f:\n",
    "    tim_dict = pickle.load(f)\n",
    "sim_dict.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 acpc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'actives_dd.sdf')\n",
    "d_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'decoys_dd.sdf')\n",
    "l_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'crystal_ligand.mol2')\n",
    "ic(a_path)\n",
    "\n",
    "# decoys = fn.Mol2MolSupplier(d_path, sanitize=True)\n",
    "dec = [m for m in Chem.SDMolSupplier(d_path, removeHs=False) if m]\n",
    "ic(len(dec))\n",
    "# Draw.MolsToGridImage(random.sample(decoys, 7), molsPerRow=7, subImgSize=(150,150), legends=[mol.GetProp('_Name') for mol in decoys[:7]],maxMols=100)\n",
    "\n",
    "act = [m for m in Chem.SDMolSupplier(a_path, removeHs=False) if m]\n",
    "ic(len(act))\n",
    "# Draw.MolsToGridImage(random.sample(actives, 7), molsPerRow=7, subImgSize=(150,150), legends=[mol.GetProp('_Name') for mol in actives[:7]],maxMols=100)\n",
    "\n",
    "lig = AllChem.MolFromMol2File(l_path, removeHs=False)\n",
    "lig = Chem.AddHs(lig, addCoords=True)\n",
    "lig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_dict = dict()\n",
    "tim_dict = dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for mol in act:\n",
    "    mol.SetProp('_Name', f'activesInChIKey{AllChem.MolToInchiKey(mol)}')\n",
    "for mol in dec:\n",
    "    mol.SetProp('_Name', f'InChIKey{AllChem.MolToInchiKey(mol)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3.1 mol2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prb_acpc_ml2_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'prb_acpc.mol2')\n",
    "ref_acpc_ml2_path = l_path\n",
    "out_acpc_ml2_path = os.path.join(base_path, 'demo-data/result/acpc', target_list[target_flag], f'{target_list[target_flag]}-ml2-result.txt')\n",
    "\n",
    "fn.MolToSpecialFormatFile(act+dec, prb_acpc_ml2_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "sim_dict['acpc-ml2'] = sm.acpc(ref_path=ref_acpc_ml2_path, prb_path=prb_acpc_ml2_path, out_path=out_acpc_ml2_path, num_core=n_jobs)\n",
    "end = timeit.default_timer()\n",
    "tim_dict['acpc-ml2'] = [len(act+dec)/(end-sat)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prb_acpc_ml2_bin_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'prb_acpc.mol2.bin')\n",
    "out_acpc_ml2_bin_path = os.path.join(base_path, 'demo-data/result/acpc', target_list[target_flag], f'{target_list[target_flag]}-ml2-bin-result.txt')\n",
    "os.system(f'acpc_codec -i {prb_acpc_ml2_path} -o {prb_acpc_ml2_bin_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "sim_dict['acpc-ml2-bin'] = sm.acpc(ref_path=ref_acpc_ml2_path, prb_path=prb_acpc_ml2_bin_path, out_path=out_acpc_ml2_bin_path, num_core=n_jobs)\n",
    "end = timeit.default_timer()\n",
    "tim_dict['acpc-ml2-bin'] = [len(act+dec)/(end-sat)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3.2 pqr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prb_acpc_pqr_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'prb_acpc.pqr')\n",
    "fn.MolToSpecialFormatFile(act+dec, prb_acpc_pqr_path)\n",
    "ref_acpc_pqr_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'ref_acpc.pqr')\n",
    "fn.MolToSpecialFormatFile([lig], ref_acpc_pqr_path)\n",
    "out_acpc_pqr_path = os.path.join(base_path, 'demo-data/result/acpc', target_list[target_flag], f'{target_list[target_flag]}-pqr-result.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "sim_dict['acpc-pqr'] = sm.acpc(ref_path=ref_acpc_pqr_path, prb_path=prb_acpc_pqr_path, out_path=out_acpc_pqr_path, num_core=n_jobs)\n",
    "end = timeit.default_timer()\n",
    "tim_dict['acpc-pqr'] = [len(act+dec)/(end-sat)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 ESP-SIM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'actives_dd.sdf')\n",
    "d_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'decoys_dd.sdf')\n",
    "l_path = os.path.join(base_path, dude_dir, target_list[target_flag], 'crystal_ligand.mol2')\n",
    "ic(a_path)\n",
    "ic(d_path)\n",
    "ic(l_path)\n",
    "\n",
    "# decoys = fn.Mol2MolSupplier(d_path, sanitize=True)\n",
    "dec = [m for m in Chem.SDMolSupplier(d_path, removeHs=False) if m]\n",
    "ic(len(dec))\n",
    "# Draw.MolsToGridImage(random.sample(decoys, 7), molsPerRow=7, subImgSize=(150,150), legends=[mol.GetProp('_Name') for mol in decoys[:7]],maxMols=100)\n",
    "\n",
    "act = [m for m in Chem.SDMolSupplier(a_path, removeHs=False) if m]\n",
    "ic(len(act))\n",
    "# Draw.MolsToGridImage(random.sample(actives, 7), molsPerRow=7, subImgSize=(150,150), legends=[mol.GetProp('_Name') for mol in actives[:7]],maxMols=100)\n",
    "\n",
    "lig = AllChem.MolFromMol2File(l_path, removeHs=False)\n",
    "lig = Chem.AddHs(lig, addCoords=True)\n",
    "lig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "act[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len_a = len(act)\n",
    "len_d = len(dec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_dict = dict()\n",
    "tim_dict = dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ali_o3a_time = 0\n",
    "sta = timeit.default_timer()\n",
    "act_o3a, align_score =al.rdkit_o3a(act, lig, n_jobs=n_jobs, verbose=1)\n",
    "dec_o3a, align_score =al.rdkit_o3a(dec, lig, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "ali_o3a_time = end - sta\n",
    "ic(ali_o3a_time)\n",
    "\n",
    "ali_cri_time = 0\n",
    "sta = timeit.default_timer()\n",
    "act_cri, align_score =al.rdkit_crippeno3a(act, lig, n_jobs=n_jobs, verbose=1)\n",
    "dec_cri, align_score =al.rdkit_crippeno3a(dec, lig, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "ali_cri_time = end - sta\n",
    "ic(ali_cri_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "act_sim_list = sm.esp_sim(act_cri, lig, n_jobs=n_jobs, verbose=1)\n",
    "dec_sim_list = sm.esp_sim(dec_cri, lig, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "sim_dict['esp-sim-crippeno3a-tanimoto-mmff-tanimoto'] = [(sim, 1) for sim in act_sim_list] + [(sim, 0) for sim in dec_sim_list]\n",
    "tim_dict['esp-sim-crippeno3a-tanimoto-mmff-tanimoto'] = [(len_a + len_d) / (end - sat + ali_cri_time)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "act_sim_list = sm.esp_sim(act_o3a, lig, n_jobs=n_jobs, verbose=1)\n",
    "dec_sim_list = sm.esp_sim(dec_o3a, lig, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "sim_dict['esp-sim-o3a-tanimoto-mmff-tanimoto'] = [(sim, 1) for sim in act_sim_list] + [(sim, 0) for sim in dec_sim_list]\n",
    "tim_dict['esp-sim-o3a-tanimoto-mmff-tanimoto'] = [(len_a + len_d) / (end - sat + ali_o3a_time)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "act_sim_list = sm.esp_sim(act_o3a, lig, shape_sim='protrude', n_jobs=n_jobs, verbose=1)\n",
    "dec_sim_list = sm.esp_sim(dec_o3a, lig, shape_sim='protrude', n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "sim_dict['esp-sim-o3a-protrude-mmff-tanimoto'] = [(sim, 1) for sim in act_sim_list] + [(sim, 0) for sim in dec_sim_list]\n",
    "tim_dict['esp-sim-o3a-protrude-mmff-tanimoto'] = [(len_a + len_d) / (end - sat + ali_o3a_time)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "act_sim_list = sm.esp_sim(act_o3a, lig, shape_sim='tversky', n_jobs=n_jobs, verbose=1)\n",
    "dec_sim_list = sm.esp_sim(dec_o3a, lig, shape_sim='tversky', n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "sim_dict['esp-sim-o3a-tversky-mmff-tanimoto'] = [(sim, 1) for sim in act_sim_list] + [(sim, 0) for sim in dec_sim_list]\n",
    "tim_dict['esp-sim-o3a-tversky-mmff-tanimoto'] = [(len_a + len_d) / (end - sat + ali_o3a_time)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "act_sim_list = sm.esp_sim(act_o3a, lig, systems='gasteiger', n_jobs=n_jobs, verbose=1)\n",
    "dec_sim_list = sm.esp_sim(dec_o3a, lig, systems='gasteiger', n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "sim_dict['esp-sim-o3a-tanimoto-gasteiger-tanimoto'] = [(sim, 1) for sim in act_sim_list] + [(sim, 0) for sim in dec_sim_list]\n",
    "tim_dict['esp-sim-o3a-tanimoto-gasteiger-tanimoto'] = [(len_a + len_d) / (end - sat + ali_o3a_time)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "prb_charge = mlCharges(act+dec)\n",
    "ref_charge = mlCharges([lig])[0]\n",
    "end = timeit.default_timer()\n",
    "ml_char_time = end - sat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "act_sim_list = sm.esp_sim(act_o3a, lig, systems='ml', prb_charge=prb_charge[:len_a], ref_charge=ref_charge, n_jobs=n_jobs, verbose=1)\n",
    "dec_sim_list = sm.esp_sim(dec_o3a, lig, systems='ml', prb_charge=prb_charge[len_a:], ref_charge=ref_charge, n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "sim_dict['esp-sim-o3a-tanimoto-ml-tanimoto'] = [(sim, 1) for sim in act_sim_list] + [(sim, 0) for sim in dec_sim_list]\n",
    "tim_dict['esp-sim-o3a-tanimoto-ml-tanimoto'] = [(len_a + len_d) / (end - sat + ali_o3a_time + ml_char_time)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat = timeit.default_timer()\n",
    "act_sim_list = sm.esp_sim(act_o3a, lig, metric='carbo', n_jobs=n_jobs, verbose=1)\n",
    "dec_sim_list = sm.esp_sim(dec_o3a, lig, metric='carbo', n_jobs=n_jobs, verbose=1)\n",
    "end = timeit.default_timer()\n",
    "sim_dict['esp-sim-o3a-tanimoto-mmff-carbo'] = [(sim, 1) for sim in act_sim_list] + [(sim, 0) for sim in dec_sim_list]\n",
    "tim_dict['esp-sim-o3a-tanimoto-mmff-carbo'] = [(len_a + len_d) / (end - sat + ali_o3a_time)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_dict.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = os.path.join(base_path, f'demo-data/result/espsim/{target_list[target_flag]}/')\n",
    "if not os.path.exists(save_dir):\n",
    "    os.system(f\"mkdir -p {save_dir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-sim_dict.pickle'), 'wb') as f:\n",
    "    pickle.dump(sim_dict, f)\n",
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-tim_dict.pickle'), 'wb') as f:\n",
    "    pickle.dump(tim_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ic(save_dir)\n",
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-sim_dict.pickle'), 'rb') as f:\n",
    "    sim_dict = pickle.load(f)\n",
    "with open(os.path.join(save_dir, f'{target_list[target_flag]}-tim_dict.pickle'), 'rb') as f:\n",
    "    tim_dict = pickle.load(f)\n",
    "sim_dict.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 评价指标"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 1\n",
    "for name, score in sim_dict.items():\n",
    "    plt.figure(figsize=(5, 5), dpi=200)\n",
    "    \n",
    "    score_list = sorted(score, key=lambda x: x[0])\n",
    "\n",
    "    y_lable = np.asarray([ele[1] for ele in score_list])\n",
    "    y_score = np.asarray([ele[0] for ele in score_list])\n",
    "\n",
    "    fpr, tpr, _ = metrics.roc_curve(y_lable, y_score)\n",
    "    roc_auc = metrics.auc(fpr, tpr)\n",
    "\n",
    "    plt.plot(\n",
    "        fpr,\n",
    "        tpr,\n",
    "        color=\"darkorange\",\n",
    "        lw=2,\n",
    "        label=\"ROC curve (area = %0.2f)\" % roc_auc,\n",
    "    )\n",
    "    plt.plot([0, 1], [0, 1], color=\"navy\", lw=2, linestyle=\"--\")\n",
    "    plt.xlim([0.0, 1.0])\n",
    "    plt.ylim([0.0, 1.05])\n",
    "    plt.xlabel(\"False Positive Rate\")\n",
    "    plt.ylabel(\"True Positive Rate\")\n",
    "    plt.title(f\"ROC {name}\")\n",
    "    plt.legend(loc=\"lower right\")\n",
    "    plt.savefig(os.path.join(save_dir, f'{i :0>2}-roc-{name}.png'))\n",
    "    plt.show()\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 1\n",
    "for name, score in sim_dict.items():\n",
    "    fig, ax = plt.subplots(figsize=(5, 5), dpi=200)\n",
    "    \n",
    "    score_list = sorted(score, key=lambda x: x[0])\n",
    "\n",
    "    y_lable = np.asarray([ele[1] for ele in score_list])\n",
    "    y_score = np.asarray([ele[0] for ele in score_list])\n",
    "\n",
    "    fpr, tpr, _ = metrics.roc_curve(y_lable, y_score)\n",
    "    \n",
    "    log_min = 0.001\n",
    "    log_max = 1.0\n",
    "\n",
    "    fpr = fpr.clip(log_min)\n",
    "    idx = (fpr <= log_max)\n",
    "    log_fpr = 1 - np.log10(fpr[idx]) / np.log10(log_min)\n",
    "    \n",
    "    log_roc_auc = metrics.auc(log_fpr, tpr[idx])\n",
    "\n",
    "\n",
    "    plt.plot(\n",
    "        fpr[idx],\n",
    "        tpr[idx],\n",
    "        color=\"darkorange\",\n",
    "        lw=2,\n",
    "        label=\"logROC curve (area = %0.2f)\" % log_roc_auc,\n",
    "    )\n",
    "    x = np.linspace(0.001, 1, 200)\n",
    "    y = x\n",
    "    plt.plot(x, y, color=\"navy\", lw=2, linestyle=\"--\")\n",
    "\n",
    "    plt.ylim([0.0, 1.05])\n",
    "    plt.xscale('log')\n",
    "    plt.xlim([0.001, 1])\n",
    "    ax.set_xticks([0.001, 0.01, 0.1, 1.0])\n",
    "    plt.xlabel(\"False Positive Rate\")\n",
    "    plt.ylabel(\"True Positive Rate\")\n",
    "    plt.title(f\"logROC {name}\")\n",
    "    plt.legend(loc=\"lower right\")\n",
    "    plt.savefig(os.path.join(save_dir, f'{i :0>2}-log_roc-{name}.png'))\n",
    "    plt.show()\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0.001, 1, 200)\n",
    "y = x\n",
    "x = 1 - np.log10(x) / np.log10(log_min)\n",
    "plt.plot(x, y, color=\"navy\", lw=2, linestyle=\"--\")\n",
    "# plt.xscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0.001, 1, 200)\n",
    "y = x\n",
    "plt.plot(x, y, color=\"navy\", lw=2, linestyle=\"--\")\n",
    "plt.xscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.asarray([1,3,4,5,6,7]).clip(4)\n",
    "sorted([1,2,5,6,7], key=lambda x: -x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_list = []\n",
    "for name, score in sim_dict.items():\n",
    "    score_list = sorted(score, key=lambda x: -x[0])\n",
    "\n",
    "    y_score = np.asarray([ele[0] for ele in score_list])\n",
    "    y_lable = np.asarray([ele[1] for ele in score_list])\n",
    "    \n",
    "    metrics_list = mt.calculate_metrics(y_lable=y_lable, y_score=y_score)\n",
    "\n",
    "    print(f'--- 🌟 {name} 🌟 ---')\n",
    "    tmp_dict = {\n",
    "        'name' : name,\n",
    "        'time' : mean(tim_dict[name])\n",
    "    }\n",
    "    print(f'time: {mean(tim_dict[name])}')\n",
    "    for m_name, m_value in metrics_list.items():\n",
    "        tmp_dict[m_name] = m_value\n",
    "        print(f'{m_name:<30}: {m_value}')\n",
    "    result_list.append(tmp_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(result_list).to_excel(os.path.join(save_dir, f'{target_list[target_flag]}_results.xlsx'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "molcalx = '/home/jovyan/work-home/molecule-3d-similarity/pde5_score.csv'\n",
    "df = pd.read_csv(molcalx)\n",
    "df_sorted = df.sort_values(by=['score'],ascending=True)\n",
    "score = df_sorted['score']\n",
    "x = np.array(score)*(-1)\n",
    "y = df_sorted['label']\n",
    "label_map = {'active': 1, 'decoy': 0}\n",
    "labels = [label_map[y_true] for y_true in y]\n",
    "y = np.array(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auc = m.roc_auc(y, x, pos_label=1, ascending_score=False)\n",
    "print(\"ROC AUC = \",auc)\n",
    "bedroc = m.bedroc(y, x, alpha=20.0, pos_label=1)\n",
    "print(\"alpha=20.0 BEDROC = \",bedroc)\n",
    "logauc = m.roc_log_auc(y, x, pos_label=1, ascending_score=False, log_min=0.001, log_max=1.0)\n",
    "print(\"logAUC [0.1%,100%]  = \",logauc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_string = 'ROCSComboscore 0.598 ± 0.152 0.681 ± 0.166 0.674 ± 0.115 0.727 ± 0.128 ROCSColorscore 0.620 ± 0.139 0.712 ± 0.159 0.677 ± 0.117 0.752 ± 0.136 ROCSShapeTanimoto 0.547 ± 0.138 0.611 ± 0.163 0.618 ± 0.105 0.667 ± 0.127 Phase Shape_Mmod 0.677 ± 0.143 0.686 ± 0.145 0.772 ± 0.105 0.769 ± 0.108 Phase Shape_Ele 0.674 ± 0.153 0.688 ± 0.158 0.753 ± 0.105 0.750 ± 0.111 Phase Shape_Pharm 0.692 ± 0.160 0.694 ± 0.168 0.761 ± 0.145 0.767 ± 0.143 Shape-it 0.541 ± 0.133 0.590 ± 0.141 0.612 ± 0.110 0.639 ± 0.115 Align-it 0.659 ± 0.137 0.680 ± 0.157 0.729 ± 0.132 0.746 ± 0.139 ShaEPbest 0.658 ± 0.122 0.660 ± 0.139 0.709 ± 0.099 0.699 ± 0.120 ShaEPshape 0.625 ± 0.139 0.632 ± 0.143 0.681 ± 0.105 0.676 ± 0.116 ShaEPESP 0.606 ± 0.109 0.591 ± 0.117 0.627 ± 0.105 0.585 ± 0.125 SHAFTS 0.733 ± 0.144 0.731 ± 0.157 0.792 ± 0.129 0.782 ± 0.135 WEGA 0.645 ± 0.143 0.659 ± 0.154 0.716 ± 0.107 0.716 ± 0.125 LIGSIFT 0.718 ± 0.133 0.755 ± 0.143 0.758 ± 0.117 0.784 ± 0.120 LS-align 0.699 ± 0.126 0.759 ± 0.119 0.773 ± 0.098 0.786 ± 0.096'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_string.split(' ')[13]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_string.split(' ')[14]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cry_single = [0.598, 0.620, ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pccs = np.corrcoef(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed_ligands = {\n",
    "    'pde5a': [\n",
    "        'CCCc1nn(c2c1nc([nH]c2=O)c1cc(ccc1OCC)S(=O)(=O)n1ccn(cc1)C)C',\n",
    "        'CCCc1nn(c2c1nc([nH]c2=O)c1cc(ccc1OCC)S(=O)(=O)N1CCN(CC1)Cc1ccc2c(c1)OCO2)C',\n",
    "    ],\n",
    "    'akt1': [\n",
    "        'Clc1c[nH]c2c1c(ncn2)N1CCc2c(C1)[nH]cn2',\n",
    "        '[NH3+]C(Cc1c[nH]c2c1cccc2)COc1cncc(c1)c1ccc2c(c1)C(C(=O)N2)C(Cc1ccccn1)(C)C'\n",
    "    ],\n",
    "    'ada': [\n",
    "        'CC(C(n1cnc(c1)C(=O)N)CCc1cccc2c1cccc2)O',\n",
    "        'CC(C(n1cnc2c1ccnc2N)CCCCCC)O'\n",
    "    ],\n",
    "    'andr': [\n",
    "        'OC1CCC2(C(=C1)CCC1C2CCC2(C1CCC2O)C)C',\n",
    "        'N#Cc1ccc(cc1C(F)(F)F)N1C(=O)C2C(C1=O)C1(OC2(C)CN(C1)c1cccc(c1)C(=O)N)C'\n",
    "    ],\n",
    "    'def': [\n",
    "        'CCCCCC(C(=O)NC(C(=O)N1CCCC1CO)C(C)C)CC(=O)NO',\n",
    "        'O=CN(CC(C(=O)NC(C(C)(C)C)C(=O)c1ccc(cc1)F)CC1CCCC1)O'\n",
    "    ],\n",
    "    'gria2': [\n",
    "        'O=c1[nH]c2cc(c(cc2n(c1=O)CP(=O)(O)O)N1CCOCC1)C(F)(F)F',\n",
    "        '[O-][N+](=O)c1cc2[nH]c(=O)c(nc2cc1n1ccc(c1)C(=O)[O-])[O-]',\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cry_smiles = Chem.MolToSmiles(Chem.MolFromSmiles(seed_ligands[target_list[target_flag]][0]))\n",
    "clu_smiles = Chem.MolToSmiles(Chem.MolFromSmiles(seed_ligands[target_list[target_flag]][1]))\n",
    "cry_inchi = Chem.MolToInchiKey(Chem.MolFromSmiles(cry_smiles))\n",
    "clu_inchi = Chem.MolToInchiKey(Chem.MolFromSmiles(clu_smiles))\n",
    "\n",
    "cry_mol = Chem.MolFromSmiles(cry_smiles)\n",
    "cry_mol = AllChem.AddHs(cry_mol, addCoords=True)\n",
    "clu_mol = Chem.MolFromSmiles(clu_smiles)\n",
    "clu_mol = AllChem.AddHs(clu_mol, addCoords=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in actives:\n",
    "    tmp_smiles = Chem.MolToSmiles(m)\n",
    "    tmp_inchi = Chem.MolToInchiKey(m)\n",
    "\n",
    "    if cry_smiles == tmp_smiles:\n",
    "        print('cry find')\n",
    "    \n",
    "    if clu_inchi == tmp_inchi:\n",
    "        print('clu find')   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp_list = []\n",
    "for m in actives:\n",
    "    fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(m,2,useFeatures=True)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(ligand,2,useFeatures=True)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(cry_mol,2,useFeatures=True)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(clu_mol,2,useFeatures=True)))\n",
    "\n",
    "tsne = TSNE(n_components=2, random_state=2022)\n",
    "z = tsne.fit_transform(fp_list)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5, 5), dpi=200)\n",
    "plt.scatter(z[:-2, 0], z[:-2, 1])\n",
    "plt.scatter(z[-3][0], z[-3][1])\n",
    "# plt.scatter(z[-2][0], z[-2][1], color='beige')\n",
    "# plt.scatter(z[67][0], z[67][1], color='beige')\n",
    "# plt.scatter(z[-1][0], z[-1][1], color='springgreen')\n",
    "# plt.scatter(z[285][0], z[285][1], color='springgreen')\n",
    "# plt.scatter(z[215][0], z[215][1], color='lightcyan')\n",
    "plt.scatter(z[291][0], z[291][1], color='lightcyan')\n",
    "# plt.scatter(z[259][0], z[259][1], color='lightskyblue')\n",
    "# plt.scatter(z[119][0], z[119][1], color='lightskyblue')\n",
    "# plt.scatter(z[26][0], z[26][1], color='lightcoral')\n",
    "# plt.scatter(z[46][0], z[46][1], color='lightcoral')\n",
    "# seed = random.sample(list(range(len(actives))), 1)\n",
    "# ic(seed)\n",
    "# plt.scatter(z[seed[0]][0], z[seed[0]][1], color='beige')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit.DataStructs import TanimotoSimilarity, TverskySimilarity\n",
    "fp1 = AllChem.GetMorganFingerprintAsBitVect(actives[-97],2,useFeatures=True)\n",
    "fp2 = AllChem.GetMorganFingerprintAsBitVect(ligand,2,useFeatures=True)\n",
    "TanimotoSimilarity(fp1, fp2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 285\n",
    "# -1\n",
    "\n",
    "# 67\n",
    "# -2\n",
    "\n",
    "# 215\n",
    "# 291\n",
    "\n",
    "# 259\n",
    "# 119\n",
    "\n",
    "# 26\n",
    "# 46"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(actives)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives[26].RemoveAllConformers()\n",
    "actives[26]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives[-12].RemoveAllConformers()\n",
    "actives[-12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives[-15].RemoveAllConformers()\n",
    "actives[-15]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives[-97].RemoveAllConformers()\n",
    "actives[-97]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives[-99].RemoveAllConformers()\n",
    "actives[-99]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives[-5].RemoveAllConformers()\n",
    "actives[-5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clu_mol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cry_mol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ligand.RemoveAllConformers()\n",
    "ligand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actives[-8].RemoveAllConformers()\n",
    "actives[-8]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp_list = []\n",
    "for m in actives:\n",
    "    fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(m,2)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(ligand,2)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(cry_mol,2)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(clu_mol,2)))\n",
    "\n",
    "tsne = TSNE(n_components=2, random_state=2022)\n",
    "z = tsne.fit_transform(fp_list)\n",
    "\n",
    "plt.figure(figsize=(5, 5), dpi=200)\n",
    "plt.scatter(z[:-2, 0], z[:-2, 1])\n",
    "plt.scatter(z[-3][0], z[-3][1])\n",
    "plt.scatter(z[-2][0], z[-2][1])\n",
    "plt.scatter(z[-1][0], z[-1][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp_list = []\n",
    "for m in actives:\n",
    "    fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(m,2,useFeatures=True)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(ligand,2,useFeatures=True)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(cry_mol,2,useFeatures=True)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(clu_mol,2,useFeatures=True)))\n",
    "\n",
    "tsne = PCA(n_components=2, random_state=2022)\n",
    "z = tsne.fit_transform(fp_list)\n",
    "\n",
    "plt.figure(figsize=(5, 5), dpi=200)\n",
    "plt.scatter(z[:-2, 0], z[:-2, 1])\n",
    "plt.scatter(z[-3][0], z[-3][1])\n",
    "plt.scatter(z[-2][0], z[-2][1])\n",
    "plt.scatter(z[-1][0], z[-1][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp_list = []\n",
    "for m in actives:\n",
    "    fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(m,2)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(ligand,2)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(cry_mol,2)))\n",
    "fp_list.append(np.asarray(AllChem.GetMorganFingerprintAsBitVect(clu_mol,2)))\n",
    "\n",
    "tsne = PCA(n_components=2, random_state=2022)\n",
    "z = tsne.fit_transform(fp_list)\n",
    "\n",
    "plt.figure(figsize=(5, 5), dpi=200)\n",
    "plt.scatter(z[:-2, 0], z[:-2, 1])\n",
    "plt.scatter(z[-3][0], z[-3][1])\n",
    "plt.scatter(z[-2][0], z[-2][1])\n",
    "plt.scatter(z[-1][0], z[-1][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = Chem.MolFromSmiles('c1ccccc(C(N)=O)1')\n",
    "mol\n",
    "mol = Chem.AddHs(mol)\n",
    "AllChem.EmbedMolecule(mol,randomSeed=0xf00d)\n",
    "AllChem.MMFFOptimizeMolecule(mol)\n",
    "mol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ic(Chem.MolToMolBlock(mol))\n",
    "for atom in mol.GetAtoms():\n",
    "    ic(atom.GetAtomicNum())\n",
    "    ic(atom.GetIdx())\n",
    "ic(dir(mol.GetConformer(0)))\n",
    "ic(mol.GetConformer(0).GetPositions())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for atm in mol.GetAtoms():\n",
    "    ic(dir(atm))\n",
    "    ic(atm.GetAtomicNum())\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Chem.GetAdjacencyMatrix(mol)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[a.GetAtomicNum() for a in mol.GetAtoms()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[a.GetBonds() for a in mol.GetAtoms()][6][].GetBondType()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an rdkit mol object\n",
    "# Iterate over the atoms\n",
    "for i, atom in enumerate(mol.GetAtoms()):\n",
    "    # For each atom, set the property \"molAtomMapNumber\" to a custom number, let's say, the index of the atom in the molecule\n",
    "    atom.SetProp(\"molAtomMapNumber\", str(atom.GetIdx()+1))\n",
    "\n",
    "mol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_list = []\n",
    "for root, dirs, files in os.walk(os.path.join(base_path, dude_dir), topdown=False):\n",
    "    tmp_dict = {'act_ism': None, 'dec_ism': None}\n",
    "    for name in files:\n",
    "        if   'actives_final.ism' in name:\n",
    "            tmp_dict['act_ism'] = os.path.join(root, name)\n",
    "        elif 'decoys_final.ism' in name:\n",
    "            tmp_dict['dec_ism'] = os.path.join(root, name)\n",
    "    tmp_list.append(tmp_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ism_path_df = pd.DataFrame(tmp_list).dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pandarallel.initialize(progress_bar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_info_from_smiles(series: pd.Series):\n",
    "    print(f\"🌟  {series.name+1}-{series['act_ism'].split('/')[-2]}  🌟\")\n",
    "    act = pd.read_csv(series['act_ism'], sep=' ',  header=None, names=['smiles', '-', 'chemblid'])\n",
    "    dec = pd.read_csv(series['dec_ism'], sep=' ',  header=None, names=['smiles', '-', 'chemblid'])\n",
    "\n",
    "    def get_charges(series: pd.Series):\n",
    "        mol = AllChem.MolFromSmiles(series['smiles'])\n",
    "        if not mol:\n",
    "            return []\n",
    "        return [a.GetAtomicNum() for a in mol.GetAtoms()]\n",
    "    \n",
    "    ic(act.head(5))\n",
    "    tqdm.pandas()\n",
    "    act['charges'] = act.progress_apply(get_charges, axis=1)\n",
    "    dec['charges'] = dec.progress_apply(get_charges, axis=1)\n",
    "\n",
    "    return [dict(Counter([cha for ele in act['charges'].tolist() for cha in ele if ele])), dict(Counter([cha for ele in dec['charges'].tolist() for cha in ele if ele]))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extract_info_from_smiles(ism_path_df.iloc[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = Parallel(n_jobs=n_jobs, verbose=1)(delayed(extract_info_from_smiles)(row) for _, row in ism_path_df.iterrows()) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sum_dict(a,b):\n",
    "    temp = dict()\n",
    "    for key in a.keys()| b.keys():\n",
    "        temp[key] = sum([d.get(key, 0) for d in (a, b)])\n",
    "    return temp\n",
    " \n",
    "def test():\n",
    "    from functools import reduce\n",
    "    return print(reduce(sum_dict,[a,b,c]))\n",
    " \n",
    "a = {'a': 1, 'b': 2, 'c': 3} \n",
    "b = {'a':1,'b':3,'d':4}\n",
    "c = {'g':3,'f':5,'a':10}\n",
    "test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "__ATOM_LIST__ = \\\n",
    "    ['-', 'h',  'he',\n",
    "     'li', 'be', 'b',  'c',  'n',  'o',  'f',  'ne',\n",
    "     'na', 'mg', 'al', 'si', 'p',  's',  'cl', 'ar',\n",
    "     'k',  'ca', 'sc', 'ti', 'v ', 'cr', 'mn', 'fe', 'co', 'ni', 'cu',\n",
    "     'zn', 'ga', 'ge', 'as', 'se', 'br', 'kr',\n",
    "     'rb', 'sr', 'y',  'zr', 'nb', 'mo', 'tc', 'ru', 'rh', 'pd', 'ag',\n",
    "     'cd', 'in', 'sn', 'sb', 'te', 'i',  'xe',\n",
    "     'cs', 'ba', 'la', 'ce', 'pr', 'nd', 'pm', 'sm', 'eu', 'gd', 'tb', 'dy',\n",
    "     'ho', 'er', 'tm', 'yb', 'lu', 'hf', 'ta', 'w',  're', 'os', 'ir', 'pt',\n",
    "     'au', 'hg', 'tl', 'pb', 'bi', 'po', 'at', 'rn',\n",
    "     'fr', 'ra', 'ac', 'th', 'pa', 'u',  'np', 'pu']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "act_out = reduce(sum_dict, [ele[0] for ele in out])\n",
    "dec_out = reduce(sum_dict, [ele[1] for ele in out])\n",
    "all_out = reduce(sum_dict, [act_out, dec_out])\n",
    "out_key = list(zip(reduce(sum_dict, [ele[0] for ele in out]+[ele[1] for ele in out]).keys(), [__ATOM_LIST__[key] for key in reduce(sum_dict, [ele[0] for ele in out]+[ele[1] for ele in out]).keys()]))\n",
    "\n",
    "print(act_out)\n",
    "print(dec_out)\n",
    "print(all_out)\n",
    "\n",
    "print(dict([(__ATOM_LIST__[ele[0]], ele[1]) for ele in sorted(act_out.items(), key=lambda x: x[0])]))\n",
    "print(dict([(__ATOM_LIST__[ele[0]], ele[1]) for ele in sorted(dec_out.items(), key=lambda x: x[0])]))\n",
    "print(dict([(__ATOM_LIST__[ele[0]], ele[1]) for ele in sorted(all_out.items(), key=lambda x: x[0])]))\n",
    "\n",
    "print(dict(sorted(out_key, key=lambda x: x[0])))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, val in dict([(__ATOM_LIST__[ele[0]], ele[1]) for ele in sorted(all_out.items(), key=lambda x: -x[1])]).items():\n",
    "    print(f'\"{key}\": {val}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.5 EGNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5.1 Preprocess data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 获取文件路径\n",
    "tmp_list = []\n",
    "for root, dirs, files in os.walk(os.path.join(base_path, dude_dir), topdown=False):\n",
    "    tmp_dict = {'target': root.split('/')[-1], 'act_sdf': None, 'dec_sdf': None}\n",
    "    for name in files:\n",
    "        if   'actives_final.sdf' in name:\n",
    "            if 'gz' in name:\n",
    "                os.system(f\"gzip -d {os.path.join(root, name)}\")\n",
    "                tmp_dict['act_sdf'] = os.path.join(root, name[:-3])\n",
    "                continue\n",
    "            tmp_dict['act_sdf'] = os.path.join(root, name)\n",
    "        elif 'decoys_final.sdf' in name:\n",
    "            if 'gz' in name:\n",
    "                os.system(f\"gzip -d {os.path.join(root, name)}\")\n",
    "                tmp_dict['dec_sdf'] = os.path.join(root, name[:-3])\n",
    "                continue\n",
    "            tmp_dict['dec_sdf'] = os.path.join(root, name)\n",
    "    tmp_list.append(tmp_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ic| sdf_path_df.shape: (101, 3)\n"
     ]
    },
    {
     "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>target</th>\n",
       "      <th>act_sdf</th>\n",
       "      <th>dec_sdf</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>cxcr4</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>70</th>\n",
       "      <td>nos1</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>92</th>\n",
       "      <td>tgfr1</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>84</th>\n",
       "      <td>pur2</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76</th>\n",
       "      <td>pgh2</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "      <td>/home/jovyan/work-home/molecule-3d-similarity/...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   target                                            act_sdf  \\\n",
       "23  cxcr4  /home/jovyan/work-home/molecule-3d-similarity/...   \n",
       "70   nos1  /home/jovyan/work-home/molecule-3d-similarity/...   \n",
       "92  tgfr1  /home/jovyan/work-home/molecule-3d-similarity/...   \n",
       "84   pur2  /home/jovyan/work-home/molecule-3d-similarity/...   \n",
       "76   pgh2  /home/jovyan/work-home/molecule-3d-similarity/...   \n",
       "\n",
       "                                              dec_sdf  \n",
       "23  /home/jovyan/work-home/molecule-3d-similarity/...  \n",
       "70  /home/jovyan/work-home/molecule-3d-similarity/...  \n",
       "92  /home/jovyan/work-home/molecule-3d-similarity/...  \n",
       "84  /home/jovyan/work-home/molecule-3d-similarity/...  \n",
       "76  /home/jovyan/work-home/molecule-3d-similarity/...  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sdf_path_df = pd.DataFrame(tmp_list).dropna()\n",
    "ic(sdf_path_df.shape)\n",
    "sdf_path_df.sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def de_duplicate(m_list):\n",
    "    name_l = []\n",
    "    inck_l = []\n",
    "    de_m_l = []\n",
    "\n",
    "    for m in tqdm(m_list):\n",
    "        name = m.GetProp('_Name')\n",
    "        inck = Chem.MolToInchiKey(m)\n",
    "        if name in name_l:\n",
    "            # ic('name duplicate')\n",
    "            pass\n",
    "        elif inck in inck_l:\n",
    "            # ic('icnk duplicate')\n",
    "            pass\n",
    "        else:\n",
    "            name_l.append(name)\n",
    "            inck_l.append(inck)\n",
    "            de_m_l.append(m)\n",
    "    \n",
    "    return de_m_l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rdkit\n",
    "map_dict = {\n",
    "    rdkit.Chem.rdchem.BondType.SINGLE   : 1,\n",
    "    rdkit.Chem.rdchem.BondType.DOUBLE   : 2,\n",
    "    rdkit.Chem.rdchem.BondType.TRIPLE   : 3,\n",
    "    rdkit.Chem.rdchem.BondType.AROMATIC : 4,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/jovyan/work-home'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.path.join(os.path.dirname(base_path), 'egnn/data', )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_info_from_sdf(series: pd.Series):\n",
    "    print(f\"🌟  {series['target']}  🌟\")\n",
    "    act = de_duplicate([m for m in Chem.SDMolSupplier(series['act_sdf'], removeHs=False) if m])\n",
    "    dec = de_duplicate([m for m in Chem.SDMolSupplier(series['dec_sdf'], removeHs=False) if m])\n",
    "\n",
    "    def get_info(mol, name, label):\n",
    "        name        = name\n",
    "        smiles      = AllChem.MolToSmiles(mol)\n",
    "        charges     = json.dumps([a.GetAtomicNum() for a in mol.GetAtoms()])\n",
    "        positions   = json.dumps(mol.GetConformer(0).GetPositions().tolist())\n",
    "        mol_adj     = AllChem.GetAdjacencyMatrix(mol)\n",
    "        # GetBondType\n",
    "        edges = []\n",
    "        for in_i, in_a in enumerate(list(mol.GetAtoms())[:-1]):\n",
    "            in_a_bt = [ele.GetBondType() for ele in in_a.GetBonds()[sum(mol_adj[in_i][:in_i+1]):]]\n",
    "            edges.extend([(in_i, wr_i, map_dict[in_a_bt[at_i]]) for at_i,wr_i  in  enumerate([to_i for to_i,to_a in enumerate(mol_adj[in_i]) if (to_i>in_i)and(to_a==1)])])\n",
    "    \n",
    "        return {'name':name, 'smiles':smiles, 'label':label, 'charges':charges, 'positions':positions, 'edges':edges}\n",
    "    act_out = Parallel(n_jobs=n_jobs, verbose=1)(delayed(get_info)(mol, mol.GetProp('_Name'), 'act') for mol in act) \n",
    "    dec_out = Parallel(n_jobs=n_jobs, verbose=1)(delayed(get_info)(mol, mol.GetProp('_Name'), 'dec') for mol in dec)\n",
    "    \n",
    "    save_path = os.path.join(os.path.dirname(base_path), 'egnn/dude/data')\n",
    "    os.system(f'mkdir -p {save_path}')\n",
    "\n",
    "    out_df = pd.DataFrame(act_out + dec_out)\n",
    "    ic(out_df.sample(5))\n",
    "    out_df.to_csv(os.path.join(save_path, series['target']+'.csv'), index=False)"
   ]
  },
  {
   "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>name</th>\n",
       "      <th>smiles</th>\n",
       "      <th>label</th>\n",
       "      <th>charges</th>\n",
       "      <th>positions</th>\n",
       "      <th>edges</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>24337</th>\n",
       "      <td>ZINC01656377</td>\n",
       "      <td>[H]c1c([H])c([H])c(N([H])[H])c(Sc2c([H])c([H])...</td>\n",
       "      <td>dec</td>\n",
       "      <td>[6, 6, 6, 6, 6, 6, 7, 16, 6, 6, 6, 6, 6, 6, 7,...</td>\n",
       "      <td>[[4.0376, 3.9609, 4.3182], [2.916, 3.2525, 3.9...</td>\n",
       "      <td>[(0, 1, 4), (0, 5, 4), (0, 20, 1), (1, 2, 4), ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13197</th>\n",
       "      <td>ZINC65280765</td>\n",
       "      <td>[H]c1c([H])c(S(=O)(=O)N([H])[H])c([H])c([N+](=...</td>\n",
       "      <td>dec</td>\n",
       "      <td>[6, 6, 6, 7, 6, 6, 8, 6, 6, 6, 6, 6, 6, 6, 6, ...</td>\n",
       "      <td>[[-0.0187, 1.5258, 0.0104], [0.0021, -0.0041, ...</td>\n",
       "      <td>[(0, 1, 1), (0, 22, 1), (0, 23, 1), (0, 24, 1)...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9715</th>\n",
       "      <td>ZINC16401530</td>\n",
       "      <td>[H]c1c([H])c(S(=O)(=O)N([H])[H])c([H])c([H])c1...</td>\n",
       "      <td>dec</td>\n",
       "      <td>[6, 6, 6, 6, 6, 6, 6, 7, 6, 16, 7, 7, 6, 6, 6,...</td>\n",
       "      <td>[[-0.3599, 2.7788, 1.8186], [-1.0099, 3.2267, ...</td>\n",
       "      <td>[(0, 1, 4), (0, 5, 4), (0, 26, 1), (1, 2, 4), ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14859</th>\n",
       "      <td>ZINC14244778</td>\n",
       "      <td>[H]c1c([H])c(C(=O)N([H])N([H])C(=O)C([H])([H])...</td>\n",
       "      <td>dec</td>\n",
       "      <td>[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 16, 8, 8, 6,...</td>\n",
       "      <td>[[3.369, -0.1417, 3.4158], [4.4076, -0.976, 3....</td>\n",
       "      <td>[(0, 1, 4), (0, 9, 4), (0, 32, 1), (1, 2, 4), ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2653</th>\n",
       "      <td>ZINC40660176</td>\n",
       "      <td>[H]c1nc(N([H])c2c([H])c([H])c(Cl)c(C(F)(F)F)c2...</td>\n",
       "      <td>dec</td>\n",
       "      <td>[6, 6, 6, 7, 6, 6, 7, 6, 6, 7, 6, 6, 6, 6, 6, ...</td>\n",
       "      <td>[[2.4124, 1.3886, 2.5638], [2.4131, 2.0734, 1....</td>\n",
       "      <td>[(0, 1, 4), (0, 5, 4), (0, 33, 1), (1, 2, 4), ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               name                                             smiles label  \\\n",
       "24337  ZINC01656377  [H]c1c([H])c([H])c(N([H])[H])c(Sc2c([H])c([H])...   dec   \n",
       "13197  ZINC65280765  [H]c1c([H])c(S(=O)(=O)N([H])[H])c([H])c([N+](=...   dec   \n",
       "9715   ZINC16401530  [H]c1c([H])c(S(=O)(=O)N([H])[H])c([H])c([H])c1...   dec   \n",
       "14859  ZINC14244778  [H]c1c([H])c(C(=O)N([H])N([H])C(=O)C([H])([H])...   dec   \n",
       "2653   ZINC40660176  [H]c1nc(N([H])c2c([H])c([H])c(Cl)c(C(F)(F)F)c2...   dec   \n",
       "\n",
       "                                                 charges  \\\n",
       "24337  [6, 6, 6, 6, 6, 6, 7, 16, 6, 6, 6, 6, 6, 6, 7,...   \n",
       "13197  [6, 6, 6, 7, 6, 6, 8, 6, 6, 6, 6, 6, 6, 6, 6, ...   \n",
       "9715   [6, 6, 6, 6, 6, 6, 6, 7, 6, 16, 7, 7, 6, 6, 6,...   \n",
       "14859  [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 16, 8, 8, 6,...   \n",
       "2653   [6, 6, 6, 7, 6, 6, 7, 6, 6, 7, 6, 6, 6, 6, 6, ...   \n",
       "\n",
       "                                               positions  \\\n",
       "24337  [[4.0376, 3.9609, 4.3182], [2.916, 3.2525, 3.9...   \n",
       "13197  [[-0.0187, 1.5258, 0.0104], [0.0021, -0.0041, ...   \n",
       "9715   [[-0.3599, 2.7788, 1.8186], [-1.0099, 3.2267, ...   \n",
       "14859  [[3.369, -0.1417, 3.4158], [4.4076, -0.976, 3....   \n",
       "2653   [[2.4124, 1.3886, 2.5638], [2.4131, 2.0734, 1....   \n",
       "\n",
       "                                                   edges  \n",
       "24337  [(0, 1, 4), (0, 5, 4), (0, 20, 1), (1, 2, 4), ...  \n",
       "13197  [(0, 1, 1), (0, 22, 1), (0, 23, 1), (0, 24, 1)...  \n",
       "9715   [(0, 1, 4), (0, 5, 4), (0, 26, 1), (1, 2, 4), ...  \n",
       "14859  [(0, 1, 4), (0, 9, 4), (0, 32, 1), (1, 2, 4), ...  \n",
       "2653   [(0, 1, 4), (0, 5, 4), (0, 33, 1), (1, 2, 4), ...  "
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_csv('/home/jovyan/work-home/egnn/dude/data/aa2ar.csv').sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extract_info_from_sdf(sdf_path_df.iloc[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Parallel(n_jobs=n_jobs)(delayed(extract_info_from_sdf)(series) for _, series in sdf_path_df.iterrows())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.9.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "3aec0b238e0cea904811e95e3be7636f0428efc14a0842ce8d58eaeab35118c2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}