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

##############################################################################
##  pyvolve: Python platform for simulating evolutionary sequences.
##
##  Written by <NAME> (<EMAIL>) 
##############################################################################


'''
Generate the instantaneous rate matrix for Markov chain.
'''



import numpy as np
from copy import deepcopy
from genetics import *
from state_freqs import *
ZERO      = 1e-8
MOLECULES = Genetics()


class MatrixBuilder(object):
    '''
        Parent class for model instantaneous matrix generation.
        
        Child class include the following:
        1. *aminoAcid_Matrix*       : Specific functionality for empirical amino acid models (currently, JTT, WAG, and LG and SCG05).
        2. *nucleotide_Matrix*      : Specific functionality for nucleotide models (GTR and nested).
        3. *mechCodon_Matrix*       : Specific functionality for so-called mechanistic codon models, which include GY-style and MG-style models (dN/dS models)
        5. *ECM_Matrix*             : Specific functionality for the ECM (Kosiol2007) empirical codon model
        6. *mutSel_Matrix*          : Specific functionality for the mutation-selection-balance model (Halpern and Bruno 1998). Extended to work for either codon or nucleotides.


        TODO:
            Incorpprate SCG05 empirical codon model. 
    '''
    
    def __init__(self, param_dict):
        self.params = param_dict
        


    def _sanity_params(self):
        '''
            Sanity-check that all necessary parameters have been supplied to construct the matrix.
        '''
        print "Parent class method, not called."
        

    def _sanity_params_state_freqs(self):
        '''
            Sanity-check specifically state_freqs key/value in the params dictionary.
        '''
        assert( 'state_freqs' in self.params ), "A list of state frequencies must be provided in params dictionary!"
        assert( len(self.params['state_freqs']) == self._size ), "state_freqs does not contain 20 values. Are you sure these are amino acid frequencies?"



    def _sanity_params_mutation_rates(self):
        '''
            Sanity-check specifically mu key/value in params dictionary.
        '''
        
        if 'mu' in self.params:
            # Single float provided
            if type(self.params['mu']) is float:
                new_mu = {'AC':1.,  'CA':1.,  'AG':1.,  'GA':1.,  'AT':1.,  'TA':1.,  'CG':1.,  'GC':1.,  'CT':1.,  'TC':1.,  'GT':1.,  'TG':1.}
                for key in new_mu:
                    new_mu[key] *= self.params['mu']
                self.params['mu'] = new_mu
        
            # Dictionary of mu's provided. Make sure dictionary is full. Anything missing, so replace with 1.
            elif type(self.params['mu']) is dict:
                for key in ['AC', 'CA', 'AG', 'GA', 'AT', 'TA', 'CG', 'GC', 'CT', 'TC', 'GT', 'TG']:
                    if key not in self.params['mu']:
                        self.params['mu'][key] = 1.
            
            else:
                raise AssertionError("You must provide EITHER a single mutation or a dictionary of mutation rates for nucleotide pairs to the key 'mu' in the 'params' dictionary.")
        
        # Nothing specified, so simply use equal mutation rates to construct matrix
        else:
            self.params['mu'] = {'AC':1.,  'CA':1.,  'AG':1.,  'GA':1.,  'AT':1.,  'TA':1.,  'CG':1.,  'GC':1.,  'CT':1.,  'TC':1.,  'GT':1.,  'TG':1.}

        # Apply kappa as needed.        
        if 'kappa' in self.params:
            temp_mu = deepcopy( self.params['mu'] )
            self.params['mu'] = {'AC': temp_mu['AC'], 'AG': temp_mu['AG'] * float(self.params['kappa']), 'AT': temp_mu['AT'], 'CG': temp_mu['CG'], 'CT': temp_mu['CT']*float(self.params['kappa']), 'GT': temp_mu['GT'], 'CA': temp_mu['CA'], 'GA': temp_mu['GA'] * float(self.params['kappa']), 'TA': temp_mu['TA'], 'GC': temp_mu['GC'], 'TC': temp_mu['TC']*float(self.params['kappa']), 'TG': temp_mu['TG']}



 
    def _sanity_params_dNdS(self, required = True):
        '''
            Sanity-check specifically beta, alpha keys/values in params dictionary.
            Ensures that a dN has been specified. If no dS, sets to 1.
            *Required* argument means that the model really does need a dN/dS. Codon models require for biology, but ECM models may receive..
        '''
        if 'omega' in self.params:
            self.params['beta'] = self.params['omega']
        if required:
            assert( 'beta' in self.params ), "You must specify a dN value (key 'omega' or key 'beta') in the params dictionary."
        else:
            if 'beta' not in self.params:
                self.params['beta'] = 1.
        if 'alpha' not in self.params:
            self.params['alpha'] = 1.
           

        


    def __call__(self):
        ''' 
            Generate the instantaneous rate matrix.
        '''    
        self._sanity_params() # sanity check parameterss
        self.inst_matrix = np.zeros( [self._size, self._size] ) # For nucleotides, self._size = 4; amino acids, self._size = 20; codons, self._size = 61.
        for s in range(self._size):
            for t in range(self._size):
                # Non-diagonal
                rate = self._calc_instantaneous_prob( s, t )                
                self.inst_matrix[s][t] = rate
                
            # Fill in the diagonal position so the row sums to 0, but ensure it doesn't became -0
            self.inst_matrix[s][s]= -1. * np.sum( self.inst_matrix[s] )
            if self.inst_matrix[s][s] == -0.:
                self.inst_matrix[s][s] = 0.
            assert ( abs(np.sum(self.inst_matrix[s])) < ZERO ), "Row in instantaneous matrix does not sum to 0."
        self._scale_matrix()
        return self.inst_matrix

