run_trinity.py

#!/usr/bin/env python2

'''
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
'''

# Import modules
import sys
import os
from glob import glob
from argparse import ArgumentParser

# Get Logging
this_path = os.path.realpath(__file__)
this_dir = os.path.dirname(this_path)
sys.path.append(this_dir)
from set_logging import set_logging
from import_config import import_config

# Parameters
max_memory = '10G'
program_name = 'trinity'


# Main function
def main(argv):
    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"
    )
    parser.add_argument(
        "-l", "--log_dir", nargs='?', default='logs',
        help="Log directory"
    )
    parser.add_argument(
        "-c", "--num_cores", nargs='?', default=1, type=int,
        help="Number of cores to be used"
    )
    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')
    global logger_time, logger_txt
    logger_time, logger_txt = set_logging(log_file)

    # 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
    )


# Define functions
def import_file(input_file):
    with open(input_file) as f_in:
        txt = (line.rstrip() for line in f_in)
        txt = list(line for line in txt if line)
    return txt


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

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

    log_output_dir = os.path.join(log_dir, program_name)
    if not os.path.exists(log_output_dir):
        os.mkdir(log_output_dir)


def run_trinity(
    bam_files, output_dir, log_dir, num_cores, max_intron, jaccard_clip_flag
):
    D_conf = import_config(this_dir)
    trinity_bin = D_conf['TRINITY_PATH']

    # Trinity --genome_guided_bam rnaseq_alignments.csorted.bam
    # --max_memory 50G --genome_guided_max_intron 2000 --CPU 6
    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 {}'.format(prefix))
        if not os.path.exists(new_output):
            log_file = os.path.join(
                log_dir, program_name, '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] {}'.format(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(
                'Running Trinity has already been finished {}'.format(prefix)
            )

        logger_time.debug('DONE : Trinity for {}'.format(prefix))


if __name__ == "__main__":
    main(sys.argv[1:])