def setUp(self): """Setup Tests""" eugene.Config.VAR = {'x' : np.arange(0, 5)} eugene.Config.truth = eugene.Config.VAR['x'] + (10 * abs(-4)) self.tree = Tree( Node('n_add', Node('x'), Node('n_mul', Node(10), Node('n_abs', Node(-4) ) ) ) )
class TreeTests(unittest.TestCase): """ Test the Tree Module. """ def setUp(self): """Setup Tests""" eugene.Config.VAR = {'x' : np.arange(0, 5)} eugene.Config.truth = eugene.Config.VAR['x'] + (10 * abs(-4)) self.tree = Tree( Node('n_add', Node('x'), Node('n_mul', Node(10), Node('n_abs', Node(-4) ) ) ) ) def test_13_tree_string(self): """Check tree string""" correct_string = '[0:0] n_add (3) - {6|15}\n |-[1:1] x (0) - {1|1}\n \\-[1:2] n_mul (2) - {4|8}\n |-[2:3] 10 (0) - {1|1}\n \\-[2:4] n_abs (1) - {2|3}\n \\-[3:5] -4 (0) - {1|1}' tree_string = self.tree.display() self.assertEqual(repr(self.tree), str(self.tree)) self.assertEqual(tree_string, correct_string) def test_14_node_list(self): """Check tree node list""" correct_nodes = "['n_add', 'x', 'n_mul', 10, 'n_abs', -4]" tree_nodes = str(self.tree.list_nodes()) self.assertEqual(str(tree_nodes), correct_nodes) def test_15_edge_list(self): """Check tree edge list""" correct_edges = "[('n_add', 'x'), ('n_add', 'n_mul'), ('n_mul', 10), ('n_mul', 'n_abs'), ('n_abs', -4)]" tree_edges = str(self.tree.list_edges()) self.assertEqual(str(tree_edges), correct_edges) def test_16_tree_attributes(self): """Check tree attributes""" correct_attributes = { # Summary of children 'height' : 4, 'node_num' : 6, 'leaf_num' : 3, 'edge_num' : 5, # Sum of subtrees 'complexity' : 15 } attributes = { # Summary of children 'height' : self.tree.height, 'node_num' : self.tree.node_num, 'leaf_num' : self.tree.leaf_num, 'edge_num' : self.tree.edge_num, # Sum of subtrees 'complexity' : self.tree.complexity } self.assertEqual(attributes, correct_attributes) def test_17_evaluate(self): """Evaluate Tree expression""" result = self.tree.evaluate() self.assertEqual(result.tolist(), eugene.Config.truth.tolist()) def test_18_evaluate(self): """Evaluate Tree expression with nan in variable""" eugene.Config.VAR['x'] = np.NaN result = self.tree.evaluate() self.assertTrue(np.isnan(result)) def test_19_evaluate(self): """Evaluate Tree expression with no set variable""" eugene.Config.VAR = {} result = self.tree.evaluate() self.assertTrue(np.isnan(result)) def test_20_tree_get_node(self): """Check tree get node""" node = self.tree.get_node(3) self.assertEqual(node.value, 10) def test_21_tree_set_node(self): """Check tree set node""" self.tree.set_node(3, Node(123)) correct_string = '[0:0] n_add (3) - {6|15}\n |-[1:1] x (0) - {1|1}\n \\-[1:2] n_mul (2) - {4|8}\n |-[2:3] 123 (0) - {1|1}\n \\-[2:4] n_abs (1) - {2|3}\n \\-[3:5] -4 (0) - {1|1}' tree_string = self.tree.display() self.assertEqual(tree_string, correct_string) def test_22_tree_pruning(self): """Check tree pruning""" self.tree.prune()