| import numpy as np |
| from tqdm import tqdm |
| import os, argparse, datetime |
|
|
| import torch |
| import torch_geometric |
| from torch_geometric.loader import DataLoader |
|
|
| from feature_extraction.ProtTrans import get_ProtTrans |
| from feature_extraction.process_structure import get_pdb_xyz, process_dssp, match_dssp |
| from utils import * |
| from model import * |
|
|
|
|
| |
| ProtTrans_path = os.environ.get("PROTTRANS_PATH", "/data/user/yuanqm/tools/Prot-T5-XL-U50") |
| |
|
|
| script_path = os.path.split(os.path.realpath(__file__))[0] + "/" |
| model_path = os.path.dirname(script_path[0:-1]) + "/model/" |
|
|
|
|
| def extract_feat(ID_list, seq_list, outpath, gpu): |
| max_len = max([len(seq) for seq in seq_list]) |
| chunk_size = 32 if max_len > 1000 else 64 |
|
|
| esmfold_cmd = "python {}/feature_extraction/esmfold.py -i {} -o {} --chunk-size {}".format(script_path, outpath + "test_seq.fa", outpath + "pdb/", chunk_size) |
| esmfold_model_dir = os.environ.get("ESMFOLD_HUB_DIR") |
| if esmfold_model_dir: |
| esmfold_cmd += " -m {}".format(esmfold_model_dir) |
| if not gpu: |
| esmfold_cmd += " --cpu-only" |
| else: |
| esmfold_cmd = "CUDA_VISIBLE_DEVICES=" + gpu + " " + esmfold_cmd |
| os.system(esmfold_cmd + " | tee {}/esmfold_pred.log".format(outpath)) |
|
|
| Min_protrans = torch.tensor(np.load(script_path + "feature_extraction/Min_ProtTrans_repr.npy"), dtype = torch.float32) |
| Max_protrans = torch.tensor(np.load(script_path + "feature_extraction/Max_ProtTrans_repr.npy"), dtype = torch.float32) |
| get_ProtTrans(ID_list, seq_list, Min_protrans, Max_protrans, ProtTrans_path, outpath, gpu) |
|
|
| print("Processing PDB files...") |
| for ID in tqdm(ID_list): |
| with open(outpath + "pdb/" + ID + ".pdb", "r") as f: |
| X = get_pdb_xyz(f.readlines()) |
| torch.save(torch.tensor(X, dtype = torch.float32), outpath + "pdb/" + ID + '.tensor') |
|
|
| print("Extracting DSSP features...") |
| for i in tqdm(range(len(ID_list))): |
| ID = ID_list[i] |
| seq = seq_list[i] |
|
|
| os.system("{}/feature_extraction/mkdssp -i {}/pdb/{}.pdb -o {}/DSSP/{}.dssp".format(script_path, outpath, ID, outpath, ID)) |
| dssp_seq, dssp_matrix = process_dssp("{}/DSSP/{}.dssp".format(outpath, ID)) |
| if dssp_seq != seq: |
| dssp_matrix = match_dssp(dssp_seq, dssp_matrix, seq) |
|
|
| torch.save(torch.tensor(np.array(dssp_matrix), dtype = torch.float32), "{}/DSSP/{}.tensor".format(outpath, ID)) |
| os.system("rm {}/DSSP/{}.dssp".format(outpath, ID)) |
|
|
|
|
| def predict(ID_list, outpath, batch, gpu): |
| device = torch.device('cuda:' + gpu if torch.cuda.is_available() and gpu else 'cpu') |
|
|
| node_input_dim = nn_config['node_input_dim'] |
| edge_input_dim = nn_config['edge_input_dim'] |
| hidden_dim = nn_config['hidden_dim'] |
| layer = nn_config['layer'] |
| augment_eps = nn_config['augment_eps'] |
| dropout = nn_config['dropout'] |
|
|
| task_list = ["PRO", "PEP", "DNA", "RNA", "ZN", "CA", "MG", "MN", "ATP", "HEME"] |
|
|
| |
| test_dataset = ProteinGraphDataset(ID_list, outpath) |
| test_dataloader = DataLoader(test_dataset, batch_size = batch, shuffle=False, drop_last=False, num_workers=8, prefetch_factor=2) |
|
|
| models = [] |
| for fold in range(5): |
| state_dict = torch.load(model_path + 'fold%s.ckpt'%fold, device) |
| model = GPSite(node_input_dim, edge_input_dim, hidden_dim, layer, augment_eps, dropout, task_list).to(device) |
| model.load_state_dict(state_dict) |
| model.eval() |
| models.append(model) |
|
|
| test_pred_dict = {} |
| for data in tqdm(test_dataloader): |
| data = data.to(device) |
|
|
| with torch.no_grad(): |
| outputs = [model(data.X, data.node_feat, data.edge_index, data.batch).sigmoid() for model in models] |
| outputs = torch.stack(outputs,0).mean(0) |
|
|
| IDs = data.name |
| outputs_split = torch_geometric.utils.unbatch(outputs, data.batch) |
| for i, ID in enumerate(IDs): |
| test_pred_dict[ID] = [] |
| for j in range(len(task_list)): |
| test_pred_dict[ID].append(list(outputs_split[i][:,j].detach().cpu().numpy())) |
|
|
| return test_pred_dict |
|
|
|
|
| def main(seq_info, outpath, batch, gpu): |
| ID_list, seq_list = seq_info |
| for dir_name in ["pdb", "ProtTrans", "DSSP", "pred"]: |
| os.makedirs(outpath + dir_name, exist_ok = True) |
|
|
| print("\n######## Feature extraction begins at {}. ########\n".format(datetime.datetime.now().strftime("%m-%d %H:%M"))) |
|
|
| extract_feat(ID_list, seq_list, outpath, gpu) |
|
|
| print("\n######## Feature extraction is done at {}. ########\n".format(datetime.datetime.now().strftime("%m-%d %H:%M"))) |
|
|
| print("\n######## Prediction begins at {}. ########\n".format(datetime.datetime.now().strftime("%m-%d %H:%M"))) |
|
|
| predictions = predict(ID_list, outpath, batch, gpu) |
|
|
| print("\n######## Prediction is done at {}. ########\n".format(datetime.datetime.now().strftime("%m-%d %H:%M"))) |
|
|
| export_predictions(predictions, seq_list, outpath) |
|
|
| print("\n######## Results are saved in {} ########\n".format(outpath + "pred/")) |
|
|
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument("-i", "--fasta", type = str, help = "Input fasta file", required=True) |
| parser.add_argument("-o", "--outpath", type = str, help = "Output path to save intermediate files and final predictions", required=True) |
| parser.add_argument("-b", "--batch", type = int, default = 4, help = "Batch size for GPSite prediction") |
| parser.add_argument("--gpu", type = str, default = None, help = "The GPU id used for feature extraction and binding site prediction") |
|
|
| args = parser.parse_args() |
| run_id = args.fasta.split("/")[-1].split(".")[0].replace(" ", "_") |
| outpath = args.outpath + "/" + run_id + "/" |
| os.makedirs(outpath, exist_ok = True) |
|
|
| seq_info = process_fasta(args.fasta, outpath) |
|
|
| if seq_info == -1: |
| print("The format of your input fasta file is incorrect! Please check!") |
| elif seq_info == 1: |
| print("Too much sequences! Up to {} sequences are supported each time!".format(MAX_INPUT_SEQ)) |
| else: |
| main(seq_info, outpath, args.batch, args.gpu) |
|
|
