orf_scanner

#!/usr/bin/env python

# GY190919

''' orf_scanner: robust CDS prediction '''


from argparse import ArgumentParser, FileType, SUPPRESS
from collections import defaultdict as dd
from copy import deepcopy
from itertools import product
import logging
from math import ceil, log
from random import sample
import re
from signal import signal, SIGPIPE, SIG_DFL
from sys import stdin, stdout
from typing import Iterator, NamedTuple, TextIO

signal(SIGPIPE, SIG_DFL)  # Gracefully handle downstream PIPE closure


# Globals #####################################################################


NUCL_COMP = str.maketrans('ATUCGXNRYWSKMDVHB', 'TAAGCXNYRWSMKHBDV')

CODON_MAP = {
    'TTT': 'F', 'TTC': 'F', 'TCT': 'S', 'TCC': 'S',
    'TAT': 'Y', 'TAC': 'Y', 'TGT': 'C', 'TGC': 'C',
    'TTA': 'L', 'TCA': 'S', 'TAA': '*', 'TGA': '*',
    'TTG': 'L', 'TCG': 'S', 'TAG': '*', 'TGG': 'W',
    'CCC': 'P', 'CTT': 'L', 'CTC': 'L', 'CCT': 'P',
    'CCC': 'P', 'CAC': 'H', 'CGC': 'R', 'CAT': 'H',
    'CAC': 'H', 'CGT': 'R', 'CGC': 'R', 'CCA': 'P',
    'CCG': 'P', 'CTA': 'L', 'CTG': 'L', 'CCA': 'P',
    'CCG': 'P', 'CAA': 'Q', 'CAG': 'Q', 'CGA': 'R',
    'CGG': 'R', 'CAA': 'Q', 'CAG': 'Q', 'CGA': 'R',
    'CGG': 'R', 'ACC': 'T', 'ATT': 'I', 'ATC': 'I',
    'ACT': 'T', 'ACC': 'T', 'AAC': 'N', 'AGC': 'S',
    'AAT': 'N', 'AAC': 'N', 'AGT': 'S', 'AGC': 'S',
    'ACA': 'T', 'AAA': 'K', 'AGA': 'R', 'ATA': 'I',
    'ACA': 'T', 'AAA': 'K', 'AGA': 'R', 'ACG': 'T',
    'AAG': 'K', 'AGG': 'R', 'ATG': 'M', 'ACG': 'T',
    'AAG': 'K', 'AGG': 'R', 'GCC': 'A', 'GTT': 'V',
    'GTC': 'V', 'GCT': 'A', 'GCC': 'A', 'GAC': 'D',
    'GGC': 'G', 'GAT': 'D', 'GAC': 'D', 'GGT': 'G',
    'GGC': 'G', 'GCA': 'A', 'GCG': 'A', 'GTA': 'V',
    'GTG': 'V', 'GCA': 'A', 'GCG': 'A', 'GAA': 'E',
    'GAG': 'E', 'GGA': 'G', 'GGG': 'G', 'GAA': 'E',
    'GAG': 'E', 'GGA': 'G', 'GGG': 'G'}


# Classes #####################################################################


class SeqRecord(NamedTuple):
    ''' A class for holding sequence data '''
    seqid: str
    seq: str
    quals: str = ''

    def __len__(self):
        return len(self.seq)

    def print(self, wrap: int=0, fileobj: TextIO=stdout) -> None:
        ''' Print a sequence as FASTA/Q '''
        if self.quals:
            print(f'@{self.seqid}\n{self.seq}\n+\n{self.quals}', file=fileobj)
        elif wrap:
            print(f'>{self.seqid}', file=fileobj)
            for g in grouper(self.seq, wrap):
                print(g, file=fileobj)
        else:
            print(f'>{self.seqid}\n{self.seq}', file=fileobj)


class MultiReader(object):
    ''' A class for reading FASTA/Q records '''

    def __init__(self, inputstream: TextIO, replaceU: bool=False):
        self._fileobj = inputstream
        self._buffer = ' '
        self.replaceU = replaceU

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        pass

    def __iter__(self):
        return self

    def __next__(self) -> SeqRecord:
        while self._buffer and self._buffer[0] not in ('>', '@'):
            self._buffer = self._fileobj.readline()
        if not self._buffer:
            raise StopIteration
        if self._buffer.startswith('>'):
            seqid = self._buffer[1:].strip('\r\n')
            self._buffer = self._fileobj.readline()
            seqlines = []
            while self._buffer and self._buffer[0] not in ('>', '@'):
                seqlines.append(self._buffer.strip('\r\n'))
                self._buffer = self._fileobj.readline()
            if not (seq := ''.join(seqlines)):
                raise EOFError
            seq = seq.upper()
            if self.replaceU:
                seq = seq.replace('U', 'T')
            return SeqRecord(seqid, seq)
        else:  # self._buffer.startswith('@')
            seqid = self._buffer[1:].strip('\r\n')
            seq = self._fileobj.readline().strip('\r\n')
            self._fileobj.readline()
            quals = self._fileobj.readline().strip('\r\n')
            self._buffer = self._fileobj.readline()
            if not (seq and quals):
                raise EOFError
            seq = seq.upper()
            if self.replaceU:
                seq = seq.replace('U', 'T')
            return SeqRecord(seqid, seq, quals)


