File size: 4,803 Bytes
feb41b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | import json
import os
import subprocess
import argparse
from tqdm import tqdm
from requests import get
from time import sleep
def process_pdb(pdb_file, output_dir):
file_name = pdb_file.split('.')[0].split('/')[-1]
if os.path.exists(f'{output_dir}/{file_name}.fasta'):
print(f'>>> {file_name} already exists')
return None
# submit a new job and get the ticket
# repeat = True
# try_times = 0
# while repeat:
# result = subprocess.run(
# [
# "curl", "-X", "POST", "-F", f"q=@{pdb_file}",
# "-F", "mode=3diaa", "-F", "database[]=afdb50",
# "-F", "database[]=afdb-proteome", "-F", "database[]=cath50",
# "-F", "database[]=mgnify_esm30", "-F", "database[]=pdb100",
# "-F", "database[]=gmgcl_id", "-F", "database[]=afdb-swissprot",
# "-F", "database[]=bfvd", "-F", "database[]=bfmd",
# "https://search.foldseek.com/api/ticket"
# ],
# stdout=subprocess.PIPE,
# stderr=subprocess.PIPE,
# text=True
# )
# try:
# result = result.stdout
# ticket = json.loads(result)
# repeat = ticket['status'] != 'COMPLETE'
# except:
# sleep(1)
# try_times += 1
# print('>>> Try again for the ' + str(try_times) + ' time')
# continue
# print('>>> Ticket:', ticket)
# result = get('https://search.foldseek.com/api/result/' + ticket['id'] + '/0').json()
structure_aln_dict = json.load(open(f'{output_dir}/{file_name}.json'))[0]
# structure_aln_dict = result
# with open('result.json', 'w') as f:
# json.dump(result, f)
results = structure_aln_dict['results']
query_seq = structure_aln_dict['queries'][0]['sequence']
# with open(f'residue_sequence/{file_name}.fasta', 'r') as f:
# seq = f.readlines()[1].strip()
# if query_seq != seq:
# print(f'Error with {file_name}')
# return None
alignment_dict = {}
for result_db in results:
# alignments is a dict with string keys, each value is a list of alignment dicts
if not result_db['alignments'] or len(result_db['alignments']) == 0:
continue
# Iterate over all alignment lists in the alignments dict
for alignments_list in result_db['alignments'].values():
for target_info in alignments_list:
name = f"{target_info['target']}/prob_{target_info['prob']}/eval_{target_info['eval']}/score_{target_info['score']}/{target_info['qStartPos']}-{target_info['qEndPos']}"
qaln = target_info['qAln']
dbaln = target_info['dbAln']
try:
# get the index list of '-' in qaln
qaln = list(qaln)
miss_index = [i for i in range(len(qaln)) if qaln[i] == '-']
# remove the residues in dbaln according to the missing index list in qaln
dbaln = list(dbaln)
dbaln = ''.join([dbaln[i] for i in range(len(dbaln)) if i not in miss_index])
except Exception as e:
print(e)
print(name)
# fill '-' to the left and right of dbaln to make it the same length as query_seq
left = target_info['qStartPos'] - 1
right = len(query_seq) - target_info['qEndPos']
dbaln = '-' * left + dbaln + '-' * right
alignment_dict[name] = dbaln
if alignment_dict == {}:
print(f'>>> {file_name} has no alignment')
with open(f'{output_dir}/{file_name}.fasta', 'w') as f:
f.write('\n')
return None
seqs = []
with open(f'{output_dir}/{file_name}.fasta', 'w') as f:
for key, value in alignment_dict.items():
if value not in seqs:
seqs.append(value)
else:
continue
f.write(f'>{key}\n{value}\n')
print(f'>>> {file_name} done')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--pdb_dir', type=str)
parser.add_argument('--pdb_file', type=str)
parser.add_argument('--output_dir', type=str, required=True)
args = parser.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
if args.pdb_dir is not None:
pdbs = sorted(os.listdir(args.pdb_dir))
process_bar = tqdm(pdbs)
for pdb in process_bar:
process_pdb(f'{args.pdb_dir}/{pdb}', args.output_dir)
process_bar.set_description(pdb)
elif args.pdb_file is not None:
process_pdb(args.pdb_file, args.output_dir) |