def get_particle_label(particle, representation, label_opts, fancy=True):
"""Return string for particle label to be displayed on graph.
Parameters
----------
particle : Particle
Particle under consideration
representation : {"NODE", "EDGE"}
Particle representation
label_opts : dict
Dict of labels for different representations and fancy/plain
fancy : bool
If True, will use HTML/unicode in labels
Returns
-------
str
Particle label string
Raises
------
RuntimeError
If representation is not one of "NODE", "EDGE"
"""
check_representation_str(representation)
style_key = "fancy" if fancy else "plain"
label = label_opts[representation.lower()][style_key].format(**particle.__dict__)
if fancy:
label = label.replace("inf", "∞")
return label
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 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 __init__(self, description, parser, default_representation,
file_extension):
"""Basic class to hold info about a parser and associated fields.
Parameters
----------
description : str
Brief description about parser
parser : class
The Parser class
default_representation : {'NODE', 'EDGE'}
Default particle representation of the parser.
file_extension : str, optional
Optional file extension to associate with this parser. (no preceeding .)
"""
self.description = description
self.parser = parser
self.file_extension = file_extension
check_representation_str(default_representation,
"default_representation")
self.default_representation = default_representation