quantify_orfs.py

#! /usr/bin/env python

import argparse
import os
import sys
from time import strftime
import pysam
from hashed_read_genome_array import HashedReadBAMGenomeArray, ReadKeyMapFactory, read_length_nmis  #, get_hashed_counts
from plastid.genomics.roitools import SegmentChain, positionlist_to_segments
import multiprocessing as mp
from scipy.optimize import nnls
import numpy as np
import pandas as pd

parser = argparse.ArgumentParser(description='Use linear regression to quantify expression of the ORFs identified by ORF-RATER. Reported values are '
                                             'in reads per nucleotide; any additional normalization(s) (such as for read depth must be performed in '
                                             'post-processing. The number of nucleotides used to quantify each ORF is also included in the output. '
                                             'ORFs may receive NaN values if they are too short (and therefore masked out by STARTMASK and STOPMASK) '
                                             'or if they are indistinguishable from another ORF after those regions have been masked.')
parser.add_argument('bamfiles', nargs='+', help='Path to transcriptome-aligned BAM file(s) for read data. Expression values will be quantified from '
                                                'each independently.')
parser.add_argument('--names', nargs='+', help='Names to use to refer to BAMFILES, e.g. if the filenames themselves are inconveniently long or '
                                               'insufficiently descriptive. (Default: inferred from BAMFILES)')
parser.add_argument('--subdir', default=os.path.curdir,
                    help='Convenience argument when dealing with multiple datasets. In such a case, set SUBDIR to an appropriate name (e.g. HARR, '
                         'CHX) to avoid file conflicts. (Default: current directory)')
parser.add_argument('--inbed', default='transcripts.bed', help='Transcriptome BED-file (Default: transcripts.bed)')
parser.add_argument('--offsetfile', default='offsets.txt',
                    help='Path to 2-column tab-delimited file with 5\' offsets for variable P-site mappings. First column indicates read length, '
                         'second column indicates offset to apply. Read lengths are calculated after trimming up to MAX5MIS 5\' mismatches. Accepted '
                         'read lengths are defined by those present in the first column of this file. If SUBDIR is set, this file is assumed to be '
                         'in that directory. (Default: offsets.txt)')
parser.add_argument('--max5mis', type=int, default=1, help='Maximum 5\' mismatches to trim. Reads with more than this number will be excluded. '
                                                           '(Default: 1)')
parser.add_argument('--startmask', type=int, nargs=2, default=[1, 2],
                    help='Region around start codons (in codons) to exclude from quantification. (Default: 1 2, meaning one full codon before the '
                         'start is excluded, as are the start codon and the codon following it).')
parser.add_argument('--stopmask', type=int, nargs=2, default=[3, 0],
                    help='Region around stop codons (in codons) to exclude from quantification. (Default: 3 0, meaning three codons before and '
                         'including the stop are excluded, but none after).')
parser.add_argument('--metagenefile', default='metagene.txt',
                    help='File to be used as the metagene, generated by regress_orfs.py. If SUBDIR is set, this file is assumed to be in that '
                         'directory. (Default: metagene.txt)')
parser.add_argument('--ratingsfile', default='orfratings.h5',
                    help='Path to pandas HDF store containing ORF ratings; generated by rate_regression_output.py (Default: orfratings.h5)')
parser.add_argument('--minrating', type=float, default=0.8, help='Minimum ORF rating to require for an ORF to be quantified (Default: 0.8)')
parser.add_argument('--minlen', type=int, default=0, help='Minimum ORF length (in amino acids) to be included in the BED file (Default: 0)')
parser.add_argument('--quantfile', default='quant.h5',
                    help='Filename to which to output the table of quantified translation values for each ORF. Formatted as pandas HDF; table name '
                         'is "quant". If SUBDIR is set, this file will be placed in that directory. (Default: quant.h5)')
parser.add_argument('--CSV', help='If included, also write output in CSV format to the provided filename.')
parser.add_argument('-v', '--verbose', action='count', help='Output a log of progress and timing (to stdout). Repeat for higher verbosity level.')
parser.add_argument('-p', '--numproc', type=int, default=1, help='Number of processes to run. Defaults to 1 but more recommended if available.')
parser.add_argument('-f', '--force', action='store_true', help='Force file overwrite')
opts = parser.parse_args()

