python——提取fasta文件中10G的序列

#!/usr/bin/python
# -*- coding: utf-8 -*-
# conda activate python3
"""
    作者：徐诗芬
    内容：截取原始序列的10G,读一行写一行，当序列长度叠加到10G时终止读写 
    日期：2021.1.11
"""
import sys
import time
import gzip

def usage():
    print('Usage: python cut_sequence_10G.py [infile.fq.gz] [outfile.fq.gz]')


def main():

    global name
    inf = gzip.open(sys.argv[1], 'rt')
    ouf = gzip.open(sys.argv[2], 'wt')

    fq = {}
    i = 0
    size = 0
    Gb = 10**9
    for line in inf:
        line = line.strip()
        i += 1
        if line.startswith('@'):
            name = line
            ouf.write(line + '\n')  # 写@所在行
            fq[name] = []  # 将剩余三行放入列表中
        else:
            fq[name].append(line)
            ouf.write(line + '\n')  # 依次写入剩下三行
        if i % 4 == 2:
            size += len(line)  # 如果是序列行，叠加序列长度
        if size > 10*Gb:
            break
    inf.close()
    ouf.close()

try:
    main()
except IndexError:
    usage()

end_time = time.process_time()
print('Running time: {:.2f} Seconds'.format(end_time))

上面的脚本只能一个文件输入和一个文件的输出，适用于单端测序数据。对于双端测序数据，通常需要两个文件的reads数目是一致的，因此，我们需要同步输入和输出两个文件，用python3里的zip函数连接，zip函数可以把两个或者两个以上的迭代器封装成生成器，这种zip生成器会从每个迭代器中获取该迭代器的下一个值，然后把这些值组装成元组（tuple）。这样，zip函数就实现了平行地遍历多个迭代器。脚本如下：注意4个输入输出文件的顺序！！

import sys
import time
import gzip

def usage():
    print('Usage: python cut_sequence_10G.py [input_file1.gz] [input_file2.gz] [outfile1.gz] [outfile2.gz]')


def main():
    global name
    inf1 = gzip.open(sys.argv[1], 'rt')
    ouf1 = gzip.open(sys.argv[3], 'wt')
    inf2 = gzip.open(sys.argv[2], 'rt')
    ouf2 = gzip.open(sys.argv[4], 'wt')
    
    i = 0       ##记录行号
    size = 0    ##序列总长度
    Gb = 10**9  
    for line1, line2 in zip(inf1, inf2):
        line1 = line1.strip()
        line2 = line2.strip()
        i += 1
        if line1.startswith('@'):
            ouf1.write(line1 + '\n')  # 写@所在行
        else:
            ouf1.write(line1 + '\n')  # 依次写入剩下三行
        if line2.startswith('@'):
            ouf2.write(line2 + '\n')  # 写@所在行
        else:
            ouf2.write(line2 + '\n')
        if i % 4 == 2:
            size += len(line1)  # 如果是序列行，叠加序列长度
        #这里只用一个file的序列长度作为终止操作的条件
        if size > 10*Gb:
                break
    inf1.close()
    inf2.close()
    ouf1.close()
    ouf2.close()

try:
    main()
except IndexError:
    usage()
end_time = time.process_time()
print('Running time: {:.2f} seconds'.format(end_time))