barcode_splitter.py

#!/usr/bin/env python
# from https://bitbucket.org/lance_parsons/paired_sequence_utils/src/355e838b92d0/barcode_splitter.py

""" Split one or more fastq files based on barcode sequence.
"""
from __future__ import division
import gzip
import optparse
import os
import re
import sys
import mapping_functions
#from fastq_utils import fastq_utils

__version__ = "0.1"
__author__ = "Lance Parsons"
__author_email__ = "lparsons@princeton.edu"
__copyright__ = "Copyright 2011, Lance Parsons"
__license__ = "BSD 2-Clause License http://www.opensource.org/licenses/BSD-2-Clause"

UNMATCHED = 'unmatched'

def main (argv=None):
    print "barcode_splitter starting run"
    if argv is None:
        argv = sys.argv
    
    usage = "Usage: %prog [options] fastq_read1 [fastq_read2] [fastq_read3]"
    parser = optparse.OptionParser(usage=usage, version='%prog version ' + globals()['__version__'], description=globals()['__doc__'])
    parser.add_option ('-v', '--verbose', action='store_true', default=False, help='verbose output')
    
    parser.add_option ('--bcfile', help='Tab delimited file with barcodes and sample ids (REQUIRED)')
    
    
    #When used in MSG (Multiplexed shotgun sequencing) the illumina indexes will be parsed out first
    #and those results will be passed through the MSG parser to find the MSG barcodes.  Since each index
    #will have the same barcodes, the MSG parser renames the results with the index.  
    #This option group creates a new barcodes file for MSG to use downstream from parsing with the updated
    #names.
    msg_group = optparse.OptionGroup(parser, "MSG Options")
    msg_group.add_option ('--make_indexed_msg_barcodes_file',help='If you have a MSG barcodes file, this creates a new version of it with indexes',
        action='store_true', default=False)
    msg_group.add_option ('--msg_barcodes', help="Tab delimited seq/label list of MSG barcodes", default='')
    parser.add_option_group(msg_group)
    
    output_group = optparse.OptionGroup(parser, "Output Options")
    output_group.add_option ('--prefix', default='', help='Prefix for output files')
    output_group.add_option ('--suffix', default='', help='Suffix for output files')
    parser.add_option_group(output_group)
    
    barcode_location_group = optparse.OptionGroup(parser, "Barcode Location")
    barcode_location_group.add_option ('--idxread', default=1, type='int', help='Indicate in which read to search for the barcode sequence (1 based) (default: 1)')
    barcode_location_group.add_option ('--barcodes_at_end', action='store_true', default=False, help='Barcodes are at the end of the index read (default is at the beginning)')
    parser.add_option_group(barcode_location_group)
    
    matching_group = optparse.OptionGroup(parser, "Matching")
    matching_group.add_option ('--mismatches', default=0, type='int',  help='Number of mismatches allowed in barcode matching')
    parser.add_option_group(matching_group)
    
    input_group = optparse.OptionGroup(parser, "Input format")
    input_group.add_option ('--format', default='fastq', help='Specify format for sequence files (fasta or fastq)')
    input_group.add_option('--gzip', action='store_true', default=False, help='Force gzip format for input and output (default is auto based on input file extension)')
    parser.add_option_group(input_group)
    
    
    try:
        (options, args) = parser.parse_args(argv[1:])
        if (not options.make_indexed_msg_barcodes_file) and len(args) < 1:
            parser.error('Must specify at least one sequence file')
        if not options.bcfile:
            parser.error('Must specify a barcodes file with "--bcfile" option')
    except SystemExit: # Prevent exit when calling as function
        return 2
    
    if options.make_indexed_msg_barcodes_file and options.msg_barcodes:
        return make_msg_barcodes_file(options.msg_barcodes, options.bcfile)
    
    # Read barcodes files into dict
    barcode_dict = read_barcodes(options.bcfile)
    total_read_count = 0
    counts = {UNMATCHED: 0}
    for barcode in barcode_dict:
        counts[barcode] = 0
    # TODO Verbose: print barcode_dict
    
    # Determine if we should use gzip for input/output
    basename, extension = os.path.splitext(args[0])
    if extension == '.gz':
        options.gzip = True

    # Open filehandles for each read
    inputs = {}
    outputs = {}
    for i in xrange(0,len(args)):
        if options.gzip:
            inputs[i] = gzip.open(args[i], 'rb')
        else:
            inputs[i] = open(args[i], 'rb')
        first_line = inputs[i].readline().strip()
        id_format = determine_id_format(first_line[1:])
        # TODO print "\nFile %s: %s\nId format: %s" % (i, args[i], id_format)
        inputs[i].seek(0)
    
    # Open filehandles for each barcode (and unmatched), for each read
    for barcode in barcode_dict:
        outputs[barcode] = {}
        for i in inputs:
            if options.gzip:
                outputs[barcode][i] = gzip.open('%s%s_read_%s%s' % (options.prefix, barcode_dict[barcode], i+1, options.suffix), 'wb')
            else:
                outputs[barcode][i] = open('%s%s_read_%s%s' % (options.prefix, barcode_dict[barcode], i+1, options.suffix), 'wb')
    outputs[UNMATCHED] = {}
    for i in inputs:
        if options.gzip:
            outputs[UNMATCHED][i] = gzip.open('%s%s_read_%s%s' % (options.prefix, UNMATCHED, i+1, options.suffix), 'wb')  
        else:
            outputs[UNMATCHED][i] = open('%s%s_read_%s%s' % (options.prefix, UNMATCHED, i+1, options.suffix), 'wb')    
            
    # For each input line in index read, get index sequence
    for index_read in read_fastq(inputs[options.idxread-1]):
        total_read_count += 1
        barcode_length = len(barcode_dict.keys()[0])
        if options.barcodes_at_end:
            index_seq = index_read['seq'][-barcode_length:]
        else:
            index_seq = index_read['seq'][0:barcode_length]
            
        # Get matching barcode(s), if more than one, warn and set to unmatched
        best_match_barcodes = match_barcodes(index_seq, barcode_dict.keys(), options.mismatches)
        if (len(best_match_barcodes) == 1):
            barcode_match = best_match_barcodes[0]
        else:
            barcode_match = UNMATCHED
            if (len(best_match_barcodes) > 1) :
                sys.stderr.write('More than one barcode matches for %s, moving to %s category\n' % (index_read['seq_id'], UNMATCHED))
        counts[barcode_match] += 1
        
        # Get sequence record from each other read, assert id matches
        for readnum in xrange(0,len(inputs)):
            if readnum != options.idxread-1:
                read = read_fastq(inputs[readnum]).next()
                try:
                    assert(match_id(index_read['seq_id'], read['seq_id'], id_format))
                except AssertionError:
                    sys.stderr.write("Id mismatch: %s does not match %s" %(index_read['seq_id'], read['seq_id']))
            else:
                read = index_read
            # Output sequences into barcode/read file    
            outputs[barcode_match][readnum].write(fastq_string(read))
    print "Sample\tBarcode\tCount\tPercent"
    for barcode in sorted(barcode_dict, key=barcode_dict.get):
        print "%s\t%s\t%s\t%.2f%%" % (barcode_dict[barcode], barcode, counts[barcode], (counts[barcode]/total_read_count)*100 )
    print "%s\t%s\t%s\t%.2f%%" % (UNMATCHED, None, counts[UNMATCHED], (counts[UNMATCHED]/total_read_count)*100 )
    return 0