offsetfilename = os.path.join(opts.subdir, opts.offsetfile)
metafilename = os.path.join(opts.subdir, opts.metagenefile)
quantfilename = os.path.join(opts.subdir, opts.quantfile)

if not opts.force:
    if os.path.exists(quantfilename):
        raise IOError('%s exists; use --force to overwrite' % quantfilename)
    if opts.CSV and os.path.exists(opts.CSV):
        raise IOError('%s exists; use --force to overwrite' % opts.CSV)

if opts.names:
    if len(opts.bamfiles) != len(opts.names):
        raise ValueError('Precisely one name must be provided for each BAMFILE')
    colnames = opts.names
else:
    colnames = [os.path.splitext(os.path.basename(bamfile))[0] for bamfile in opts.bamfiles]  # '/path/to/myfile.bam' -> 'myfile'

if opts.verbose:
    sys.stdout.write(' '.join(sys.argv) + '\n')

    def logprint(nextstr):
        sys.stdout.write('[%s] %s\n' % (strftime('%Y-%m-%d %H:%M:%S'), nextstr))
        sys.stdout.flush()

    log_lock = mp.Lock()

rdlens = []
Pdict = {}
with open(offsetfilename, 'rU') as infile:
    for line in infile:
        ls = line.strip().split()
        rdlen = int(ls[0])
        for nmis in range(opts.max5mis+1):
            Pdict[(rdlen, nmis)] = int(ls[1])+nmis  # e.g. if nmis == 1, offset as though the read were missing that base entirely
        rdlens.append(rdlen)
rdlens.sort()

# hash transcripts by ID for easy reference later
with open(opts.inbed, 'rU') as inbed:
    bedlinedict = {line.split()[3]: line for line in inbed}

with pd.HDFStore(opts.ratingsfile, mode='r') as ratingstore:
    chroms = ratingstore.select('orfratings/meta/chrom/meta').values  # because saved as categorical, this is the list of all chromosomes

if opts.verbose:
    logprint('Loading metagene')

metagene = pd.read_csv(metafilename, sep='\t').set_index(['region', 'position'])
metagene.columns = metagene.columns.astype(int)  # they are read lengths
assert (metagene.columns == rdlens).all()
cdsprof = metagene.loc['CDS']
assert len(cdsprof) == 3
cdsprof = cdsprof.values.sum(1)  # quantification as implemented here is not readlength-sensitive
cdsprof /= cdsprof.mean()  # divide by mean so it sums to 3 - so values will be per nt, not per codon

startmask = (-abs(opts.startmask[0])*3, abs(opts.startmask[1])*3)  # force <=0 and >= 0 for the mask
stopmask = (-abs(opts.stopmask[0])*3, abs(opts.stopmask[1])*3)


