def test_simplify_shower():
# this should do nothing to an empty shower
empty = FormShower.Shower([], [], [], [])
empty.simplify_shower()
assert len(empty) == 0
# now a simple shower where A -> B, C
particle_idxs = [1, 2, 3]
parents = [[], [1], [1]]
children = [[2, 3], [], []]
labels = ['A', 'B', 'C']
amalgam = False
simple = FormShower.Shower(particle_idxs, parents, children, labels,
amalgam)
simple.simplify_shower()
assert len(simple) == 3
# now a shower where there is a link that could be removed
particle_idxs.append(4)
parents.append([2])
children[1].append(4)
children.append([])
labels.append('D')
simple = FormShower.Shower(particle_idxs, parents, children, labels,
amalgam)
simple.simplify_shower()
assert len(simple) == 3
def test_Shower():
# try making an empty shower
empty = FormShower.Shower([], [], [], [])
assert len(empty) == 0
assert empty.n_particles == 0
assert len(empty.root_idxs) == 0
empty.find_ranks()
assert len(empty.ranks) == 0
empty_graph = empty.graph()
assert isinstance(empty_graph, DrawTrees.DotGraph)
with pytest.raises(AttributeError):
empty.outside_connections
with pytest.raises(AttributeError):
empty.roots
assert len(empty.outside_connection_idxs) == 0
assert len(empty.ends) == 0
assert len(empty.flavour) == 0
# now a simple shower where A -> B, C
particle_idxs = [1, 2, 3]
parents = [[], [1], [1]]
children = [[2, 3], [], []]
labels = ['A', 'B', 'C']
amalgam = True
simple = FormShower.Shower(particle_idxs, parents, children, labels,
amalgam)
assert len(simple) == 3
assert simple.n_particles == 3
tst.assert_allclose(simple.root_idxs, [1])
simple.find_ranks()
tst.assert_allclose(simple.ranks, [0, 1, 1])
simple_graph = simple.graph()
assert isinstance(simple_graph, DrawTrees.DotGraph)
assert len(simple.outside_connection_idxs) == 0
tst.assert_allclose(simple.ends, [2, 3])
particle_idxs = [1, 2, 3, 4]
parents = [[], [1], [1], [5]]
children = [[2, 3], [6], [], []]
labels = ['A', 'B', 'C', 'D']
amalgam = True
double_root = FormShower.Shower(particle_idxs, parents, children, labels,
amalgam)
tst.assert_allclose(sorted(double_root.root_idxs), [1, 4])
tst.assert_allclose(sorted(double_root.outside_connection_idxs), [4])
with pytest.raises(AssertionError):
amalgam = False
double_root = FormShower.Shower(particle_idxs, parents, children,
labels, amalgam)
particle_idxs = [2, 6, 7]
parents = [[1], [2], [6]]
children = [[6], [7], []]
labels = ['B', 'E', 'F']
amalgam = False
addition = FormShower.Shower(particle_idxs, parents, children, labels,
amalgam)
double_root.amalgamate(addition)
assert len(double_root) == 6
tst.assert_allclose(sorted(double_root.root_idxs), [1, 4])
tst.assert_allclose(sorted(double_root.outside_connection_idxs), [4])
def test_from_shower():
# will need an eventwise with Parents, Children, MCPID
# layer -1 0 1 1 -1 2 2 3 3 3 -1
# idx 0 1 2 3 4 5 6 7 8 9 10
children = [[], [0, 2, 3], [5], [6, 5, 4], [], [], [7, 8, 9], [], [], []]
parents = [[1], [], [1], [1], [3], [2, 3], [3], [6], [6], [6]]
mcpid = [4, 5, 5, 3, 2, 1, -5, -1, 7, 11]
n_nodes = len(children)
shower = FormShower.Shower(list(range(n_nodes)), parents, children, mcpid)
dot = DrawTrees.DotGraph(shower)
with TempTestDir("tst") as dir_name:
path = os.path.join(dir_name, "graph.dot")
with open(path, 'w') as graph_file:
graph_file.write(str(dot))
graph = read_dot(path)
assert len(graph.nodes) == n_nodes
assert len(graph.edges) == 10