def read_barcodes(filename):
    '''Read barcodes file into dictionary'''
    barcode_dict = {}
    linenum = 0
    if filename.lower().endswith('.gz'):
        filehandle = gzip.open(filename, 'rb')
    else:
        filehandle = open(filename, 'rb')
    for line in filehandle:
        linenum += 1
        line = line.strip()
        if line[0] != '#':
            (sample_id, barcode_sequence) = line.split('\t')
            if (sample_id is not None) and (barcode_sequence is not None):
                barcode_dict[barcode_sequence] = sample_id
            else:
                raise Exception("Unable to read barcode from line %s: '%s'" % (linenum, line))
    return barcode_dict

def make_msg_barcodes_file(msg_barcodes, bcfile):
    """
    When used in MSG (Multiplexed shotgun sequencing) the illumina indexes will be parsed out first
    and those results will be passed through the MSG parser to find the MSG barcodes.  Since each index
    will have the same barcodes, the MSG parser renames the results with the index.  
    This function creates a new barcodes file for MSG to use downstream from parsing with the updated
    names.
    """
    new_barcodes = open(msg_barcodes + '.after.index.parsing','w')
    for index_id in read_barcodes(bcfile).values():
        for ind in mapping_functions.read_barcodes(msg_barcodes):
            #For reference, MSG parsing output file_name is composed of: 'indiv' + ind[1] + '_' + ind[0]
            new_ind = ind[:]
            new_ind[1] = new_ind[1] + index_id
            new_barcodes.write('\t'.join(new_ind) + '\n')
    new_barcodes.close()  
    return 0

