generate_counterexample_typeslines.py

#!/usr/bin/env python3

'''
This script will generate the lines for a new types file with the iterative poses generated from counterexample_generation_jobs.py

!!WARNING!!
Part of this process is to determine which newly generated poses are NOT REDUNDANT with the previously generated ones.
This requires an O(n^2) calculation to calculate the RMSD between every pose...
Ergo, this calculation depending on the number of poses in a given pocket could take a very long time.
This script also works on all ligands present in the pocket, so there is the potential for multiple O(n^2) calculations to take place.

We have done our best to avoid needless calculations, but this is why we generate the lines for each pocket independently

ASSUMPTIONS:
	i) Poses with <2 RMSD to the crystal pose will be labeled as positive poses
	ii) you have obrms installed, and can run it from your commandline
	iii) the jobfile provided as input contains the full PATH to the files specified.
	iv) the gninatypes files (generated by gninatyper) for the poses in args.input have ALREADY BEEN generated.
	v) The crystal ligand files are formatted PDBid_LignameLIGSUFFIX
	vi) The OLD sdf file with the unique poses is named LignameOLDUNIQUESUFFIX

INPUT:
	i) The path to the pocket you are working on
	ii) the threshold RMSD to determine if they are the same pose
	iii) the name for the txt file that contains the lines to write (will be written in the POCKET DIRECTORY)
	iv) the suffix of the NEW sdf file that contains all of the unique poses
	v) the commands file generated from counterexample_generation_jobs.py
	vi) --OPTIONAL-- the suffix of the OLD sdf file that contains all of the unique poses

OUTPUT:
	==Normal==
	i) the typesfile lines to add to generate the new types file
	ii) A SDF file containing all of the unique poses for a given ligand -- named LignameUNIQUE_SUFFIX
	iii) a ___.sdf file which will be the working file for obrms.
'''

import argparse, re, subprocess, os, sys
import pandas as pd
from rdkit.Chem import AllChem as Chem

def check_exists(filename):
	if os.path.isfile(filename) and os.path.getsize(filename)>0:
		return True
	else:
		return False

def get_pocket_lines(filename,pocket):
	'''
	This function reads the lines from filename, and returns only the lines which contain pocket in them.
	'''
	all_lines=open(filename).readlines()
	lines=[x for x in all_lines if pocket in x]
	return lines

def calc_ligand_dic(lines,ligand_suffix):
	'''
	This function will parse the input list of lines and construct 2 dictionaries
		 1) ligand name -> [docked files with that ligand]
		 2) docked_filename -> crystal_file for that pose
	'''
	data={}
	docked_lookup={}
	for line in lines:
		#1) Getting the crystal ligand file
		ligfile=re.split('--autobox_ligand ',line)[1].split()[0]

		#2) Getting the name of the ligand ** here we assume the ligfile is PATH/<PDBid>_<ligname><LIGSUFFIX>
		ligname=ligfile.split('/')[-1].split(ligand_suffix)[0].split('_')[1]

		#3) Check if ligname in data
		if ligname not in data:
			data[ligname]=[]

		#4) grabbing the docked files
		outfile=re.split('-o ',line)[1].split()[0]

		#5) Adding these files to their corresponding places in the dictionary
		data[ligname].append(outfile)
		docked_lookup[outfile]=ligfile

	return data, docked_lookup

def run_obrms(ligand_file,crystal_file):
	'''
	This function returns a list of rmsds of the docked ligand file to the crystal file. The list is in the order of the poses.
	'''

	rmsds=subprocess.check_output(f'obrms {ligand_file} {crystal_file}',shell=True)
	rmsds=str(rmsds,'utf-8').rstrip().split('\n')
	rmsds=[float(x.split()[-1]) for x in rmsds]
	return rmsds

