OneScience-Group/TemStaPro-main

1#!/usr/bin/env python32 3# Program that makes thermostability predictions4 5from optparse import OptionParser6from datetime import datetime7import sys8import os9import numpy10from torch.utils.data import DataLoader11from torch.utils.data import TensorDataset12import torch13 14PARAMETERS = {15	"PT_MODEL_PATH": "Rostlab/prot_t5_xl_half_uniref50-enc",16	"DATASET": "major",17	"EMB_TYPE": "mean",18	"CLASSIFIER_TYPE": "imbal",19	"THRESHOLDS": {20		":(40-65]:": ["40", "45", "50", "55", "60", "65"],21		":(40-80]:": ["40", "45", "50", "55", "60", "65", "70", "75", "80"],22	},23	"SEEDS": ["1", "2", "3", "4", "5"],24	"INPUT_SIZE": 1024,25	"HIDDEN_LAYER_SIZES": [256, 128],26	"DEVICE": torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),27	"THRESHOLDS_RANGE": ":(40-65]:",28	"TEMPERATURE_RANGES": {29		":(40-65]:": ["<40", "[40-45)", "[45-50)", "[50-55)", "[55-60)", 30			"[60-65)", "65<="],31		":(40-80]:": ["<40", "[40-45)", "[45-50)", "[50-55)", "[55-60)", 32			"[60-65)", "[65-70)", "[70-75)", "[75-80)", "80<="],33	},34	"THERMOPHILICITY_LABELS": {35		"mesophilic": ["<40", "[40-45)", "<45"],36		"thermophilic": ["[45-50)", "[50-55)", "[55-60)", "[60-65)", 37			"65<=", "[65-70)", "[70-75)", "<75"],38		"hyperthermophilic": ["[75-80)", "80<="]39	},40	"PRINT_THERMOPHILICITY": {41		":(40-65]:": False,42		":(40-80]:": True43	}44}45 46SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))47PROJECT_DIR = os.path.dirname(SCRIPT_DIR)48 49parser = OptionParser()50 51parser.add_option("--input-fasta", "-f", dest="fasta",52	default=None, help="path to the input FASTA file.")53 54parser.add_option("--embeddings-dir", "-e", dest="emb_dir",55	default=None, help="path to the directory to which embeddings "+\56	"files will be saved (cache).")57 58parser.add_option("--PT-directory", "-d", dest="pt_dir",59	default=None, help="path to the directory of ProtTrans model.")60 61parser.add_option("--temstapro-directory", "-t", dest="tsp_dir",62	default=PROJECT_DIR, help="path to the directory of TemStaPro program "+\63	"with its dependencies.")64 65parser.add_option("--more-thresholds", dest="more_thresholds",66	action="store_true", help="option for the mode that outputs "+\67	"additional predictions for upper temperature thresholds and the "+\68	"thremophilicity label")69 70parser.add_option("--mean-output", dest="mean_out",71	default=None, help="path to the output TSV file with mean predictions. "+\72	"Predictions made from the mean embeddings are always printed to STDOUT."+\73	" If this option is given, the output is directed to the given file")74 75parser.add_option("--per-res-output", dest="per_res_out",76	default=None, help="path to the output TSV file with per-residue "+\77	"predictions.")78 79parser.add_option("--per-segment-output", dest="per_segment_out",80	default=None, help="path to the output TSV file with per-residue "+\81	"predictions made for each segment of the sequence.")82 83parser.add_option("--segment-size", dest="segment_size",84	default=41, help="option to set the window size for average smoothening "+\85	"of per residue embeddings ('per-segment-output' option). Default: 41.")86 87parser.add_option("--window-size-predictions", "-w", 88	dest="window_size_predictions",89	default=81, help="option to set the window size for average smoothening "+\90	"of per residue predictions for plotting (option for 'per-res-output' "+\91	"and 'per-segment-output'). Default: 81.")92 93parser.add_option("--per-residue-plot-dir", "-p", dest="plot_dir",94	default=None, help="path to the directory to which inferences "+\95	"plots will be saved (option for 'per-res-output' and "+\96	"'per-res-segment-output' modes. Default: './'.")97 98parser.add_option("--curve-smoothening", "-c", dest="curve_smoothening",99	default=False, action="store_true", 100	help="option for 'per-segment-output' run mode, which adjusts the "+\101	"plot by making an additional smoothening of the curve.")102 103parser.add_option("--portion-size", dest="portion_size",104	default=1000, 105	help="option to set the portions', into which to divide the input "+\106	"of sequences, maximum size. If no division is needed, set the "+\107	"option to 0. Default: 1000.")108 109parser.add_option("--version", "-v", dest="version",110	default=False, action="store_true",111	help="print version of the program and exit.")112 113(options, args) = parser.parse_args()114 115if(options.version):116	print(f"TemStaPro 0.2.{int(os.popen('git rev-list --count HEAD').read().strip())-61}")117	exit()118 119options.window_size_predictions = int(options.window_size_predictions)120options.segment_size = int(options.segment_size)121if(options.more_thresholds): PARAMETERS['THRESHOLDS_RANGE'] = ":(40-80]:"122 123try:124	assert (options.fasta != None), f"{sys.argv[0]}: a FASTA file is required."125except AssertionError as message:126	print(message, file=sys.stderr)127	exit()128 129try:130	assert (options.pt_dir != None), (131		f"{sys.argv[0]}: a path to the ProtTrans model location is required."132	)133except AssertionError as message:134	print(message, file=sys.stderr)135	exit()136 137temstapro_dir = os.path.abspath(options.tsp_dir)138PARAMETERS["CLASSIFIERS_DIR"] = os.path.join(temstapro_dir, "weight")139 140# Importing local modules141 142sys.path.append(os.path.join(temstapro_dir, "scripts"))143sys.path.append(os.path.join(temstapro_dir, "model"))144 145import prottrans_models146import data_process147import model_flow148import results149 150# Standardization of the FASTA file151(sequences, orig_headers, orig_seqs) = prottrans_models.process_FASTA(options.fasta)152 153# Loading the ProtTrans model154print("%s: beginning to load the model " % datetime.now(), file=sys.stderr)155 156pt_model, tokenizer = prottrans_models.load_model_and_tokenizer(options.pt_dir, 157	PARAMETERS["PT_MODEL_PATH"])158 159print("%s: finished loading the model" % datetime.now(), file=sys.stderr)160 161# Dividing sequences into portions162options.portion_size = int(options.portion_size)163if(options.portion_size == 0): options.portion_size = len(sequences)164 165per_res_mode = (options.per_res_out or options.per_segment_out)166 167for i in range(0, len(list(sequences.keys())), options.portion_size):168	portion_keys = list(sequences.keys())[i:i+options.portion_size]169	170	sequences_portion = {}171	for key in portion_keys:172		sequences_portion[key] = sequences[key]173 174	# Check which sequences do not have embeddings generated175	if(options.emb_dir and os.path.exists(options.emb_dir)):176		seqs_wo_emb_portion = data_process.get_sequences_without_embeddings(177			sequences_portion, options.emb_dir, per_res=per_res_mode)178	else:179		seqs_wo_emb_portion = sequences_portion180 181	embeddings = {}182	per_res_dataset = {}183	per_res_sequences_portion = {}184 185	if(len(seqs_wo_emb_portion)):186		gen_emb_start = datetime.now()187		print(f"{datetime.now()}: beginning to generate embeddings", file=sys.stderr)188 189		# Generating embeddings190		embeddings = prottrans_models.get_embeddings(pt_model, tokenizer, 191			seqs_wo_emb_portion, 192			per_residue=per_res_mode, 193			per_protein=True)194 195		gen_emb_end = datetime.now()196		197		# If cache given, save embeddings198		if(options.emb_dir and os.path.exists(options.emb_dir)):199			if(per_res_mode):200				prottrans_models.save_embeddings(seqs_wo_emb_portion, embeddings,201					options.emb_dir, "per_res")202			prottrans_models.save_embeddings(seqs_wo_emb_portion, embeddings, 203				options.emb_dir, "mean")204		elif(options.emb_dir and not os.path.exists(options.emb_dir)):205			print("The given directory (option -e) does not exist, "+\206				"embeddings' PT files will not be saved.", file=sys.stderr)207 208		try:209			prottrans_models.print_embeddings_generation_stats(i, 210				options.portion_size, embeddings, seqs_wo_emb_portion, 211				gen_emb_start, gen_emb_end)212		except ZeroDivisionError:213			print(f"{sys.argv[0]}: no embeddings were generated.", file=sys.stderr)214			sys.exit(1)215 216	# Collecting the required type of embeddings217	dataset = data_process.collect_mean_embeddings(sequences_portion, 218		embeddings=embeddings, emb_dir=options.emb_dir, 219		input_size=PARAMETERS["INPUT_SIZE"])220 221	if(options.per_res_out):222		per_res_dataset = data_process.collect_per_res_embeddings(sequences_portion, 223			orig_seqs, embeddings=embeddings, emb_dir=options.emb_dir, 224			input_size=PARAMETERS["INPUT_SIZE"])225		per_res_sequences_portion = per_res_dataset["z_test"]226	elif(options.per_segment_out):227		per_res_dataset = data_process.collect_per_res_embeddings(sequences_portion, 228			orig_seqs, embeddings=embeddings,229			emb_dir=options.emb_dir, input_size=PARAMETERS["INPUT_SIZE"], smoothen=True, 230			window_size=options.segment_size)231		per_res_sequences_portion = per_res_dataset["z_test"]232 233	test_loader, per_res_test_loader = model_flow.prepare_data_loaders([234		dataset, per_res_dataset], 'test')235 236	print("%s: beginning to make inferences" % datetime.now(), 237		file=sys.stderr)238 239	averaged_inferences, binary_inferences, labels, clashes = model_flow.make_inferences(240		sequences_portion, per_res_sequences_portion, test_loader, 241		per_res_test_loader, PARAMETERS, PARAMETERS["THRESHOLDS_RANGE"])242 243	print("%s: finished making inferences" % datetime.now(), file=sys.stderr)244			245	# Processing results246	for j, loader in enumerate([test_loader, per_res_test_loader]):247		if(loader is None): break248		for seq in averaged_inferences[j].keys():249			labels[j][seq].append(results.get_temperature_label(250				averaged_inferences[j][seq], 251				PARAMETERS["TEMPERATURE_RANGES"][PARAMETERS["THRESHOLDS_RANGE"]], left_hand=True))252			labels[j][seq].append(results.get_temperature_label(253				averaged_inferences[j][seq],254				PARAMETERS["TEMPERATURE_RANGES"][PARAMETERS["THRESHOLDS_RANGE"]], left_hand=False))255			clashes[j][seq].append(results.detect_clash(averaged_inferences[j][seq],256				left_hand=True))257 258	# Processing printing of mean predictions259	if(options.mean_out):260		os.system(f"mkdir -p {os.path.dirname(options.mean_out)}")261		f_mean = open(options.mean_out, "w") if i == 0 else open(options.mean_out, "a")262	else:263		f_mean = sys.stdout264 265	if(i == 0): results.print_inferences_header(f_mean, 266		PARAMETERS["THRESHOLDS"][PARAMETERS["THRESHOLDS_RANGE"]], 267		PARAMETERS["PRINT_THERMOPHILICITY"][PARAMETERS["THRESHOLDS_RANGE"]])268 269	results.print_inferences(averaged_inferences[0], binary_inferences[0],270		orig_headers, labels[0], clashes[0], 271		PARAMETERS["THERMOPHILICITY_LABELS"], f_mean, orig_seqs,272		"mean", PARAMETERS["PRINT_THERMOPHILICITY"][PARAMETERS["THRESHOLDS_RANGE"]])273 274	# Printing per-residue inferences275	if(options.per_res_out):276		os.system(f"mkdir -p {os.path.dirname(options.per_res_out)}")277		f_per_res = open(options.per_res_out, "w") if i == 0 else open(options.per_res_out, "a")278		if(i == 0): results.print_inferences_header(f_per_res, 279			PARAMETERS["THRESHOLDS"][PARAMETERS["THRESHOLDS_RANGE"]], 280			PARAMETERS["PRINT_THERMOPHILICITY"][PARAMETERS["THRESHOLDS_RANGE"]])281 282		results.print_inferences(averaged_inferences[1], binary_inferences[1], 283			orig_headers, labels[1],284			clashes[1], PARAMETERS["THERMOPHILICITY_LABELS"],285			f_per_res, per_res_sequences_portion, "per-res", 286			PARAMETERS["PRINT_THERMOPHILICITY"][PARAMETERS["THRESHOLDS_RANGE"]])287	elif(options.per_segment_out):288		os.system(f"mkdir -p {os.path.dirname(options.per_segment_out)}")289		f_per_res = open(options.per_segment_out, "w") if i == 0 else open(options.per_segment_out, "a")290		if(i == 0): results.print_inferences_header(f_per_res, 291			PARAMETERS["THRESHOLDS"][PARAMETERS["THRESHOLDS_RANGE"]],292			PARAMETERS["PRINT_THERMOPHILICITY"][PARAMETERS["THRESHOLDS_RANGE"]])293 294		results.print_inferences(averaged_inferences[1], binary_inferences[1], 295			orig_headers, labels[1],296			clashes[1], PARAMETERS["THERMOPHILICITY_LABELS"], f_per_res, 297			per_res_sequences_portion, "per-segment",298			PARAMETERS["PRINT_THERMOPHILICITY"][PARAMETERS["THRESHOLDS_RANGE"]])299 300	# Plotting inferences301	if(options.plot_dir):302		os.system(f"mkdir -p {options.plot_dir}")303		results.plot_inferences(304			options.per_res_out, options.per_segment_out,305			averaged_inferences[1],306			PARAMETERS["THRESHOLDS"][PARAMETERS["THRESHOLDS_RANGE"]], options.plot_dir,307			options.window_size_predictions, options.segment_size, 308			options.curve_smoothening)309