<commit_before>#!/usr/bin/env ccp4-python

'''
09.11.2015

@author: hlfsimko
'''

import collections
import logging
import os
import tempfile

from ample.parsers import tmscore_parser
from ample.util import ample_util
from ample.util import pdb_edit

try:
    from ample.parsers import alignment_parser
    _BIOPYTHON = True
except ImportError:
    _BIOPYTHON = False


TMScoreModel = collections.namedtuple("TMScoreModel", 
                                      ["name", "pdbin", "TMSCORE_log", "structure", 
                                      "tm", "maxsub", "gdtts", "gdtha", "rmsd", 
                                      "nrResiduesCommon"])


def tmscoreAvail():
    """Check if TMscore binary is available"""
    try:
        ample_util.find_exe("TMscore")
    except:
        return False
    return True


class TMscorer(object):
    
    def __init__(self, structure, tmscore_exe, wdir=None):
        self.logger = logging.getLogger()
        self.pickle_file = None
        self.structure = structure
        self.tmscore_exe = tmscore_exe
        self.workingDIR = wdir
        return

    def main(self, pdb_list_file, keep_modified_structures=False, identical_sequences=False):
        """Wrapper function for compare_to_structure
        
        :pdb_list_file:            file containing a list of models
        :keep_modified_structures: do not delete any intermediate, modified structure files
        :sequence_identical:       avoid any modification of files due to sequence identity
        """
        
        # Read all pdb files from list
        pdbs = self._read_list(pdb_list_file)

        # Compare each pdb in list against structure defined
        self.entries = self.compare_to_structure(pdbs, 
                                                 keep_modified_structures=keep_modified_structures, 
                                                 identical_sequences=identical_sequences)

        return

    def compare_to_structure(self, pdb_list, keep_modified_structures=False, identical_sequences=False):
        """Compare a list of structures to a reference structure
        
        :pdb_list:                 list containing paths to model files
        :keep_modified_structures: do not delete any intermediate, modified structure files
        :sequence_identical:       avoid any modification of files due to sequence identity
        
        :returns:                  list of TMscore data per model
        """
        
        entries = []                                                    # For data storage
        
        self.logger.info('-------Evaluating decoys/models-------')

        structure = os.path.abspath(self.structure)                        # Path to reference structure
        structure_name = os.path.splitext(os.path.basename(structure))[0]  # Filename
        pdb_list_abs = [ os.path.abspath(model) for model in pdb_list ]    # Full paths to models

        for pdbin in pdb_list_abs:
            self.logger.debug("Working on %s" % pdbin)
            pdbin_name = os.path.splitext(os.path.basename(pdbin))[0]  # Filename        
            if not os.path.exists(pdbin): 
                self.logger.warning("Cannot find {0}".format(pdbin))
                continue
            
            # Modify structures to be identical as required by TMscore binary
            if not identical_sequences:
                pdbin_mod     = os.path.join(self.workingDIR, pdbin_name + "_mod.pdb")
                structure_mod = os.path.join(self.workingDIR, pdbin_name + "_" + 
                                             structure_name + "_mod.pdb")
                self.mod_structures(pdbin, pdbin_mod, structure, structure_mod)    
            model     = pdbin_mod     if not identical_sequences else pdbin
            reference = structure_mod if not identical_sequences else structure
    
            log = os.path.join(self.workingDIR, pdbin_name + "_tmscore.log")
            self.execute_comparison(model, reference, log)

            # Delete the modified structures if not wanted       
            if not keep_modified_structures and not identical_sequences:
                os.remove(pdbin_mod)
                os.remove(structure_mod)
            
            # Do the try clause here to allow anything that is required from here to throw
            # exceptions. In that case we revert to the TMscoreLogParser default values of
            # 0.0 for every score.
            pt = tmscore_parser.TMscoreLogParser()
            
            try:
                pt.parse(log)
            except Exception as e:
                self.logger.critical(e.msg)
                log = "None"
                
            entry = self._store(pdbin_name, pdbin, log, self.structure, pt)
            entries.append(entry)
                
        return entries

    def execute_comparison(self, model, reference, log=None):
        # Create a command list and execute TMscore
        cmd = [ self.tmscore_exe, model, reference ]
        p = ample_util.run_command(cmd, logfile=log, directory=self.workingDIR)
        return p
    
    def dump_csv(self, csv_file):
        if not len(self.entries): return
        import csv
        with open(csv_file, 'w') as f:
            fieldnames = self.entries[0]._asdict().keys()
            dw = csv.DictWriter(f, fieldnames=fieldnames)
            dw.writeheader()
            for e in self.entries:
                dw.writerow(e._asdict())
        print "Wrote csvfile: {0}".format(os.path.abspath(csv_file))
        return

    def mod_structures(self, pdbin, pdbin_mod, structure, structure_mod):
        """Make sure the decoy and the xtal pdb align to get an accurate TM-score"""
        
        if not _BIOPYTHON:
            raise ImportError
        
        # Disable the info logger to not spam the user with which chain of native extracted.
        # Happens for every model + native below
        # http://stackoverflow.com/questions/2266646/how-to-i-disable-and-re-enable-console-logging-in-python
        logging.disable( logging.INFO )

        pdbin_seq     = pdb_edit.sequence(pdbin).values()[0]
        structure_seq = pdb_edit.sequence(structure).values()[0]

        # Align the sequences to see how much of the predicted decoys are in the xtal
        aligned_seq_list = alignment_parser.AlignmentParser().align_sequences(pdbin_seq, 
                                                                              structure_seq)<commit_msg>LOGGER made module variable; variable renamed<commit_after>#!/usr/bin/env ccp4-python

'''
09.11.2015

@author: hlfsimko
'''