def get_lines_towrite(crystal_lookup,list_of_docked,affinity_lookup,crystal_suffix,good_pose_thresh,bad_pose_thresh):
	'''
	This function will calculate the RMSD of every input pose, to the provided crystal pose.

	returns a dictionary of lines -->  'docked pose filename':[lines to write]
	'''
	lines={}

	for docked in list_of_docked:
		#Figure out affinity.
		affinity=0.0
		crystal=crystal_lookup[docked]
		cr_lookup=crystal.split(crystal_suffix)[0]
		if cr_lookup in affinity_lookup:
			affinity=affinity_lookup
		print(docked,crystal)
		rmsds=run_obrms(docked,crystal)
		counter=0
		lines[docked]=[]
		for r in rmsds:
			if r < good_pose_thresh:
				label='1'
				neg_aff=''
			elif r >= bad_pose_thresh:
				label='0'
				neg_aff='-'
			else:
				#if the pose is in (good_pose_thesh, bad_pose_thresh) it will be skipped
				continue

			rec_gninatypes=docked.split('rec')[0]+'rec_0.gninatypes'
			lig_gninatypes=docked.replace('.sdf','_'+str(counter)+'.gninatypes')
			lines[docked].append(f'{label} {neg_aff}{affinity} {r} {rec_gninatypes} {lig_gninatypes}\n')
			counter+=1
	return lines

def run_obrms_cross(filename):
	'''
	This function returns a pandas dataframe of the RMSD between every pose and every other pose, which is generated using obrms -x
	'''

	csv=subprocess.check_output('obrms -x '+filename,shell=True)
	csv=str(csv,'utf-8').rstrip().split('\n')
	data=pd.DataFrame([x.split(',')[1:] for x in csv],dtype=float)
	return data


parser=argparse.ArgumentParser(description='Create lines to add to types files from counterexample generation. Assumes data file structure is ROOT/POCKET/FILES.')
parser.add_argument('-p','--pocket',type=str,required=True,help='Name of the pocket that you will be generating the lines for.')
parser.add_argument('-r','--root',type=str,required=True,help='PATH to the ROOT of the pockets.')
parser.add_argument('-i','--input',type=str,required=True,help='File that is output from counterexample_generation_jobs.py')
parser.add_argument('-cs','--crystal_suffix',default='_lig.pdb',help='Expresssion to glob the crystal ligand PDB. Defaults to _lig.pdb. Needs to match what was used with counterexample_generation_jobs.py')
parser.add_argument('--old_unique_suffix',type=str,default=None,help='Suffix for the unique ligand sdf file from a previous run. If set we will load that in and add to it. Default behavior is to generate it from provided input file.')
parser.add_argument('-us','--unique_suffix',type=str,default='_it1___.sdf',help='Suffix for the unique ligand sdf file for this run. Defaults to _it1___.sdf. One will be created for each ligand in the pocket.')
parser.add_argument('--unique_threshold',default=0.25,help='RMSD threshold for unique poses. IE poses with RMSD > thresh are considered unique. Defaults to 0.25.')
parser.add_argument('--lower_confusing_threshold',default=0.5,help='CNNscore threshold for identifying confusing good poses. Score < thresh & under 2RMSD is kept and labelled 1. 0<thresh<1. Default 0.5')
parser.add_argument('--upper_confusing_threshold',default=0.9,help='CNNscore threshold for identifying confusing poor poses. If CNNscore > thresh & over 2RMSD pose is kept and labelled 0. lower<thresh<1. Default 0.9')
parser.add_argument('--good_pose_thresh',default=2.0,help='RMSD threshold to identify a good pose. If ligand RMSD to crystal < this value, the pose is labeled good. Defaults to 2.0')
parser.add_argument('--bad_pose_thresh',default=2.0,help='RMSD threshold to identify a bad pose. If the ligand RMSD to crystal >= this value, the pose is labelled as bad. Defaults to 2.0')
parser.add_argument('-o','--outname',type=str,required=True,help='Name of the text file to write the new lines in. DO NOT WRITE THE FULL PATH!')
parser.add_argument('-a','--affinity_lookup',default='pdbbind2017_affs.txt',help='File mapping the PDBid and ligname of the ligand to its pK value. Assmes space delimited "PDBid ligname pK". Defaults to pdbbind2017_affs.txt')
args=parser.parse_args()

#Setting the myroot and root remove variable for use in the script
myroot=os.path.join(args.root,args.pocket,'')
root_remove=os.path.join(args.root,'')


#sanity check threshold
assert args.unique_threshold > 0, "Unique RMSD threshold needs to be positive"
assert 0<args.lower_confusing_threshold <1, "Lower_confusing_threshold needs to be in (0,1)"
assert args.lower_confusing_threshold<args.upper_confusing_threshold<1, "Upper_confusing_threshold needs to be in (lower_confusing_threshold,1)"

