docking_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Docking Analysis\n",
    "\n",
    "For both RDKit (rdkit_250_etkdg) and DMCG (dmcg_250) we generate these small ensembles:\n",
    " * 1 unminimized (single_unranked_unmin)\n",
    " * 1 unbiased sampled minimized (same as above - single_unranked_min)\n",
    " * 1 minimal energy (single_eranked_min)\n",
    " * 5 unbiased sampled energy minimized (5_unranked_min)\n",
    " * 5 sorted by energy and filtered by RMSDs: 0, 0.5, 1.0, 2.0 (5_filtered_R_min)\n",
    " \n",
    "Each of these is docked.  The single minimal energy pose is also docked with 5x the exhaustiveness as a reference for the 5 sized ensembles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re, sys, glob\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "from rdkit.Chem import AllChem as Chem\n",
    "import gzip\n",
    "from matplotlib import cm\n",
    "from matplotlib.lines import Line2D\n",
    "from matplotlib.patches import Rectangle\n",
    "from matplotlib.legend_handler import HandlerLineCollection, HandlerTuple\n",
    "import matplotlib.transforms as transforms\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dkind = 'docked'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "'CNNscore'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-a54e137f739d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     11\u001b[0m         \u001b[0msupp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mChem\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mForwardSDMolSupplier\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msdffile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mdkind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'docked'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m             \u001b[0mmols\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mmol\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGetProp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'CNNscore'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mmol\u001b[0m \u001b[0;32min\u001b[0m \u001b[0msupp\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m             \u001b[0mmols\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmol\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGetProp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'minimizedAffinity'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mmol\u001b[0m \u001b[0;32min\u001b[0m \u001b[0msupp\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-3-a54e137f739d>\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     11\u001b[0m         \u001b[0msupp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mChem\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mForwardSDMolSupplier\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msdffile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mdkind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'docked'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m             \u001b[0mmols\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mmol\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGetProp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'CNNscore'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mmol\u001b[0m \u001b[0;32min\u001b[0m \u001b[0msupp\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m             \u001b[0mmols\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmol\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGetProp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'minimizedAffinity'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mmol\u001b[0m \u001b[0;32min\u001b[0m \u001b[0msupp\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: 'CNNscore'"
     ]
    }
   ],
   "source": [
    "data = []\n",
    "errors = []\n",
    "for fname in glob.glob(f'wierbowski_cd/*/*_{dkind}_*.rms.txt') + glob.glob(f'refined-set/*/*_{dkind}.rms.txt'):\n",
    "    if 'tmp' in fname:\n",
    "        continue\n",
    "    m = re.search(r'(\\S+)/(\\S+)/(\\S+)_(rdkit|dmcg)_250_(etkdg_|dg_)?2dock_(5|single)_(eranked|unranked|filtered_\\S+)_(min|unmin)_(e5_|s5_)?s?docked(_\\S+)?.rms.txt',fname)\n",
    "    bench,target,prefix,generator,rdgen,size,ranking,minimized,exhaustiveness,pdb = m.groups()\n",
    "    sdf = fname.replace('.rms.txt','.sdf.gz')\n",
    "    num = 0\n",
    "    with gzip.open(sdf) as sdffile:\n",
    "        supp = Chem.ForwardSDMolSupplier(sdffile)\n",
    "        if dkind == 'docked':\n",
    "            mols = [mol.GetProp('CNNscore') for mol in supp]\n",
    "        else:\n",
    "            mols = [-float(mol.GetProp('minimizedAffinity')) for mol in supp]\n",
    "        if len(mols) == 0:\n",
    "            errors.append(fname)\n",
    "            continue\n",
    "        idx = np.argmax(mols)\n",
    "        num = len(mols)\n",
    "    \n",
    "    rmsds = pd.read_csv(fname,delim_whitespace=True,usecols=(2,),names=['rmsd'])\n",
    "    if len(rmsds) == 0:\n",
    "        errors.append(fname)\n",
    "        continue\n",
    "    rmsd = rmsds.iloc[idx].rmsd\n",
    "\n",
    "        \n",
    "    thresh = 0\n",
    "    m = re.search(r'filtered_(\\S+)',ranking)\n",
    "    if m:\n",
    "        thresh = float(m.group(1))\n",
    "        \n",
    "    if 'single' in size:\n",
    "        size = 1\n",
    "    else:\n",
    "        size = int(size)\n",
    "        \n",
    "    if exhaustiveness == 'e5_':\n",
    "        exhaustiveness = 40\n",
    "    elif exhaustiveness == 's5_':\n",
    "        size = 5\n",
    "        exhaustiveness = 8\n",
    "        ranking = 'same'\n",
    "    elif exhaustiveness == None:\n",
    "        exhaustiveness = 8\n",
    "    else:\n",
    "        assert 'Unknown exhaustiveness'\n",
    "        \n",
    "    minimized = minimized == 'min'\n",
    "    \n",
    "    if pdb:\n",
    "        pdb = pdb.lstrip('_')\n",
    "        \n",
    "    data.append([bench,target,pdb,prefix,generator,size,ranking,minimized,exhaustiveness,thresh,rmsd,idx,num])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = pd.DataFrame(data,columns=('Benchmark','Target','Receptor','Ligand','Generator','Size','Ranking','Minimized','Exhaustiveness','Threshold','RMSD','BestIndex','NumConfs'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate the distribution of docked pairs across targets.  Most are 100, so I think we are fine combining these all without worry too much about target bias."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = pd.read_csv('wierbowski_cd/ds_cd_input_pairs.txt',delim_whitespace=True,names=('rec','lig','lig2','out'))\n",
    "lines['targ'] = lines.rec.str.extract(r'\\S+/(\\S+)/\\S+_PRO.pdb')\n",
    "plt.hist(lines.groupby('targ').count().out);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(lines.groupby('targ').count()['out'] ==100)/92"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#colors from bioactivity analysis\n",
    "\n",
    "subset_sizes = [1,5,10,25,50,100,250]\n",
    "\n",
    "colormaps = {\n",
    "    'rdkitdg': cm.get_cmap('Greys'),\n",
    "    'rdkitetkdg': cm.get_cmap('autumn_r'),\n",
    "    'rdkit': cm.get_cmap('autumn_r'),\n",
    "    'dmcg': cm.get_cmap('winter_r'),\n",
    "    'auto3d': cm.get_cmap('Purples')\n",
    "}\n",
    "\n",
    "colors = {\n",
    "    'rdkitdg': {N:colormaps['rdkitdg']((i+25)/(len(subset_sizes)+25)) for i,N in enumerate(subset_sizes)},\n",
    "    'rdkitetkdg': {N:colormaps['rdkitetkdg'](i/len(subset_sizes)) for i,N in enumerate(subset_sizes)},\n",
    "    'rdkit': {N:colormaps['rdkitetkdg'](i/len(subset_sizes)) for i,N in enumerate(subset_sizes)},\n",
    "    'dmcg': {N:colormaps['dmcg'](i/len(subset_sizes)) for i,N in enumerate(subset_sizes)},\n",
    "    'auto3d': {N:colormaps['auto3d']((i+2)/(len(subset_sizes)+2)) for i,N in enumerate(subset_sizes)}\n",
    "}\n",
    "\n",
    "names = {\n",
    "    'platinum2017': \"Platinum 2017\",\n",
    "    'refined-set': \"PDBbind 2020 Refined\",\n",
    "    'rdkitdg': \"RDKit DG\",\n",
    "    'rdkitetkdg': \"RDKit ETKDGv3\",\n",
    "    'rdkit': \"RDKit ETKDGv3\",\n",
    "    'dmcg': \"DMCG\",\n",
    "    'auto3d':\"Auto3D\",\n",
    "    'wierbowski_cd': 'Wierbowski CrossDock'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D.groupby(['Benchmark','Generator','Size','Ranking','Minimized','Exhaustiveness','Threshold']).count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider generating a single conformer with RDKit or DMCG and with or without minimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "        \n",
    "for b in ['wierbowski_cd','refined-set']:\n",
    "    plt.figure(figsize=(6,4),dpi=600)\n",
    "    plt.axvline(2,color='gray',linestyle='--',linewidth=1)    \n",
    "    bench = D[(D.Benchmark == b) & (D.Ranking == 'unranked') & (D.Exhaustiveness == 8) & (D.Size == 1)] \n",
    "    for a in ['rdkit','dmcg']:        \n",
    "        for M in [False,True]:\n",
    "                df = bench[(bench.Minimized==M) & (bench.Generator == a)]\n",
    "                r = df.sort_values(by='RMSD').RMSD\n",
    "                r = np.array(r)\n",
    "                plt.plot([0]+list(r),np.linspace(0,1,len(r)+1),\n",
    "                         linestyle='-' if M else '--',\n",
    "                         alpha=0.5 if a == 'rdkit' else 1.0,\n",
    "                         label=f'{names[a]}' if M else None, color=colors[a][1])\n",
    "                \n",
    "    plt.gca().add_line(Line2D([0], [0], color='black',label='Minimized'))\n",
    "    plt.gca().add_line(Line2D([0], [0], linestyle='--', color='black',label='Unminimized'))\n",
    "    plt.xlim(0,8)\n",
    "    plt.ylim(0,1)\n",
    "    plt.xlabel('RMSD of First Docking Pose')\n",
    "    plt.ylabel(f'Fraction of {names[b]}')\n",
    "    plt.legend();\n",
    "    plt.savefig(f'{dkind}_single_unranked_{b}.pdf',bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Best of 250 versus single sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "        \n",
    "for b in ['wierbowski_cd','refined-set']:\n",
    "    plt.figure(figsize=(6,4),dpi=600)\n",
    "    plt.axvline(2,color='gray',linestyle='--',linewidth=1)    \n",
    "    bench = D[(D.Benchmark == b) & (D.Exhaustiveness == 8) & (D.Minimized) & (D.Size == 1)] \n",
    "    for a in ['rdkit','dmcg']:        \n",
    "        for G in ['unranked','eranked']:\n",
    "                df = bench[(bench.Ranking==G) & (bench.Generator == a)]\n",
    "                r = df.sort_values(by='RMSD').RMSD\n",
    "                r = np.array(r)\n",
    "                plt.plot([0]+list(r),np.linspace(0,1,len(r)+1),                         \n",
    "                         label=f'{names[a]} {\"Single\" if G == \"unranked\" else \"Lowest Energy (250)\"}', color=colors[a][1 if G == 'unranked' else 250])\n",
    "                \n",
    "    plt.xlim(0,8)\n",
    "    plt.ylim(0,1)\n",
    "    plt.xlabel('RMSD of First Docking Pose')\n",
    "    plt.ylabel(f'Fraction of {names[b]}')\n",
    "    plt.legend();\n",
    "    plt.savefig(f'{dkind}_single_eranked_{b}.pdf',bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "        \n",
    "for b in ['wierbowski_cd','refined-set']:\n",
    "    for a in ['rdkit','dmcg']:        \n",
    "\n",
    "        plt.figure(figsize=(6,4),dpi=600)\n",
    "        plt.axvline(2,color='gray',linestyle='--',linewidth=1)    \n",
    "        bench = D[(D.Benchmark == b) & (D.Generator == a) & (D.Minimized)]\n",
    "        \n",
    "        df = bench[(bench.Size == 1) & (bench.Ranking == 'eranked') & (bench.Exhaustiveness == 8)]\n",
    "        r = df.sort_values(by='RMSD').RMSD\n",
    "        r = np.array(r)\n",
    "        plt.plot([0]+list(r),np.linspace(0,1,len(r)+1),   \n",
    "                 linewidth=1,\n",
    "                label=f'Lowest Energy Conformer', color=colors[a][1])\n",
    "\n",
    "#        df = bench[(bench.Size == 5) & (bench.Ranking == 'same') & (bench.Exhaustiveness == 8)]\n",
    "#        r = df.sort_values(by='RMSD').RMSD\n",
    "#        r = np.array(r)\n",
    "#        plt.plot([0]+list(r),np.linspace(0,1,len(r)+1),   \n",
    "#                label=f'Lowest Energy (5x)', color='k')\n",
    "        \n",
    "        df = bench[(bench.Size == 1) & (bench.Ranking == 'eranked') & (bench.Exhaustiveness == 40)]\n",
    "        r = df.sort_values(by='RMSD').RMSD\n",
    "        r = np.array(r)\n",
    "        plt.plot([0]+list(r),np.linspace(0,1,len(r)+1),   \n",
    "                label=f'Lowest Energy (5x Sampling)', color='purple',alpha=0.8)        \n",
    "        \n",
    "        for G in sorted(pd.unique(bench.Ranking)):\n",
    "            \n",
    "            df = bench[(bench.Ranking==G) & (bench.Size == 5)]\n",
    "            if len(df) == 0:\n",
    "                continue\n",
    "            r = df.sort_values(by='RMSD').RMSD\n",
    "            r = np.array(r)\n",
    "            t = 0\n",
    "            if 'filtered_0.5' in G:\n",
    "                linestyle='dotted'\n",
    "                t = 0.5\n",
    "            elif 'filtered_1.0' in G:\n",
    "                linestyle='dashdot'\n",
    "                t = 1.0\n",
    "            elif 'filtered_2.0' in G:\n",
    "                linestyle = 'dashed'\n",
    "                t = 2.0\n",
    "            elif 'filtered' in G:\n",
    "                linestyle = '-'\n",
    "                continue\n",
    "            else:         \n",
    "                continue\n",
    "                linestyle = '-'\n",
    "            plt.plot([0]+list(r),np.linspace(0,1,len(r)+1),       \n",
    "                     linewidth=1, color=colors[a][5], linestyle=linestyle,\n",
    "                     label=f'{t} RMSD Cutoff (5 Conformers)')\n",
    "\n",
    "        plt.xlim(0,8)\n",
    "        plt.ylim(0,1)\n",
    "        plt.title(f'{names[a]}')\n",
    "        plt.xlabel('RMSD of First Docking Pose')\n",
    "        plt.ylabel(f'Fraction of {names[b]}')\n",
    "        plt.legend();\n",
    "        plt.savefig(f'{dkind}_multi_eranked_{a}_{b}.pdf',bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = bench[(bench.Size == 1) & (bench.Ranking == 'eranked') & (bench.Exhaustiveness == 40)]\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D.groupby(['Benchmark','Generator','Size','Ranking','Minimized','Exhaustiveness','Threshold']).count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider just the 2A level of accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench = D[(D.Minimized)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench.groupby(['Benchmark','Generator','Size','Ranking','Minimized','Exhaustiveness','Threshold']).count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench = bench.copy()\n",
    "bench.loc[:,'combo'] =  bench[['Size','Ranking','Exhaustiveness']].astype(str).apply(' '.join, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate bootstrap statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "bench.loc[:,'top1'] = (bench.RMSD <= 2).astype(int)\n",
    "\n",
    "p = sns.catplot(data=bench, col_order=['refined-set','wierbowski_cd'], hue='Generator',col='Benchmark',\\\n",
    "                x='combo',y='top1',\\\n",
    "                order=['1 unranked 8', '1 eranked 8', '1 eranked 40', '5 unranked 8', \\\n",
    "                       '5 filtered_0 8','5 filtered_0.5 8', '5 filtered_1.0 8', '5 filtered_2.0 8'],\\\n",
    "                kind='bar',palette={'rdkit': colors['rdkit'][5], 'dmcg': colors['dmcg'][5]},\n",
    "               legend=None, n_boot=1000, errwidth=2,linestyle='-',edgecolor='k',linewidth=0.5)\n",
    "\n",
    "\n",
    "p.fig.set_figwidth(8)\n",
    "p.fig.set_figheight(4)\n",
    "p.fig.set_dpi(300)\n",
    "p.axes[0,0].set_ylabel('Fraction <2Å RMSD',fontsize=14)\n",
    "p.axes[0,0].set_xlabel(None)\n",
    "p.axes[0,1].set_xlabel(None)\n",
    "\n",
    "p.axes[0,0].set_title(names[p.axes[0,0].get_title().split()[-1]])\n",
    "p.axes[0,1].set_title(names[p.axes[0,1].get_title().split()[-1]])\n",
    "\n",
    "p.axes[0,0].errorbar\n",
    "\n",
    "for i in range(3):\n",
    "    for j in range(2):\n",
    "        p.axes[0,j].patches[i+8].set_facecolor(colors['dmcg'][1])\n",
    "        p.axes[0,j].patches[i].set_facecolor(colors['rdkit'][1])\n",
    "\n",
    "p.axes[0,1].legend(handles=p.axes[0,1].get_legend_handles_labels()[0], \n",
    "                   labels=[names[l] for l in p.axes[0,1].get_legend_handles_labels()[1]],\n",
    "                   loc='upper center',ncol=2)\n",
    "p.set_xticklabels(['Random', 'Lowest\\nEnergy', 'Lowest\\nEnergy (5x)', 'Random', \\\n",
    "                       'Lowest\\nEnergy','RMSD\\nCutoff 0.5', 'RMSD\\nCutoff 1.0', 'RMSD\\nCutoff 2.0'],rotation=90)\n",
    "\n",
    "\n",
    "#p.axes[0,0].set_clip_on(False)\n",
    "\n",
    "for i in range(2):\n",
    "    annotation = p.axes[0,i].annotate('Single Conformer',annotation_clip=False,xy=(1,-0.225),xytext=(1,-0.24),\n",
    "                                  va='top', \n",
    "                                  ha='center',arrowprops={'arrowstyle':'-[, widthB=4, lengthB=0.5,angleB=0'})\n",
    "    \n",
    "    p.axes[0,i].annotate('5 Conformers',annotation_clip=False,xy=(5,-0.225),xytext=(5,-0.24),\n",
    "                                  va='top',\n",
    "                                  ha='center',arrowprops={'arrowstyle':'-[, widthB=6, lengthB=0.5,angleB=0'})\n",
    "\n",
    "    \n",
    "plt.savefig(f'{dkind}_2A.pdf',bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = sns.catplot(data=bench[bench.combo=='1 eranked 8'], x='Benchmark', hue='Generator',y='top1',\n",
    "                kind='bar',palette={'rdkit': '#003594', 'dmcg': '#FFB81C'},\n",
    "                legend=None,n_boot=1000, errwidth=2,linestyle='-',edgecolor='k',linewidth=0.5)\n",
    "\n",
    "p.fig.set_figwidth(4)\n",
    "p.fig.set_figheight(4)\n",
    "p.fig.set_dpi(300)\n",
    "a = p.axes[0,0]\n",
    "a.set_ylabel('Fraction <2Å RMSD',fontsize=16)\n",
    "a.set_xlabel(None)\n",
    "a.set_xticklabels(['Crossdocking','Redocking'],fontsize=14)\n",
    "a.legend(handles=[h[0] for h in a.get_legend_handles_labels()[0]], \n",
    "                   labels=['RDKit','DMCG'],\n",
    "                   loc='upper left',fontsize=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.get_legend_handles_labels()[0][1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grouped = bench.groupby(['Benchmark','Generator','combo'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import ttest_ind\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pvals = np.zeros((len(grouped),len(grouped)))\n",
    "for i,g in enumerate(grouped.groups):\n",
    "    for j,g2 in enumerate(grouped.groups):\n",
    "        pvals[i,j] = ttest_ind(grouped.get_group(g)['top1'],grouped.get_group(g2)['top1'])[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.colors import LogNorm\n",
    "\n",
    "plt.figure(figsize=(10,10))\n",
    "sns.heatmap(pvals,xticklabels=grouped.groups,yticklabels=grouped.groups,norm=LogNorm(vmin=pvals.min()+0.00000001, vmax=pvals.max()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for g2 in grouped.groups:\n",
    "    g = (g2[0],g2[1],'1 eranked 40')\n",
    "    print(g2,ttest_ind(grouped.get_group(g)['top1'],grouped.get_group(g2)['top1'])[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}