Mercurial > repos > davidvanzessen > imgt_locus_split

comparison imgt_locus_split.py @ 1:418b7dbc8947 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Uploaded
author davidvanzessen
date Mon, 17 Jul 2017 08:54:02 -0400
parents b00c257f0a67
children 4bb8f6523130
comparison
equal deleted inserted replaced
0:b00c257f0a67 1:418b7dbc8947
42 file = os.path.join(check, file) 42 file = os.path.join(check, file)
43 shutil.move(file, new_file) 43 shutil.move(file, new_file)
44 shutil.rmtree(check) 44 shutil.rmtree(check)
45 45
46 46
47 def filter_tabular_file(old_file, new_file, column, regex): 47 def filter_imgt_file(old_file, new_file, column, fltr):
48 logging.debug("Filtering {0} with {1}".format(old_file, regex.pattern)) 48 logging.debug("Filtering {0} with {1}".format(old_file, fltr))
49 first = True 49 first = True
50 total = 0 50 total = 0
51 remain = 0 51 remain = 0
52 with open(old_file, 'r') as of, open(new_file, 'w') as nf: 52 with open(old_file, 'r') as of, open(new_file, 'w') as nf:
53 column_index = -1 53 column_index = -1
57 column_index = splt.index(column) 57 column_index = splt.index(column)
58 first = False 58 first = False
59 nf.write(line) 59 nf.write(line)
60 continue 60 continue
61 total += 1 61 total += 1
62 if len(splt) >= column_index and regex.search(splt[column_index]): 62 if len(splt) > column_index and splt[column_index].find(fltr) != -1:
63 remain += 1 63 remain += 1
64 nf.write(line) 64 nf.write(line)
65 return total, remain 65 return total, remain
66 66
67 67
68 def all_same_in_list(l): 68 def all_same_in_list(l):
69 return all(l[0] == x for x in l[1:]) 69 return all(l[0] == x for x in l[1:])
70 70
71 71
72 def filter_imgt_dir(imgt_dir, loci): 72 def filter_imgt_dir(imgt_dir, locus):
73 logging.info("Filtering {0} with {1}".format(imgt_dir, loci)) 73 logging.info("Working on {0}".format(locus))
74 imgt_files = [f for f in os.listdir(imgt_dir) if imgt_file_regex.match(f)] 74 imgt_files = [f for f in os.listdir(imgt_dir) if imgt_file_regex.match(f)]
75 tmp_file = os.path.join(imgt_dir, "tmp.txt") 75 tmp_file = os.path.join(imgt_dir, "tmp.txt")
76 totals = [] 76 totals = []
77 remains = [] 77 remains = []
78 loci_regex = re.compile("|".join(loci))
79 for imgt_file in imgt_files: 78 for imgt_file in imgt_files:
80 imgt_file = os.path.join(imgt_dir, imgt_file) 79 imgt_file = os.path.join(imgt_dir, imgt_file)
81 total, remain = filter_tabular_file(imgt_file, tmp_file, "V-GENE and allele", loci_regex) 80 total, remain = filter_imgt_file(imgt_file, tmp_file, "V-GENE and allele", locus)
82 totals.append(total) 81 totals.append(total)
83 remains.append(remain) 82 remains.append(remain)
84 logging.debug("{0} rows, {1} after filtering".format(total, remain)) 83 logging.debug("{0} rows, {1} after filtering".format(total, remain))
85 shutil.move(tmp_file, imgt_file) 84 shutil.move(tmp_file, imgt_file)
86 if not (all_same_in_list(totals) and all_same_in_list(remains)): 85 if not (all_same_in_list(totals) and all_same_in_list(remains)):
87 logging.warning("Not all files had the same number of sequences remaining for {0}: {1}".format(imgt_dir, remains)) 86 logging.warning("Not all files had the same number of sequences remaining for {0}".format(imgt_dir))
88 return totals[0], remains[0] 87 return totals[0], remains[0]
89 88
90 89
91 def make_new_xz_file(input_dir, output_file): 90 def make_new_xz_file(input_dir, output_file):
92 logging.info("Creating new IMGT zip for {0} at {1}".format(input_dir, output_file)) 91 logging.info("Creating new IMGT zip for {0} at {1}".format(input_dir, output_file))
98 out.add(imgt_file, arcname=os.path.basename(imgt_file)) 97 out.add(imgt_file, arcname=os.path.basename(imgt_file))
99 98
100 99
101 def main(): 100 def main():
102 parser = argparse.ArgumentParser() 101 parser = argparse.ArgumentParser()
103 parser.add_argument("-i", "--input", help="The input IMGT file", required=True) 102 parser.add_argument("--input", help="The input IMGT file", required=True)
104 parser.add_argument("-l", "--loci", help="The Loci to filter on", required=True) 103 parser.add_argument("--output-ig", help="The output file for new IMGT ZIP with just IG sequences", default="None")
105 parser.add_argument("-o", "--output", help="The output file for the new IMGT zip with just the filtered sequences", required=True) 104 parser.add_argument("--output-igh", help="The output file for new IMGT ZIP with just IGH sequences", default="None")
105 parser.add_argument("--output-igk", help="The output file for new IMGT ZIP with just IGK sequences", default="None")
106 parser.add_argument("--output-igl", help="The output file for new IMGT ZIP with just IGL sequences", default="None")
107 parser.add_argument("--output-tr", help="The output file for new IMGT ZIP with just TR sequences", default="None")
108 parser.add_argument("--output-tra", help="The output file for new IMGT ZIP with just TRA sequences", default="None")
109 parser.add_argument("--output-trb", help="The output file for new IMGT ZIP with just TRB sequences", default="None")
110 parser.add_argument("--output-trd", help="The output file for new IMGT ZIP with just TRD sequences", default="None")
111 parser.add_argument("--output-trg", help="The output file for new IMGT ZIP with just TRG sequences", default="None")
106 112
107 logging.basicConfig(filename="./log.html", level=logging.DEBUG, format="%(asctime)s:&emsp;%(message)s <br />", 113 logging.basicConfig(filename="./log.html", level=logging.DEBUG, format="%(asctime)s:&emsp;%(message)s <br />",
108 datefmt='%Y/%m/%d %H:%M:%S') 114 datefmt='%Y/%m/%d %H:%M:%S')
109 logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) 115 logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
110 logging.info("Started IMGT locus split") 116 logging.info("Started IMGT locus split")
111 117
112 args = parser.parse_args() 118 args = parser.parse_args()
113 input_file = args.input 119 input_file = args.input
114 loci = args.loci.split(",") 120 output_ig = args.output_ig
115 output_file = args.output 121 output_igh = args.output_igh
122 output_igk = args.output_igk
123 output_igl = args.output_igl
124 output_tr = args.output_tr
125 output_tra = args.output_tra
126 output_trb = args.output_trb
127 output_trd = args.output_trd
128 output_trg = args.output_trg
129
130 loci = {
131 "IG": output_ig,
132 "IGH": output_igh,
133 "IGK": output_igk,
134 "IGL": output_igl,
135 "TR": output_tr,
136 "TRA": output_tra,
137 "TRB": output_trb,
138 "TRD": output_trd,
139 "TRG": output_trg
140 }
141
142 loci_to_filter = {}
116 143
117 logging.debug("All Parameters:") 144 logging.debug("All Parameters:")
118 logging.debug("Input: {0}".format(input_file)) 145 logging.debug("Input: {0}".format(input_file))
119 logging.debug("Loci: {0}".format(loci)) 146 for locus, path in loci.items():
120 logging.debug("Output: {0}".format(output_file)) 147 logging.debug("{0}: {1}".format(locus, path))
148 if path != "None" and os.path.isdir(os.path.split(path)[0]):
149 loci_to_filter[locus] = path
121 150
122 if len(loci) == 0: 151 if len(loci_to_filter) == 0:
123 raise Exception("No locus selected, nothing to do") 152 raise Exception("No locus selected, nothing to do")
153
154 logging.info("Parameters:")
155 for locus, path in loci_to_filter.items():
156 logging.info("{0}: {1}".format(locus, path))
124 157
125 work_dir = tempfile.mkdtemp() 158 work_dir = tempfile.mkdtemp()
126 original_files_dir = os.path.join(work_dir, "original") 159 original_files_dir = os.path.join(work_dir, "original")
127 os.mkdir(original_files_dir) 160 os.mkdir(original_files_dir)
128 161
129 unpack_imgt_zip(input_file, original_files_dir) 162 unpack_imgt_zip(input_file, original_files_dir)
130 163
131 total, remain = filter_imgt_dir(original_files_dir, loci) 164 for locus, path in loci_to_filter.items():
132 logging.info("{0}\t{1}".format(total, remain)) 165 locus_dir = os.path.join(work_dir, locus)
166 shutil.copytree(original_files_dir, locus_dir)
167 total, remain = filter_imgt_dir(locus_dir, locus)
168 logging.info("{0}\t{1}\t{2}\t{3}".format(locus, path, total, remain))
133 169
134 make_new_xz_file(original_files_dir, output_file) 170 make_new_xz_file(locus_dir, loci_to_filter[locus])
135 171
136 172
137 if __name__ == "__main__": 173 if __name__ == "__main__":
138 main() 174 main()