2b_first_step_ligMPNN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import re\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import math\n",
    "import glob\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import subprocess\n",
    "from shutil import copyfile\n",
    "import tempfile\n",
    "import shutil\n",
    "from collections import defaultdict\n",
    "# import xarray\n",
    "import gzip\n",
    "import time\n",
    "import json\n",
    "from Bio import SeqIO\n",
    "\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "from sklearn import metrics\n",
    "from scipy.cluster.hierarchy import dendrogram\n",
    "import scipy as sp\n",
    "import sklearn as sk\n",
    "\n",
    "from sklearn.cluster import KMeans\n",
    "import sklearn.metrics as sm\n",
    "%matplotlib inline\n",
    "pd.options.display.max_colwidth = 1000\n",
    "\n",
    "sys.path.append('./software/')\n",
    "from get_rmsds_functions import TMalign\n",
    "from get_rmsds_functions import get_RMSD\n",
    "import check_interactions_to_lig\n",
    "\n",
    "# from pyrosetta import *\n",
    "# param_f = '/home/linnaan/ligands/4KL_xanthurenic_acid/4KL_paper2/ready_to_use/good_to_use/4KL_ref2015.params'\n",
    "# init(f'-beta_nov16 -corrections:beta_nov16 -extra_res {param_f}')\n",
    "\n",
    "sys.path.append('./lib/')\n",
    "from libSlurm import make_submit_file\n",
    "from libSlurm import make_dist_plots\n",
    "from libSlurm import sed_inplace\n",
    "from libSlurm import grep\n",
    "from libSlurm import get_flags\n",
    "from libCommonJupyterFunc import get_total_scores\n",
    "from check_interactions_to_lig import get_lig_num\n",
    "from check_interactions_to_lig import get_hb2lig\n",
    "sys.path.append('./software/')\n",
    "import silent_tools # from git@github.com:bcov77/silent_tools.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "home_dir = '/home/linnaan/software/Pseudocycle_small_molecule_binder/' # change to your script directory\n",
    "scratch_dir = '/net/scratch/linnaan/New_AF2_scaffolds/220214_dock/' # change to your temp output directory\n",
    "\n",
    "#ligands\n",
    "# format: {name,[param, [rotamers],requirements,required_atms]}\n",
    "# required_atms should be adjusted based on rifgen rif size results. too small rif\n",
    "# hurt rifdock significantly\n",
    "ligands = {           'T44':['/home/linnaan/ligands/T44/sel_low_energy/good_to_use/T44_0001_relax_charged_genpot.params',\\\n",
    "                  '/home/linnaan/ligands/T44/sel_low_energy/good_to_use/T44_ref2015.params',\\\n",
    "                  [i for i in glob.glob('/home/linnaan/ligands/T44/sel_low_energy/good_to_use/T44_*_relax_charged_ref2015.pdb') if 'linker' not in i],\\\n",
    "                  '2,3,5,6,8,9','N1,H8,H9,H11,O3,O4,O1,O2,H10'],\n",
    "          } \n",
    "\n",
    "# Check burial of tail\n",
    "# edit manually\n",
    "# format: ligand_name: [check_burial_atom, linker_added_ligand_params,burial_cutoff (adjust this),\n",
    "#                     lig_heavy_atm,lig_aromatic_amds]\n",
    "ligands_tails = {   'T44':['C16',\n",
    "                        '/home/linnaan/ligands/T44/sel_low_energy/T44_0005_relax_charged_ref2015_linker.params',\n",
    "                        7,('O1','O2','O3','O4','H10','N1','H8','H9','H11'),\n",
    "                        ('O1','O2','O3','O4','H10','N1')],\n",
    "                } #add tail ligands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this contains 1 scaffold from each cluster, no filter on scaffold quality.\n",
    "# can change with more knowledge\n",
    "scaffold_df = pd.read_csv(f'{home_dir}finalist_21k_info_v220106_clustered_hullpos_forDock_cleaned_155aa.csv')\n",
    "print(len(scaffold_df))\n",
    "scaffold_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_hb_filters(atm_list,scorefxn):\n",
    "    filters,protocols = [],[]\n",
    "    for atm in atm_list:\n",
    "        filters.append(f'<SimpleHbondsToAtomFilter name=\"hb_to_{atm}\" n_partners=\"1\" hb_e_cutoff=\"-0.3\" target_atom_name=\"{atm}\" res_num=\"%%ligand_res_number%%\" scorefxn=\"{scorefxn}\" confidence=\"0\"/>')\n",
    "        protocols.append(f'<Add filter=\"hb_to_{atm}\"/>')\n",
    "    return filters,protocols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot2pandas(df1,df2,features,legend1='sel',legend2='all',ncols=3):\n",
    "    nrows = math.ceil(len(features) / ncols)\n",
    "    (fig, axs) = plt.subplots(ncols=ncols, nrows=nrows, figsize=[4*ncols,3*nrows])\n",
    "    axs = axs.reshape(-1)\n",
    "    for (i, metric) in enumerate(features):\n",
    "        sns.distplot(df1[metric].dropna(), ax=axs[i],label=legend1).legend()\n",
    "        sns.distplot(df2[metric].dropna(), ax=axs[i],label=legend2).legend()\n",
    "    sns.despine()\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manually edit this\n",
    "# feature:[quantile,ascending=T/F,actual_value] #should take top quantile\n",
    "FINAL_CHOSE_QUANTILE = {\n",
    "                    'T44':{'total_hb_to_lig':[5,False,5],\n",
    "                         'contact_molecular_surface':[50,False,367],\n",
    "                         'ddg2':[50,True,-26.6], #this one visually cannot distinguish\n",
    "                         'dsasa':[70,False,0.976],\n",
    "                         'holes_around_lig':[30,True,1.58],\n",
    "                         #'total_atom_depth':[50,True,6.16] #don't use this feature for CHD since checked before\n",
    "                        }\n",
    "                 }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Use ligand MPNN to perform designs from step 6 input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generate commands for ligMPNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Generate commands for ligMPNN, ligMPNN scrtips take silent file\n",
    "#PYTHON = '/home/aivan/.conda/envs/SE3-a6000/bin/python'\n",
    "#PYTHON = '/software/conda/envs/SE3-nvidia/bin/python'\n",
    "PYTHON = '/software/containers/mpnn_binder_design.sif' # this docker is good for py\n",
    "LIGMPNN = f'{home_dir}toolkits/ligMPNN_FR_silent_in.py'\n",
    "ligMPNN_versions = {'v2Noise':\n",
    "    [32,f'{home_dir}ligandMPNN_checkpoint/nov1_biounits_207.pt']}\n",
    "bias_fs = {f'{home_dir}ligandMPNN_bias/bias_AA.jsonl':'_nobias',\\\n",
    "          f'{home_dir}ligandMPNN_bias/polar_bias.json':'_pbias'}\n",
    "use_noise_option = [True]\n",
    "n_mpnn_FR_cycle = 3 # a naive number\n",
    "OUTPUT_NUM = 1 \n",
    "\n",
    "for n,lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):\n",
    "        miss_input = 0\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        os.makedirs(design_dir,exist_ok=True)\n",
    "# option 1: use the finished designs from step 7 final rosetta design\n",
    "# 4KL and CHD used all chosen docks, in the future, use all actual designs which is more accurate\n",
    "#         input_df= pd.read_csv(f'{scratch_dir}6_redo_ref2015_predictor/{lig}_prediction/chosen_scaffolds.csv')\n",
    "        input_df= pd.read_csv(f'{scratch_dir}7_final_design/{lig}_design/sel_for_af2.csv')\n",
    "# option 2: use the finished silent files if you already got many at this step\n",
    "#         predictor_sc = f'{scratch_dir}3_prediction_ref2015_oldmethod/{lig}_prediction/chosen_scaffolds_to_extract.csv' #f'{design_dir}all_predictor.sc'\n",
    "# option 3: use the finished silent files at step 6\n",
    "#         predictor_sc = f'{scratch_dir}6_redo_ref2015_predictor/{lig}_prediction/chosen_scaffolds.csv'\n",
    "#         input_df= pd.read_csv(predictor_sc)\n",
    "        os.chdir(design_dir)\n",
    "        cmds = []\n",
    "        for rotamer in ligands[lig][2]:\n",
    "            rotamer_name = os.path.basename(rotamer).replace('.pdb','')\n",
    "            print(f'For {rotamer_name}:________________________________')\n",
    "            rotamer_sel_df = input_df[input_df['rotamer']==rotamer_name]\n",
    "            for i,cluster in enumerate(list(set(rotamer_sel_df.cluster))):\n",
    "                if i >=0:\n",
    "                    lig_num = int(scaffold_df.loc[scaffold_df['cluster']==cluster,'nres'].values[0])+1\n",
    "                    indi_cluster_dir = scaffold_df.loc[scaffold_df['cluster']==cluster,'indi_cluster_dir'].values[0]\n",
    "                    pos_res_f = scaffold_df.loc[scaffold_df['cluster']==cluster,'posfile_by_hull'].values[0]\n",
    "                    pos_res = []\n",
    "                    with open(pos_res_f,'r') as f:\n",
    "                        for line in f:\n",
    "                            pos_res.append(line.rstrip())\n",
    "                    pos_res = ','.join(pos_res)\n",
    "                    rotamer_cluster_sel = rotamer_sel_df[rotamer_sel_df['cluster']==cluster]\n",
    "                    os.makedirs(f'{rotamer_name}/{indi_cluster_dir}',exist_ok=True)\n",
    "                    for pdb_in in list(set(rotamer_cluster_sel.pdb_path)):\n",
    "                        close_hb_list,close_2aro_list,sat_lig_atm,buttress_list = get_hb2lig(pdb_in,\\\n",
    "                                                        lig, ligands_tails[lig][3], ligands_tails[lig][4])\n",
    "                        print(f'for design {pdb_in} from {cluster}:\\n close hb {close_hb_list}\\n',\n",
    "                              f'close aroma{close_2aro_list}\\nsatisfied atm{sat_lig_atm}\\n',\n",
    "                              f'buttress res{buttress_list}')\n",
    "                        fix_res = ','.join(str(i) for i in list(set(close_hb_list+close_2aro_list+buttress_list)))\n",
    "                        for bias_f in bias_fs:\n",
    "                            for use_noise in use_noise_option:\n",
    "                                for ligMPNNver in ligMPNN_versions:\n",
    "                                    cmd = f'{PYTHON} {ligMPNN} '+\\\n",
    "                                      f'--pdb_path {pdb_in} ' +\\\n",
    "                                      f'--freeze_res_1st_list {fix_res} '+\\\n",
    "                                      f'--target_atm_for_cst '+ ligands[lig][4] +\\\n",
    "                                      f' --ligand_params_path '+ ligands[lig][1]+\\\n",
    "                                      f' --repackable_res {pos_res}'+\\\n",
    "                                      f' --out_folder {rotamer_name}/{indi_cluster_dir}'+\\\n",
    "                                      f' --use_ligand 1 --num_seq_per_target 1 --omit_AAs C '+\\\n",
    "                                      f'--sampling_temp 0.2 --ligand_genpot_params_path '+ ligands[lig][1]+\\\n",
    "                                      f' --out_folder {rotamer_name}/{indi_cluster_dir}'+\\\n",
    "                                      f' --bias_AA_jsonl {bias_f} '+\\\n",
    "                                      f' --suffix {bias_fs[bias_f]}_cyc{n_mpnn_FR_cycle}_noise'+\\\n",
    "                                      f' --target_hb_atms '+ligands[lig][4] + \\\n",
    "                                      f' --n_mpnn_FR_cycle {n_mpnn_FR_cycle}' +\\\n",
    "                                      f' --ligand_res_number {lig_num} --use_ligmpnn_version {ligMPNNver}'+\\\n",
    "                                      f' --checkpoint_path {ligMPNN_versions[ligMPNNver][1]}'+\\\n",
    "                                      f' --save_probs 1 --num_connections {ligMPNN_versions[ligMPNNver][0]}'+\\\n",
    "                                      f' --out_name ligmpnn{bias_fs[bias_f]}'\n",
    "                                    cmds.append(cmd)\n",
    "                      \n",
    "        cmd_f = design_dir + f'{lig}_ligMPNN_cmd'\n",
    "        with open(cmd_f,'w') as fp:\n",
    "            fp.write('\\n'.join(cmds))\n",
    "        make_submit_file(cmds=cmd_f,submitfile=f'{cmd_f}.submit.sh',time='12:00:00',group_size=230,\n",
    "                logsfolder='logs',cpus=1,mem='1800mb') ## SSpred usage save cache which will create error if visited too often. \n",
    "        print(f'write {len(cmds)} design commands, miss input {miss_input}')\n",
    "        SUBMIT = True\n",
    "        if SUBMIT:\n",
    "            subprocess.check_call(f'sbatch {cmd_f}.submit.sh',shell=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rotamer_cluster_sel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. collect all ligMPNN results, FR them and score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## collect all results, FR them and score\n",
    "features = {'cms':True,'ddg':False,'res_totalscore':False,'totalscore':False}\n",
    "for n,lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):\n",
    "        miss_input,total_pdb = 0,0\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        input_df= pd.read_csv(f'{scratch_dir}6_redo_ref2015_predictor/{lig}_prediction/chosen_scaffolds.csv')\n",
    "        os.chdir(design_dir)\n",
    "        df_list = []\n",
    "        for rotamer in ligands[lig][2]:\n",
    "            rotamer_name = os.path.basename(rotamer).replace('.pdb','')\n",
    "            print(f'For {rotamer_name}:________________________________')\n",
    "            rotamer_sel_df = input_df[input_df['rotamer']==rotamer_name]\n",
    "            for i,cluster in enumerate(list(set(rotamer_sel_df.cluster))):\n",
    "                if i >= 0:\n",
    "                    indi_cluster_dir = scaffold_df.loc[scaffold_df['cluster']==cluster,\n",
    "                                                       'indi_cluster_dir'].values[0]\n",
    "                    rotamer_cluster_sel = rotamer_sel_df[rotamer_sel_df['cluster']==cluster]\n",
    "                    ligMPNN_desigs = glob.glob(f'{rotamer_name}/{indi_cluster_dir}*.pdb')\n",
    "                    total_pdb += len(ligMPNN_desigs)\n",
    "                    pdbs = [i for i in ligMPNN_desigs if 'cyc3' in i]\n",
    "                    if len(pdbs)>0:\n",
    "                        for pdb in pdbs:\n",
    "                            scores = {}\n",
    "                            with open(pdb,'r') as f:\n",
    "                                for line in f:\n",
    "                                    if line.rstrip() != '':\n",
    "                                        if line.split()[0] in features:\n",
    "                                              scores[line.split()[0]] = [float(line.split()[-1])]\n",
    "                            indi_df = pd.DataFrame.from_dict(data=scores)\n",
    "                            indi_df['n_mpnn_FR_cycle'] = 3\n",
    "                            indi_df['pdb_path'] = design_dir+pdb \n",
    "                            indi_df['parent'] = pdb.split('_FR0')[0]\n",
    "                            indi_df['bias'] = pdb.split('_')[-3]\n",
    "                            indi_df['indi_cluster_dir'] = '/'.join(pdb.split('/')[-3:-1])+'/'\n",
    "                            cluster = scaffold_df.loc[scaffold_df['indi_cluster_dir']==indi_cluster_dir,\n",
    "                                                      'cluster'].values[0]\n",
    "                            indi_df['cluster'] = cluster\n",
    "                            if 'noise' in pdb:\n",
    "                                indi_df['ligMPNN_group'] = 'noise'\n",
    "                            else:\n",
    "                                indi_df['ligMPNN_group'] = 'no_noise'\n",
    "                            df_list.append(indi_df)\n",
    "            print(f'{rotamer_name} {len(df_list)}')\n",
    "        all_df = pd.concat(df_list)\n",
    "        all_df = all_df.reset_index()\n",
    "        all_df['rotamer'] = all_df['parent'].str.split('/').str[0]\n",
    "        all_df.to_csv(f'{design_dir}all_ligmpnn_scores.csv')\n",
    "        print(f'{lig} got {len(all_df)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Final FR with all to make scores comparable\n",
    "## Redo design for superimposed pdbs, but only do ProteinMPNN+FR\n",
    "# 149sec/design\n",
    "#PYTHON = '/software/conda/envs/pyrosetta/bin/python'\n",
    "PYTHON = '/software/containers/mpnn_binder_design.sif'\n",
    "# this design script count the close residues and freeze them in fd1\n",
    "SCRIPTS = 'toolkits/design_ligand_full_noHBNet_fromDesign_onlyFR_noFD.py'\n",
    "for n,lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):#(lig=='LLC') or \n",
    "        cmds = []\n",
    "        MPNN_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        all_df = pd.read_csv(f'{MPNN_dir}all_ligmpnn_scores.csv')\n",
    "        os.chdir(MPNN_dir)\n",
    "        SUM = 0\n",
    "        for i,(j,row) in enumerate(all_df.iterrows()):\n",
    "            if i >= 0:\n",
    "                pdb = row['pdb_path']\n",
    "                if os.path.isfile(pdb[:-4]+'_FR.pdb'):\n",
    "                    continue\n",
    "                else:\n",
    "                    cluster = row['cluster']  \n",
    "                    name = os.path.basename(pdb).replace('.pdb','')\n",
    "                    indi_dir = row['indi_cluster_dir']\n",
    "                    rotamer = row['rotamer']\n",
    "                    lig_res_num = int(scaffold_df.loc[scaffold_df['cluster']==cluster,'nres'])+1\n",
    "                    cmds.append(f'cd {MPNN_dir}{rotamer}/{indi_dir}; {PYTHON} {SCRIPTS} '+\\\n",
    "                                f'--pdb {pdb} --params ' + ligands[lig][1] + ' ' + \\\n",
    "                                f'--heavy_atms ' + ligands[lig][4] + \\\n",
    "                                f' --ligand_res_num {int(lig_res_num)} --suffix _FR '+\\\n",
    "                                f'--save_to_pdb True')\n",
    "        if len(cmds) > 0:\n",
    "            print(f'for {lig}...')\n",
    "            cmd_f = f'{lig}_ligmpnn_onlyFR'\n",
    "            with open(cmd_f,'w') as f:\n",
    "                f.write('\\n'.join(cmds))\n",
    "            make_submit_file(cmds=cmd_f,submitfile=f'{cmd_f}.submit.sh',time=\"6:00:00\",\n",
    "            group_size=100,logsfolder='logs',cpus=1,mem='1500mb')\n",
    "            SUBMIT = True\n",
    "            if SUBMIT:\n",
    "                subprocess.check_call(f'sbatch {cmd_f}.submit.sh',shell=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Collect FR results and filter with already validated Rosetta Filter (generation step from ProteinMPNN parts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Collect results and filter with Already validated Rosetta Filter\n",
    "## collect all results, FR them and score\n",
    "features = {'cms':True,'ddg':False,'res_totalscore':False,'totalscore':False}\n",
    "for n,lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):\n",
    "        print(f'working on {lig}...')\n",
    "        miss_input,total_pdb = 0,0\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        input_df= pd.read_csv(f'{scratch_dir}6_redo_ref2015_predictor/{lig}_prediction/chosen_scaffolds.csv')\n",
    "        os.chdir(design_dir)\n",
    "        df_list = []\n",
    "        for rotamer in ligands[lig][2]:\n",
    "            rotamer_name = os.path.basename(rotamer).replace('.pdb','')\n",
    "            print(f'For {rotamer_name}:________________________________')\n",
    "            rotamer_sel_df = input_df[input_df['rotamer']==rotamer_name]\n",
    "            for i,cluster in enumerate(list(set(rotamer_sel_df.cluster))):\n",
    "                if i >= 0:\n",
    "                    indi_cluster_dir = scaffold_df.loc[scaffold_df['cluster']==cluster,\n",
    "                                                       'indi_cluster_dir'].values[0]\n",
    "                    ligMPNN_desigs = glob.glob(f'{rotamer_name}/{indi_cluster_dir}*_FR.pdb')\n",
    "                    total_pdb += len(ligMPNN_desigs)\n",
    "                    pdbs = [i for i in ligMPNN_desigs if 'cyc3' in i]\n",
    "                    print(f'got {len(pdbs)} pdbs...')\n",
    "                    if len(pdbs)>0:\n",
    "                        for pdb in pdbs:\n",
    "                            sc = pdb[:-4]+'.sc'\n",
    "                            try:\n",
    "                                indi_df = pd.read_csv(sc)\n",
    "                                indi_df['n_mpnn_FR_cycle'] = 3\n",
    "                                indi_df['pdb_path'] = design_dir+pdb \n",
    "                                indi_df['parent'] = pdb.split('_FR0')[0]\n",
    "                                indi_df['bias'] = pdb.split('_')[-5]\n",
    "                                indi_df['indi_cluster_dir'] = indi_cluster_dir\n",
    "                                indi_df['cluster'] = cluster\n",
    "                                indi_df['rotamer'] = rotamer_name\n",
    "                                if 'noise' in pdb:\n",
    "                                    indi_df['ligMPNN_group'] = 'noise'\n",
    "                                else:\n",
    "                                    indi_df['ligMPNN_group'] = 'no_noise'\n",
    "                                df_list.append(indi_df)\n",
    "                            except:\n",
    "                                print(f'error {sc}')\n",
    "            print(len(df_list))\n",
    "        all_ligmpnn_df = pd.concat(df_list)\n",
    "        all_ligmpnn_df['total_hb_to_lig'] = 0\n",
    "        for feature in list(all_ligmpnn_df.columns):\n",
    "            if 'hb_to' in feature:\n",
    "                all_ligmpnn_df['total_hb_to_lig'] += all_ligmpnn_df[feature]\n",
    "        all_ligmpnn_df = all_ligmpnn_df.reset_index()\n",
    "        all_ligmpnn_df.to_csv(f'{design_dir}all_ligmpnn_FR_scores.csv')\n",
    "        print(len(all_ligmpnn_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finalized_SEL = True\n",
    "for n, lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        all_ligmpnn_df = pd.read_csv(f'{design_dir}all_ligmpnn_FR_scores.csv')\n",
    "        chose_df = all_ligmpnn_df.copy()\n",
    "        if 'total_atom_depth' not in list(chose_df.columns):\n",
    "            chose_df['total_atom_depth'] = 0\n",
    "        print(f'{lig} start from {len(chose_df)}')\n",
    "        for feature in FINAL_CHOSE_QUANTILE[lig]:\n",
    "            chose_df = chose_df.sort_values(by=feature,ascending=FINAL_CHOSE_QUANTILE[lig][feature][1])\n",
    "            value = FINAL_CHOSE_QUANTILE[lig][feature][2]\n",
    "            if FINAL_CHOSE_QUANTILE[lig][feature][1]: #ascending=True\n",
    "                chose_df = chose_df[chose_df[feature] <= value]\n",
    "                print(f'{feature} at quantile {FINAL_CHOSE_QUANTILE[lig][feature][0]} less than value {value},',\n",
    "                      f'left design {len(chose_df)}')\n",
    "            else:\n",
    "                chose_df = chose_df[chose_df[feature] >= value]\n",
    "                print(f'{feature} at quantile {FINAL_CHOSE_QUANTILE[lig][feature][0]} greater than value {value},',\n",
    "                      f'left design {len(chose_df)}')\n",
    "        print(f'Rosetta filter success rate {100*len(chose_df)/len(all_ligmpnn_df)}%')\n",
    "        print(f'chose scaffolds '+str(len(set(list(chose_df['cluster'])))))\n",
    "        ncols = 3\n",
    "        nrows = math.ceil(len(FINAL_CHOSE_QUANTILE[lig].keys()) / ncols)\n",
    "        (fig, axs) = plt.subplots(ncols=ncols, nrows=nrows, figsize=[20,3*nrows])\n",
    "        axs = axs.reshape(-1)\n",
    "\n",
    "        for (i, metric) in enumerate(FINAL_CHOSE_QUANTILE[lig].keys()):\n",
    "            sns.distplot(chose_df[metric].dropna(), ax=axs[i],label='chose')\n",
    "            sns.distplot(all_ligmpnn_df[metric].dropna(), ax=axs[i],label='all')\n",
    "            axs[i].legend()\n",
    "        sns.despine()\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "        \n",
    "        if finalized_SEL:\n",
    "            chose_df.to_csv(f'{design_dir}sel_ligmpnn_FR_scores_for_af2.csv')\n",
    "            #design these chosen docks for accurate MPNN resampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## (optional) check designs from selected list\n",
    "# for n, lig in enumerate(ligands):\n",
    "#     if lig == '4KL':\n",
    "#         design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "#         os.chdir(design_dir)\n",
    "#         chose_df = pd.read_csv(f'{design_dir}sel_ligmpnn_FR_scores_for_af2.csv')\n",
    "#         check_df = chose_df.sample(n = 50)\n",
    "#         for i,(j,row) in enumerate(check_df.iterrows()):\n",
    "#             pdb = ('/'.join(row['pdb_path'].split('/')[-4:]))\n",
    "#             subprocess.check_call(f'ln -s ../{pdb} check_sample/',shell=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Final check with AF2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Finally check with AF2\n",
    "## if gpu: 8min 16 pdb\n",
    "AF2 = '/software/containers/mlfold.sif /home/linnaan/software/af2/run_af2.py'\n",
    "BATCH = 30 # request 32g, 200 aa need 2 g, can load 32/2 = 16 proteins at maximum\n",
    "for n, lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):\n",
    "        cmds = []\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        os.chdir(design_dir)\n",
    "        chA_only = f'check_af2/chA_only/'\n",
    "        AF2_models = f'check_af2/af2_results/'\n",
    "        os.makedirs(chA_only,exist_ok=True)\n",
    "        os.makedirs(AF2_models,exist_ok=True)\n",
    "        chose_df = pd.read_csv(f'{design_dir}sel_ligmpnn_FR_scores_for_af2.csv')\n",
    "        for rotamer in ligands[lig][2]:\n",
    "            rotamer_name = os.path.basename(rotamer).replace('.pdb','')\n",
    "            print(f'For {rotamer_name}:________________________________')\n",
    "            rotamer_sel_df = chose_df[chose_df['rotamer']==rotamer_name]\n",
    "            for cluster in list(set(rotamer_sel_df.cluster)):\n",
    "                cluster_chose_df = rotamer_sel_df[rotamer_sel_df['cluster']==cluster]\n",
    "                indi_cluster_dir = scaffold_df.loc[scaffold_df['cluster']==cluster,'indi_cluster_dir'].values[0]\n",
    "                if len(cluster_chose_df) > 0:\n",
    "                    print(f'pdbs {len(cluster_chose_df)}')\n",
    "                    length = int(scaffold_df.loc[scaffold_df['cluster']==cluster,'nres'].values[0])\n",
    "                    pdbs = list(cluster_chose_df['pdb_path'])\n",
    "                    af2_indi_dir = f'{chA_only}{rotamer_name}/{indi_cluster_dir}'\n",
    "                    af2_result_dir = f'{AF2_models}{rotamer_name}/{indi_cluster_dir}'\n",
    "                    os.makedirs(af2_indi_dir,exist_ok=True)\n",
    "                    os.makedirs(af2_result_dir,exist_ok=True)\n",
    "                    fastas = []\n",
    "                    for i,pdb in enumerate(pdbs): \n",
    "                        fasta_f = f'{af2_indi_dir}'+os.path.basename(pdb)[:-4]+'.fasta'\n",
    "                        af2_model = glob.glob(af2_result_dir+'/'+os.path.basename(pdb)[:-4]+\\\n",
    "                                              '*_model_*.pdb')\n",
    "                        if len(af2_model) > 0:\n",
    "                            subprocess.check_call(f'rm -f {fasta_f}',shell=True)\n",
    "                        else:\n",
    "                            try:\n",
    "                                record_gen = SeqIO.parse(pdb,\"pdb-atom\")\n",
    "                                for record in record_gen:\n",
    "                                    seq = str(record.seq)\n",
    "                                with open(fasta_f,'w') as f:\n",
    "                                    f.write('>'+os.path.basename(pdb)[:-4]+'\\n'+seq)\n",
    "                                fastas.append(fasta_f)\n",
    "                            except:\n",
    "                                print(f'[ERROR] {pdb} ')\n",
    "                    if len(fastas) > 0:\n",
    "                        fasta_list = [fastas[i:i + BATCH] for i in range(0, len(fastas), BATCH)]\n",
    "                        for fasta_fs in fasta_list:\n",
    "                            batch_size = len(fasta_fs)\n",
    "                            fasta_fs = ','.join(fasta_fs)\n",
    "                            cmds.append(f'{AF2} --out_dir {af2_result_dir} '+\\\n",
    "                                        f' --min_length {length} --num_models 5 '+\\\n",
    "                                        f' --batch_size {batch_size} --fasta_fs {fasta_fs}')\n",
    "        if len(cmds) > 0:\n",
    "            cmd_f = f'{lig}_check_af2_cmd'\n",
    "            with open(cmd_f,'w') as f:\n",
    "                f.write('\\n'.join(cmds))\n",
    "            make_submit_file(cmds=cmd_f,submitfile=f'{cmd_f}.submit.sh',time=\"2-0:00:00\",\n",
    "            group_size=(int(len(cmds)/160)+1),logsfolder='logs',cpus=1,mem='6g',needs_gpu=True,\n",
    "                    gpu_type='rtx2080')\n",
    "            SUBMIT = True\n",
    "            if SUBMIT:\n",
    "                subprocess.check_call(f'sbatch {cmd_f}.submit.sh',shell=True)\n",
    "        else:\n",
    "            print(f'{lig} all done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmd_f = f'{lig}_check_af2_cmd'\n",
    "with open(cmd_f,'w') as f:\n",
    "    f.write('\\n'.join(cmds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add AF2 scores\n",
    "## Finally check with AF2\n",
    "## if gpu: 8min 16 pdb\n",
    "## AF2 171min1.74 \n",
    "for n, lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):# or (lig == 'MOA'):\n",
    "        print(f'working on {lig}...')\n",
    "        cmds,df_list = [],[]\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        os.chdir(design_dir)\n",
    "        chA_only = f'check_af2/chA_only/'\n",
    "        AF2_models = f'check_af2/af2_results/'\n",
    "        chose_df = pd.read_csv(f'{design_dir}sel_ligmpnn_FR_scores_for_af2.csv')\n",
    "        for feature in chose_df.columns: # prevent repeakmerge\n",
    "            if 'model' in feature:\n",
    "                chose_df = chose_df.drop(columns=[feature])\n",
    "        for rotamer in ligands[lig][2]:\n",
    "            rotamer_name = os.path.basename(rotamer).replace('.pdb','')\n",
    "            print(f'For {rotamer_name}:________________________________')\n",
    "            rotamer_sel_df = chose_df[chose_df['rotamer']==rotamer_name]\n",
    "            for cluster in list(set(rotamer_sel_df.cluster)):\n",
    "                cluster_chose_df = rotamer_sel_df[rotamer_sel_df['cluster']==cluster]\n",
    "                indi_cluster_dir = scaffold_df.loc[scaffold_df['cluster']==cluster,\\\n",
    "                                                   'indi_cluster_dir'].values[0]\n",
    "                if len(cluster_chose_df) > 0:\n",
    "                    pdbs = list(cluster_chose_df['pdb_path'])\n",
    "                    af2_indi_dir = f'{chA_only}{rotamer_name}/{indi_cluster_dir}'\n",
    "                    af2_result_dir = f'{AF2_models}{rotamer_name}/{indi_cluster_dir}'\n",
    "                    for i,pdb in enumerate(pdbs):\n",
    "                        if i >= 0:\n",
    "                            pdb_info = {}\n",
    "                            af2_models = glob.glob(af2_result_dir+\\\n",
    "                                                   os.path.basename(pdb)[:-4]+'*_model_*.pdb')\n",
    "                            if len(af2_models) > 0:\n",
    "                                chA = f'{af2_indi_dir}'+os.path.basename(pdb)\n",
    "                                subprocess.check_call(f'grep \" A \" {pdb} > {chA}',shell=True)\n",
    "                                for af2_model in af2_models:\n",
    "                                    with open(af2_model,'r') as f:\n",
    "                                        for line in f:\n",
    "                                            if 'plddt' in line:\n",
    "                                                pdb_info[line.split()[0]] = [line.split()[-1]]\n",
    "                                    if 'model_4' in af2_model:\n",
    "                                        pdb_info['unrelaxed_model_4'] = [af2_model]\n",
    "                                        rmsd = get_RMSD(TMalign(chA,af2_model))\n",
    "                                        pdb_info['rmsd_model4'] = rmsd\n",
    "                                pdb_info['pdb_path'] = [pdb]\n",
    "                                indi_df = pd.DataFrame.from_dict(pdb_info)\n",
    "                                df_list.append(indi_df)\n",
    "                            else:\n",
    "                                continue\n",
    "        all_plddt_df = pd.concat(df_list)\n",
    "        chose_df = chose_df.merge(all_plddt_df,how='left',on='pdb_path')\n",
    "        chose_df.to_csv(f'{design_dir}sel_ligmpnn_FR_scores_for_af2.csv')\n",
    "chose_df  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pLDDT_CUTOFF = 85\n",
    "RMSD_CUTOFF = 1.5\n",
    "cutoff_final = True\n",
    "for n,lig in enumerate(ligands):\n",
    "    if (lig == 'T44'):# or (lig == 'MOA'):\n",
    "        print(f'working on {lig}...')\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        chose_df = pd.read_csv(f'{design_dir}sel_ligmpnn_FR_scores_for_af2.csv')\n",
    "        print(f'{lig} start {len(chose_df)}')\n",
    "        sel_df = chose_df[((chose_df['model_1'] > pLDDT_CUTOFF) | \\\n",
    "                          (chose_df['model_2'] > pLDDT_CUTOFF) | \\\n",
    "                          (chose_df['model_3'] > pLDDT_CUTOFF) | \\\n",
    "                          (chose_df['model_4'] > pLDDT_CUTOFF) | \\\n",
    "                          (chose_df['model_5'] > pLDDT_CUTOFF)) & \\\n",
    "                          (chose_df['rmsd_model4'] <= RMSD_CUTOFF) ]\n",
    "        print(f'obtained {len(sel_df)}, success rate {100*len(sel_df)/len(chose_df)}%')\n",
    "        features = ['model_4','rmsd_model4']\n",
    "        plot2pandas(sel_df,chose_df,features,legend1='pass_rosetta_af2',legend2='pass_rosetta',ncols=2)\n",
    "\n",
    "        order_dir = f'{scratch_dir}9_for_order/{lig}_design/'\n",
    "        ligmpnn_design_dir = f'{order_dir}sel_ligMPNN_designs/'\n",
    "        os.makedirs(ligmpnn_design_dir,exist_ok=True)\n",
    "        plot2pandas(sel_df,chose_df,FINAL_CHOSE_QUANTILE[lig],legend1='pass_rosetta_af2',\n",
    "                    legend2='pass_rosetta',ncols=3)\n",
    "        \n",
    "        if cutoff_final:\n",
    "            sel_df['order_name'] = ''\n",
    "            for i,(j,row) in enumerate(sel_df.iterrows()):\n",
    "                pdb = row['pdb_path']\n",
    "                rotamer = row['rotamer']\n",
    "                cluster = row['cluster']\n",
    "                short_name = scaffold_df.loc[scaffold_df['cluster']==cluster,'short_name'].values[0]\n",
    "                original_name = os.path.basename(scaffold_df.loc[scaffold_df['cluster']==cluster,\\\n",
    "                                                                 'path'].values[0]).replace('.pdb','')\n",
    "                new_pdb_name = os.path.basename(pdb)[:-4].replace(original_name,short_name)+f'_{rotamer}.pdb'\n",
    "                cmd = f'cp -n {pdb} {ligmpnn_design_dir}{new_pdb_name}'\n",
    "                subprocess.check_call(cmd,shell=True)\n",
    "                with open(f'{ligmpnn_design_dir}{new_pdb_name}','a+') as f:\n",
    "                    f.write(f'original_pdb {pdb}')\n",
    "                sel_df.loc[j,'order_name'] = f'{ligmpnn_design_dir}{new_pdb_name}'\n",
    "            sel_df.to_csv(f'{order_dir}sel_ligMPNN_designs.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-----------------------------------------------------------\n",
    "## Old scripts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Correlation before FR intermediate results to final scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_score_lines(pdb):\n",
    "    features = {'cms':-1,'ddg':9999,'res_totalscore':9999,'totalscore':9999}\n",
    "    with open(pdb,'r') as fp:\n",
    "        for line in fp:\n",
    "            if line.rstrip() != '':\n",
    "                if line.split()[0] in features:\n",
    "                    features[line.split()[0]] = float(line.split()[1])\n",
    "                else:\n",
    "                    continue\n",
    "            else:\n",
    "                continue\n",
    "    features['description'] = pdb\n",
    "    df = pd.DataFrame.from_dict(features,orient='index')\n",
    "    df_T = df.transpose()\n",
    "    return df_T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Collect results and filter with Already validated Rosetta Filter\n",
    "## collect all results, FR them and score\n",
    "features = {'cms':True,'ddg':False,'res_totalscore':False,'totalscore':False}\n",
    "for n,lig in enumerate(ligands):\n",
    "    if lig == '4KL':\n",
    "        miss_input,total_pdb = 0,0\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/'\n",
    "        input_df= pd.read_csv(f'{scratch_dir}6_redo_ref2015_predictor/4KL_prediction/chosen_scaffolds.csv')\n",
    "        os.chdir(design_dir)\n",
    "        df_list = []\n",
    "        for rotamer in ligands[lig][2]:\n",
    "            rotamer_name = os.path.basename(rotamer).replace('.pdb','')\n",
    "            print(f'For {rotamer_name}:________________________________')\n",
    "            rotamer_sel_df = input_df[input_df['rotamer']==rotamer_name]\n",
    "            for i,cluster in enumerate(list(set(rotamer_sel_df.cluster))):\n",
    "                if i >= 0:\n",
    "                    indi_cluster_dir = scaffold_df.loc[scaffold_df['cluster']==cluster,\n",
    "                                                       'indi_cluster_dir'].values[0]\n",
    "                    ligMPNN_desigs = glob.glob(f'{rotamer_name}/{indi_cluster_dir}*_cyc3_0.pdb')\n",
    "                    total_pdb += len(ligMPNN_desigs)\n",
    "                    pdbs = [i for i in ligMPNN_desigs if 'cyc3' in i]\n",
    "                    if len(pdbs)>0:\n",
    "                        for pdb in pdbs:\n",
    "                            indi_df = get_score_lines(pdb)\n",
    "                            df_list.append(indi_df)\n",
    "            print(len(df_list))\n",
    "        all_ligmpnn_df = pd.concat(df_list)\n",
    "        all_ligmpnn_df = all_ligmpnn_df.reset_index()\n",
    "        #all_ligmpnn_df.to_csv(f'{design_dir}all_ligmpnn_quickFR_scores.csv')\n",
    "        print(len(all_ligmpnn_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_ligmpnn_df.to_csv(f'{design_dir}all_ligmpnn_quickFR_scores.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "design_dir = f'{scratch_dir}7_ligandMPNN_design/4KL_ligMPNN/'\n",
    "all_ligmpnn_df = pd.read_csv(f'{design_dir}all_ligmpnn_quickFR_scores.csv')\n",
    "df_real = pd.read_csv(f'{design_dir}all_ligmpnn_FR_scores.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_real['before_FR'] = df_real['pdb_path'].str.replace('_FR.pdb','.pdb').str.split('/').str[-4:].str.join('/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_real_quick = all_ligmpnn_df.merge(df_real[['contact_molecular_surface','ddg2','total','score_per_res','before_FR']],\\\n",
    "                left_on='description',right_on='before_FR',how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_real_quick = df_real_quick.rename(columns={\"cms\": \"cms_quick\", \"ddg\": \"ddg_quick\",\n",
    "                                              \"res_totalscore\":\"score_per_res_quick\",\n",
    "                                              \"totalscore\":\"total_score_quick\",\n",
    "                                              \"contact_molecular_surface\":\"cms_real\",\n",
    "                                              \"ddg2\":\"ddg_real\",\"total\":\"total_score_real\",\n",
    "                                              \"score_per_res\":\"score_per_res_real\"})\n",
    "df_real_quick"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score  \n",
    "from sklearn.metrics import roc_curve, auc \n",
    "from sklearn.metrics import roc_auc_score\n",
    "sns.set_context('notebook')  \n",
    "\n",
    "# Find relevant features  \n",
    "# Features to draw ROC for, you should have true and predict value for all of them \n",
    "\n",
    "features = ['cms','ddg','total_score','score_per_res']  #,'total_score','score_per_res']\n",
    "higher_better = ['cms']  \n",
    "method = '_quick'\n",
    "df_feature = df_real_quick.copy()\n",
    "df_feature = df_feature.drop(columns=['index'])\n",
    "df_feature = df_feature.dropna()\n",
    "# Prep axes; draw figures \n",
    "ncols = 4  \n",
    "nrows = math.ceil(len(features) / ncols)*2 # *2 because we will also be plotting ROCs  \n",
    "(fig, axs) = plt.subplots(nrows=nrows, ncols=ncols, figsize=[4*ncols,3*nrows])  \n",
    "axs = axs.reshape(-1)  \n",
    "\n",
    "# Plot  \n",
    "for i,f in enumerate(features):\n",
    "    print(f)\n",
    "    #-------- First we plot the scatter ---------  \n",
    "    # f+'_min': predict value; f+'_full': true value \n",
    "    df_feature[f+method] = df_feature[f+method].astype('float')\n",
    "    sns.scatterplot(data=df_feature,x=f+method, y=f+'_real', ax=axs[i*2]) \n",
    "#    plt.plot(df_feature[f+method], df_feature[f+'_real'], ax=axs[i*2]) \n",
    "    #-------- Next we make roc plots ------------\n",
    "    ps = [99,95,85,70,50,25,1] if f in higher_better else [5,15,30,50,75,100]\n",
    "    print(ps)\n",
    "    #ps = [1] if f in higher_better else [5,15,30,50,75,100]  \n",
    "    cutoffs = np.percentile(df_feature[f+'_real'], ps) \n",
    "    print(f'{f} cutoffs {cutoffs}')\n",
    "    dftmp = pd.DataFrame.from_dict({'percentiles':ps, 'cutoffs': cutoffs}).drop_duplicates(subset='cutoffs') #remove duplicates in case after some percentile, the cutoffs are the same \n",
    "    ps = list(dftmp['percentiles'])  \n",
    "    cutoffs = list(dftmp['cutoffs'])  \n",
    "    alt_cuts = list(set(df_feature[f+'_real'])) \n",
    "    print(len(alt_cuts))\n",
    "    if len(alt_cuts)<=5: # if actual data number less or equal to cutoff number \n",
    "        cutoffs = alt_cuts  \n",
    "        ps = len(cutoffs)*['_']            \n",
    "    for percentile, cutoff in zip(ps, cutoffs):  \n",
    "        print(percentile, cutoff)\n",
    "        truth = (df_feature[f+'_real'] >= cutoff) if f in higher_better else (df_feature[f+'_real'] <= cutoff)  \n",
    "        if sum(truth)==0 or sum(truth)==len(truth): # if no training set in range \n",
    "            continue  \n",
    "        pred = df_feature[f+method] if f in higher_better else - df_feature[f+method]  \n",
    "        print(len(truth), len(pred))\n",
    "        y_true = truth\n",
    "        y_score = pred\n",
    "        AUC = roc_auc_score(y_true, y_score)\n",
    "        print(AUC)\n",
    "        fpr, tpr, _ = roc_curve(truth, pred) # https://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_curve.html; fpr: false possitive rate; tpr; true possitive rate \n",
    "#         print(fpr, tpr)\n",
    "        sns.lineplot(fpr,tpr,ax=axs[i*2+1], label=f'{percentile}_c={np.round(cutoff,0)} AUC={np.round(AUC,2)}')#   \n",
    "        axs[i*2+1].set_xlabel('tpr')  \n",
    "        axs[i*2+1].set_ylabel('fpr')  \n",
    "        axs[i*2+1].set_title(f'{f}{method}')  \n",
    "\n",
    "plt.tight_layout()  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmark for ligMPNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Benchmark for ligMPNN\n",
    "PYTHON = '/home/aivan/.conda/envs/SE3-a6000/bin/python'\n",
    "ligMPNN = '/home/linnaan/scripts/scripts/pseudocycle_scripts/GRL_dl_design_scripts/ligMPNN_FR.py'\n",
    "bias_fs = {'/home/drhicks1/scripts/proteinmpnn/design_scripts/bias_AA.jsonl':'_nobias',\\\n",
    "          '/home/linnaan/software/ml_mpnn/polar_bias.json':'_pbias'}\n",
    "for n,lig in enumerate(ligands):\n",
    "    if lig == '4KL':\n",
    "        miss_input = 0\n",
    "        design_dir = f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/benchmark/'\n",
    "        input_df= pd.read_csv(f'{scratch_dir}6_redo_ref2015_predictor/4KL_prediction/chosen_scaffolds.csv')\n",
    "        os.chdir(design_dir)\n",
    "        cmds = []\n",
    "        for rotamer in ligands[lig][2]:\n",
    "            rotamer_name = os.path.basename(rotamer).replace('.pdb','')\n",
    "            print(f'For {rotamer_name}:________________________________')\n",
    "            rotamer_sel_df = input_df[input_df['rotamer']==rotamer_name]\n",
    "            random_df = rotamer_sel_df.sample(25)\n",
    "            for i,cluster in enumerate(list(set(random_df.cluster))):\n",
    "                if i >= 0:\n",
    "                    indi_cluster_dir = scaffold_df.loc[scaffold_df['cluster']==cluster,'indi_cluster_dir'].values[0]\n",
    "                    pos_res_f = scaffold_df.loc[scaffold_df['cluster']==cluster,'posfile_by_hull'].values[0]\n",
    "                    pos_res = []\n",
    "                    with open(pos_res_f,'r') as f:\n",
    "                        for line in f:\n",
    "                            pos_res.append(line.rstrip())\n",
    "                    pos_res = ','.join(pos_res)\n",
    "                    rotamer_cluster_sel = random_df[random_df['cluster']==cluster]\n",
    "                    os.makedirs(f'{rotamer_name}/{indi_cluster_dir}',exist_ok=True)\n",
    "                    for silent_f in list(set(rotamer_cluster_sel.orig_silent_f)):\n",
    "                        tags = list(rotamer_cluster_sel[rotamer_cluster_sel['orig_silent_f']==silent_f].description)\n",
    "                        tag = tags[0]\n",
    "                        if os.path.isfile(f'../{rotamer_name}/{indi_cluster_dir}{tag}.pdb'):\n",
    "                            for bias_f in bias_fs:\n",
    "                                for n_mpnn_FR_cycle in range(1,11):\n",
    "                                    cmds.append(f'{PYTHON} {ligMPNN} '+\\\n",
    "                                          f'--pdb_path ../{rotamer_name}/{indi_cluster_dir}{tag}.pdb '+\\\n",
    "                                          f'--target_atm_for_cst '+ ligands[lig][4] +\\\n",
    "                                          f' --ligand_params_path '+ ligands[lig][1]+\\\n",
    "                                          f' --repackable_res {pos_res} --omit_AAs C'+\\\n",
    "                                          f' --out_folder {rotamer_name}/{indi_cluster_dir}'+\\\n",
    "                                          f' --use_ligand 1 --num_seq_per_target 4 '+\\\n",
    "                                          f'--sampling_temp 0.2 --ligand_genpot_params_path '+ ligands[lig][1]+\\\n",
    "                                          f' --out_folder {rotamer_name}/{indi_cluster_dir}'+\\\n",
    "                                          f' --bias_AA_jsonl {bias_f} '+\\\n",
    "                                          f' --suffix {bias_fs[bias_f]}_cyc{n_mpnn_FR_cycle}'+\\\n",
    "                                          f' --target_hb_atms '+ligands[lig][4] + \\\n",
    "                                          f' --n_mpnn_FR_cycle {n_mpnn_FR_cycle}')\n",
    "                        else:\n",
    "                            miss_input += 1\n",
    "        cmd_f = design_dir + 'ligMPNN_cmd'\n",
    "        with open(cmd_f,'w') as fp:\n",
    "            fp.write('\\n'.join(cmds))\n",
    "        make_submit_file(cmds=cmd_f,\n",
    "                submitfile=f'{cmd_f}.submit.sh',\n",
    "                queue='medium',\n",
    "                group_size=50,\n",
    "                logsfolder='logs',\n",
    "                cpus=1,\n",
    "                mem='3g',\n",
    "                limit_job_to100=False,\n",
    "                needs_gpu=False) ## SSpred usage save cache which will create error if visited too often. \n",
    "        print(f'write {len(cmds)} design commands, miss input {miss_input}')\n",
    "        SUBMIT = False\n",
    "        if SUBMIT:\n",
    "            subprocess.check_call(f'sbatch {cmd_f}.submit.sh',shell=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## collect results to check correlation between ddg v.s. n_mpnn_FR_cycle numbers, \n",
    "## also eyeball to check which is better\n",
    "lig = '4KL'\n",
    "os.chdir(f'{scratch_dir}7_ligandMPNN_design/{lig}_ligMPNN/benchmark')\n",
    "pdbs = glob.glob('*/*/*/*.pdb')\n",
    "print(len(pdbs))\n",
    "df_list = []\n",
    "features = {'cms':True,'ddg':False,'res_totalscore':False,'totalscore':False}\n",
    "for pdb in pdbs:\n",
    "    scores = {}\n",
    "    with open(pdb,'r') as f:\n",
    "        for line in f:\n",
    "            if line.rstrip() != '':\n",
    "                if line.split()[0] in features:\n",
    "                      scores[line.split()[0]] = [float(line.split()[-1])]\n",
    "    indi_df = pd.DataFrame.from_dict(data=scores)\n",
    "    indi_df['n_mpnn_FR_cycle'] = int(pdb.split('_')[-2].replace('cyc',''))\n",
    "    indi_df['pdb_path'] = pdb \n",
    "    indi_df['parent'] = pdb.split('_FR0')[0]\n",
    "    indi_df['bias'] = pdb.split('_')[-3]\n",
    "    df_list.append(indi_df)\n",
    "all_df = pd.concat(df_list)\n",
    "df_list = []\n",
    "for parent in list(all_df.parent):\n",
    "    parent_df = all_df[all_df['parent']==parent]\n",
    "    for feature in features:\n",
    "        max_v, min_v = parent_df[feature].max(),parent_df[feature].min()\n",
    "        if features[feature]:\n",
    "            parent_df[feature+'_relative'] = (parent_df[feature]-min_v)/(max_v-min_v)\n",
    "        else:\n",
    "            parent_df[feature+'_relative'] = (parent_df[feature]-max_v)/(min_v-max_v)\n",
    "    df_list.append(parent_df)\n",
    "all_df = pd.concat(df_list)\n",
    "all_df = all_df.drop_duplicates(subset = 'pdb_path')\n",
    "all_df = all_df.sort_values(by='ddg',ascending=False).reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_features = [i+'_relative' for i in features]\n",
    "#g1= sns.distplot(data=all_df,x='ddg_relative',y='cms_relative',hue='n_mpnn_FR_cycle',style='bias',palette=\"deep\")\n",
    "g1= sns.displot(data=all_df,x='ddg',hue='n_mpnn_FR_cycle',palette=\"deep\")#,style='bias',palette=\"deep\")\n",
    "g2= sns.displot(data=all_df,x='cms',hue='n_mpnn_FR_cycle',palette=\"deep\")#,style='bias',palette=\"deep\")\n",
    "g3 = sns.scatterplot(data=all_df,x='ddg_relative',y='cms_relative',hue='n_mpnn_FR_cycle',palette=\"deep\")#,style='bias',palette=\"deep\")\n",
    "all_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g3 = sns.scatterplot(data=all_df,x='ddg_relative',y='cms_relative',hue='n_mpnn_FR_cycle',style='bias',palette=\"deep\")\n",
    "plt.legend(bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "common_use",
   "language": "python",
   "name": "common_use"
  },
  "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.6.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}