run_trinity.py

#!/usr/bin/env python3

'''
Run Trinity for transcripts assembly

Trinity \
  --genome_guided_bam <BAM_FILE> \
  --genome_guided_max_intron 2000 \
  --max_memory 10G \
  --CPU <NUMBER_OF_CORES> \
  --output <OUTPUT_DIR> \
  --jaccard_clip

<BAM_FILE> is generated by Hisat2 running.
--genome_guided_max_intron: it is set to 2000 for fungal genomes, but can be
    modified with --max_intron parameter of this script.
--jaccard_clip: set if you have paired reads and you expect high gene
    density with UTR overlap.

Input: BAM file generated by Hisat
Output: transcript assembly file in FASTA
Last updated: Aug 12, 2020
'''

import os
import sys
from argparse import ArgumentParser
from glob import glob

from import_config import import_config
from set_logging import set_logging

# Parameters
D_CONF = import_config()
MAX_MEMORY = '10G'


def main():
    '''Main function'''
    argparse_usage = (
        'run_trinity.py -b <bam_files> -o <output_dir> -l <log_dir> '
        '-c <num_cores> -m <max_intron>'
    )
    parser = ArgumentParser(usage=argparse_usage)
    parser.add_argument(
        '-b', '--bam_files', nargs='+', required=True,
        help='Sorted BAM files generated by HISAT2'
    )
    parser.add_argument(
        '-o', '--output_dir', nargs='?', default='trinity_out',
        help='Output directory (default: trinity_out)'
    )
    parser.add_argument(
        '-l', '--log_dir', nargs='?', default='logs',
        help='Log directory (default: logs)'
    )
    parser.add_argument(
        '-c', '--num_cores', nargs='?', default=1, type=int,
        help='Number of cores to be used (default: 1)'
    )
    parser.add_argument(
        '-m', '--max_intron', nargs='?', default=2000, type=int,
        help='Max intron length (Default: 2000 bp)'
    )
    parser.add_argument(
        '--jaccard_clip', action='store_true',
        help='--jaccard_clip flag in Trinity'
    )

    args = parser.parse_args()
    output_dir = os.path.abspath(args.output_dir)
    log_dir = os.path.abspath(args.log_dir)
    bam_files = [os.path.abspath(x) for x in args.bam_files]
    num_cores = args.num_cores
    max_intron = args.max_intron

    if args.jaccard_clip:
        jaccard_clip_flag = '--jaccard_clip'
    else:
        jaccard_clip_flag = ''

    # Create necessary dirs
    create_dir(output_dir, log_dir)

    # Set logging
    log_file = os.path.join(log_dir, 'run_trinity.log')
    logger = set_logging(log_file)
    logger_txt = logger[1]

    # Check bamfile
    bam_files = [x for x in bam_files if glob(x)]
    if not bam_files:
        logger_txt.debug('[ERROR] You provided wrong BAM FILES. Please check')
        sys.exit(2)

    # Run functions :)
    run_trinity(
        bam_files, output_dir, log_dir, num_cores, max_intron,
        jaccard_clip_flag, logger
    )


def import_file(input_file):
    '''Import file'''
    with open(input_file) as f_in:
        txt = list(line.rstrip() for line in f_in)
    return txt


def create_dir(output_dir, log_dir):
    '''Create directories'''
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    if not os.path.exists(log_dir):
        os.mkdir(log_dir)


def run_trinity(
        bam_files, output_dir, log_dir, num_cores, max_intron,
        jaccard_clip_flag, logger):
    '''Run Trinity'''
    trinity_bin = D_CONF['TRINITY_PATH']
    # Trinity --genome_guided_bam rnaseq_alignments.csorted.bam
    # --max_memory 50G --genome_guided_max_intron 2000 --CPU 6
    logger_time, logger_txt = logger
    for bam_file in bam_files:
        prefix = (
            os.path.splitext(os.path.basename(bam_file))[0]
        )
        outdir = os.path.join(output_dir, 'trinity_{}'.format(prefix))

        new_output = os.path.join(outdir, 'Trinity_{}.fasta'.format(prefix))
        logger_time.debug('START: Trinity for %s', prefix)
        if not os.path.exists(new_output):
            log_file = os.path.join(log_dir, 'trinity_{}.log'.format(prefix))
            command = (
                '{} {} --genome_guided_bam {} --genome_guided_max_intron {} '
                '--max_memory {} --CPU {} --output {} > {} 2>&1'.format(
                    trinity_bin, jaccard_clip_flag, bam_file, max_intron,
                    MAX_MEMORY, num_cores, outdir, log_file
                )
            )
            logger_txt.debug('[Run] %s', command)
            os.system(command)

            # Rename the file
            trinity_output = os.path.join(outdir, 'Trinity-GG.fasta')
            os.rename(trinity_output, new_output)

        else:
            logger_txt.debug(
                '[Note] Running Trinity has already been finished %s', prefix
            )

        logger_time.debug('DONE : Trinity for %s', prefix)


if __name__ == '__main__':
    main()