def match_barcodes(sequence, barcodes, mismatches):
    '''Find closest match(es) in barcodes to specified sequence with max number of mismatches'''
    best_distance = mismatches
    results = []
    for barcode in barcodes:
        if mismatches == 0:
            if (sequence == barcode):
                results.append(barcode)
        else:
            distance = hamming_distance(sequence, barcode)
            if (distance <= best_distance):
                best_distance = distance
                results.append(barcode)
    return results
    
def hamming_distance(s1, s2):
    assert len(s1) == len(s2)
    return sum(ch1 != ch2 for ch1, ch2 in zip(s1, s2))    

''' Supported types of Fastq IDs '''
ILLUMINA = 'illumina' # CASVA 1.8+, match up to space
STRIPONE = 'stripone' # Illumina CASAVA 1.7 and lower (/1 /2) and NCBI SRA (/f /r), match all but last character
OTHER = 'other'       # Other, match exactly

def determine_id_format(seq_id):
    '''Determine if the id is new illumina, old illumina (/1 /2 ...), sanger (/f /r), or other'''
    
    id_format = None
    # Illumina CASAVA 1.8+ fastq headers use new format
    read_id_regex = re.compile(r'(?P<instrument>[a-zA-Z0-9_-]+):(?P<run_number>[0-9]+):(?P<flowcell_id>[a-zA-Z0-9]+):(?P<lane>[0-9]+):(?P<tile>[0-9]+):(?P<x_pos>[0-9]+):(?P<y_pos>[0-9]+) (?P<read>[0-9]+):(?P<is_filtered>[YN]):(?P<control_number>[0-9]+):(?P<index_sequence>[ACGT]+){0,1}')
    # Old illumina and sanger reads use /1 /2 or /f /r
    strip_one_endings = ['/1', '/2', '/3', '/f', '/r']
    
    if read_id_regex.match(seq_id):
        id_format = ILLUMINA
    elif (seq_id[-2:] in strip_one_endings):
        id_format = STRIPONE
    else:
        id_format = OTHER
    return id_format

def strip_read_from_id(seq_id, id_format=None):
    new_id = seq_id
    if not id_format:
        id_format = determine_id_format(seq_id)
    elif id_format == STRIPONE:
        new_id = seq_id[0:-1]
    elif id_format == ILLUMINA:
        new_id = seq_id.split(' ')[0]
    return new_id

def strip_read_from_id_stripone(seq_id):
    return seq_id[0:-1]

def strip_read_from_id_illumina(seq_id):
    return seq_id.split(' ')[0]

def match_id(id1, id2, id_format=OTHER):
    ''' Return true if id's match using rules for specified format '''
    if id_format == STRIPONE:
        if id1[0:-1] == id2[0:-1]:
            return True
        else:
            return False
    elif id_format == ILLUMINA:
        if (id1.split(' ')[0] == id2.split(' ')[0]):
            return True
        else:
            return False
    elif id1 == id2:
        return True
    else:
        return False

def read_fastq(filehandle):
    ''' Return dictionary with 'seq_id', 'seq', 'qual_id', and 'qual' '''
    record_line = 0
    read_number = 0
    fastq_record = dict()
    for line in filehandle:
        record_line += 1
        if record_line == 1:
            fastq_record['seq_id'] = line.strip()
        elif record_line == 2:
            fastq_record['seq'] = line.strip()
        elif record_line == 3:
            fastq_record['qual_id'] = line.strip()
        elif record_line == 4:
            record_line = 0
            fastq_record['qual'] = line.strip()
            read_number += 1
            yield fastq_record
            

def fastq_string(record):
    return "%s\n%s\n%s\n%s\n" % (record['seq_id'], record['seq'], record['qual_id'], record['qual'])

if __name__ == '__main__':
    sys.exit(main())