Mercurial > repos > estrain > seqsero2_v1_0_1

annotate SeqSero2_update_kmer_database.py @ 14:6d28f83c35b1

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Uploaded
author estrain
date Sat, 07 Sep 2019 11:35:02 -0400
parents bf80acab3d66
children
rev   line source
estrain@2 1 #!/usr/bin/env python3
estrain@2 2
estrain@2 3 import argparse
estrain@2 4 import os,subprocess
estrain@2 5 import pickle
estrain@2 6
estrain@2 7 ### SeqSero Kmer
estrain@2 8 def parse_args():
estrain@2 9 "Parse the input arguments, use '-h' for help."
estrain@2 10 parser = argparse.ArgumentParser(usage='Just type "SeqSero2_update_kmer_database.py", it will update kmer database automatically')
estrain@2 11 return parser.parse_args()
estrain@2 12
estrain@2 13 def reverse_complement(sequence):
estrain@2 14 complement = {
estrain@2 15 'A': 'T',
estrain@2 16 'C': 'G',
estrain@2 17 'G': 'C',
estrain@2 18 'T': 'A',
estrain@2 19 'N': 'N',
estrain@2 20 'M': 'K',
estrain@2 21 'R': 'Y',
estrain@2 22 'W': 'W',
estrain@2 23 'S': 'S',
estrain@2 24 'Y': 'R',
estrain@2 25 'K': 'M',
estrain@2 26 'V': 'B',
estrain@2 27 'H': 'D',
estrain@2 28 'D': 'H',
estrain@2 29 'B': 'V'
estrain@2 30 }
estrain@2 31 return "".join(complement[base] for base in reversed(sequence))
estrain@2 32
estrain@2 33 def multifasta_dict(multifasta):
estrain@2 34 multifasta_list = [
estrain@2 35 line.strip() for line in open(multifasta, 'r') if len(line.strip()) > 0
estrain@2 36 ]
estrain@2 37 headers = [i for i in multifasta_list if i[0] == '>']
estrain@2 38 multifasta_dict = {}
estrain@2 39 for h in headers:
estrain@2 40 start = multifasta_list.index(h)
estrain@2 41 for element in multifasta_list[start + 1:]:
estrain@2 42 if element[0] == '>':
estrain@2 43 break
estrain@2 44 else:
estrain@2 45 if h[1:] in multifasta_dict:
estrain@2 46 multifasta_dict[h[1:]] += element
estrain@2 47 else:
estrain@2 48 multifasta_dict[h[1:]] = element
estrain@2 49 return multifasta_dict
estrain@2 50
estrain@2 51 def createKmerDict_reads(list_of_strings, kmer):
estrain@2 52 kmer_table = {}
estrain@2 53 for string in list_of_strings:
estrain@2 54 sequence = string.strip('\n')
estrain@2 55 for i in range(len(sequence) - kmer + 1):
estrain@2 56 new_mer = sequence[i:i + kmer].upper()
estrain@2 57 new_mer_rc = reverse_complement(new_mer)
estrain@2 58 if new_mer in kmer_table:
estrain@2 59 kmer_table[new_mer.upper()] += 1
estrain@2 60 else:
estrain@2 61 kmer_table[new_mer.upper()] = 1
estrain@2 62 if new_mer_rc in kmer_table:
estrain@2 63 kmer_table[new_mer_rc.upper()] += 1
estrain@2 64 else:
estrain@2 65 kmer_table[new_mer_rc.upper()] = 1
estrain@2 66 return kmer_table
estrain@2 67
estrain@2 68 def multifasta_to_kmers_dict(multifasta):
estrain@2 69 multi_seq_dict = multifasta_dict(multifasta)
estrain@2 70 lib_dict = {}
estrain@2 71 for h in multi_seq_dict:
estrain@2 72 lib_dict[h] = set(
estrain@2 73 [k for k in createKmerDict_reads([multi_seq_dict[h]], 27)])
estrain@2 74 return lib_dict
estrain@2 75
estrain@2 76 def get_salmid_invA_database(ex_dir):
estrain@2 77 # read invA kmer and return it
estrain@2 78 a = open(ex_dir + '/invA_mers_dict', 'rb')
estrain@2 79 invA_dict = pickle.load(a)
estrain@2 80 try:
estrain@2 81 del invA_dict['version']
estrain@2 82 except:
estrain@2 83 pass
estrain@2 84 return invA_dict
estrain@2 85
estrain@2 86 def get_salmid_rpoB_database(ex_dir):
estrain@2 87 # read invA kmer and return it
estrain@2 88 a = open(ex_dir + '/rpoB_mers_dict', 'rb')
estrain@2 89 rpoB_dict = pickle.load(a)
estrain@2 90 try:
estrain@2 91 del rpoB_dict['version']
estrain@2 92 except:
estrain@2 93 pass
estrain@2 94 return rpoB_dict
estrain@2 95
estrain@2 96 def main():
estrain@2 97 args = parse_args()
estrain@2 98 ex_dir = os.path.dirname(os.path.realpath(__file__))
estrain@2 99 lib_dict = multifasta_to_kmers_dict(ex_dir + '/H_and_O_and_specific_genes.fasta')
estrain@2 100 invA_dict=get_salmid_invA_database(ex_dir)
estrain@2 101 #rpoB_dict=get_salmid_rpoB_database(ex_dir)
estrain@2 102 lib_dict_new = lib_dict.copy()
estrain@2 103 #print(len(lib_dict_new))
estrain@2 104 lib_dict_new.update(invA_dict)
estrain@2 105 #print(len(lib_dict_new))
estrain@2 106 #lib_dict_new.update(rpoB_dict)
estrain@2 107 #print(len(lib_dict_new))
estrain@2 108 f = open(ex_dir + '/antigens.pickle', "wb")
estrain@2 109 pickle.dump(lib_dict_new, f)
estrain@2 110 f.close()
estrain@2 111
estrain@2 112 if __name__ == '__main__':
estrain@2 113 main()