def test_topological_order_for_tree(self):
D = Leaf(scope=[0])
E = Leaf(scope=[0])
F = Leaf(scope=[0])
B = 0.5 * D + 0.5 * E
C = 0.5 * E + 0.5 * F
A = 0.5 * B + 0.5 * C
A.aname = "A"
B.aname = "B"
C.aname = "C"
D.aname = "D"
E.aname = "E"
F.aname = "F"
result = get_topological_order(A)
self.assertEqual(result[0], D)
self.assertEqual(result[1], E)
self.assertEqual(result[2], F)
self.assertEqual(result[3], B)
self.assertEqual(result[4], C)
self.assertEqual(result[5], A)
self.assertEqual(len(result), 6)
def test_topological_order_for_non_tree(self):
D = Leaf(scope=[0])
E = Leaf(scope=[0])
F = Leaf(scope=[0])
B = 0.5 * D + 0.5 * E
C = 0.5 * E + 0.5 * F
H = 0.5 * D + 0.5 * E
I = 0.5 * D + 0.5 * E
G = 0.5 * H + 0.5 * I
A = 0.5 * B + 0.5 * C
Z = 0.5 * A + 0.5 * G
Z.aname = "Z"
A.aname = "A"
B.aname = "B"
C.aname = "C"
D.aname = "D"
E.aname = "E"
F.aname = "F"
G.aname = "G"
H.aname = "H"
I.aname = "I"
result = get_topological_order(Z)
self.assertEqual(result[0], D)
self.assertEqual(result[1], E)
self.assertEqual(result[2], F)
self.assertEqual(result[3], B)
self.assertEqual(result[4], H)
self.assertEqual(result[5], I)
self.assertEqual(result[6], C)
self.assertEqual(result[7], G)
self.assertEqual(result[8], A)
self.assertEqual(result[9], Z)
self.assertEqual(len(result), 10)
layers = get_topological_order_layers(Z)
self.assertEqual(set(layers[0]), set([D, E, F]))
self.assertEqual(set(layers[1]), set([I, H, B, C]))
self.assertEqual(set(layers[2]), set([G, A]))
self.assertEqual(set(layers[3]), set([Z]))