<s> '''<STR_LIT>''' <EOL> import numpy as np <EOL> from scipy import linalg <EOL> from copy import deepcopy <EOL> from genetics import * <EOL> from state_freqs import * <EOL> ZERO = <NUM_LIT> <EOL> MOLECULES = Genetics ( ) <EOL> class MatrixBuilder ( object ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , model_type , parameters ) : <EOL> '''<STR_LIT>''' <EOL> self . model_type = model_type . lower ( ) <EOL> self . params = parameters <EOL> def _build_matrix ( self , parameters = None ) : <EOL> '''<STR_LIT>''' <EOL> if parameters is None : <EOL> parameters = self . params <EOL> matrix = np . zeros ( [ self . _size , self . _size ] ) <EOL> for s in range ( self . _size ) : <EOL> for t in range ( self . _size ) : <EOL> rate = self . _calc_instantaneous_prob ( s , t , parameters ) <EOL> matrix [ s ] [ t ] = rate <EOL> matrix [ s ] [ s ] = - <NUM_LIT:1.> * np . sum ( matrix [ s ] ) <EOL> if matrix [ s ] [ s ] == - <NUM_LIT:0.> : <EOL> matrix [ s ] [ s ] = <NUM_LIT:0.> <EOL> assert ( abs ( np . sum ( matrix [ s ] ) ) < ZERO ) , "<STR_LIT>" <EOL> return matrix <EOL> def _scaling_factor_from_matrix ( self , frequencies , matrix ) : <EOL> '''<STR_LIT>''' <EOL> scaling_factor = <NUM_LIT:0.> <EOL> for i in range ( self . _size ) : <EOL> scaling_factor += ( matrix [ i ] [ i ] * frequencies [ i ] ) <EOL> return scaling_factor <EOL> def _obtain_scaling_factor ( self ) : <EOL> '''<STR_LIT>''' <EOL> if self . scale_matrix == "<STR_LIT>" : <EOL> factor = self . _scaling_factor_from_matrix ( self . params [ "<STR_LIT>" ] , self . inst_matrix ) <EOL> else : <EOL> frequencies , matrix = self . _build_scaling_matrix ( ) <EOL> factor = self . _scaling_factor_from_matrix ( frequencies , matrix ) <EOL> return factor <EOL> def _build_scaling_matrix ( self ) : <EOL> '''<STR_LIT>''' <EOL> params = self . _build_scaling_params ( ) <EOL> matrix = self . _build_matrix ( parameters = params ) <EOL> return params [ "<STR_LIT>" ] , matrix <EOL> def __call__ ( self ) : <EOL> '''<STR_LIT>''' <EOL> self . inst_matrix = self . _build_matrix ( ) <EOL> scaling_factor = self . _obtain_scaling_factor ( ) <EOL> self . inst_matrix /= - <NUM_LIT:1.> * scaling_factor <EOL> return self . inst_matrix <EOL> def _init_empirical_matrix ( self , name ) : <EOL> '''<STR_LIT>''' <EOL> from . import empirical_matrices as em <EOL> try : <EOL> self . emp_matrix = eval ( "<STR_LIT>" + name + "<STR_LIT>" ) <EOL> except : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> def _is_TI ( self , source , target ) : <EOL> '''<STR_LIT>''' <EOL> ti_pyrim = source in MOLECULES . pyrims and target in MOLECULES . pyrims <EOL> ti_purine = source in MOLECULES . purines and target in MOLECULES . purines <EOL> if ti_pyrim or ti_purine : <EOL> return True <EOL> else : <EOL> return False <EOL> def _is_syn ( self , source , target ) : <EOL> '''<STR_LIT>''' <EOL> source_codon = MOLECULES . codons [ source ] <EOL> target_codon = MOLECULES . codons [ target ] <EOL> if ( MOLECULES . codon_dict [ source_codon ] == MOLECULES . codon_dict [ target_codon ] ) : <EOL> return True <EOL> else : <EOL> return False <EOL> def _get_nucleotide_diff ( self , source , target ) : <EOL> '''<STR_LIT>''' <EOL> source_codon = MOLECULES . codons [ source ] <EOL> target_codon = MOLECULES . codons [ target ] <EOL> return "<STR_LIT>" . join ( [ source_codon [ i ] + target_codon [ i ] for i in range ( len ( source_codon ) ) if source_codon [ i ] != target_codon [ i ] ] ) <EOL> def _calc_instantaneous_prob ( self , source , target , params = None ) : <EOL> '''<STR_LIT>''' <EOL> print ( "<STR_LIT>" ) <EOL> class AminoAcid_Matrix ( MatrixBuilder ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , * args ) : <EOL> super ( AminoAcid_Matrix , self ) . __init__ ( * args ) <EOL> self . _size = <NUM_LIT:20> <EOL> self . code = MOLECULES . amino_acids <EOL> self . scale_matrix = "<STR_LIT>" <EOL> self . _init_empirical_matrix ( self . model_type ) <EOL> def _calc_instantaneous_prob ( self , source , target , parameters = None ) : <EOL> '''<STR_LIT>''' <EOL> if parameters is None : <EOL> parameters = self . params <EOL> return self . emp_matrix [ source ] [ target ] * parameters [ '<STR_LIT>' ] [ target ] <EOL> class Nucleotide_Matrix ( MatrixBuilder ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , * args ) : <EOL> super ( Nucleotide_Matrix , self ) . __init__ ( * args ) <EOL> self . _size = <NUM_LIT:4> <EOL> self . _code = MOLECULES . nucleotides <EOL> self . scale_matrix = "<STR_LIT>" <EOL> def _calc_instantaneous_prob ( self , source , target , parameters = None ) : <EOL> '''<STR_LIT>''' <EOL> if parameters is None : <EOL> parameters = self . params <EOL> source_nuc = self . _code [ source ] <EOL> target_nuc = self . _code [ target ] <EOL> if source_nuc == target_nuc : <EOL> return <NUM_LIT:0.> <EOL> else : <EOL> return parameters [ '<STR_LIT>' ] [ target ] * parameters [ '<STR_LIT>' ] [ "<STR_LIT>" . join ( sorted ( source_nuc + target_nuc ) ) ] <EOL> class MechCodon_Matrix ( MatrixBuilder ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , * args ) : <EOL> super ( MechCodon_Matrix , self ) . __init__ ( * args ) <EOL> self . _size = <NUM_LIT> <EOL> self . _code = MOLECULES . codons <EOL> if "<STR_LIT>" in self . params : <EOL> self . scale_matrix = "<STR_LIT>" <EOL> else : <EOL> self . scale_matrix = "<STR_LIT>" <EOL> def _calc_prob ( self , target_codon , target_nuc , nuc_pair , subrate ) : <EOL> '''<STR_LIT>''' <EOL> prob = self . params [ '<STR_LIT>' ] [ nuc_pair ] * subrate <EOL> if self . model_type == '<STR_LIT>' : <EOL> prob *= self . params [ '<STR_LIT>' ] [ target_codon ] <EOL> else : <EOL> prob *= self . params [ "<STR_LIT>" ] [ MOLECULES . nucleotides . index ( target_nuc ) ] <EOL> return prob <EOL> def _calc_instantaneous_prob ( self , source , target , parameters = None ) : <EOL> '''<STR_LIT>''' <EOL> if parameters is None : <EOL> parameters = self . params <EOL> nuc_diff = self . _get_nucleotide_diff ( source , target ) <EOL> if len ( nuc_diff ) != <NUM_LIT:2> : <EOL> return <NUM_LIT:0.> <EOL> else : <EOL> nuc_pair = "<STR_LIT>" . join ( sorted ( nuc_diff [ <NUM_LIT:0> ] + nuc_diff [ <NUM_LIT:1> ] ) ) <EOL> if self . _is_syn ( source , target ) : <EOL> return self . _calc_prob ( target , nuc_diff [ <NUM_LIT:1> ] , nuc_pair , parameters [ '<STR_LIT>' ] ) <EOL> else : <EOL> return self . _calc_prob ( target , nuc_diff [ <NUM_LIT:1> ] , nuc_pair , parameters [ '<STR_LIT>' ] ) <EOL> def _build_scaling_params ( self ) : <EOL> '''<STR_LIT>''' <EOL> new_parameters = { "<STR_LIT>" : <NUM_LIT:1.> , "<STR_LIT>" : self . params [ "<STR_LIT>" ] , "<STR_LIT>" : self . params [ "<STR_LIT>" ] } <EOL> if self . scale_matrix == "<STR_LIT>" : <EOL> try : <EOL> new_parameters [ "<STR_LIT>" ] = self . params [ "<STR_LIT>" ] <EOL> except : <EOL> raise KeyError ( "<STR_LIT>" ) <EOL> elif self . scale_matrix == "<STR_LIT>" : <EOL> new_parameters [ "<STR_LIT>" ] = <NUM_LIT:1.> <EOL> else : <EOL> raise AssertionError ( "<STR_LIT>" ) <EOL> return new_parameters <EOL> class MutSel_Matrix ( MatrixBuilder ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , * args ) : <EOL> super ( MutSel_Matrix , self ) . __init__ ( * args ) <EOL> try : <EOL> self . _size = len ( self . params [ "<STR_LIT>" ] ) <EOL> except : <EOL> self . _size = len ( self . params [ "<STR_LIT>" ] ) <EOL> self . scale_matrix = "<STR_LIT>" <EOL> if self . _size == <NUM_LIT:4> : <EOL> self . _code = MOLECULES . nucleotides <EOL> elif self . _size == <NUM_LIT> : <EOL> self . _code = MOLECULES . codons <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> def _calc_fixrate_state_freqs ( self , source , target , nuc_diff , parameters ) : <EOL> '''<STR_LIT>''' <EOL> pi_i = parameters [ '<STR_LIT>' ] [ source ] <EOL> pi_j = parameters [ '<STR_LIT>' ] [ target ] <EOL> mu_ij = parameters [ "<STR_LIT>" ] [ nuc_diff ] <EOL> mu_ji = parameters [ "<STR_LIT>" ] [ nuc_diff [ <NUM_LIT:1> ] + nuc_diff [ <NUM_LIT:0> ] ] <EOL> if abs ( pi_i ) <= ZERO or abs ( pi_j ) <= ZERO : <EOL> fixation_rate = <NUM_LIT:0.> <EOL> else : <EOL> pi_mu = ( mu_ji * pi_j ) / ( mu_ij * pi_i ) <EOL> if abs ( <NUM_LIT:1.> - pi_mu ) <= ZERO : <EOL> fixation_rate = <NUM_LIT:1.> <EOL> else : <EOL> fixation_rate = np . log ( pi_mu ) / ( <NUM_LIT:1.> - <NUM_LIT:1.> / pi_mu ) <EOL> return fixation_rate <EOL> def _calc_fixrate_fitness ( self , source , target , parameters ) : <EOL> '''<STR_LIT>''' <EOL> sij = parameters [ '<STR_LIT>' ] [ target ] - parameters [ '<STR_LIT>' ] [ source ] <EOL> if abs ( sij ) <= ZERO : <EOL> fixation_rate = <NUM_LIT:1.> <EOL> else : <EOL> fixation_rate = ( sij ) / ( <NUM_LIT:1> - np . exp ( - <NUM_LIT:1.> * sij ) ) <EOL> return fixation_rate <EOL> def _calc_instantaneous_prob ( self , source , target , parameters = None ) : <EOL> '''<STR_LIT>''' <EOL> if parameters is None : <EOL> parameters = self . params <EOL> nuc_diff = self . _get_nucleotide_diff ( source , target ) <EOL> if len ( nuc_diff ) != <NUM_LIT:2> : <EOL> return <NUM_LIT:0.> <EOL> else : <EOL> if self . params [ "<STR_LIT>" ] : <EOL> fixation_rate = self . _calc_fixrate_state_freqs ( source , target , nuc_diff , parameters ) <EOL> else : <EOL> fixation_rate = self . _calc_fixrate_fitness ( source , target , parameters ) <EOL> return fixation_rate * parameters [ '<STR_LIT>' ] [ nuc_diff ] <EOL> def _build_scaling_params ( self ) : <EOL> '''<STR_LIT>''' <EOL> assert ( self . scale_matrix == "<STR_LIT>" ) , "<STR_LIT>" <EOL> new_parameters = { "<STR_LIT>" : self . params [ "<STR_LIT>" ] } <EOL> new_parameters [ "<STR_LIT>" ] = np . repeat ( <NUM_LIT:1.> / self . _size , self . _size ) <EOL> new_parameters [ "<STR_LIT>" ] = np . ones ( self . _size ) <EOL> return new_parameters <EOL> class ECM_Matrix ( MatrixBuilder ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , * args ) : <EOL> super ( ECM_Matrix , self ) . __init__ ( * args ) <EOL> if self . model_type == '<STR_LIT>' : <EOL> self . restricted = True <EOL> elif self . model_type == '<STR_LIT>' : <EOL> self . restricted = False <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . _size = <NUM_LIT> <EOL> self . _code = MOLECULES . codons <EOL> self . scale_matrix = "<STR_LIT>" <EOL> self . _init_empirical_matrix ( self . model_type ) <EOL> def _set_kappa_param ( self , nuc_diff ) : <EOL> '''<STR_LIT>''' <EOL> num_ti = <NUM_LIT:0> <EOL> num_tv = <NUM_LIT:0> <EOL> for i in range ( <NUM_LIT:0> , len ( nuc_diff ) , <NUM_LIT:2> ) : <EOL> source_nuc = nuc_diff [ i ] <EOL> target_nuc = nuc_diff [ i + <NUM_LIT:1> ] <EOL> if self . _is_TI ( source_nuc , target_nuc ) : <EOL> num_ti += <NUM_LIT:1> <EOL> else : <EOL> num_tv += <NUM_LIT:1> <EOL> return self . params [ '<STR_LIT>' ] ** num_ti * self . params [ '<STR_LIT>' ] ** num_tv <EOL> def _calc_instantaneous_prob ( self , source , target , parameters = None ) : <EOL> '''<STR_LIT>''' <EOL> if parameters is None : <EOL> parameters = self . params <EOL> nuc_diff = self . _get_nucleotide_diff ( source , target ) <EOL> if len ( nuc_diff ) == <NUM_LIT:0> or ( self . restricted and len ( nuc_diff ) != <NUM_LIT:2> ) : <EOL> return <NUM_LIT:0.> <EOL> else : <EOL> kappa_param = self . _set_kappa_param ( nuc_diff ) <EOL> if self . _is_syn ( source , target ) : <EOL> return self . emp_matrix [ source ] [ target ] * parameters [ '<STR_LIT>' ] [ target ] * parameters [ '<STR_LIT>' ] * kappa_param <EOL> else : <EOL> return self . emp_matrix [ source ] [ target ] * parameters [ '<STR_LIT>' ] [ target ] * parameters [ '<STR_LIT>' ] * kappa_param </s>
