def annotate_tree_depths(tree: CassiopeiaTree) -> None: """Annotates tree depth at every node. Adds two attributes to the tree: how far away each node is from the root of the tree and how many triplets are rooted at that node. Modifies the tree in place. Args: tree: An ete3 Tree Returns: A dictionary mapping depth to the list of nodes at that depth. """ depth_to_nodes = defaultdict(list) for n in tree.depth_first_traverse_nodes(source=tree.root, postorder=False): if tree.is_root(n): tree.set_attribute(n, "depth", 0) else: tree.set_attribute(n, "depth", tree.get_attribute(tree.parent(n), "depth") + 1) depth_to_nodes[tree.get_attribute(n, "depth")].append(n) number_of_leaves = 0 correction = 0 for child in tree.children(n): number_of_leaves += len(tree.leaves_in_subtree(child)) correction += nCr(len(tree.leaves_in_subtree(child)), 3) tree.set_attribute(n, "number_of_triplets", nCr(number_of_leaves, 3) - correction) return depth_to_nodes
def compute_expansion_pvalues( tree: CassiopeiaTree, min_clade_size: int = 10, min_depth: int = 1, copy: bool = False, ) -> Union[CassiopeiaTree, None]: """Call expansion pvalues on a tree. Uses the methodology described in Yang, Jones et al, BioRxiv (2021) to assess the expansion probability of a given subclade of a phylogeny. Mathematical treatment of the coalescent probability is described in Griffiths and Tavare, Stochastic Models (1998). The probability computed corresponds to the probability that, under a simple neutral coalescent model, a given subclade contains the observed number of cells; in other words, a one-sided p-value. Often, if the probability is less than some threshold (e.g., 0.05), this might indicate that there exists some subclade under this node that to which this expansion probability can be attributed (i.e. the null hypothesis that the subclade is undergoing neutral drift can be rejected). This function will add an attribute "expansion_pvalue" to the tree, and return None unless :param:`copy` is set to True. On a typical balanced tree, this function will perform in O(n log n) time, but can be up to O(n^3) on highly unbalanced trees. A future endeavor may be to impelement the function in O(n) time. Args: tree: CassiopeiaTree min_clade_size: Minimum number of leaves in a subtree to be considered. min_depth: Minimum depth of clade to be considered. Depth is measured in number of nodes from the root, not branch lengths. copy: Return copy. Returns: If copy is set to False, returns the tree with attributes added in place. Else, returns a new CassiopeiaTree. """ tree = tree.copy() if copy else tree # instantiate attributes _depths = {} for node in tree.depth_first_traverse_nodes(postorder=False): tree.set_attribute(node, "expansion_pvalue", 1.0) if tree.is_root(node): _depths[node] = 0 else: _depths[node] = _depths[tree.parent(node)] + 1 for node in tree.depth_first_traverse_nodes(postorder=False): n = len(tree.leaves_in_subtree(node)) k = len(tree.children(node)) for c in tree.children(node): if len(tree.leaves_in_subtree(c)) < min_clade_size: continue depth = _depths[c] if depth < min_depth: continue b = len(tree.leaves_in_subtree(c)) # this value below is a simplification of the quantity: # sum[simple_coalescent_probability(n, b2, k) for \ # b2 in range(b, n - k + 2)] p = nCk(n - b, k - 1) / nCk(n - 1, k - 1) tree.set_attribute(c, "expansion_pvalue", p) return tree if copy else None
def overlay_data(self, tree: CassiopeiaTree): """Overlays Cas9-based lineage tracing data onto the CassiopeiaTree. Args: tree: Input CassiopeiaTree """ if self.random_seed is not None: np.random.seed(self.random_seed) # create state priors if they don't exist. # This will set the instance's variable for mutation priors and will # use this for all future simulations. if self.mutation_priors is None: self.mutation_priors = {} probabilites = [ self.state_generating_distribution() for _ in range(self.number_of_states) ] Z = np.sum(probabilites) for i in range(self.number_of_states): self.mutation_priors[i + 1] = probabilites[i] / Z number_of_characters = self.number_of_cassettes * self.size_of_cassette # initialize character states character_matrix = {} for node in tree.nodes: character_matrix[node] = [-1] * number_of_characters for node in tree.depth_first_traverse_nodes(tree.root, postorder=False): if tree.is_root(node): character_matrix[node] = [0] * number_of_characters continue parent = tree.parent(node) life_time = tree.get_time(node) - tree.get_time(parent) character_array = character_matrix[parent] open_sites = [ c for c in range(len(character_array)) if character_array[c] == 0 ] new_cuts = [] for site in open_sites: mutation_rate = self.mutation_rate_per_character[site] mutation_probability = 1 - (np.exp(-life_time * mutation_rate)) if np.random.uniform() < mutation_probability: new_cuts.append(site) # collapse cuts that are on the same cassette cuts_remaining = new_cuts if self.collapse_sites_on_cassette and self.size_of_cassette > 1: character_array, cuts_remaining = self.collapse_sites( character_array, new_cuts ) # introduce new states at cut sites character_array = self.introduce_states( character_array, cuts_remaining ) # silence cassettes silencing_probability = 1 - ( np.exp(-life_time * self.heritable_silencing_rate) ) character_array = self.silence_cassettes( character_array, silencing_probability, self.heritable_missing_data_state, ) character_matrix[node] = character_array # apply stochastic silencing for leaf in tree.leaves: character_matrix[leaf] = self.silence_cassettes( character_matrix[leaf], self.stochastic_silencing_rate, self.stochastic_missing_data_state, ) tree.set_all_character_states(character_matrix)