<s> '''<STR_LIT>''' <EOL> from model import * <EOL> class Partition ( ) : <EOL> def __init__ ( self , ** kwargs ) : <EOL> '''<STR_LIT>''' <EOL> self . size = kwargs . get ( '<STR_LIT:size>' , None ) <EOL> self . MRCA = kwargs . get ( '<STR_LIT>' , None ) <EOL> self . models = kwargs . get ( '<STR_LIT>' , None ) <EOL> if self . models is None : <EOL> self . model = kwargs . get ( '<STR_LIT>' , None ) <EOL> self . root_model_name = kwargs . get ( '<STR_LIT>' , None ) <EOL> self . _shuffle = False <EOL> self . _root_model = None <EOL> self . _partition_sanity ( ) <EOL> self . _size_MRCA_sanity ( ) <EOL> if self . MRCA is None : <EOL> self . _divvy_partition_size ( ) <EOL> def _size_MRCA_sanity ( self ) : <EOL> '''<STR_LIT>''' <EOL> assert ( self . size is not None or self . MRCA is not None ) , "<STR_LIT>" <EOL> if self . MRCA is not None : <EOL> assert ( type ( self . MRCA ) is str ) , "<STR_LIT>" <EOL> if self . size is not None : <EOL> print ( "<STR_LIT>" ) <EOL> code_step = len ( self . _root_model . code [ <NUM_LIT:0> ] ) <EOL> self . size = [ len ( self . MRCA ) / code_step ] <EOL> for model in self . models : <EOL> model . rate_probs = np . array ( [ <NUM_LIT:1.> ] ) <EOL> model . rate_factors = np . array ( [ <NUM_LIT:1.> ] ) <EOL> def _partition_sanity ( self ) : <EOL> '''<STR_LIT>''' <EOL> assert ( self . models is not None ) , "<STR_LIT>" <EOL> if type ( self . models ) is not list : <EOL> self . models = [ self . models ] <EOL> code1 = self . models [ <NUM_LIT:0> ] . code <EOL> if len ( self . models ) > <NUM_LIT:1> : <EOL> for m in self . models [ <NUM_LIT:1> : ] : <EOL> assert ( m . code == code1 ) , "<STR_LIT>" <EOL> if self . branch_het ( ) : <EOL> for m in self . models : <EOL> if m . name == self . root_model_name : <EOL> self . _root_model = m <EOL> break <EOL> else : <EOL> self . _root_model = self . models [ <NUM_LIT:0> ] <EOL> assert ( self . _root_model != None ) , "<STR_LIT>" <EOL> if self . site_het ( ) : <EOL> self . _shuffle = True <EOL> for model in self . models : <EOL> assert ( len ( model . rate_probs ) == len ( self . _root_model . rate_probs ) ) , "<STR_LIT>" <EOL> def _divvy_partition_size ( self ) : <EOL> '''<STR_LIT>''' <EOL> nc = self . _root_model . num_classes ( ) <EOL> rate_occurrences = np . random . choice ( nc , int ( self . size ) , p = self . _root_model . rate_probs ) <EOL> new_size = np . bincount ( rate_occurrences , minlength = nc ) <EOL> assert ( sum ( new_size ) == self . size ) , "<STR_LIT>" <EOL> assert ( len ( new_size ) == nc ) , "<STR_LIT>" <EOL> self . size = list ( new_size ) <EOL> def branch_het ( self ) : <EOL> '''<STR_LIT>''' <EOL> if isinstance ( self . models , Model ) or len ( self . models ) == <NUM_LIT:1> : <EOL> return False <EOL> elif len ( self . models ) > <NUM_LIT:1> : <EOL> return True <EOL> else : <EOL> raise AssertionError ( "<STR_LIT>" ) <EOL> def site_het ( self ) : <EOL> '''<STR_LIT>''' <EOL> if self . MRCA is not None : <EOL> return False <EOL> if self . models [ <NUM_LIT:0> ] . num_classes ( ) > <NUM_LIT:1> : <EOL> return True <EOL> else : <EOL> return False <EOL> def is_codon_model ( self ) : <EOL> '''<STR_LIT>''' <EOL> return self . models [ <NUM_LIT:0> ] . is_codon_model ( ) </s>
