def _calc_full_traversal_lnL_for_model(self, model): _LOG.debug("_calc_full_traversal_lnL_for_model called for %s" % str(model)) assert(self._scheduler is None) from pytbeaglehon.op_scheduling import TogglePartialScheduler es = model.eigen_soln self._scheduler = TogglePartialScheduler(self, model) try: for node in self._tree.postorder_internal_node_iter(): self._scheduler.add_internal_node_to_partial_calc(node) finally: self._scheduler.end_partial_calculations() prev_sched = self._scheduler self._scheduler = None root_partials = prev_sched._LCE_last_queued_dest model.transmit_category_weights() model.transmit_state_freq() assert (es is model.eigen_soln) return self._LCE.integrate_likelihood(model, root_partials)
class TogglePartialTreeScorer(TreeScorer): '''Assumes that there are: - enough partials for every internal node to have two - enough prob mats for every edge to have two. This enables an efficient API of repeatedly: 1. propose-new-state (by calling ...param_changed methods) followed by, 2. "accept" or "reject" call Note that when reject is called, the client must then call: 1. start_revert() 2. parameter changing moves to return the tree to the previous state, 3. end_revert() So that the changes to restore the tree do not result in "dirty-ing" of partials. ''' def __init__(self, like_calc_env, tree): TreeScorer.__init__(self, like_calc_env, tree) _LOG.debug("TogglePartialTreeScorer.__init__.entire_tree_is_dirty = %s" % str(self.entire_tree_is_dirty)) def initialize_tree(self): LCE = self._LCE tips = [] for nd in self._tree.postorder_internal_node_iter(): if nd.parent is None: nd._LCE_prob_mat_stored = None nd._LCE_prob_mat_scratch = None nd._LCE_edge_len_stored = None nd._LCE_edge_len_scratch = None else: nd._LCE_prob_mat_stored = {} nd._LCE_prob_mat_scratch = {} for mod in self.model_list: nd._LCE_prob_mat_stored[mod] = [LCE._prob_mat_cache.get_writable_object() for i in range(mod.num_rate_categories)] nd._LCE_prob_mat_scratch[mod] = [LCE._prob_mat_cache.get_writable_object() for i in range(mod.num_rate_categories)] nd._LCE_edge_len_stored = nd.edge_length nd._LCE_edge_len_scratch = nd.edge_length nd._LCE_edge_len_curr = None nd._LCE_prob_mat_curr = {} nd._LCE_partial_curr = {} nd._LCE_is_internal = True nd._LCE_partial_stored = {} nd._LCE_partial_scratch = {} for mod in self.model_list: nd._LCE_partial_stored[mod] = [LCE._partial_cache.get_writable_object() for i in range(mod.num_rate_categories)] nd._LCE_partial_scratch[mod] = [LCE._partial_cache.get_writable_object() for i in range(mod.num_rate_categories)] for c in nd.children: if not bool(c.children): tips.append(c) leaf_inds = set() for nd in tips: if nd.leaf_index in leaf_inds: raise ValueError("The leaf_index %s repeated!" % str(nd.leaf_index)) leaf_inds.add(nd.leaf_index) nd._LCE_prob_mat_stored = {} nd._LCE_prob_mat_scratch = {} for mod in self.model_list: nd._LCE_prob_mat_stored[mod] = [LCE._prob_mat_cache.get_writable_object() for i in range(mod.num_rate_categories)] nd._LCE_prob_mat_scratch[mod] = [LCE._prob_mat_cache.get_writable_object() for i in range(mod.num_rate_categories)] nd._LCE_prob_mat_curr = {} nd._LCE_edge_len_stored = nd.edge_length nd._LCE_edge_len_scratch = nd.edge_length nd._LCE_edge_len_curr = None nd._LCE_is_internal = False nd._LCE_buffer_index = nd.leaf_index nd._LCE_state_codes = LCE._wrap_state_code_array[nd.leaf_index] if not nd._LCE_state_codes.is_calculated: raise ValueError("State codes (data for the leaf of a tree) has not been specified for leaf_index=%d" % nd.leaf_index) tips.sort(cmp=lambda x, y: cmp(x.leaf_index, y.leaf_index)) max_leaf_ind = tips[-1].leaf_index self._tips = [None] *(1+max_leaf_ind) for nd in tips: self._tips[nd.leaf_index] = nd self._scheduler = None self.get_ln_L() self.accept() def accept(self): for nd in self._tree.postorder_internal_node_iter(): if nd._LCE_edge_len_curr is nd._LCE_edge_len_scratch: nd._LCE_edge_len_stored, nd._LCE_edge_len_scratch = nd._LCE_edge_len_scratch, nd._LCE_edge_len_stored for mod in self.model_list: if nd.parent is not None: if nd._LCE_prob_mat_curr[mod] is nd._LCE_prob_mat_scratch[mod]: pel = nd._LCE_prob_mat_stored[mod] nd._LCE_prob_mat_stored[mod], nd._LCE_prob_mat_scratch[mod] = nd._LCE_prob_mat_scratch[mod], nd._LCE_prob_mat_stored[mod] pel[0] = None if nd._LCE_partial_curr[mod] is nd._LCE_partial_scratch[mod]: sel = nd._LCE_partial_stored[mod] nd._LCE_partial_stored[mod], nd._LCE_partial_scratch[mod] = nd._LCE_partial_scratch[mod], sel sel[0] = None def revert(self): for nd in self._tree.postorder_internal_node_iter(): nd._LCE_edge_len_curr = nd._LCE_edge_len_stored nd._LCE_prob_mat_curr = dict(nd._LCE_prob_mat_stored) nd._LCE_partial_curr = dict(nd._LCE_partial_stored) def _calc_full_traversal_lnL_for_model(self, model): _LOG.debug("_calc_full_traversal_lnL_for_model called for %s" % str(model)) assert(self._scheduler is None) from pytbeaglehon.op_scheduling import TogglePartialScheduler es = model.eigen_soln self._scheduler = TogglePartialScheduler(self, model) try: for node in self._tree.postorder_internal_node_iter(): self._scheduler.add_internal_node_to_partial_calc(node) finally: self._scheduler.end_partial_calculations() prev_sched = self._scheduler self._scheduler = None root_partials = prev_sched._LCE_last_queued_dest model.transmit_category_weights() model.transmit_state_freq() assert (es is model.eigen_soln) return self._LCE.integrate_likelihood(model, root_partials)