|
jpayne@0
|
1 from bz2 import open as bzopen
|
|
jpayne@0
|
2 from gzip import open as gzopen
|
|
jpayne@0
|
3
|
|
jpayne@0
|
4 from contextlib import ExitStack
|
|
jpayne@0
|
5 from itertools import zip_longest
|
|
jpayne@0
|
6 from pathlib import Path
|
|
jpayne@0
|
7 from sys import argv
|
|
jpayne@0
|
8
|
|
jpayne@0
|
9 import random
|
|
jpayne@0
|
10
|
|
jpayne@0
|
11
|
|
jpayne@0
|
12 usage = """
|
|
jpayne@0
|
13
|
|
jpayne@0
|
14 """
|
|
jpayne@0
|
15
|
|
jpayne@0
|
16 def grouper(iterable, n, fillvalue=None):
|
|
jpayne@0
|
17 "Collect data into fixed-length chunks or blocks"
|
|
jpayne@0
|
18 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
|
|
jpayne@0
|
19 args = [iter(iterable)] * n
|
|
jpayne@0
|
20 return zip_longest(*args, fillvalue=fillvalue)
|
|
jpayne@0
|
21
|
|
jpayne@0
|
22 # file compression signatures
|
|
jpayne@0
|
23 magics = {
|
|
jpayne@0
|
24 b'\x1f\x8b\x08':gzopen,
|
|
jpayne@0
|
25 b'\x42\x5a\x68':bzopen,
|
|
jpayne@0
|
26 }
|
|
jpayne@0
|
27
|
|
jpayne@0
|
28 def sniff(path):
|
|
jpayne@0
|
29 "Sniff first three bytes of the file to determine format based on the magic number."
|
|
jpayne@0
|
30 with open(path, 'rb') as fp:
|
|
jpayne@0
|
31 magic = fp.read(3)
|
|
jpayne@0
|
32 return magics.get(magic, open)
|
|
jpayne@0
|
33
|
|
jpayne@0
|
34
|
|
jpayne@0
|
35 def coverage(collection, genome_size):
|
|
jpayne@0
|
36 "Collection of 1 or 2 tuples, whose 2nd item is the read string"
|
|
jpayne@0
|
37 return sum((len(read[0][1]) for read in collection)) / genome_size # reverse read pair doesn't contribute to coverage so we can ignore it
|
|
jpayne@0
|
38
|
|
jpayne@0
|
39
|
|
jpayne@0
|
40 try:
|
|
jpayne@0
|
41 fin, rin, fout, rout, cov, gen_size, *opts = argv[1:]
|
|
jpayne@0
|
42 ins = [fin, rin]
|
|
jpayne@0
|
43 outs = [fout, rout]
|
|
jpayne@0
|
44 except ValueError: # not enough values to unpack
|
|
jpayne@0
|
45 try:
|
|
jpayne@5
|
46 fin, fout, cov, gen_size, *opts = argv[1:]
|
|
jpayne@0
|
47 ins = [fin]
|
|
jpayne@0
|
48 outs = [fout]
|
|
jpayne@0
|
49 except ValueError:
|
|
jpayne@0
|
50 print(usage)
|
|
jpayne@0
|
51 quit(1)
|
|
jpayne@0
|
52 try:
|
|
jpayne@0
|
53 cov = float(cov)
|
|
jpayne@0
|
54 gen_size = int(gen_size)
|
|
jpayne@0
|
55 except ValueError:
|
|
jpayne@0
|
56 print("Desired coverage and assumed genome size should be numbers")
|
|
jpayne@0
|
57 print(usage)
|
|
jpayne@0
|
58 quit(1)
|
|
jpayne@0
|
59
|
|
jpayne@0
|
60 seed = "ed2b99d842cddc1ac81d7c01a0bf0555"
|
|
jpayne@0
|
61 if opts:
|
|
jpayne@0
|
62 seed = opts[0]
|
|
jpayne@0
|
63 random.seed(seed)
|
|
jpayne@0
|
64
|
|
jpayne@0
|
65 assert len(ins) == len(outs)
|
|
jpayne@0
|
66 file_openers = [sniff(path) for path in ins] # output format determined by input format
|
|
jpayne@0
|
67 with ExitStack() as stack:
|
|
jpayne@0
|
68 ins = [stack.enter_context(openn(path, 'r')) for openn, path in zip(file_openers, ins)] # opened input files
|
|
jpayne@0
|
69 inns = [iter(grouper(inn, 4)) for inn in ins] # stateful 4-ply iterator over lines in the input
|
|
jpayne@0
|
70 outs = [stack.enter_context(openn(path, 'w')) for openn, path in zip(file_openers, outs)] # opened output files
|
|
jpayne@0
|
71
|
|
jpayne@4
|
72 for file in ins:
|
|
jpayne@5
|
73 if hasattr(file, "name"):
|
|
jpayne@5
|
74 print(file.name)
|
|
jpayne@4
|
75
|
|
jpayne@0
|
76 # https://en.m.wikipedia.org/wiki/Reservoir_sampling
|
|
jpayne@0
|
77
|
|
jpayne@0
|
78 reservoir = []
|
|
jpayne@0
|
79 # this is going to be 1 or 2-tuples of 4-tuples representing the 4 lines of the fastq file
|
|
jpayne@0
|
80 # we determine its current coverage (and thus its reservoir size) to fill it, which consumes reads
|
|
jpayne@0
|
81 # from the open files
|
|
jpayne@4
|
82 reads = 0
|
|
jpayne@4
|
83 for i, readpair in enumerate(zip(*inns)):
|
|
jpayne@4
|
84 reads += len(readpair[0][1])
|
|
jpayne@0
|
85 reservoir.append(readpair)
|
|
jpayne@4
|
86 if reads / gen_size > cov:
|
|
jpayne@0
|
87 break
|
|
jpayne@0
|
88
|
|
jpayne@0
|
89 k = len(reservoir) # this is about how big the reservoir needs to be to get cov coverage
|
|
jpayne@0
|
90 #W = exp(log(random.random()) / k)
|
|
jpayne@0
|
91
|
|
jpayne@0
|
92 random.shuffle(reservoir)
|
|
jpayne@0
|
93
|
|
jpayne@0
|
94 print(f"{k} reads selected to achieve {coverage(reservoir, gen_size):.3f}X coverage.")
|
|
jpayne@0
|
95
|
|
jpayne@0
|
96 # if the number of reads is too few to meet the coverage cutoff, then the iterators
|
|
jpayne@0
|
97 # should be exhausted and this won't run
|
|
jpayne@0
|
98 # this is essentially Algorithm L, as I understand it
|
|
jpayne@0
|
99 for i, readpair in enumerate(zip(*inns)):
|
|
jpayne@0
|
100 r = random.randint(0, i)
|
|
jpayne@0
|
101 if r < k:
|
|
jpayne@0
|
102 reservoir[r] = readpair
|
|
jpayne@0
|
103
|
|
jpayne@0
|
104 for readpair in reservoir: # output the sampled reads
|
|
jpayne@0
|
105 for read, file in zip(readpair, outs):
|
|
jpayne@0
|
106 defline, read, spacer, quals = read
|
|
jpayne@0
|
107 file.write(defline)
|
|
jpayne@0
|
108 file.write(read)
|
|
jpayne@0
|
109 file.write(spacer)
|
|
jpayne@0
|
110 file.write(quals)
|
|
jpayne@0
|
111
|
|
jpayne@0
|
112 # [fp.close() for fp in ins]
|
|
jpayne@0
|
113 # [fp.close() for fp in outs]
|
