reseda/FastqSplitOnSequenceLength.py at master · EDS-Bioinformatics-Laboratory/reseda · GitHub

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
from __future__ import print_function
import sys
import os
import gzip
from random import randrange
from Bio import SeqIO
import argparse


def splitFastq(f, seq_length):
    '''
    Description: split fastq file based on sequence length
    In: f (file name, str), seq_length (int)
    Out: new files will be written to disk, count_short (int), count_long (int), count_total (int)
    '''
    try:
        fh = gzip.open(f, "rt")
    except:
        sys.exit("cannot open file:" + f)

    f = f.split("/")
    fastq_short = f[:-1] + ["short" + str(seq_length) + "-" + f[-1]]
    fastq_long = f[:-1] + ["long" + str(seq_length) + "-" + f[-1]]
    fastq_short = "/".join(fastq_short)
    fastq_long = "/".join(fastq_long)

    try:
        fh_short = gzip.open(fastq_short, "wt")
        fh_long = gzip.open(fastq_long, "wt")
    except:
        sys.exit("cannot write to file: " + fastq_short + " or " + fastq_long)

    count_total = 0
    count_short = 0
    count_long = 0
    for record in SeqIO.parse(fh, "fastq"):
        if len(record.seq) < seq_length:
            SeqIO.write(record, fh_short, "fastq")
            count_short += 1
        else:
            SeqIO.write(record, fh_long, "fastq")
            count_long += 1
        count_total += 1

    fh.close()
    fh_short.close()
    fh_long.close()
    print("Wrote", fastq_short, "to disk")
    print("Wrote", fastq_long, "to disk")
    return(fastq_short, fastq_long, count_short, count_long, count_total)


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Split fastq files based on sequence length. Threshold will be used as minimum for the longer sequences (greater than or equal to)')
    parser.add_argument('-l', '--length', default=270, type=int, help='Threshold: sequence length (default: %(default)s)')
    parser.add_argument("fastq_files", type=str, nargs='+', help='Path(s) to fastq file(s)')
    args = parser.parse_args()

    randnumber = randrange(1000000)

    fh_report = open("report-SEQLENGTH-" + str(randnumber) + ".txt", "w")
    fh_samples_short = open("SAMPLES_short", "w")
    fh_samples_long = open("SAMPLES_long", "w")

    print("File Total Short Long Perc_short Perc_long Threshold", file=fh_report)
    for fastq_file in args.fastq_files:
        (fastq_short, fastq_long, count_short, count_long, count_total) = splitFastq(fastq_file, args.length)
        print(fastq_short, file=fh_samples_short)
        print(fastq_long, file=fh_samples_long)
        perc_short = 100 * count_short / float(count_total)
        perc_long = 100 * count_long / float(count_total)
        print(fastq_file, count_total, count_short, count_long, perc_short, perc_long, args.length, file=fh_report)
    fh_report.close()
    fh_samples_long.close()
    fh_samples_short.close()
    print("Wrote report-SEQLENGTH-" + str(randnumber) + ".txt to disk")
    print("Wrote SAMPLES_short to disk")
    print("Wrote SAMPLES_long to disk")