mb_barcode_count.py

'''
Mission Bio Barcode and Trim
Ben Demaree - 10.4.2019

A simple script for pre-processing of Mission Bio data.

'''

import os
import subprocess
import sys
import argparse
from multiprocessing import Process
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
plt.rcParams['axes.unicode_minus'] = False
import pandas as pd

# add mission bio cell processing script
sys.path.append(os.path.join(sys.path[0], 'resources'))
import mb_resources_2

def wait(processes):
    # waits for processes to finish
    return [process.communicate() for process in processes]

def file_summary(samples):
    # print summary of all samples identified

    for sample in samples:

        s = vars(sample)

        for item in s:
            print item,': ' ,s[item]
        print '\n'

    print '%d samples identified.\n' % len(samples)

def generate_samples(R1_files, R2_files, output_folder, sample_label):
    # store sample filenames in Sample objects

    samples = []        # list of sample objects

    for i in range(len(R1_files)):
        # assign samples to objects

        r1 = R1_files[i]
        r2 = R2_files[i]

        sample_num = i + 1
        sample_label = sample_label + '-' + str(sample_num)

        samples.append(mb_resources_2.TapestriSample(sample_num, r1, r2, output_folder))

    return samples


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='''
    
    Mission Bio Barcode Count
    Ben Demaree - 12.10.2020
    
    A simple script for counting cell barcodes from Mission Bio data.

    ''', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('sample_label', type=str,
                        help='Label for this sample')

    parser.add_argument('input_fastq_folder', type=str,
                        help='Input folder containing raw FASTQ files (must have fastq.gz extension)')

    parser.add_argument('output_folder', type=str,
                        help='Output folder for processed files (must exist)')

    parser.add_argument('--chem_version', default='V2', choices=['V1', 'V2'],
                        help='Chemistry version (V1 or V2) (default: V2)')

    parser.add_argument('--single', action='store_true', default=False,
                        help='Option to process single-end sequencing run (default: paired-end)')

    args = parser.parse_args()  # parse arguments

    sample_label = args.sample_label
    output_folder = os.path.abspath(os.path.expanduser(args.output_folder))
    input_fastq_folder = os.path.abspath(os.path.expanduser(args.input_fastq_folder))
    chem_version = args.chem_version
    single = args.single
    if single:
        paired_end = False
    else:
        paired_end = True

    # check that folders exist
    if not os.path.isdir(output_folder):
        raise OSError(output_folder + ' is not a directory. Exiting...')

    if not os.path.isdir(input_fastq_folder):
        raise OSError(input_fastq_folder + ' is not a directory. Exiting...')

    fastq_files = [input_fastq_folder + '/' + f for f in os.listdir(input_fastq_folder) if '.fastq.gz' in f or '.fastq' in f]
    fastq_files.sort()

    ### experiment properties

    barcode_counts = output_folder + '/' + sample_label + '_barcode_counts.tsv'
    cell_barcode_cutadapt = output_folder + '/' + sample_label + '_cell_barcode_cutadapt.txt'

    # filtering parameters for cutadapt based on chemistry version
    if chem_version == 'V1':
        r1_start = 'CGATGACG'
        r1_end = 'CTGTCTCTTATACACATCT'
        r2_end = 'CGTCATCG'
        bar_ind_1, bar_ind_2 = range(8), range(-8, 0)
        cell_barcode_csv = sys.path[0] + '/resources/v1_barcodes.csv'

    elif chem_version == 'V2':
        r1_start = 'GTACTCGCAGTAGTC'
        r1_end = 'CTGTCTCTTATACACATCT'
        r2_end = 'GACTACTGCGAGTAC'
        bar_ind_1, bar_ind_2 = range(9), range(-9, 0)
        cell_barcode_csv = sys.path[0] + '/resources/v2_barcodes.csv'

    print '''
####################################################################################
# Step 1: get input file names and store in TapestriSample objects
####################################################################################
    '''

    if paired_end:

        # get all fastq filenames
        fastq_files.sort()
        R1_files = [fastq_files[i] for i in range(len(fastq_files)) if i % 2 == 0]
        R2_files = [fastq_files[i] for i in range(len(fastq_files)) if i % 2 == 1]

        # R1_files = [f for f in R1_files if '_R1_' in f]
        # R2_files = [f for f in R2_files if '_R2_' in f]

        assert len(R1_files) == len(R2_files), 'Number of R1 files does not match number of R2 files!'

        # # check filenames
        # for i in range(len(R1_files)):
        #     assert '_R1_' in R1_files[i], 'Bad R1 filename: %s' % R1_files[i]
        #     assert '_R2_' in R2_files[i], 'Bad R2 filename: %s' % R2_files[i]
        #     assert R1_files[i].split('_')[0] == R2_files[i].split('_')[0], 'Filename mismatch!'

        # store sample info in Sample objects
        samples = generate_samples(R1_files,
                                   R2_files,
                                   output_folder,
                                   sample_label)

        # display sample summary
        file_summary(samples)

    else:
        R1_files = [fastq_files[i] for i in range(len(fastq_files))]

        # check filenames
        for i in range(len(R1_files)):
            assert '_R1_' in R1_files[i], 'Bad R1 filename: %s' % R1_files[i]

        # store sample info in Sample objects
        samples = generate_samples(R1_files,
                                   R1_files,
                                   output_folder,
                                   sample_label)

        # display sample summary
        file_summary(samples)

    print '''
####################################################################################
# Step 2: filter reads for cell barcode, perform error correction, trim reads
####################################################################################
    '''

    # load mission bio barcode csv file
    barcodes = mb_resources_2.load_barcodes(cell_barcode_csv, 1, False)

    # generate hamming dictionary for error correction
    barcodes = mb_resources_2.generate_hamming_dict(barcodes)

    print 'Barcode sequences loaded into dictionary.\n'

    # for panel reads, filter reads with valid barcode structure and export to new fastq
    print 'Extracting barcodes from raw fastq files...\n'

    # cut adapters from reads and add barcodes to header

    barcode_samples = []

    for sample in samples:
        p = Process(
            target=sample.count_barcodes,
            args=(r1_start,
                  bar_ind_1,
                  bar_ind_2,
                  barcodes))
        barcode_samples.append(p)
        p.start()

    # wait for processes to finish
    for p in barcode_samples:
        p.join()

    # combine barcode count files
    bc = [s.barcode_counts for s in samples]
    wait([subprocess.Popen('cat %s > %s' % (' '.join(bc), barcode_counts), shell=True)])

    # plot barcode kneeplot
    all_df = pd.read_csv(filepath_or_buffer=barcode_counts, sep='\t', index_col=0)
    reads_per_cell = list(all_df.iloc[:, 0])

    # plot log-log reads per cell vs cells
    reads_per_cell.sort(reverse=True)
    ax = plt.figure(figsize=(7, 7))
    plt.loglog(range(1, len(reads_per_cell) + 1), reads_per_cell, color='k', linewidth=1.5)
    ax = plt.axes()
    ax.grid()
    plt.xlabel('Cell Barcode #', fontsize=18, labelpad=12)
    plt.ylabel('Reads per Cell Barcode', fontsize=18, labelpad=12)
    plt.xticks(fontsize=16)
    plt.yticks(fontsize=16)

    plt.title(sample_label)
    plt.tight_layout()
    plt.savefig(output_folder + '/' + sample_label + '.' + chem_version + '.kneeplot.png', dpi=300)

    # delete temporary files
    [os.remove(s.barcode_counts) for s in samples]