-
Notifications
You must be signed in to change notification settings - Fork 1
/
extract-hashvals-by-sample.py
executable file
·103 lines (81 loc) · 3.25 KB
/
extract-hashvals-by-sample.py
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
#! /usr/bin/env python
"""
Build a reverse index of hashes-to-samples.
"""
import sys, os
import sourmash_lib, sourmash_lib.signature
import argparse
from pickle import dump, load
from collections import defaultdict
def traverse_find_sigs(dirnames):
for dirname in dirnames:
for root, dirs, files in os.walk(dirname):
for name in files:
if name.endswith('.sig') or name.endswith('.sbt'):
fullname = os.path.join(root, name)
yield fullname
def main():
p = argparse.ArgumentParser()
p.add_argument('savename')
p.add_argument('sigs', nargs='+')
p.add_argument('-k', '--ksize', default=31, type=int)
p.add_argument('--scaled', default=10000, type=int)
p.add_argument('--traverse-directory', action='store_true')
args = p.parse_args()
# track hashval -> list of signature IDs
hashval_to_sigids = defaultdict(list)
# track hashval -> list of abundances
hashval_to_abunds = defaultdict(list)
# track sigIDs -> (filename, signature md5)
signum = 1
sigid_to_siginfo = {}
# for every minhash in every signature, link it to its NCBI taxonomic ID.
if args.traverse_directory:
inp_files = list(traverse_find_sigs(args.sigs))
else:
inp_files = list(args.sigs)
print('ksize={} scaled={}'.format(args.ksize, args.scaled))
print('loading signatures & traversing hashes')
bad_input = 0
for n, filename in enumerate(inp_files):
if n % 100 == 0:
print('... loading file #', n, 'of', len(inp_files), end='\r')
try:
sig = sourmash_lib.signature.load_one_signature(filename,
ksize=args.ksize)
except (FileNotFoundError, ValueError):
if not args.traverse_directory:
raise
bad_input += 1
continue
# record which file
md5sum = sig.md5sum()
sigid_to_siginfo[signum] = (filename, md5sum)
this_signum = signum
signum += 1
if sig.minhash.scaled < args.scaled:
sig.minhash = sig.minhash.downsample_scaled(args.scaled)
# add hashval -> signature info
mins = sig.minhash.get_mins(with_abundance=True)
try:
mins.keys
except AttributeError: # not a dictionary
mins = dict([ (m, 1) for m in mins ])
for m, abund in mins.items():
hashval_to_sigids[m].append(this_signum)
hashval_to_abunds[m].append(abund)
print('\n...done. {} hashvals, {} signatures'.format(len(hashval_to_sigids), len(sigid_to_siginfo)))
if bad_input:
print('failed to load {} of {} files found'.format(bad_input,
len(inp_files)))
hashval_picklefile = args.savename + '.hashvals'
sig_picklefile = args.savename + '.siginfo'
abund_picklefile = args.savename + '.abunds'
with open(hashval_picklefile, 'wb') as hashval_fp:
dump(hashval_to_sigids, hashval_fp)
with open(sig_picklefile, 'wb') as sig_fp:
dump(sigid_to_siginfo, sig_fp)
with open(abund_picklefile, 'wb') as abund_fp:
dump(hashval_to_abunds, abund_fp)
if __name__ == '__main__':
main()