def test_get_leaves(self): # # 1D # node = Tree(2) forest = Forest([node]) leaves = forest.get_leaves() # Only a ROOT node, it should be the only LEAF self.assertEqual(leaves, [node], 'Cell should be its own leaf.') # # Split cell and L child - find leaves # node.split() l_child = node.get_child(0) l_child.split() leaves = forest.get_leaves() self.assertEqual(len(leaves), 3, 'Cell should have 3 leaves.') # # Depth first order # addresses_depth_first = [[0, 0, 0], [0, 0, 1], [0, 1]] leaves = forest.get_leaves(mode='depth-first') for i in range(len(leaves)): leaf = leaves[i] self.assertEqual(leaf.get_node_address(), addresses_depth_first[i], 'Incorrect order, depth first search.') # # Breadth first order # addresses_breadth_first = [[0, 1], [0, 0, 0], [0, 0, 1]] leaves = node.get_leaves(mode='breadth-first') for i in range(len(leaves)): leaf = leaves[i] self.assertEqual(leaf.get_node_address(), addresses_breadth_first[i], 'Incorrect order, breadth first search.') node.get_child(0).get_child(0).mark('1') node.get_child(1).mark('1') node.make_rooted_subtree('1') leaves = node.get_leaves(subtree_flag='1') self.assertEqual(len(leaves),2, \ 'There should only be 2 flagged leaves') # # 2D # node = Tree(4) forest = Forest([node]) # # Split cell and SW child - find leaves # node.split() sw_child = node.get_child(0) sw_child.split() leaves = node.get_leaves() self.assertEqual(len(leaves), 7, 'Node should have 7 leaves.') # # Nested traversal # leaves = node.get_leaves() self.assertEqual(leaves[0].get_node_address(),[0,1], \ 'The first leaf in the nested enumeration should have address [1]') leaves = node.get_leaves(mode='depth-first') self.assertEqual(leaves[0].get_node_address(), [0,0,0], \ 'First leaf in un-nested enumeration should be [0,0].') # # Merge SW child - find leaves # sw_child.delete_children() leaves = node.get_leaves() self.assertEqual(len(leaves), 4, 'Node should have 4 leaves.') # # Marked Leaves # node = Tree(4) node.mark(1) forest = Forest([node]) self.assertTrue(node in forest.get_leaves(flag=1), \ 'Node should be a marked leaf node.') self.assertTrue(node in forest.get_leaves(), \ 'Node should be a marked leaf node.') node.split() sw_child = node.get_child(0) sw_child.split() sw_child.mark(1) self.assertEqual(node.get_leaves(subtree_flag=1), \ [sw_child], 'SW child should be only marked leaf') sw_child.remove() self.assertEqual(forest.get_leaves(subforest_flag=1), \ [node], 'node should be only marked leaf') # # Nested traversal # node = Tree(4) node.split() forest = Forest([node]) for child in node.get_children(): child.split() node.get_child(1).mark(1, recursive=True) node.get_child(3).mark(1) forest.root_subtrees(1) leaves = forest.get_leaves(subforest_flag=1) self.assertEqual(len(leaves), 7, 'This tree has 7 flagged LEAF nodes.') self.assertEqual(leaves[0], node.get_child(0), 'The first leaf should be the NE child.') self.assertEqual(leaves[3], node.get_child(1).get_child(0), '4th flagged leaf should be SE-NW grandchild.')
from mesh import Tree import matplotlib.pyplot as plt import numpy as np tree = Tree(n_children=2, regular=True) tree.split() for c in tree.get_children(): c.split() for cc in c.get_children(): cc.split() print(tree.get_depth()) leaves = tree.get_leaves() leaves.pop() w = 0 max_depth = 0 for leaf in leaves: ld = leaf.get_depth() w += 2**(-leaf.get_depth()) if ld > max_depth: max_depth = ld h = 2**(-max_depth) fig, ax = plt.subplots() for node in tree.traverse(mode='breadth-first'):