def simulate_tree(self, ) -> CassiopeiaTree:
"""Simulates a complete binary tree.
Returns:
A CassiopeiaTree with the tree topology initialized with the
simulated tree
"""
def node_name_generator() -> Generator[str, None, None]:
"""Generates unique node names for the tree."""
i = 0
while True:
yield str(i)
i += 1
names = node_name_generator()
tree = nx.balanced_tree(2, self.depth, create_using=nx.DiGraph)
mapping = {"root": next(names)}
mapping.update({node: next(names) for node in tree.nodes})
# Add root, which indicates the initiating cell
tree.add_edge("root", 0)
nx.relabel_nodes(tree, mapping, copy=False)
cassiopeia_tree = CassiopeiaTree(tree=tree)
# Initialize branch lengths
time_dict = {
node: cassiopeia_tree.get_time(node) / (self.depth + 1)
for node in cassiopeia_tree.nodes
}
cassiopeia_tree.set_times(time_dict)
return cassiopeia_tree
def extract_tree_statistics(
tree: CassiopeiaTree, ) -> Tuple[List[float], int, bool]:
"""A helper function for testing simulated trees.
Outputs the total lived time for each extant lineage, the number of extant
lineages, and whether the tree has the expected node degrees (to ensure
unifurcations were collapsed).
Args:
tree: The tree to test
Returns:
The total time lived for each leaf, the number of leaves, and if the
degrees only have degree 0 or 2
"""
times = []
out_degrees = []
for i in tree.nodes:
if tree.is_leaf(i):
times.append(tree.get_time(i))
out_degrees.append(len(tree.children(i)))
out_degrees.pop(0)
correct_degrees = all(x == 2 or x == 0 for x in out_degrees)
return times, len(times), correct_degrees