#generating our affinity lookup dictionary
affinity_lookup={}
with open(args.affinity_lookup) as infile:
	for line in infile:
		items=line.split()
		key=items[0]+'_'+items[1]
		val=items[2]
		affinity_lookup[key]=val

#first we will generate the dictionary for the ligand - poses we will use.
tocheck=get_pocket_lines(args.input, args.pocket)
datadic, docked_to_crystal_lookup=calc_ligand_dic(tocheck,args.crystal_suffix)

#main loop of the script
with open(myroot+args.outname,'w') as outfile:
	#loop over the ligands
	for cr_name, list_o_ligs in datadic.items():
		if cr_name!='iqz':
			continue
		#0)  Make sure that the working sdf is free.
		sdf_name=myroot+'___.sdf'
		sdf_tmp=myroot+'___tmp.sdf'
		#if this "___sdf" file already exists, we need to delete it and make a new one.
		if check_exists(sdf_name):
			os.remove(sdf_name)

		#1) Figure out ALL of the lines to write
		line_dic=get_lines_towrite(crystal_lookup=docked_to_crystal_lookup,list_of_docked=list_o_ligs,affinity_lookup=affinity_lookup,crystal_suffix=args.crystal_suffix, good_pose_thresh=args.good_pose_thresh, bad_pose_thresh=args.pad_pose_thresh)

		#2) Set up the 'working sdf' for the obrms -x calculations, consisting of the confusing examples + any possible previously generated examples
		# i) iterate over the possible lines for this ligand, keep only the confusing ones,
		#         and write the confusing poses into the working sdf file.

		w=Chem.SDWriter(sdf_name)
		keys=list(line_dic.keys())
		for key in keys:
			kept_lines=[]
			supply=Chem.SDMolSupplier(key,sanitize=False)
			for i,mol in enumerate(supply):
				curr_line=line_dic[key][i]
				score=mol.GetProp('CNNscore')
				label=curr_line.split()[0]
				#if scored "well", but was a bad pose
				if float(score) > args.upper_confusing_threshold and label=='0':
					kept_lines.append(curr_line)
					w.write(mol)
				#or if scored "poor", but was a good pose
				elif float(score) < args.lower_confusing_threshold and label=='1':
					kept_lines.append(curr_line)
					w.write(mol)
			#after the lines have been checked, we overwrite and only store the lines we kept.
			line_dic[key]=kept_lines
		w=None

		# ii) Prepend ___.sdf with the previously existing unique poses sdf
		offset=0

		if args.old_unique_suffix:
			print('Prepending existing similarity sdf to working sdf file')
			old_sdfname=myroot+cr_name+args.old_unique_suffix
			supply=Chem.SDMolSupplier(old_sdfname,sanitize=False)
			offset=len(supply)
			subprocess.check_call('mv %s %s'%(sdf_name,sdf_tmp),shell=True)
			subprocess.check_call('cat %s %s > %s'%(old_sdfname,sdf_tmp,sdf_name),shell=True)

		#3) run obrms -x working_sdf to calculate the rmsd between each pose. This is the O(n^2) calculation
		unique_data=run_obrms_cross(sdf_name)

		#4) determine the newly found "unique" poses
		assignments={}
		for (r,row) in unique_data.iterrows():
			if r not in assignments:
				for simi in row[row<args.unique_threshold].index:
					if simi not in assignments:
						assignments[simi]=r

		to_remove=set([k for (k,v) in assignments.items() if k!=v])
		#5) write the remaining lines for the newly found "unique" poses.
		counter=offset
		for key in keys:
			for line in line_dic[key]:
				if counter not in to_remove:
					outfile.write(line.replace(root_remove,''))
				counter+=1

		#6) Write out the new "uniques" sdf file to allow for easier future generation
		new_unique_sdfname=myroot+cr_name+args.unique_suffix
		w=Chem.SDWriter(new_unique_sdfname)
		supply=Chem.SDMolSupplier(sdf_name,sanitize=False)
		for i,mol in enumerate(supply):
			if i not in to_remove:
				w.write(mol)