<commit_msg>option to scale me matrix or not<commit_after>#! /usr/bin/env python

##############################################################################
##  pyvolve: Python platform for simulating evolutionary sequences.
##
##  Written by <NAME> (<EMAIL>) 
##############################################################################


'''
Generate the instantaneous rate matrix for Markov chain.
'''



import numpy as np
from copy import deepcopy
from genetics import *
from state_freqs import *
ZERO      = 1e-8
MOLECULES = Genetics()


class MatrixBuilder(object):
    '''
        Parent class for model instantaneous matrix generation.
        
        Child class include the following:
        1. *aminoAcid_Matrix*       : Specific functionality for empirical amino acid models (currently, JTT, WAG, and LG and SCG05).
        2. *nucleotide_Matrix*      : Specific functionality for nucleotide models (GTR and nested).
        3. *mechCodon_Matrix*       : Specific functionality for so-called mechanistic codon models, which include GY-style and MG-style models (dN/dS models)
        5. *ECM_Matrix*             : Specific functionality for the ECM (Kosiol2007) empirical codon model
        6. *mutSel_Matrix*          : Specific functionality for the mutation-selection-balance model (Halpern and Bruno 1998). Extended to work for either codon or nucleotides.


        TODO:
            Incorpprate SCG05 empirical codon model. 
    '''
    
    def __init__(self, param_dict, scale_matrix = True):
        self.params = param_dict
        self.rescale = scale_matrix
        


    def _sanity_params(self):
        '''
            Sanity-check that all necessary parameters have been supplied to construct the matrix.
        '''
        print "Parent class method, not called."
        

    def _sanity_params_state_freqs(self):
        '''
            Sanity-check specifically state_freqs key/value in the params dictionary.
        '''
        assert( 'state_freqs' in self.params ), "A list of state frequencies must be provided in params dictionary!"
        assert( len(self.params['state_freqs']) == self._size ), "state_freqs does not contain 20 values. Are you sure these are amino acid frequencies?"



    def _sanity_params_mutation_rates(self):
        '''
            Sanity-check specifically mu key/value in params dictionary.
        '''
        
        if 'mu' in self.params:
            # Single float provided
            if type(self.params['mu']) is float:
                new_mu = {'AC':1.,  'CA':1.,  'AG':1.,  'GA':1.,  'AT':1.,  'TA':1.,  'CG':1.,  'GC':1.,  'CT':1.,  'TC':1.,  'GT':1.,  'TG':1.}
                for key in new_mu:
                    new_mu[key] *= self.params['mu']
                self.params['mu'] = new_mu
        
            # Dictionary of mu's provided. Make sure dictionary is full. Anything missing, so replace with 1.
            elif type(self.params['mu']) is dict:
                for key in ['AC', 'CA', 'AG', 'GA', 'AT', 'TA', 'CG', 'GC', 'CT', 'TC', 'GT', 'TG']:
                    if key not in self.params['mu']:
                        self.params['mu'][key] = 1.
            
            else:
                raise AssertionError("You must provide EITHER a single mutation or a dictionary of mutation rates for nucleotide pairs to the key 'mu' in the 'params' dictionary.")
        
        # Nothing specified, so simply use equal mutation rates to construct matrix
        else:
            self.params['mu'] = {'AC':1.,  'CA':1.,  'AG':1.,  'GA':1.,  'AT':1.,  'TA':1.,  'CG':1.,  'GC':1.,  'CT':1.,  'TC':1.,  'GT':1.,  'TG':1.}

        # Apply kappa as needed.        
        if 'kappa' in self.params:
            temp_mu = deepcopy( self.params['mu'] )
            self.params['mu'] = {'AC': temp_mu['AC'], 'AG': temp_mu['AG'] * float(self.params['kappa']), 'AT': temp_mu['AT'], 'CG': temp_mu['CG'], 'CT': temp_mu['CT']*float(self.params['kappa']), 'GT': temp_mu['GT'], 'CA': temp_mu['CA'], 'GA': temp_mu['GA'] * float(self.params['kappa']), 'TA': temp_mu['TA'], 'GC': temp_mu['GC'], 'TC': temp_mu['TC']*float(self.params['kappa']), 'TG': temp_mu['TG']}



 
    def _sanity_params_dNdS(self, required = True):
        '''
            Sanity-check specifically beta, alpha keys/values in params dictionary.
            Ensures that a dN has been specified. If no dS, sets to 1.
            *Required* argument means that the model really does need a dN/dS. Codon models require for biology, but ECM models may receive..
        '''
        if 'omega' in self.params:
            self.params['beta'] = self.params['omega']
        if required:
            assert( 'beta' in self.params ), "You must specify a dN value (key 'omega' or key 'beta') in the params dictionary."
        else:
            if 'beta' not in self.params:
                self.params['beta'] = 1.
        if 'alpha' not in self.params:
            self.params['alpha'] = 1.
           

        


    def __call__(self):
        ''' 
            Generate the instantaneous rate matrix.
        '''    
        self._sanity_params() # sanity check parameterss
        self.inst_matrix = np.zeros( [self._size, self._size] ) # For nucleotides, self._size = 4; amino acids, self._size = 20; codons, self._size = 61.
        for s in range(self._size):
            for t in range(self._size):
                # Non-diagonal
                rate = self._calc_instantaneous_prob( s, t )                
                self.inst_matrix[s][t] = rate
                
            # Fill in the diagonal position so the row sums to 0, but ensure it doesn't became -0
            self.inst_matrix[s][s]= -1. * np.sum( self.inst_matrix[s] )
            if self.inst_matrix[s][s] == -0.:
                self.inst_matrix[s][s] = 0.
            assert ( abs(np.sum(self.inst_matrix[s])) < ZERO ), "Row in instantaneous matrix does not sum to 0."
        if self.rescale:
            self._scale_matrix()
        return self.inst_matrix


{'fraction_non_alphanumeric': 0.1518646350702921, 'fraction_numerical': 0.010270224766394477, 'mean_word_length': 2.2256312788487644, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1033373', 'n_tokens_mistral': 3579, 'n_tokens_neox': 3333, 'n_words': 1096}