def _quantify_tfam(orf_set, gnds):
    """Performs non-negative least squares regression to quantify all of the ORFs in a transcript family, using a simplified profile consisting of
    the same three numbers tiled across each ORF. All readlengths are treated identically. Regions around start and stop codons are masked in
    accordance with startmask and stopmask"""
    strand = orf_set['strand'].iat[0]
    chrom = orf_set['chrom'].iat[0]
    tids = orf_set['tid'].drop_duplicates().tolist()
    all_tfam_genpos = set()
    tid_genpos = {}
    tlens = {}
    for (i, tid) in enumerate(tids):
        currtrans = SegmentChain.from_bed(bedlinedict[tid])
        curr_pos_set = currtrans.get_position_set()
        tlens[tid] = len(curr_pos_set)
        tid_genpos[tid] = curr_pos_set
        all_tfam_genpos.update(curr_pos_set)
    all_tfam_genpos = np.array(sorted(all_tfam_genpos))
    if strand == '-':
        all_tfam_genpos = all_tfam_genpos[::-1]
    nnt = len(all_tfam_genpos)
    tid_indices = {tid: np.flatnonzero(np.in1d(all_tfam_genpos, list(curr_tid_genpos), assume_unique=True))
                   for (tid, curr_tid_genpos) in tid_genpos.iteritems()}
    orf_matrix = np.zeros((nnt, len(orf_set)))
    ignore_coords = []
    for (orf_num, (tid, tcoord, tstop, AAlen)) in enumerate(orf_set[['tid', 'tcoord', 'tstop', 'AAlen']].itertuples(False)):
        orf_matrix[tid_indices[tid][tcoord:tstop], orf_num] = np.tile(cdsprof, AAlen + 1)
        ignore_coords.append(tid_indices[tid][max(tcoord+startmask[0], 0):tcoord+startmask[1]])
        ignore_coords.append(tid_indices[tid][max(tstop+stopmask[0], 0):tstop+stopmask[1]])
    ignore_coords = np.unique(np.concatenate(ignore_coords))
    orf_matrix[ignore_coords, :] = 0  # mask out all positions within the mask region around starts and stops
    valid_orfs = np.array([(orf_matrix[:, i] > 0).any() and (orf_matrix.T[i, :] != orf_matrix.T[:i, :]).any(1).all() for i in xrange(len(orf_set))])
    # require at least one valid position, and if >1 ORFs are identical, only include one of them
    orf_matrix[:, ~valid_orfs] = 0  # completely ignore these positions
    valid_nts = (orf_matrix > 0).any(1)  # only bother checking nucleotides where there is a valid ORF
    orf_res = orf_set.copy()
    if valid_nts.any():
        orf_matrix = orf_matrix[valid_nts, :]
        valid_nt_segs = SegmentChain(*positionlist_to_segments(chrom, strand, list(all_tfam_genpos[valid_nts])))
        orf_res['nts_quantified'] = (orf_matrix > 0).sum(0)  # the number of nucleotides included in the quantification
        for colname, gnd in zip(colnames, gnds):
            orf_res[colname] = nnls(orf_matrix, valid_nt_segs.get_counts(gnd))[0]
            # gnd is a HashedReadBAMGenomeArray, but it still works with get_counts(), which will collapse all read lengths to a single array
        return orf_res
    else:
        orf_res['nts_quantified'] = 0
        for colname in colnames:
            orf_res[colname] = 0.
        return orf_res


def _quantify_chrom(chrom_to_do):
    """Applies _quantify_tfam() to all of the transcript families on a chromosome"""
    chrom_orfs = pd.read_hdf(opts.ratingsfile, 'orfratings', mode='r',
                             where="chrom == %r and orfrating >= %f and AAlen >= %d" % (chrom_to_do, opts.minrating, opts.minlen),
                             columns=['orfname', 'tfam', 'tid', 'tcoord', 'tstop', 'AAlen', 'chrom', 'gcoord', 'gstop', 'strand',
                                      'codon', 'orftype', 'annot_start', 'annot_stop', 'orfrating'])
    if chrom_orfs.empty:
        if opts.verbose > 1:
            logprint('No ORFs found on %s' % chrom_to_do)
        return pd.DataFrame()

    inbams = [pysam.Samfile(infile, 'rb') for infile in opts.bamfiles]
    gnds = [HashedReadBAMGenomeArray([inbam], ReadKeyMapFactory(Pdict, read_length_nmis)) for inbam in inbams]

    res = pd.concat([_quantify_tfam(tfam_set, gnds) for (tfam, tfam_set) in chrom_orfs.groupby('tfam')])

    for inbam in inbams:
        inbam.close()

    if opts.verbose > 1:
        logprint('%s complete' % chrom_to_do)

    return res

if opts.verbose:
    logprint('Quantifying ORFs by chromosome')
workers = mp.Pool(opts.numproc)
quant = pd.concat(workers.map(_quantify_chrom, chroms))
workers.close()

if opts.verbose:
    logprint('Saving results')

for catfield in ['chrom', 'strand', 'codon', 'orftype']:
    quant[catfield] = quant[catfield].astype('category')  # saves disk space and read/write time

quant.to_hdf(quantfilename, 'quant', format='t', data_columns=True)
if opts.CSV:
    quant.to_csv(opts.CSV, index=False)

if opts.verbose:
    logprint('Tasks complete')