Benchmark

#!/usr/bin/env python
'''
    Benchmark program accepts two input files in UniProt-GOA GAF format at
    two distinct time points t1 and t2. The input files can be downloaded
    from ftp://ftp.ebi.ac.uk/pub/databases/GO/goa/. The simplest way to run
    this program is:

    python Benchmark -I1=input_file_at_t1 -I2=input_file_at_t2

    Running this program creates SIX benchmark files: THREE limited-knolwedge
    (LK) benchmark files - one in each of the three ontolgy categories (MFO, 
    BPO, and CCO) and THREE no-knowledge (NK) benchmark files - one in each 
    ontology category.
   
    Complete usage directions of this program can be obtained through the
    following command:

    python Benchmark --help
'''
import os
import sys
import shutil
import subprocess
from os.path import basename 
#from Bio.UniProt import GOA
import GOAParser as GOA

import ArgParser_Benchmark as ap
import Config
import CreateBenchmark as cb
import FormatChecker as fc
import GOAParser_cafa as gc
import LocateDataset as ld
import PaperTermFrequency as ptf

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

# Default configuration file name:
config_filename = '.cafarc' 

class Benchmark:
    def __init__(self):
        # Obtain user supplied argument values in a dictionary: 
        self.parsed_dict = ap.parse_args('benchmark')
        # Collect config file entries:
        self.ConfigParam = Config.read_config(config_filename) 
        # Retreive file name at time t1:
        t1 = self.parsed_dict['t1'] 
        # Retreive file name at time t2:
        t2 = self.parsed_dict['t2'] 
        if self.parsed_dict['t1'] == self.parsed_dict['t2']:
            print 'Both input files are from the same time point. ' + \
                'This will not create a valid benchmark set.'
            print 'Program quiting ...'
            sys.exit(1)

        # Retreive output file name:
        outfile_basename = basename(self.parsed_dict['outfile'])
        # Retreive work directory name:
        self.work_dir = (self.ConfigParam['workdir']).rstrip('/')

        # Create work direcoty, if it does not exist:
        if not os.path.exists(self.work_dir):
            os.makedirs(self.work_dir)
        # Locate t1 file:
        self.t1_input_file = ld.locate_GOAfile(t1, self.work_dir) 
        # Locate t2 file:
        self.t2_input_file = ld.locate_GOAfile(t2, self.work_dir) 

        # Names for SIX ouput files: bpo, cco, and mfo for 
        # LK and NK benchmark types:
        self.output_filename_LK_bpo = self.create_outfilename('LK_bpo')
        self.output_filename_LK_cco = self.create_outfilename('LK_cco')
        self.output_filename_LK_mfo = self.create_outfilename('LK_mfo')
        self.output_filename_NK_bpo = self.create_outfilename('NK_bpo')
        self.output_filename_NK_cco = self.create_outfilename('NK_cco')
        self.output_filename_NK_mfo = self.create_outfilename('NK_mfo')

        # Names for THREE files to store non-EXP and EXP type entries:
        # These files will be deleted once the calculation is done

        # File name for entries in t1 file with non-EXP evidence codes:
        self.t1_iea_name = self.t1_input_file + '.iea'

        # File name for entries in t1 file with EXP evidence codes:
        self.t1_exp_name = self.t1_input_file + '.exp'

        # File name for entries in t2 file with EXP evidence codes:
        self.t2_exp_name = self.t2_input_file + '.exp'

        # Name for GO ID frequency per pubmed id for t2 file:
        # This file will be deleted once the calculations are done
        self.t2_ptf_file =  self.t2_input_file + \
                            '_with_annotations_per_paper.txt'
        
        # Names for SIX intermediate benchmark files:
        self.bmfile_LK_bpo = self.t2_exp_name + '.bpo_LK_bench.txt'
        self.bmfile_LK_cco = self.t2_exp_name + '.cco_LK_bench.txt'
        self.bmfile_LK_mfo = self.t2_exp_name + '.mfo_LK_bench.txt'
        self.bmfile_NK_bpo = self.t2_exp_name + '.bpo_NK_bench.txt'
        self.bmfile_NK_cco = self.t2_exp_name + '.cco_NK_bench.txt'
        self.bmfile_NK_mfo = self.t2_exp_name + '.mfo_NK_bench.txt'

    def create_outfilename(self, ontType):
        """
        This method creates an output filename according to the following
        rules: 
            (1) When the user supplies the optional output filename prefix,
                this method looks for the latest version of the related file
                name in the workspace (See the while loop). It creates a new
                file name for the subsequent version:

                output_filename = self.work_dir + '/' + ob + '.' + str(index)
            
                This ensures that multiple runs of Benchmark program with the 
                same arguments creates new version of output files.

            (2) When the user does not supply the optional output filename 
                prefix, this method creates a new prefix based on the two 
                input data file names (See code in else block). Then the 
                program looks for the latest version of the related file 
                name in the workspace (See the while loop). Then, it creates
                a new file name for the subsequent version:

                output_filename = self.work_dir + '/' + ob + '.' + str(index)

                Here again, this ensures that multiple runs of Benchmark 
                program with the same arguments creates new version of output
                files.
        At the end, the method returns the newly created filename.
        """

        if not self.parsed_dict['outfile'] == '':
            ob = basename(self.parsed_dict['outfile']) + \
                 '.benchmark' + '_' + ontType
        else:

            if bool(self.parsed_dict['Taxon_ID']):
                ob = basename(self.parsed_dict['t2']) + '-' + \
                    ((basename(self.parsed_dict['t1'])).split('.'))[-1] + \
                    '.' + str((list(self.parsed_dict['Taxon_ID']))[0]) + \
                    '.benchmark' + '_' + ontType
            else: 
                ob = basename(self.parsed_dict['t2']) + '-' + \
                    ((basename(self.parsed_dict['t1'])).split('.'))[-1] + \
                    '.benchmark' + '_' + ontType
        index = 1
        while os.path.exists(self.work_dir + '/' + ob + '.' + str(index)):
            index = index + 1
        output_filename = self.work_dir + '/' + ob + '.' + str(index)
        return output_filename

    def create_outfilename_old(self, ontType):
        """
        This method creates an output filename according to the following
        rules: 
            (1) When the user supplies the optional output filename prefix,
                this method looks for the latest version of the related file
                name in the workspace (See the while loop). It creates a new
                file name for the subsequent version:

                output_filename = self.work_dir + '/' + ob + '.' + str(index)
            
                This ensures that multiple runs of Benchmark program with the 
                same arguments creates new version of output files.

            (2) When the user does not supply the optional output filename 
                prefix, this method creates a new prefix based on the two 
                input data file names (See code in else block). Then the 
                program looks for the latest version of the related file 
                name in the workspace (See the while loop). Then, it creates
                a new file name for the subsequent version:

                output_filename = self.work_dir + '/' + ob + '.' + str(index)

                Here again, this ensures that multiple runs of Benchmark 
                program with the same arguments creates new version of output
                files.
        At the end, the method returns the newly created filename.
        """
        if not self.parsed_dict['outfile'] == '':
            ob = basename(self.parsed_dict['outfile']) + \
                 '.benchmark' + '_' + ontType
        else:
            ob = basename(self.parsed_dict['t2']) + '-' + \
                 ((basename(self.parsed_dict['t1'])).split('.'))[-1] + \
                '.benchmark' + '_' + ontType
        index = 1
        while os.path.exists(self.work_dir + '/' + ob + '.' + str(index)):
            index = index + 1
        output_filename = self.work_dir + '/' + ob + '.' + str(index)
        return output_filename

    def create_iterator(self, infile):
        """
        This method creates an iterator object for the input UniProt-GOA file
        and returns it along with a list of all field names contained in the
        UniProt-GOA file. The UniProt-GOA file can either be in GAF 1.0 or
        GAF 2.0 file format.
        """
        infile_handle = open(infile, 'r')
        iter_handle = GOA.gafiterator(infile_handle)
        for ingen in iter_handle:
            if len(ingen) == 17:
                GAFFIELDS = GOA.GAF20FIELDS
                break
            else:
                GAFFIELDS = GOA.GAF10FIELDS
                break
        infile_handle = open(infile, 'r')
        iter_handle = GOA.gafiterator(infile_handle)
        return iter_handle, GAFFIELDS

    def remove_redundant_benchmarks(self):
        if os.stat(self.bmfile_LK_bpo).st_size == 0:
            print('Your limited-knowledge benchmark set for ' + \
                  'Biological Process Ontology is empty.')
            os.system('cp ' + self.bmfile_LK_bpo + ' ' + \
                              self.output_filename_LK_bpo)
        else:
            os.system('sort ' + self.bmfile_LK_bpo + ' | ' + \
                   'uniq >' + self.output_filename_LK_bpo)
        if os.stat(self.bmfile_LK_cco).st_size == 0:
            print('Your limited-knowledge benchmark set for ' + \
                  'Cellular Component Process Ontology is empty.')
            os.system('cp ' + self.bmfile_LK_cco + ' ' + \
                              self.output_filename_LK_cco)
        else:
            os.system('sort ' + self.bmfile_LK_cco + ' | ' + \
                   'uniq > ' + self.output_filename_LK_cco)
        if os.stat(self.bmfile_LK_mfo).st_size == 0:
            print('Your limited-knowledge benchmark set for '+ \
                  'Molecular Function Ontology is empty.')
            os.system('cp ' + self.bmfile_LK_mfo + ' ' + \
                              self.output_filename_LK_mfo)
        else:
            os.system('sort ' + self.bmfile_LK_mfo + ' | ' + \
                   'uniq > ' + self.output_filename_LK_mfo)
        if os.stat(self.bmfile_NK_bpo).st_size == 0:
            print('Your no-knowledge benchmark set for ' + \
                  'Biological Process Ontology is empty.')
            os.system('cp ' + self.bmfile_NK_bpo + ' ' + \
                              self.output_filename_NK_bpo)
        else:
            os.system('sort ' + self.bmfile_NK_bpo + ' | ' + \
                   'uniq > ' + self.output_filename_NK_bpo)
        if os.stat(self.bmfile_NK_cco).st_size == 0:
            print('Your no-knowledge benchmark set for ' + \
                  'Cellular Component Process Ontology is empty.')
            os.system('cp ' + self.bmfile_NK_cco + ' ' + \
                              self.output_filename_NK_cco)
        else:
            os.system('sort ' + self.bmfile_NK_cco + ' | ' + \
                   'uniq > ' + self.output_filename_NK_cco)
        if os.stat(self.bmfile_NK_mfo).st_size == 0:
            print('Your no-knowledge benchmark set for ' + \
                  'Molecular Function Ontology is empty.')
            os.system('cp ' + self.bmfile_NK_mfo + ' ' + \
                              self.output_filename_NK_mfo)
        else:
            os.system('sort ' + self.bmfile_NK_mfo + ' | ' + \
                   'uniq > ' + self.output_filename_NK_mfo)
        return None

    def create_intermediate_files(self):
        """
        This method creates all the necessary intermediate files 
        that are needed to create the desired benchmark sets.
        """
        # Create paper-term freq file for t2 file:
        ann_conf = ptf.paper_term_freq( open(self.t2_input_file,'r'),
                                        open(self.t2_ptf_file,'w'),
                                        self.parsed_dict)
        # Create an iterator object for filtering t2 file:
        iter_handle, GAFFIELDS = self.create_iterator(self.t2_input_file)

        # Create tax_id_name_mapping for filtering t2 file:
        tax_id_name_mapping = gc.parse_tax_file(self.ConfigParam['tax_file'])
        #print(tax_id_name_mapping)
        #sys.exit(0)
        # Create t2_exp_name file:
        # Filter t2 file for all proteins with EXP evidence:
        print 'Parsing t2 file: ' + basename(self.t2_input_file) + ' ...'
        t2_exp_handle = open(self.t2_exp_name, 'w')
        # Iterate through entries of the input file at time t2:
        for ingen in iter_handle: 
            retval = gc.record_has_forBenchmark(ingen,
                                                ann_conf,
                                                self.parsed_dict,
                                                tax_id_name_mapping,
                                                self.ConfigParam['exp_eec'],
                                                GAFFIELDS)
            # If retval is TRUE, write out the record to the file t2_exp_name:
            if retval:
                GOA.writerec(ingen, t2_exp_handle, GAFFIELDS)
        t2_exp_handle.close()

        # If t2.exp is empty, program quits:
        if os.stat(self.t2_exp_name).st_size == 0:
            print('Empty intermediate file: ' + basename(self.t2_exp_name))
            print('Your benchmark set will be empty with the ' + \
                  'parameters provided.')
            print('Program will quit after deleting the intermediate files ...')
            # Delete exp_name files:
            print('Deleting intermediate file: '+ basename(self.t2_exp_name))
            #os.remove(self.t2_exp_name)
            # Delete paper term frequency file for t2 file:
            print('Deleting intermediate file: '+ basename(self.t2_ptf_file))
            #os.remove(self.t2_ptf_file)
            print('Quiting ...')
            sys.exit(1)
       
        # Create t1.iea_name and t1.exp_name files:

        # Create an iterator handle for t1_input_file:
        iter_handle, GAFFIELDS = self.create_iterator(self.t1_input_file)
        print 'Parsing t1 file: ' + basename(self.t1_input_file) + ' ...'
        # Filter t1 file and create files t1.iea_name and t1.exp_name:
        gc.t1_filter(iter_handle, self.t1_iea_name, self.t1_exp_name,
                    self.t2_exp_name, GAFFIELDS, self.ConfigParam['exp_eec'])
        return None

    def delete_intermediate_files(self):
        print 'Cleaning working directory ...'
        # Delete SIX intermediate benchmark files:
        os.remove(self.bmfile_LK_bpo)
        os.remove(self.bmfile_LK_cco)
        os.remove(self.bmfile_LK_mfo)
        os.remove(self.bmfile_NK_bpo)
        os.remove(self.bmfile_NK_cco)
        os.remove(self.bmfile_NK_mfo)
        # Delete t1.iea_name, t1.exp_name, and t2.exp_name files:
        os.remove(self.t1_iea_name)
        os.remove(self.t1_exp_name)
        os.remove(self.t2_exp_name)
        # Delete paper term frequency file for t2 file:
        os.remove(self.t2_ptf_file)
        # Delete any empty files from the workspace (subdirectories included):
        for root, dirs, files in os.walk(self.work_dir):
            for fname in files:
                if os.path.getsize(root + '/' + fname) == 0:
                    os.remove(root + '/' + fname)
            break
        return None

    def check_gaf_format(self, goa_fname):
        """
        This method exits the Benchmark program on any of the
        following conditions:
            Case 1: if the file is empty
            Case 2: if the file is NOT in GAF format. To check this
                    it invokes check_gaf_format method of
                    FormatChecker module.
        """
        if os.stat(goa_fname).st_size == 0:
            print bcolors.WARNING + 'You submitted an empty file: ' + goa_fname + \
                  bcolors.ENDC
            sys.exit(1)
        elif not fc.check_gaf_format(open(goa_fname, 'r')):
            print bcolors.WARNING + 'File format error: ' + \
                  basename(goa_fname) + bcolors.ENDC
            print bcolors.WARNING + 'File must be in GAF 1.0 or GAF 2.0 ' + \
                'format' + bcolors.ENDC
            sys.exit(1)

    def print_prolog(self):
        print "*************************************************"
        print "Running Benchmark Creation Tool !!!!!"
        print 'Following is a list of user supplied inputs:'
        for arg in self.parsed_dict:
            print arg + ': ' + str(self.parsed_dict[arg])
        print '*********************************************\n'
        return None

    def print_epilog(self):
        print(bcolors.OKGREEN + 'The following benchmark files ' + \
                                'are created:' + bcolors.ENDC)
        if os.path.exists(self.output_filename_LK_bpo):
            print basename(self.output_filename_LK_bpo)
        if os.path.exists(self.output_filename_LK_cco):
            print basename(self.output_filename_LK_cco)
        if os.path.exists(self.output_filename_LK_mfo):
            print basename(self.output_filename_LK_mfo)
        if os.path.exists(self.output_filename_NK_bpo):
            print basename(self.output_filename_NK_bpo)
        if os.path.exists(self.output_filename_NK_cco):
            print basename(self.output_filename_NK_cco)
        if os.path.exists(self.output_filename_NK_mfo):
            print basename(self.output_filename_NK_mfo)
        print(bcolors.OKGREEN + 'Thank you for using Benchmark ' + \
                                'Creation Tool' + bcolors.ENDC)
        return None
    
    def process_data(self): 
        """ 
        This method processes user data, creates necessary intermediate files, 
        creates benchmark sets, and afterwards deletes the intermediate files.
        """
        # Print the welcome message and argument list:
        self.print_prolog()
        # File format check for t1 file:
        self.check_gaf_format(self.t1_input_file)
        # File format check for t2 file:
        self.check_gaf_format(self.t2_input_file)

        # Create necessary intermediate files:
        self.create_intermediate_files()
        # Populate benchmark files:
        cb.create_benchmarks(open(self.t1_iea_name, 'r'),
                             open(self.t1_exp_name, 'r'),
                             open(self.t2_exp_name, 'r'),
                             open(self.bmfile_LK_bpo, 'w'),
                             open(self.bmfile_LK_cco, 'w'),
                             open(self.bmfile_LK_mfo, 'w'),
                             open(self.bmfile_NK_bpo, 'w'),
                             open(self.bmfile_NK_cco, 'w'),
                             open(self.bmfile_NK_mfo, 'w'))
        # Remove redundant benchmark entries:
        self.remove_redundant_benchmarks()
        # Delete intermediate files:
        #self.delete_intermediate_files()
        # Print summary of running this program:
        self.print_epilog()
        return None

if __name__ == '__main__':
    if len(sys.argv) == 1:
        print (sys.argv[0] + ':')
        print(__doc__)
    else:
        # Create an instance of Benchmark class:
        bm = Benchmark()
        # Process data and create benchmark sets:
        bm.process_data()
    sys.exit(0)