Count and classify novel splice junctions for diffSplice

novel_junc_count.py

import sys
from collections import Counter
from functools import partial
import numpy as np
import pandas as pd
import pysam
import click


def load_bams(bam_fns):
    return [pysam.AlignmentFile(b) for b in bam_fns]


BED_CONVERTERS = {
    'chrom': str,
    'start': int,
    'end': int,
    'gene_id': str,
    'is_reverse': lambda x: x == '-',
    'exon_sizes': partial(np.fromstring, sep=',', dtype='i8'),
    'exon_starts': partial(np.fromstring, sep=',', dtype='i8')
}

def load_bed(bed_fn):
    bed = pd.read_table(
        bed_fn,
        sep='\t',
        names=['chrom', 'start', 'end', 'gene_id', '_', 'is_reverse',
               '_', '_', '_', '_', 'exon_sizes', 'exon_starts'],
        converters=BED_CONVERTERS,
        usecols=[0, 1, 2, 3, 5, 10, 11],
        index_col='gene_id'
    )
    bed['exon_starts'] = bed['exon_starts'] + bed['start']
    bed['exon_ends'] = bed['exon_starts'] + bed['exon_sizes']
    return bed.drop('exon_sizes', axis=1)


def check_orientation(aln, gene_reverse):
    # correct for reversely stranded libraries
    if gene_reverse ^ bool(aln.is_reverse):
        return bool(aln.is_read1)
    else:
        return bool(aln.is_read2)


def extract_sj_from_aln(aln, start, end):
    leftmost = aln.reference_start
    curr_pos = 0
    splice_juncts = []
    for i, (cig_type, n_bases) in enumerate(aln.cigar):
        if cig_type in (0, 2):
            curr_pos += n_bases
        elif cig_type == 3:
            if aln.reference_start > start and aln.reference_end < end:
                splice_juncts.append((leftmost + curr_pos,
                                      leftmost + curr_pos + n_bases))
            curr_pos += n_bases
    return splice_juncts


def fetch_splice_juncs(bam, chrom, start, end, gene_reverse):
    aligned_reads = bam.fetch(chrom, start, end)
    splice_juncts = []
    for aln in aligned_reads:
        if 'N' in aln.cigarstring and check_orientation(aln, gene_reverse):
            splice_juncts += extract_sj_from_aln(aln, start, end)
    return Counter(splice_juncts)


def assign_splice_junc_classes(splice_juncs, exon_starts, exon_ends):
    sj_classes = {}
    for sj_start, sj_end in splice_juncs:
        donor_canon = sj_start in exon_ends
        acceptor_canon = sj_end in exon_starts
        # check if normal donor/acceptor
        if donor_canon and acceptor_canon:
            # check if exon skipping
            if np.any((exon_starts > sj_start) & (exon_ends < sj_end)):
                sj_classes[(sj_start, sj_end)] = 'skipping'
            else:
                sj_classes[(sj_start, sj_end)] = 'constitutive'
        elif donor_canon or acceptor_canon:
            # different donor or acceptor site
            sj_classes[(sj_start, sj_end)] = 'alternate'
        else:
            # check if contained within exon
            if np.any((exon_starts < sj_start) & (exon_ends > sj_end)):
                sj_classes[(sj_start, sj_end)] = 'retained/exitronic'
            else:
                sj_classes[(sj_start, sj_end)] = 'other'
    return sj_classes

@click.command()
@click.option(
    '-a', '--bed_fn',
    required=True,
    help='bed12 file containing flattened gene models'
)
@click.option(
    '-b', '--bam_fns',
    required=True,
    help='comma separated list of bam files'
)
@click.option(
    '-m', '--min-supporting', required=False, default=20,
    help='minimum supporting reads total (across all bams) to filter low abundance SJs'
)
def count_novel_splice_juncs(bed_fn, bam_fns, min_supporting):
    bam_fns = bam_fns.split(',')
    alignment_files = load_bams(bam_fns)
    bed = load_bed(bed_fn)
    sys.stdout.write('gene_id\tchrom\tstart\tend\tstrand\tsj_class\t{}\n'.format('\t'.join(bam_fns)))
    for gene_id, *inv, exon_starts, exon_ends in bed.itertuples():
        splice_juncs = [fetch_splice_juncs(bam, *inv) for bam in alignment_files]
        total_splice_juncs = sum(splice_juncs, Counter())
        splice_junc_classes = assign_splice_junc_classes(
            total_splice_juncs, exon_starts, exon_ends)
        for junc, count in total_splice_juncs.items():
            if count >= min_supporting:
                sys.stdout.write('{}\t{}\t{}\t{}\t{}\t{}\t{}\n'.format(
                    gene_id, inv[0], *junc,
                    '-' if inv[3] else '+',
                    splice_junc_classes[junc],
                    '\t'.join('{:d}'.format(counts[junc]) for counts in splice_juncs)
                ))

if __name__ == '__main__':
    count_novel_splice_juncs()