Subsampling

#!/usr/bin/env python
'''USAGE: subsample_fasta.py file.fasta --size n --total n' --pair y/n
Randomly subsamples a fasta file to the given size, handles both single and paired-end read data
'''

# It's important to know the correct number of sequences per file for this script to run properly,
# please run seq_stats.py prior to running subsample_fasta.py

import sys
import random
import argparse


parser = argparse.ArgumentParser(description='Subsample reads from a fasta file.')
parser.add_argument('newfile')
parser.add_argument('--size', default=0, help='Number of sequences to subsample (# of pairs for paired-end)')
parser.add_argument('--total', default=0, help='Total number of sequences in the fasta file')
parser.add_argument('--pair', default='n', help='Indicates if the file is interleaved, pired-end reads')
args = parser.parse_args()

if int(args.size) == 0 or int(args.total) == 0:
	print('ERROR: File or subsample value are of size 0')
	sys.exit()


# Check if the subsample is that same or larger than the total sequence count, and then
# generates a random list of positions to pick from fasta file
print 'Creating subsampling distribution...'

if args.pair == 'n':
	sample_list = range(1, int(args.total) + 1)
	if int(args.size) >= int(args.total):
		sample_list = set(sample_list)
		print 'Subsample size is greater than or equal to the total number of sequences.  Using all sequences.'
	else:
		sample_list = set(sorted(random.sample(sample_list, int(args.size) + 1)))
elif args.pair == 'y':
	sample_list = range(1, int(args.total) + 1, 2)
	if int(args.size) >= int(args.total):
		sample_list = set(sample_list)
		print 'Subsample size is greater than or equal to the total number of sequences.  Using all sequences.'
	else:
		sample_list_forward = sorted(random.sample(sample_list, int(args.size) + 1))
		sample_list_reverse = []
		for index in sample_list_forward: sample_list_reverse.append(index + 1)
		sample_list = set(sorted(sample_list_forward + sample_list_reverse))
elif args.pair not in ['y', 'n']:
	print('ERROR: Invalid input file type')
	sys.exit()


# Name and open output file
outfile_str = str(args.newfile).rstrip('fasta') + 'pick.fasta' 
outfile = open(outfile_str, 'w')

# Create logfile of subsampling
logfile_str = str(args.newfile).rstrip('fasta') + 'pick.log' 
logfile = open(logfile_str, 'w')

# Label the input as pe or se
if args.pair == 'y':
	file_type = 'Paired-end'
else:
	file_type = 'Single-end'

log_str = '''Input file name: {infile}
Input file type:  {type}
Output file name: {outfile}
Total sequences: {total}
Subsample size: {size}
'''.format(infile=str(args.newfile), type=file_type, outfile=outfile_str, total=str(args.total), size=str(int(args.size)))
logfile.write(log_str)
logfile.close()

# Open and begin looping through fasta file and picking sequences
print 'Writing subsampled fasta file...'
with open(args.newfile, 'r') as infile:

	outfile = open(outfile_str, 'w')

	seq_count = 0
	include = 0
	iteration = 0
	pair_include = 0

	for line in infile:
	
		line = line.strip()
	
		if line[0] == '>':
			
			include = 0
			seq_count += 1
			iteration += 1
			
			if iteration == 10000:
				iteration = 0
				progress_str = str(seq_count) + ' of ' + str(args.total)
				print(progress_str)

			if seq_count in sample_list:
				sample_list.discard(seq_count)
				outfile.write(line)
				outfile.write('\n')
				include = 1
				continue
				
		elif line[0] != '>' and include == 1: 
			outfile.write(line)
			outfile.write('\n')
			continue
				
	
outfile.close()			
print 'Done.'