# Functions ###################################################################


def iterkmers(iterable: Iterator, k: int, step: int=1) -> Iterator[Iterator]:
    ''' Yield all possible kmers of length `k` from a memory-mapped `iterable`
    !!! Sequence will be truncated if iterable % k != 0 !!! '''
    for i in range(0, len(iterable)-k+1, step):
        yield iterable[i:i+k]


def grouper(iterable: Iterator, k: int) -> Iterator[Iterator]:
    ''' Yield a memory-mapped `iterable` in chunks of `k`
    !!! If iterable % k != 0 the length of the last chunk will be <k !!! '''
    for i in range(len(iterable) // k + (1 if len(iterable) % k else 0)):
        yield iterable[i*k:i*k+k]


def fnv1a_64(data: str) -> int:
    ''' 64-bit Fowler-Noll_Vo 1a hash function '''
    hval = 0xcbf2_9ce4_8422_2325
    for c in data:
        hval = ((hval ^ ord(c)) * 0x100_0000_01b3) % 0x1_0000_0000_0000_0000
    return hval


def find_orfs(record, unstranded: bool=False) -> list:
    ''' Search all frames (0 -> 5) for ORFs using the provided regex '''
    matches = {}
    for m in re.finditer(orf_pattern, record.seq):
        if (k := (m.end(1), m.start(1)%3)) not in matches:
            matches[k] = SeqRecord(
                f'orf{hex(fnv1a_64(m.group(1)))[2:]} ' \
                f'{record.seqid.split()[0]}:+:{m.start(1)+1}-{m.end(1)+1}',
                m.group(1),
                record.quals[m.start(1):m.end(1)])
    if unstranded:
        for m in re.finditer(orf_pattern, record.seq[::-1].translate(NUCL_COMP)):
            if (k := (m.end(1), (m.start(1)%3)+3)) not in matches:
                matches[k] = SeqRecord(
                    f'orf{hex(fnv1a_64(m.group(1)))[2:]} ' \
                    f'{record.seqid.split()[0]}:-:{m.start(1)+1}-{m.end(1)+1}',
                    m.group(1),
                    record.quals[m.start(1):m.end(1)])
    return list(matches.values())


###############################################################################


if __name__ == '__main__':

    # Parse commandline args ##################################################

    parser = ArgumentParser(
        description='orf_scanner: robust CDS prediction',
        argument_default=SUPPRESS)
    length_group = parser.add_mutually_exclusive_group()

    parser.add_argument(
        'input', nargs='*', type=FileType('r'), default=[stdin],
        help='FASTA/Q input file(s) (use "-" or leave blank for stdin)')
    parser.add_argument(
        '-m', '--model', nargs='*', type=FileType('r'),
        help='FASTA/Q file(s) of validated, in-frame, sense-oriented CDS \
        sequences for amino acid 3mer and codon frequency modeling')
    length_group.add_argument(
        '-l', '--min_length', nargs='?', dest='cds_len', const=100, default='100', type=int,
        help='Minimum length of CDS regions (default %(default)s AAs)')
    length_group.add_argument(
        '-p', '--max_prob', nargs='?', dest='cds_len', const=0.001, type=float,
        help='Calculate the mimium length of CDS regions to report such that \
        the probability of a random DNA sequence exceeding this value is <`p` \
        (%(const)s when supplied alone)')
    parser.add_argument(
        '--unstranded', action='store_true', default=False,
        help='Consider both sense and antisense frames (default behaviour \
        considers only the sense strand)')
    parser.add_argument(
        '--longest', action='store_true', default=False,
        help='Report only the longest CDS prediction for each seqeunce')
    parser.add_argument(
        '--complete', action='store_true', default=False,
        help='Report only CDS predictions ending in a stop codon (default \
        behaviour reports predictions extending 3\' off the seqeunce)')
    parser.add_argument(
        '-w', '--wrap', nargs='?', type=int, const=80, default=0,
        help='Wrap lines at `w` when outputting FASTA (%(const)s when \
        supplied alone)')
    parser.add_argument(
        '-q', '--quiet', action='store_true', default=False,
        help='Quiet! Silence stderr progress reporting')

    args = parser.parse_args()

    # Set up logging system ###################################################

    logger = logging.getLogger()
    logstream = logging.StreamHandler()
    logstream.setFormatter(logging.Formatter(
        '%(asctime)s\t%(message)s', '%y%m%d %H:%M:%S'))
    loglevel = logging.ERROR if args.quiet else logging.INFO
    logging.basicConfig(level=loglevel, handlers=[logstream])
    logger.info('Started orf_scanner')

    # Set up CDS regex ########################################################

    if isinstance(args.cds_len, float):
        vars(args)['cds_len'] = int(ceil(log(args.cds_len/(3/64), 1-(3/64))))

    orf_pattern = re.compile('(?=(ATG(?:(?!TAA|TGA|TAG)...){' +
        '{}'.format(args.cds_len-1) + ',}(?:TAA|TGA|TAG' +
        ('' if args.complete else '|..$|.$|$') + ')))')

    logger.info(f'Reporting CDS features >={args.cds_len} residues')

    # Build hexamer and nucleotide models #####################################

    if 'model' in args:

        logger.info('Building Markov chains for CDS scoring')

        # markov chains for AA 3mer usage
        # aa_chain[XXX] = float
        aa_chain = dd(float)
        aa_chain.update({''.join(g): 1.0 for g in
            product(set(CODON_MAP.values()), repeat=3)})
        aa_nullchain = deepcopy(aa_chain)
        # markov chains for codon usage
        # codon_chain[X][NNN] = float
        codon_chain = dd(lambda: dd(float))
        for g in product('ACGT', repeat=3):
            codon_chain[CODON_MAP.get(g := ''.join(g), 'X')][g] = 1.0
        codon_nullchain = deepcopy(codon_chain)

        for f in args.model:
            logger.info(f'\t{f.name}')
            with MultiReader(f, replaceU=True) as F:
                for n_records, record in enumerate(F, 1):
                    if len(record) < args.cds_len:
                        continue
                    # Fill the scores for the true sequences
                    resis = []
                    for g in iterkmers(record.seq, 3, 3):
                            resis.append(resi := CODON_MAP.get(g, 'X'))
                            codon_chain[resi][g] += 1.0
                    for g in iterkmers(''.join(resis), 3):
                        aa_chain[g] += 1
                    # Fill the nullscores from the scrambled sequences
                    shuf = ''.join(sample(record.seq, k=len(record)))
                    resis = []
                    for g in iterkmers(shuf, 3, 3):
                            resis.append(resi := CODON_MAP.get(g, 'X'))
                            codon_nullchain[resi][g] += 1.0
                    for g in iterkmers(''.join(resis), 3):
                        aa_nullchain[g] += 1
            logger.info(f'\t\t{n_records} sequences processed')
            f.close()

        # Process the counts into loglikelihoods
        for k in [k for k in aa_chain.keys() if 'N' in k]:
            del(aa_chain[k])
        count = sum(aa_chain.values())
        for k, v in aa_chain.items():
            aa_chain[k] = log(v/count, 2)

        for k in [k for k in aa_nullchain.keys() if 'N' in k]:
            del(aa_nullchain[k])
        count = sum(aa_nullchain.values())
        for k, v in aa_nullchain.items():
            aa_nullchain[k] = log(v/count, 2)

        if 'X' in codon_chain:
            del(codon_chain['X'])
        for resi, codons in codon_chain.items():
            count = sum(codons.values())
            for k, v in codons.items():
                codon_chain[resi][k] = log(v/count, 2)

        if 'X' in codon_nullchain:
            del(codon_nullchain['X'])
        for resi, codons in codon_nullchain.items():
            count = sum(codons.values())
            for k, v in codons.items():
                codon_nullchain[resi][k] = log(v/count, 2)

    # Identify ORFs ###########################################################

    logger.info('Identifying CDS sequences')

    retained = 0
    for f in args.input:
        logger.info(f'\t{f.name}')
        with MultiReader(f, replaceU=True) as F:
            if 'model' in args:
                for n_records, record in enumerate(F, 1):
                    cds_regions = []
                    for cds in find_orfs(record, args.unstranded):
                        resis = []
                        codon_score, codon_nullscore = 0.0, 0.0
                        for g in iterkmers(cds.seq, 3, 3):
                            resis.append(resi := CODON_MAP.get(g, 'X'))
                            try:
                                codon_score += codon_chain[resi][g]
                                codon_nullscore += codon_nullchain[resi][g]
                            except KeyError:
                                pass
                        aa_score, aa_nullscore = 0.0, 0.0
                        for g in iterkmers(''.join(resis), 3):
                            try:
                                aa_score += aa_chain[g]
                                aa_nullscore += aa_nullchain[g]
                            except KeyError:
                                pass
                        codon_bitscore = codon_score - codon_nullscore
                        aa_bitscore = aa_score - aa_nullscore
                        final_bitscore = (aa_bitscore + codon_bitscore) / len(cds)
                        if final_bitscore > 0:
                            cds_regions.append(cds)
                    if args.longest and cds_regions:
                        cds_regions = [sorted(cds_regions, key=lambda x: len(x), reverse=True)[0]]
                    for cds in cds_regions:
                        retained += 1
                        cds.print(wrap=args.wrap)
            else:
                for record in F:
                    cds_regions = find_orfs(record, args.unstranded)
                    if args.longest and cds_regions:
                        cds_regions = [sorted(cds_regions, key=lambda x: len(x), reverse=True)[0]]
                    for cds in cds_regions:
                        retained += 1
                        cds.print(wrap=args.wrap)
        f.close()
        logger.info(f'\t\t{retained} CDS regions predicted from {n_records} sequences')