def test_intial_final_state(self):
"""Test whether particles marked as initial/final state correctly
Uses the 2 particles to 1 to 3 setup
(0) - p1 -> (2) (2) - p3 -> (3) (3) - p4 -> (4)
(1) - p2 -> (2) (3) - p5 -> (5)
(3) - p6 -> (6)
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=11),
vtx_out_barcode=0,
vtx_in_barcode=2)
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=-11),
vtx_out_barcode=1,
vtx_in_barcode=2)
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=22),
vtx_out_barcode=2,
vtx_in_barcode=3)
p4 = EdgeParticle(particle=Particle(barcode=4, pdgid=11),
vtx_out_barcode=3,
vtx_in_barcode=4)
p5 = EdgeParticle(particle=Particle(barcode=5, pdgid=12),
vtx_out_barcode=3,
vtx_in_barcode=5)
p6 = EdgeParticle(particle=Particle(barcode=6, pdgid=13),
vtx_out_barcode=3,
vtx_in_barcode=6)
particles = [p1, p2, p3, p4, p5, p6]
g = eg.assign_particles_edges(particles)
self.assertTrue(p1.particle.initial_state)
self.assertTrue(p2.particle.initial_state)
self.assertTrue(p4.particle.final_state)
self.assertTrue(p5.particle.final_state)
self.assertTrue(p6.particle.final_state)
def test_2_to_1_to_3(self):
"""2 particles to 1 to 3
(0) - p1 -> (2) (2) - p3 -> (3) (3) - p4 -> (4)
(1) - p2 -> (2) (3) - p5 -> (5)
(3) - p6 -> (6)
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=11),
vtx_out_barcode=0,
vtx_in_barcode=2)
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=-11),
vtx_out_barcode=1,
vtx_in_barcode=2)
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=22),
vtx_out_barcode=2,
vtx_in_barcode=3)
p4 = EdgeParticle(particle=Particle(barcode=4, pdgid=11),
vtx_out_barcode=3,
vtx_in_barcode=4)
p5 = EdgeParticle(particle=Particle(barcode=5, pdgid=12),
vtx_out_barcode=3,
vtx_in_barcode=5)
p6 = EdgeParticle(particle=Particle(barcode=6, pdgid=13),
vtx_out_barcode=3,
vtx_in_barcode=6)
particles = [p1, p2, p3, p4, p5, p6]
g = eg.assign_particles_edges(particles)
self.check_graph_particles([p.particle for p in particles], g)
edges = [(0, 2), (1, 2), (2, 3), (3, 4), (3, 5), (3, 6)]
self.check_graph_edges(edges, g)
def test_1_to_2_edge_to_node(self):
"""Very simple scenario: 1 particle decays to 2 daughters
(0) - p1 -> (1) (1) - p2 -> (2)
(1) - p3 -> (3)
becomes
p1 ->- p2
|- p3
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=11),
vtx_out_barcode="0",
vtx_in_barcode="1")
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=12),
vtx_out_barcode="1",
vtx_in_barcode="2")
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=13),
vtx_out_barcode="1",
vtx_in_barcode="3")
particles = [p1, p2, p3]
g_edge = assign_particles_edges(particles)
g_node = edge_to_node(g_edge)
n1 = Particle(barcode=1, pdgid=11)
n2 = Particle(barcode=2, pdgid=12)
n3 = Particle(barcode=3, pdgid=13)
n_particles = [n1, n2, n3]
log.debug(vars(g_node))
ideal_edges = [(1, 3), (1, 2)]
self.check_graph_edges(ideal_edges, g_node)
self.check_graph_node_particles(n_particles, g_node)
def test_1_to_2_edge_to_node(self):
"""Very simple scenario: 1 particle decays to 2 daughters
(0) - p1 -> (1) (1) - p2 -> (2)
(1) - p3 -> (3)
becomes
p1 ->- p2
|- p3
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=11),
vtx_out_barcode="0", vtx_in_barcode="1")
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=12),
vtx_out_barcode="1", vtx_in_barcode="2")
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=13),
vtx_out_barcode="1", vtx_in_barcode="3")
particles = [p1, p2, p3]
g_edge = assign_particles_edges(particles)
g_node = edge_to_node(g_edge)
n1 = Particle(barcode=1, pdgid=11)
n2 = Particle(barcode=2, pdgid=12)
n3 = Particle(barcode=3, pdgid=13)
n_particles = [n1, n2, n3]
log.debug(vars(g_node))
ideal_edges = [(1, 3), (1, 2)]
self.check_graph_edges(ideal_edges, g_node)
self.check_graph_node_particles(n_particles, g_node)
def assign_particles_to_graph(particles, default_repr, desired_repr=None,
filter_pdgid=None, filter_pdgid_final=None,
remove_redundants=True):
"""Wrapper to easily assign particles to graph, change representation,
filter out certain PDGIDs, and remove redundants.
Parameters
----------
particles : list[NodeParticle], list[EdgeParticle]
List of particles to be assigned to a graph. Must include relationship information.
default_repr : {"NODE", "EDGE"}
Particle representation for particles
desired_repr : {"NODE", "EDGE"}, optional
Desired output representation.
filter_pdgid : list[int], optional
filter_pdgid_final : list[int], optional
These two args allow for removal of particles from the graph based on PDGID,
the latter only removing final-state particles.
Note that both particles and anti-particles will be removed.
remove_redundants : bool, optional
Whether to remove redundant particles from the graph.
Returns
-------
NetworkX.MultiDiGraph
"""
check_representation_str(default_repr, "default_repr")
if default_repr == "NODE":
graph = assign_particles_nodes(particles)
elif default_repr == "EDGE":
graph = assign_particles_edges(particles)
new_repr = default_repr
if desired_repr and desired_repr != default_repr:
check_representation_str(desired_repr, "desired_repr")
new_repr = desired_repr
if (default_repr, desired_repr) == ("NODE", "EDGE"):
graph = node_to_edge(graph)
elif (default_repr, desired_repr) == ("EDGE", "NODE"):
graph = edge_to_node(graph)
filterer = remove_nodes_by_pdgid if new_repr == "NODE" else remove_edges_by_pdgid
filter_pdgid = filter_pdgid or []
for pdgid in filter_pdgid:
filterer(graph, pdgid, False)
filter_pdgid_final = filter_pdgid_final or []
for pdgid in filter_pdgid_final:
filterer(graph, pdgid, True)
if remove_redundants:
if new_repr == "NODE":
remove_redundant_nodes(graph)
elif new_repr == "EDGE":
remove_redundant_edges(graph)
return graph
def assign_particles_to_graph(particles,
default_repr,
desired_repr=None,
remove_redundants=True):
"""Wrapper to easily assign particles to graph, change representation, remove redundants.
Parameters
----------
particles : list[NodeParticle], list[EdgeParticle]
List of particles to be assigned to a graph. Must include relationship information.
default_repr : {"NODE", "EDGE"}
Particle representation for particles
desired_repr : {"NODE", "EDGE"}, optional
Desired output representation.
remove_redundants : bool, optional
Whether to remove redundant particles from the graph.
Returns
-------
NetworkX.MultiDiGraph
"""
check_representation_str(default_repr, "default_repr")
if default_repr == "NODE":
graph = assign_particles_nodes(particles)
elif default_repr == "EDGE":
graph = assign_particles_edges(particles)
remove_edges_by_pdgid(graph, 22, True)
new_repr = default_repr
if desired_repr and desired_repr != default_repr:
check_representation_str(desired_repr, "desired_repr")
new_repr = desired_repr
if (default_repr, desired_repr) == ("NODE", "EDGE"):
graph = node_to_edge(graph)
elif (default_repr, desired_repr) == ("EDGE", "NODE"):
graph = edge_to_node(graph)
if remove_redundants:
if new_repr == "NODE":
remove_redundant_nodes(graph)
elif new_repr == "EDGE":
remove_redundant_edges(graph)
return graph
def test_redundant_simple(self):
"""Check remove_redundants code with a very simple case.
(i) is node i, pj is edge/particle j
(-1)--p1(u)--(-2)--p3(u)--(-3)--p4(u)--(-5)--p5(u)--(-6)--p6(u)--(-7)
| |
(-2)--p2(g)--(-4) (-6)--p7(g)--(-8)
should simplify to
(-1)--p1(u)--(-2)----p3(u)----(-6)----p6(u)----(-7)
| |
(-2)--p2(g)--(-4) |
(-6)--p7(g)--(-8)
since p4 and p5 have the same pdgid as p3 and are therefore redundant
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=2),
vtx_out_barcode=-1,
vtx_in_barcode=-2)
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=22),
vtx_out_barcode=-2,
vtx_in_barcode=-4)
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=2),
vtx_out_barcode=-2,
vtx_in_barcode=-3)
p4 = EdgeParticle(particle=Particle(barcode=4, pdgid=2),
vtx_out_barcode=-3,
vtx_in_barcode=-5)
p5 = EdgeParticle(particle=Particle(barcode=5, pdgid=2),
vtx_out_barcode=-5,
vtx_in_barcode=-6)
p6 = EdgeParticle(particle=Particle(barcode=6, pdgid=2),
vtx_out_barcode=-6,
vtx_in_barcode=-7)
p7 = EdgeParticle(particle=Particle(barcode=7, pdgid=22),
vtx_out_barcode=-6,
vtx_in_barcode=-8)
particles = [p1, p2, p3, p4, p5, p6, p7]
graph = eg.assign_particles_edges(particles)
eg.remove_redundant_edges(graph)
edges = [(-1, -2), (-2, -4), (-2, -3), (-3, -7), (-3, -8)]
self.check_graph_edges(edges, graph)
def test_2_to_1_to_3_edge_to_node(self):
"""2 particles to 1 to 3
(0) - p1 -> (2) (2) - p3 -> (3) (3) - p4 -> (4)
(1) - p2 -> (2) (3) - p5 -> (5)
(3) - p6 -> (6)
becomes
p1 ->- p3 ->- p4
p2 -| |- p5
|- p6
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=11),
vtx_out_barcode=0,
vtx_in_barcode=2)
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=12),
vtx_out_barcode=1,
vtx_in_barcode=2)
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=13),
vtx_out_barcode=2,
vtx_in_barcode=3)
p4 = EdgeParticle(particle=Particle(barcode=4, pdgid=14),
vtx_out_barcode=3,
vtx_in_barcode=4)
p5 = EdgeParticle(particle=Particle(barcode=5, pdgid=15),
vtx_out_barcode=3,
vtx_in_barcode=5)
p6 = EdgeParticle(particle=Particle(barcode=6, pdgid=16),
vtx_out_barcode=3,
vtx_in_barcode=6)
particles = [p1, p2, p3, p4, p5, p6]
g_edge = assign_particles_edges(particles)
g_node = edge_to_node(g_edge)
n1 = Particle(barcode=1, pdgid=11)
n2 = Particle(barcode=2, pdgid=12)
n3 = Particle(barcode=3, pdgid=13)
n4 = Particle(barcode=4, pdgid=14)
n5 = Particle(barcode=5, pdgid=15)
n6 = Particle(barcode=6, pdgid=16)
self.check_graph_node_particles([n1, n2, n3, n4, n5, n6], g_node)
ideal_edges = [(1, 3), (2, 3), (3, 4), (3, 5), (3, 6)]
self.check_graph_edges(ideal_edges, g_node)
def assign_particles_to_graph(particles, default_repr, desired_repr=None, remove_redundants=True):
"""Wrapper to easily assign particles to graph, change representation, remove redundants.
Parameters
----------
particles : list[NodeParticle], list[EdgeParticle]
List of particles to be assigned to a graph. Must include relationship information.
default_repr : {"NODE", "EDGE"}
Particle representation for particles
desired_repr : {"NODE", "EDGE"}, optional
Desired output representation.
remove_redundants : bool, optional
Whether to remove redundant particles from the graph.
Returns
-------
NetworkX.MultiDiGraph
"""
check_representation_str(default_repr, "default_repr")
if default_repr == "NODE":
graph = assign_particles_nodes(particles)
elif default_repr == "EDGE":
graph = assign_particles_edges(particles)
remove_edges_by_pdgid(graph, 22, True)
new_repr = default_repr
if desired_repr and desired_repr != default_repr:
check_representation_str(desired_repr, "desired_repr")
new_repr = desired_repr
if (default_repr, desired_repr) == ("NODE", "EDGE"):
graph = node_to_edge(graph)
elif (default_repr, desired_repr) == ("EDGE", "NODE"):
graph = edge_to_node(graph)
if remove_redundants:
if new_repr == "NODE":
remove_redundant_nodes(graph)
elif new_repr == "EDGE":
remove_redundant_edges(graph)
return graph
def test_edge_removal_simple(self):
"""Test whether edge removal works in simple scenario.
(0) - p1 -> (1) - p2 -> (2) - p3 -> (3)
becomes
(0) - p1 -> (1) - p3 -> (3)
"""
p1 = EdgeParticle(particle=Particle(barcode=1),
vtx_out_barcode=0,
vtx_in_barcode=1)
p2 = EdgeParticle(particle=Particle(barcode=2),
vtx_out_barcode=1,
vtx_in_barcode=2)
p3 = EdgeParticle(particle=Particle(barcode=3),
vtx_out_barcode=2,
vtx_in_barcode=3)
g = eg.assign_particles_edges([p1, p2, p3])
eg.remove_particle_edge(g, (1, 2))
self.check_graph_edges([(0, 1), (1, 3)], g)
self.check_graph_particles([ep.particle for ep in [p1, p3]], g)
def test_1_to_2(self):
"""Very simple scenario: 1 particle decays to 2 daughters
(0) - p1 -> (1) (1) - p2 -> (2)
(1) - p3 -> (3)
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=11),
vtx_out_barcode=0,
vtx_in_barcode=1)
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=11),
vtx_out_barcode=1,
vtx_in_barcode=2)
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=11),
vtx_out_barcode=1,
vtx_in_barcode=3)
particles = [p1, p2, p3]
g = eg.assign_particles_edges(particles)
self.check_graph_particles([p.particle for p in particles], g)
edges = [(0, 1), (1, 3), (1, 2)]
self.check_graph_edges(edges, g)
def test_2_to_1_to_3_edge_to_node(self):
"""2 particles to 1 to 3
(0) - p1 -> (2) (2) - p3 -> (3) (3) - p4 -> (4)
(1) - p2 -> (2) (3) - p5 -> (5)
(3) - p6 -> (6)
becomes
p1 ->- p3 ->- p4
p2 -| |- p5
|- p6
"""
p1 = EdgeParticle(particle=Particle(barcode=1, pdgid=11),
vtx_out_barcode=0, vtx_in_barcode=2)
p2 = EdgeParticle(particle=Particle(barcode=2, pdgid=12),
vtx_out_barcode=1, vtx_in_barcode=2)
p3 = EdgeParticle(particle=Particle(barcode=3, pdgid=13),
vtx_out_barcode=2, vtx_in_barcode=3)
p4 = EdgeParticle(particle=Particle(barcode=4, pdgid=14),
vtx_out_barcode=3, vtx_in_barcode=4)
p5 = EdgeParticle(particle=Particle(barcode=5, pdgid=15),
vtx_out_barcode=3, vtx_in_barcode=5)
p6 = EdgeParticle(particle=Particle(barcode=6, pdgid=16),
vtx_out_barcode=3, vtx_in_barcode=6)
particles = [p1, p2, p3, p4, p5, p6]
g_edge = assign_particles_edges(particles)
g_node = edge_to_node(g_edge)
n1 = Particle(barcode=1, pdgid=11)
n2 = Particle(barcode=2, pdgid=12)
n3 = Particle(barcode=3, pdgid=13)
n4 = Particle(barcode=4, pdgid=14)
n5 = Particle(barcode=5, pdgid=15)
n6 = Particle(barcode=6, pdgid=16)
self.check_graph_node_particles([n1, n2, n3, n4, n5, n6], g_node)
ideal_edges = [(1, 3), (2, 3), (3, 4), (3, 5), (3, 6)]
self.check_graph_edges(ideal_edges, g_node)