digenome-run

#!/usr/bin/python

import os
from sys import argv
import pysam

def usage():
    print 'digenome-toolkit v1.2, written by Jeongbin Park'
    print 
    print 'USAGE: digenome-run [-s] [-c count_cutoff=10] [-r cutoff_ratio=20.0] [-g difference=1] [-u range_sum] [-p prefix] bam-file-path [step=1]'

def get_chrom_list(fn):
    chrom_list = []
    with pysam.Samfile(fn, "rb") as f:
        for item in f.header['SQ']:
            try:
                if item['SN'] != "*":
                    chrom_list.append(item['SN'])
            except:
                pass
    return chrom_list

def main():
    ratio_cutoff = 20.0
    difference = 1
    count_cutoff = 10
    range_sum = 1
    sum_cutoffs = False

    opts = []
    prefix = ""
    for i in range(1, len(argv)):
        if i != len(argv)-1 and argv[i] == "-p":
            prefix = argv[i+1]
        elif i != len(argv)-1 and argv[i] == "-g":
            difference = int(argv[i+1])
        elif i != len(argv)-1 and argv[i] == "-r":
            ratio_cutoff = float(argv[i+1])
        elif i != len(argv)-1 and argv[i] == "-c":
            count_cutoff = int(argv[i+1])
        elif i != len(argv)-1 and argv[i] == "-u":
            range_sum = int(argv[i+1])
        elif argv[i] == "-s":
            sum_cutoffs = True
        elif argv[i-1] != "-p" and argv[i-1] != "-g" and argv[i-1] != "-r" and argv[i-1] != "-c" and argv[i-1] != "-u":
            opts.append(argv[i])

    try:
        digenome_home = os.environ['DIGENOME_HOME']
    except:
        digenome_home = ''
    if digenome_home == '':
        print 'warning: DIGENOME_HOME environment variable is not set, using current dir...'
        print
        digenome_home = './'
    elif digenome_home[-1] != '/':
        digenome_home += '/'

    if len(opts) > 2 or len(opts) == 0:
        usage()
        return

    if len(opts) == 2:
        step = int(opts[1])
    else:
        step = 1

    chrom_list = get_chrom_list(opts[0])

    if prefix != "":
        opt_prefix = "-p " + prefix + " "
        pre_prefix = prefix + "_"
    else:
        opt_prefix = ""
        pre_prefix = ""

    if step <= 1:
        print('Running 1.find_position_bam...')
        os.system('{0}1.find_position_bam {1}{2}'.format(digenome_home, opt_prefix, opts[0]))
        print('')
        step = 2
    if step == 2:
        print('Running 2.sort.py...')
        for chrom in chrom_list:
            os.system('pypy {0}2.sort.py {1}{2}_reverse.txt'.format(digenome_home, pre_prefix, chrom))
        print('')
        step += 1
    if step == 3:
        print('Running 3.count.py...')
        for chrom in chrom_list:
            os.system('pypy {0}3.count.py -u {1} {2}{3}_forward.txt'.format(digenome_home, range_sum, pre_prefix, chrom))
            os.system('pypy {0}3.count.py -u {1} {2}{3}_reverse_sorted.txt'.format(digenome_home, range_sum, pre_prefix, chrom))
        print('')
        step += 1
    if step == 4:
        print('Running 4.cut_threshold.py...')
        for chrom in chrom_list:
            os.system('pypy {0}4.cut_threshold.py {1}{2}_forward_freq.txt'.format(digenome_home, pre_prefix, chrom))
            os.system('pypy {0}4.cut_threshold.py {1}{2}_reverse_sorted_freq.txt'.format(digenome_home, pre_prefix, chrom))
        print('')
        step += 1
    if step == 5:
        print('Running 5.coverage.py (it takes time)...')
        os.system('python {0}5.coverage.py {1} {2}'.format(digenome_home, prefix, opts[0]))
        print('')
        step += 1
    if step == 6:
        print('Running 6.get_depth.py...')
        try:
            print('Removing {0}forward.txt and {0}reverse.txt if exist...'.format(pre_prefix))
            os.remove('{0}forward.txt'.format(pre_prefix))
            os.remove('{0}reverse.txt'.format(pre_prefix))
        except:
            pass
        for chrom in chrom_list:
            os.system('pypy {0}6.get_depth.py {1} -d forward {2}{3}_forward_freq_from_5.txt {2}{3}_depth.txt'.format(digenome_home, opt_prefix, pre_prefix, chrom))
            os.system('pypy {0}6.get_depth.py {1} -d reverse {2}{3}_reverse_sorted_freq_from_5.txt {2}{3}_depth.txt'.format(digenome_home, opt_prefix, pre_prefix, chrom))
        print('')
        step += 1
    if step == 7:
        print('Running 7.ana_depth.py...')
        fns = ['{0}{1}_forward_freq_from_5.txt'.format(pre_prefix, chrom) for chrom in chrom_list]
        os.system('pypy {0}7.ana_depth.py {1} -d forward {2}'.format(digenome_home, opt_prefix, ' '.join(fns)))
        fns = ['{0}{1}_reverse_sorted_freq_from_5.txt'.format(pre_prefix, chrom) for chrom in chrom_list]
        os.system('pypy {0}7.ana_depth.py {1} -d reverse {2}'.format(digenome_home, opt_prefix, ' '.join(fns)))
        print('')
        step += 1
    if step == 8:
        print('Running 8.analysis.py...')
        for chrom in chrom_list:
            os.system('pypy {0}8.analysis.py{1} -c {2} -r {3} -g {4} {5}{6}_forward_freq_from_5_ana.txt {5}{6}_reverse_sorted_freq_from_5_ana.txt'.format(digenome_home, ' -s' if sum_cutoffs else '', count_cutoff, ratio_cutoff, difference, pre_prefix, chrom))
main()