import collections
import logging
import os
import tempfile

from ample.parsers import tmscore_parser
from ample.util import ample_util
from ample.util import pdb_edit

try:
    from ample.parsers import alignment_parser
    _BIOPYTHON = True
except ImportError:
    _BIOPYTHON = False

LOGGER = logging.getLogger(__name__)

TMScoreModel = collections.namedtuple("TMScoreModel", 
                                      ["name", "pdbin", "TMSCORE_log", "structure", 
                                      "tm", "maxsub", "gdtts", "gdtha", "rmsd", 
                                      "nrResiduesCommon"])


def tmscoreAvail():
    """Check if TMscore binary is available"""
    try: ample_util.find_exe("TMscore")
    except: return False
    return True


class TMscorer(object):
    
    def __init__(self, structure, tmscore_exe, wdir=None):
        self.pickle_file = None
        self.structure = structure
        self.tmscore_exe = tmscore_exe
        self.work_dir = wdir
        return

    def main(self, pdb_list_file, keep_modified_structures=False, identical_sequences=False):
        """Wrapper function for compare_to_structure
        
        :pdb_list_file:            file containing a list of models
        :keep_modified_structures: do not delete any intermediate, modified structure files
        :sequence_identical:       avoid any modification of files due to sequence identity
        """
        # Read all pdb files from list
        pdbs = self._read_list(pdb_list_file)
        # Compare each pdb in list against structure defined
        self.entries = self.compare_to_structure(pdbs, 
                                                 keep_modified_structures=keep_modified_structures, 
                                                 identical_sequences=identical_sequences)
        return

    def compare_to_structure(self, pdb_list, keep_modified_structures=False, identical_sequences=False):
        """Compare a list of structures to a reference structure
        
        :pdb_list:                 list containing paths to model files
        :keep_modified_structures: do not delete any intermediate, modified structure files
        :sequence_identical:       avoid any modification of files due to sequence identity
        
        :returns:                  list of TMscore data per model
        """
        LOGGER.info('-------Evaluating decoys/models-------')        
        entries = []                                                    # For data storage

        structure = os.path.abspath(self.structure)                        # Path to reference structure
        structure_name = os.path.splitext(os.path.basename(structure))[0]  # Filename
        pdb_list_abs = [ os.path.abspath(model) for model in pdb_list ]    # Full paths to models

        for pdbin in pdb_list_abs:
            LOGGER.debug("Working on %s" % pdbin)
            pdbin_name = os.path.splitext(os.path.basename(pdbin))[0]  # Filename        
            if not os.path.exists(pdbin): 
                LOGGER.warning("Cannot find {0}".format(pdbin))
                continue
            
            # Modify structures to be identical as required by TMscore binary
            if not identical_sequences:
                pdbin_mod     = os.path.join(self.work_dir, pdbin_name + "_mod.pdb")
                structure_mod = os.path.join(self.work_dir, pdbin_name + "_" + 
                                             structure_name + "_mod.pdb")
                self.mod_structures(pdbin, pdbin_mod, structure, structure_mod)    
            model     = pdbin_mod     if not identical_sequences else pdbin
            reference = structure_mod if not identical_sequences else structure
    
            log = os.path.join(self.work_dir, pdbin_name + "_tmscore.log")
            self.execute_comparison(model, reference, log)

            # Delete the modified structures if not wanted       
            if not keep_modified_structures and not identical_sequences:
                os.remove(pdbin_mod)
                os.remove(structure_mod)
            
            # Do the try clause here to allow anything that is required from here to throw
            # exceptions. In that case we revert to the TMscoreLogParser default values of
            # 0.0 for every score.
            pt = tmscore_parser.TMscoreLogParser()
            
            try:
                pt.parse(log)
            except Exception as e:
                LOGGER.critical(e.msg)
                log = "None"
                
            entry = self._store(pdbin_name, pdbin, log, self.structure, pt)
            entries.append(entry)
                
        return entries

    def execute_comparison(self, model, reference, log=None):
        # Create a command list and execute TMscore
        cmd = [ self.tmscore_exe, model, reference ]
        p = ample_util.run_command(cmd, logfile=log, directory=self.work_dir)
        return p
    
    def dump_csv(self, csv_file):
        if not len(self.entries): return
        import csv
        with open(csv_file, 'w') as f:
            fieldnames = self.entries[0]._asdict().keys()
            dw = csv.DictWriter(f, fieldnames=fieldnames)
            dw.writeheader()
            for e in self.entries:
                dw.writerow(e._asdict())
        print "Wrote csvfile: {0}".format(os.path.abspath(csv_file))
        return

    def mod_structures(self, pdbin, pdbin_mod, structure, structure_mod):
        """Make sure the decoy and the xtal pdb align to get an accurate TM-score"""
        
        if not _BIOPYTHON:
            raise ImportError
        
        # Disable the info logger to not spam the user with which chain of native extracted.
        # Happens for every model + native below
        # http://stackoverflow.com/questions/2266646/how-to-i-disable-and-re-enable-console-logging-in-python
        logging.disable( logging.INFO )

        pdbin_seq     = pdb_edit.sequence(pdbin).values()[0]
        structure_seq = pdb_edit.sequence(structure).values()[0]

        # Align the sequences to see how much of the predicted decoys are in the xtal
        aligned_seq_list = alignment_parser.AlignmentParser().align_sequences(pdbin_seq, 
                                                                              structure_seq)
{'fraction_non_alphanumeric': 0.07570831398049234, 'fraction_numerical': 0.003870568199411674, 'mean_word_length': 1.8811329170383586, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7027174', 'n_tokens_mistral': 3576, 'n_tokens_neox': 3347, 'n_words': 930}