import argparse import pandas as pd import os import numpy as np import re import subprocess import sys def parse_fasta_seq(fasta_file): try: with open(fasta_file, 'r') as file: return [line.strip() for line in file if not line.startswith('>')] except Exception as e: raise ValueError(f"Error parsing FASTA file: {fasta_file}") from e def parse_fasta_name(fasta_file): try: with open(fasta_file, 'r') as file: return [line.strip()[1:] for line in file if line.startswith('>')] except Exception as e: raise ValueError(f"Error parsing FASTA file: {fasta_file}") from e def create_tsv_frags(fastas, smis, screening_batch_size): protein_sequence_list = parse_fasta_seq(fastas) protein_name_list = parse_fasta_name(fastas) full_df = pd.read_csv(smis, delimiter='\t', names=['smiles', 'chem_id']) #splits = int(np.ceil((len(protein_sequence_list))/screening_batch_size)) #protein_sequence_list = np.array_split(protein_sequence_list, splits) #protein_name_list = np.array_split(protein_name_list, splits) try: [full_df.assign(protein_id = [protein_id]*full_df.shape[0]).assign(protein_sequence = [protein_sequence]*full_df.shape[0])[['protein_id', 'chem_id', 'protein_sequence', 'smiles']].to_csv(f"{protein_id}.tsv", sep='\t', header=False, index=False) for protein_id, protein_sequence in zip(protein_name_list,protein_sequence_list)] tsv_files = [f"{protein_id}.tsv" for protein_id in protein_name_list] return tsv_files except Exception as e: print(f"Failed to process files: {e}") raise def screen_frags(tsv_files): #[subprocess.run(["conplex-dti", "predict","--data-file", tsv_file,"--model-path", "/home/omic/ConPLex/models/ConPLex_v1_BindingDB.pt","--outfile", f"{os.path.splitext(os.path.basename(tsv_file))[0]}_results.tsv"]) for tsv_file in tsv_files] [subprocess.run(["conplex-dti", "predict","--data-file", tsv_file,"--model-path", "/home/omic/ConPLex/models/Run_best_model_epoch46.pt","--outfile", f"{os.path.splitext(os.path.basename(tsv_file))[0]}_results.tsv"]) for tsv_file in tsv_files] scores_files = [f"{os.path.splitext(os.path.basename(tsv_file))[0]}_results.tsv" for tsv_file in tsv_files] #scores_files = [] #for tsv_file in tsv_files: #base_name = os.path.splitext(os.path.basename(tsv_file))[0] #scores_file = f"{base_name}_scores.tsv" #subprocess.run([ # "conplex-dti", "predict", # "--data-file", tsv_file, # "--model-path", "/home/omic/ConPLex/models/ConPLex_v1_BindingDB.pt", # "--outfile", f"{base_name}_results.tsv" #]) #with open(scores_file, 'w') as f: # f.write("chem_id\tprotein_id\tscore\n") #with open(f"{base_name}_results.tsv", 'r') as results_file, open(f"{base_name}_temp.tsv", 'w') as temp_file: # for line in results_file: # fields = line.strip().split('\t') # temp_file.write('\t'.join(fields) + '\n') #merged_df = pd.read_csv(f"{base_name}_temp.tsv", sep='\t', names=['chem_id', 'protein_id', 'score']) #tsv_df = pd.read_csv(tsv_file, sep='\t', names=['protein_id', 'chem_id', 'sequence', 'smiles']) #merged_df = merged_df.merge(tsv_df[['chem_id', 'smiles']], on='chem_id') #merged_df.to_csv(scores_file, sep='\t', index=False) #os.remove(f"{base_name}_temp.tsv") #scores_files.append(scores_file) return scores_files def collect_frags(scores_files, fasta_file, outdir, split_write_df = 'no'): def concatenate_dataframes(file_list): df_list = [] header = None for i, file in enumerate(file_list): if os.path.getsize(file) > 0: if i == 0: df = pd.read_csv(file, sep='\t', index_col=None) if df.columns[0] == 'chem_id': header = df.columns else: df = pd.read_csv(file, sep='\t', index_col=None, header=None) df.columns = ['chem_id', 'protein_id', 'score', 'smiles'] else: df = pd.read_csv(file, sep='\t', index_col=None, header=None) if header is not None: df.columns = header df_list.append(df) else: print(f"Skipping empty file: {file}") if df_list: concatenated_df = pd.concat(df_list, ignore_index=True) return concatenated_df else: return pd.DataFrame() scores_df = concatenate_dataframes(scores_files) if not scores_df.empty: score_column_index = 2 if 'score' not in scores_df.columns else 'score' scores_df[score_column_index] = pd.to_numeric(scores_df[score_column_index], errors='coerce') scores_df = scores_df.dropna(subset=[score_column_index]) sorted_scores_df = scores_df.sort_values(by=score_column_index, ascending=False) if 'score' not in scores_df.columns: sorted_scores_df.columns = ['chem_id', 'protein_id', 'score', 'hit', 'smiles'] base_name = os.path.splitext(os.path.basename(fasta_file))[0] all_scores_file = os.path.join(outdir, f"{base_name}_all_scores.tsv") sorted_scores_df.to_csv(all_scores_file, sep='\t', index=False, header=True) else: print("No score files found or all score files are empty.") with open(fasta_file, 'r') as fasta: fasta_lines = fasta.readlines() fasta_sequence = fasta_lines[1].strip() if len(fasta_lines) > 1 else '' if not scores_df.empty: df = sorted_scores_df header = df.columns.tolist() if split_write_df == 'yes': for index, row in df.iterrows(): row_df = pd.DataFrame([row], columns=header) csv_filename = f"{row['complex_name'].replace(' ', '-')}_hit.csv" row_df.to_csv(os.path.join(outdir, csv_filename), index=False) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Process and screen fragments.") parser.add_argument("--fastas", required=True, help="Path to the FASTA file.") parser.add_argument("--smis", required=True, help="Path to the SMILES file.") parser.add_argument("--screening_batch_size", type=int, default=100000, help="Batch size for screening.") parser.add_argument("--outdir", default="output", help="Output directory.") parser.add_argument("--split_write_df", default="no", help="Write all complexes in separate files") args = parser.parse_args() tsv_files = create_tsv_frags(args.fastas, args.smis, args.screening_batch_size) scores_files = screen_frags(tsv_files) #collect_frags(scores_files, args.fastas, args.outdir, args.split_write_df)