class Unsummarize(Transform):
"""
Undo a summarization.
Useful when used in combination with other transforms,
e.g. -v{gluballs,pluck,unsummarize}
"""
_name = "Unsummarize"
_args = [
Arg("pno", int, "id of particle to unsummarize", default=None),
Arg("vno", int, "id of vertex to unsummarize", default=None)
]
def __call__(self, graph_view):
if self.options["pno"] is None and self.options["vno"] is None:
retry = False
for obj in list(graph_view.particles) + list(graph_view.vertices):
if isinstance(obj, Summary):
obj.undo_summary()
else:
if not self.options["pno"] is None:
for obj in list(graph_view.particles):
if isinstance(obj, Summary):
if self.options["pno"] in obj.particle_numbers:
vs, ve = obj.start_vertex, obj.end_vertex
if isinstance(vs, Summary):
vs.undo_summary()
if isinstance(ve, Summary):
ve.undo_summary()
obj.undo_summary()
if not self.options["vno"] is None:
for obj in list(graph_view.vertices):
if isinstance(obj, Summary):
if self.options["vno"] in obj.vertex_number:
obj.undo_summary()
class LineWidthPt(Style):
"""
Make the particle line width dependent on the transverse momentum.
"""
_name = "LineWidthPt"
_args = [
Arg("scale", float, "scale of the line effects", default=1.0),
Arg("min", float, "minimal width of a line", default=0.1),
]
def __call__(self, layout):
if isinstance(layout, FeynmanLayout):
elements = layout.edges
else:
elements = layout.nodes
for edge in layout.edges:
particle = edge.going
if hasattr(particle, "pt"):
edge.style_args["stroke-width"] = self.options[
"min"] + self.options["scale"] * ln(particle.pt +
1) * 0.1
for element in elements:
particle = element.item
if hasattr(particle, "pt"):
element.style_args["stroke-width"] = self.options[
"min"] + +self.options["scale"] * ln(particle.pt + 1) * 0.1
class Jets(Transform):
_name = "Jets"
_args = [
Arg("algorithm",
str,
"Jet Algorithm",
default="antikt",
choices=_jet_algos),
Arg("r", float, "Delta R for the Jet Algorithm", default=0.4),
Arg("n_max", int, "Maximum number of jets to form", default=5),
Arg("tracks",
Arg.bool,
"Only cluster charged particles",
default=False)
]
def __call__(self, graph_view):
from mcviz.jet import cluster_jets, JetAlgorithms
from math import hypot
track_jets = self.options["tracks"]
if track_jets:
raise Exception("Unimplemented!")
def pselect(p):
return p.final_state and not p.invisible and (p.charge != 0 if
track_jets else True)
final_state_particles = [p for p in graph_view.particles if pselect(p)]
jets = cluster_jets(final_state_particles,
getattr(JetAlgorithms, self.options["algorithm"]))
print "Converted %i final state particles into %i jets" % (
len(final_state_particles), len(jets))
def pt(jet):
return hypot(*jet.p[:2])
for i, jet in enumerate(sorted(jets, key=pt, reverse=True)):
print "Created jet: np=%2i, %r, %r" % (len(
jet.particles), jet.p, jet.e)
if i >= self.options["n_max"]:
break
jet_start_vertices = set(p.start_vertex for p in jet.particles)
jet_end_vertices = set(p.end_vertex for p in jet.particles)
vs_summary = graph_view.summarize_vertices(jet_start_vertices)
ve_summary = graph_view.summarize_vertices(jet_end_vertices)
p_summary = graph_view.summarize_particles(jet.particles)
vs_summary.tag("jet")
ve_summary.tag("jet")
p_summary.tag("jet")
p_summary.tag("jet_{0:d}".format(i))
class DualDecongestedHad(DualLayout):
"""
UNDOCUMENTED
Takes the all-to-all connections at the hadronization step and replaces it
with a all-to-one-to-all vertex labelled "A miracle occurs"
"""
_name = "DualDecongestedHad"
_args = [
Arg("label",
str,
"label for the hadronization vertex",
default="A miracle occurs")
]
def get_vertex(self, vertex, node_style=None):
if vertex.hadronization:
items = []
had_node = LayoutNode(vertex, label=self.options["label"])
had_node.width = 5.0
had_node.height = 1.0
items.append(had_node)
for particle in vertex.incoming:
items.append(LayoutEdge(vertex, particle, vertex))
for particle in vertex.outgoing:
items.append(LayoutEdge(vertex, vertex, particle))
return items
else:
return super(DualDecongestedHad,
self).get_vertex(vertex, node_style)
class FancyLines(Style):
"""
Draw particle lines with arrows on them, and draw gluons with curls
"""
_name = "FancyLines"
_args = [
Arg("scale", float, "scale of the line effects", default=1.0),
]
def __call__(self, layout):
""" set fancy line types, curly gluons, wavy photons etc."""
for edge in layout.edges:
particle = edge.item
edge.style_args["scale"] = 0.2 * self.options["scale"]
if not hasattr(particle, "gluon"):
return
# colouring
if "jet" in particle.tags:
edge.style_line_type = "jet"
elif "cluster" in particle.tags:
edge.style_line_type = "hadron"
edge.style_args["scale"] = 0.2 * self.options["scale"]
edge.style_args["stroke-width"] = 0.2
elif "cut_summary" in particle.tags:
edge.style_line_type = "cut"
try:
edge.style_args["n"] = edge.n_represented
except AttributeError:
pass # Edge has no n_represented property
elif particle.gluon:
edge.style_args["scale"] = 0.2 * self.options["scale"]
edge.style_line_type = "gluon"
elif particle.photon:
if particle.final_state:
edge.style_line_type = "final_photon"
else:
edge.style_line_type = "photon"
elif particle.invisible:
edge.style_line_type = "invisible"
elif particle.squark:
edge.style_line_type = "sfermion"
elif particle.colored:
edge.style_line_type = "fermion"
elif particle.lepton:
edge.style_line_type = "fermion"
elif particle.boson:
edge.style_line_type = "boson"
elif particle.gluino:
edge.style_args["scale"] = 0.2 * self.options["scale"]
edge.style_line_type = "gluino"
elif particle.chargino:
edge.style_line_type = "chargino"
elif particle.slepton:
edge.style_line_type = "sfermion"
else:
edge.style_line_type = "hadron"
class GraphvizEngine(LayoutEngine):
_global_args = ["dump_dot"]
_args = [Arg("extra", str, "extra graphviz options", default="")]
_base = True
def graphviz_pass(self, engine, graphviz_options, dot_data):
log.debug("dot_data hash: 0x%0X", hash(dot_data))
# Dump the whole dot file before passing it to graphviz if requested
if self.options["dump_dot"]:
log.debug("Data passed to %s:" % engine)
# TODO: flush log FD
print dot_data
# Process the DOT data with graphviz
log.verbose("Calling '%s' with options %s" %
(engine, graphviz_options))
with timer("run graphviz", log.VERBOSE):
output, errors = run_graphviz(engine, dot_data, graphviz_options)
errors = map(str.strip, errors.split("\n"))
errors = filter(lambda e: e and not "Warning: gvrender" in e, errors)
if errors:
log.warning("********* GraphViz Output **********")
for error in errors:
log.warning(error)
log.warning("************************************")
if not output.strip():
log.error("No output from %s " % engine)
log.error("There may be too many constraints on the graph.")
if self.options["dump_dot"]:
log.debug("Data received from %s:" % engine)
print output
return output
def __call__(self, layout):
engine = self._name
opts = self.options["extra"].split()
if not any(opt.startswith("-T") for opt in opts):
opts.append("-Tplain")
plain = self.graphviz_pass(engine, opts, self.dot(layout))
layout.update_from_plain(plain)
def dot(self, layout):
out = ["digraph pythia {"]
out.append('dpi=1;')
if layout.width and layout.height:
out.append('size="%s,%s!";' % (layout.width, layout.height))
if layout.ratio:
out.append("ratio=%s;" % layout.ratio)
out.append(layout.dot)
out.append("}")
return "\n".join(out)
class Highlight(Style):
"""
Colour particles matching some criteria.
e.g. -sHighlight:6:color=blue highlights top quarks in blue,
-sHighlight:param=eta:0:2.5 highlights central particles,
-sHighlight:1000000:2000015 highlights susy particles
"""
_name = "Highlight"
_args = [
Arg("start", float, "start pdgid to highlight", default=6),
Arg("end", float, "end pdgid", default=0),
Arg("color", str, "highlight color", default="red"),
Arg("param", str, "parameter", default="pdgid"),
]
def __call__(self, layout):
""" highlight all particles in rage of interest """
start = self.options["start"]
if self.options["end"] == 0:
end = start
else:
end = self.options["end"]
color = self.options["color"]
param = self.options["param"]
for edge in layout.edges:
try:
if start <= abs(getattr(edge.item, param)) <= end:
edge.style_args["stroke"] = edge.style_args["fill"] = color
except AttributeError: # looks like the item doesn't have that param
pass
for node in layout.nodes:
try:
# label in Inline; particle in Dual
if start <= abs(getattr(node.item, param)) <= end:
node.style_args["fill"] = color
except AttributeError: # looks like the item doesn't have that param
pass
class DotEngine(GraphvizEngine):
_name = "dot"
_args = [
Arg("orientation",
str,
"orientation of the graph",
"TB",
choices=["LR", "RL", "TB", "BT"])
]
def dot(self, layout):
""" tuning parameters specific to dot:
* rankdir (G) - direction of layout LR, RL, TB, BT
* nodesep (G, 0.25) - minimum in-rank separation
* ranksep (G, 0.5) - minimum separation between ranks
* aspect (G, unset) - if set can narrow layouts
* compound (G, false) - allow edges between clusters
* ordering (G) - force ordering of in/out particles per vertex
* remincross (G) - reminimize cluster (??)
* mclimit (G, 1) - tweak number of iterations
* searchsize (G, 30) - optimization parameter, bigger better
* rank (S) - same, min, max, source or sink (force to that rank)
* group (N) - if same at start and end, try to connect straight
* constraint (E, true) - edge used for ranking
* minlen (E, 1) - minumum edge length in rank difference
* weight (E, 1) - heavier is shorter, 0 is minimum
"""
out = ["digraph pythia {"]
#out.append(layout.options["extra_dot"])
out.append('rankdir="%s";' % self.options["orientation"])
#out.append('ranksep=10;')
out.append('dpi=1;')
if layout.width and layout.height:
out.append('size="%s,%s!";' % (layout.width, layout.height))
if layout.ratio:
out.append("ratio=%s;" % layout.ratio)
out.append("edge [labelangle=90]")
out.append(layout.dot)
out.append("}")
return "\n".join(out)
class ThickenColor(Style):
"""
This can be used to make a single (anti)colour line very thick, so that it
can be easily seen.
"""
_name = "ThickenColor"
_args = [Arg("color_id", int, "id of the color to thicken")]
def __call__(self, layout):
color_id = self.options["color_id"]
for edge in layout.edges:
particle = edge.item
if color_id in (particle.color, particle.anticolor):
edge.style_args["stroke-width"] = 0.5
class LabelSizePt(Style):
"""
Make the label size dependent on the transverse momentum.
"""
_name = "LabelSizePt"
_args = [
Arg("scale", float, "scale of the line effects", default=1.0),
]
def __call__(self, layout):
for element in list(layout.edges) + list(layout.nodes):
particle = element.item
if not hasattr(particle, "pt"):
continue
element.label_size = self.options["scale"] * ln(particle.pt +
1) * 0.5 + 0.5
class StdPainter(Painter):
_global_args = ("output_file", )
_args = (Arg("output_file", str, "output filename", default="mcviz.svg"), )
_base = True
def write_data(self, data_string):
output_file = self.options["output_file"]
# Dump the data to stdout if required
log.debug("data hash: 0x%0X", hash(data_string))
if output_file == "-":
print data_string
elif hasattr(output_file, "write"):
output_file.write(data_string.encode("UTF-8"))
else:
# Write the data to file otherwise
log.info('writing "%s"' % output_file)
with open(output_file, "w") as f:
f.write(data_string.encode("UTF-8"))
def recalculate_boundingbox(self):
# if we do not know the scale; (i.e. bounding box) calculate it
x_min, y_min, x_max, y_max = 0, 0, 0, 0
def get_minmax(x, y):
return min(x, x_min), min(y, y_min), max(x, x_max), max(y, y_max)
for edge in self.layout.edges:
if edge.spline:
x0, x1, y0, y1 = edge.spline.boundingbox
x_min, y_min, x_max, y_max = get_minmax(x0, y0)
x_min, y_min, x_max, y_max = get_minmax(x1, y0)
x_min, y_min, x_max, y_max = get_minmax(x0, y1)
x_min, y_min, x_max, y_max = get_minmax(x1, y1)
for node in self.layout.nodes:
if node.center:
x_min, y_min, x_max, y_max = get_minmax(
node.center.x, node.center.y)
wx = x_max - x_min
wy = y_max - y_min
if not self.layout.width:
self.layout.width = 100
self.layout.height = 100
self.layout.scale = 1
else:
self.layout.scale = min(self.width / wx, self.height / wy)
class DemoOptionSet(CommandLineOptionSet):
_name = "demo"
_args = [Arg("layout", str, "choose a layout to demo: Feynman, Dual or InlineLabels",
choices=("Feynman", "Dual", "InlineLabels"), default="Dual")]
def __call__(self, tools):
for tool_type in sorted(tool_types.keys()):
tools[tool_type] = []
setdefault(tools, "painter", "navisvg")
setdefault(tools, "layout", self.options["layout"])
setdefault(tools, "layout", "FixIni")
tools["style"].append(ToolSetting("Default"))
tools["style"].append(ToolSetting("SimpleColors"))
tools["style"].append(ToolSetting("FancyLines"))
tools["transform"].append(ToolSetting("NoKinks"))
tools["transform"].append(ToolSetting("Gluballs"))
tools["transform"].append(ToolSetting("Chainmail"))
log.info("--demo is equivalent to '-pnavisvg:mcviz.svg -sSimpleColors -sFancyLines -tNoKinks -tGluballs -tChainmail -l%s -lFixIni'" % self.options["layout"])
return super(DemoOptionSet, self).__call__(tools)
class FixedInitialLayout(BaseLayout):
"""
Place all of the initial vertices on the same rank.
"""
_name = "FixIni"
_args = [
Arg("stretch",
float,
"pull initial vertices apart, 1 is max",
default=0.8)
]
def process(self):
sg_options = self.subgraph_options.setdefault("initial", [])
# rank=source means "put these on the first rank" to graphviz
sg_options.append('rank="source"')
super(FixedInitialLayout, self).process()
initial = self.subgraphs["initial"]
initial_pairs = len(initial) // 2
for i, p in enumerate(initial):
pair = i // 2
if self.width and self.height:
# Attempt to fix the initial particles on the left and right
# of the graph.
stretch = self.options["stretch"] * self.width / 2.0
xposition = stretch + (i % 2) * (self.width - 2 * stretch)
yposition = (1 + pair) * self.height / (initial_pairs + 1)
p.dot_args["pos"] = "%s,%s!" % (xposition, yposition)
def process_node(self, obj):
if isinstance(obj.item, ViewVertex):
if obj.item.initial:
obj.subgraph = "initial"
elif obj.item.start_vertex.initial:
if obj.dot_args.get("group", "") != "particlelabels":
obj.subgraph = "initial"
return super(FixedInitialLayout, self).process_node(obj)
class DualLayout(BaseLayout, FundamentalTool):
"""
The Dual layout, so named because it is the "Dual" in the graph sense of
Feynman diagrams, shows particles as nodes.
"""
_name = "Dual"
_args = [
Arg("helper_vertices",
Arg.bool,
"add helper vertices if there are many-to-many vertices",
default=True)
]
def get_subgraph(self, particle):
if particle.initial_state:
return "initial"
def get_particle(self, particle):
lo = LayoutNode(particle)
lo.subgraph = self.get_subgraph(particle)
lo.label = particle.pdgid
lo.label_size = self.options["label_size"]
if particle.initial_state:
# Big red initial vertices
lo.width = lo.height = 1.0
elif "cluster" in particle.tags:
lo.label = "cluster (%.4g %seV)" % self.graph.units.pick_mag(
particle.pt)
elif "cut_summary" in particle.tags:
pass #lo.label = None
elif "jet" in particle.tags:
jet_id = 0
for tag in particle.tags:
if tag != 'jet' and tag[:4] == 'jet_':
jet_id = int(tag[4:]) + 1
lo.label = "jet {0:d} ({1:.4g} {2:s}eV)".format(
jet_id, *self.graph.units.pick_mag(particle.pt))
return lo
def get_vertex(self, vertex, node_style=None):
items = []
if "cut_summary" in vertex.tags and len(vertex.incoming) > 1:
return None
need_help = False
if self.options["helper_vertices"]:
if len(vertex.incoming) > 1 and len(vertex.outgoing) > 1:
need_help = True
elif vertex.vacuum:
need_help = True
if need_help:
helper_node = LayoutNode(vertex)
helper_node.width = 0.5
helper_node.height = 0.5
items.append(helper_node)
for particle in vertex.incoming:
items.append(LayoutEdge(vertex, particle, vertex))
for particle in vertex.outgoing:
items.append(LayoutEdge(vertex, vertex, particle))
return items
for particle in vertex.outgoing:
for mother in particle.mothers:
items.append(LayoutEdge(vertex, mother, particle))
return items
class PhiLayout(FeynmanLayout):
_name = "Phi"
_args = [Arg("scale", float, "length scale", default=1.0)]
def process(self):
super(PhiLayout, self).process()
incoming_counter = 0
cx, cy = 0, 0
DR = 50
# create a tree that is used to generate the visuals
coming = {}
going = {}
for edge in self.edges:
coming.setdefault(edge.coming, []).append(edge)
going.setdefault(edge.going, []).append(edge)
nodes_remaining = list(self.nodes)
item_nodes = {}
for node in self.nodes:
item_nodes[node.item] = node
if not node.item in going:
xposition = 10 * DR * 1 if incoming_counter % 2 == 0 else -1
yposition = cy # - 10 * DR * 1 if incoming_counter % 2 == 0 else -1
incoming_counter += 1
node.center = Point2D(xposition, yposition)
node.dot_args["pos"] = "%s,%s!" % node.center.tuple()
node.dot_args["pin"] = "true"
nodes_remaining.remove(node)
while nodes_remaining:
for node in list(nodes_remaining):
pinned = [
edge for edge in going[node.item]
if item_nodes[edge.coming].center
]
if pinned:
nongluons = [p for p in pinned if not p.item.gluon]
if nongluons:
edge = nongluons[0]
force_scale = None
else:
edge = pinned[0]
force_scale = 0.1
phi = edge.item.phi
px, py, pz = edge.item.p
pt = edge.item.pt
e = edge.item.e
signum = 1 if sin(phi) * px + cos(phi) * py > 0 else -1
theta = atan2(pz, pt * signum)
#phi = 0
#phi = theta
#scale = ln(edge.item.e)
scale = e / 1000.0 + pt * 10
if scale < 0.1:
scale = 0.1
elif scale > 100:
scale = 100
if force_scale:
scale = force_scale
scale *= self.options["scale"]
delta = Point2D(DR * sin(phi), DR * cos(phi)) * scale
node.center = item_nodes[edge.coming].center + delta
node.dot_args["pos"] = "%s,%s!" % node.center.tuple()
node.dot_args["pin"] = "true"
nodes_remaining.remove(node)
break
class CircleLayout(FeynmanLayout):
_name = "Circle"
_args = [
Arg("scale", float, "length scale", default=1.0),
Arg("view",
str,
"view on the event",
choices=["front", "side", "eta_pt"],
default="side"),
Arg("phi", float, "rotate the view in phi [radian]", default=0.0),
Arg("pt",
float,
"minimum pT for a particle to be fixed [GeV]",
default=0.0),
]
def process(self):
super(CircleLayout, self).process()
# create a tree that is used to generate the visuals
coming = {}
going = {}
for edge in self.edges:
coming.setdefault(edge.coming, []).append(edge)
going.setdefault(edge.going, []).append(edge)
nodes_remaining = list(self.nodes)
item_nodes = {}
for node in self.nodes:
item_nodes[node.item] = node
rotate_phi = self.options["phi"]
min_fix_pt = self.options["pt"]
D = self.options["scale"]
def rotate_phi_proj(momentum, dims):
p0 = cos(rotate_phi) * momentum[0] + sin(rotate_phi) * momentum[1]
p1 = -sin(rotate_phi) * momentum[0] + cos(rotate_phi) * momentum[1]
p = p0, p1, momentum[2]
return Point2D(p[dims[0]], p[dims[1]])
if self.options["view"] == "front":
pfunc = lambda m: rotate_phi_proj(m, (0, 1))
mfunc = lambda m: rotate_phi_proj(m, (0, 1))
elif self.options["view"] == "side":
pfunc = lambda m: rotate_phi_proj(m, (2, 1))
mfunc = lambda m: rotate_phi_proj(m, (2, 1))
elif self.options["view"] == "eta_pt":
mfunc = lambda m: rotate_phi_proj(m, (2, 1))
def pfunc(momentum):
p0, p1 = rotate_phi_proj(momentum, (0, 1)).tuple()
pt = sqrt(p0**2 + p1**2)
y = pt * (1 if p0 > 0 else -1)
if all(momentum[i] == 0 for i in (0, 1, 2)):
return 0, 0
if pt == 0:
x = -momentum[2] / abs(momentum[2]) * 5
else:
x = -ln(tan(atan2(pt, momentum[2]) / 2.))
return Point2D(x, y)
for node in self.nodes:
recursive_fix = all(p.final_state and p.pt < min_fix_pt
for p in node.item.outgoing)
nofix = node.item.final and list(
node.item.incoming)[0].pt < min_fix_pt
if not nofix and (node.item.final or recursive_fix):
p = list(node.item.incoming)[0]
phi, pt, e = p.phi, p.pt, p.e
scale = pt
if scale < 1:
scale = 1
elif scale > 100:
scale = 100
scale *= self.options["scale"]
momentum = mfunc(p.p) * scale
sv = p.start_vertex
sv_p = [
sum(par.p[i] for par in sv.outgoing) for i in (0, 1, 2)
]
svn = item_nodes[sv]
dlvl = get_max_descendant_levels(sv) + 1 - (1 if recursive_fix
else 0)
if sv_p[0] == 0 and sv_p[1] == 0:
continue
stretch = 1.0
assert 0 < stretch <= 1
pos = pfunc(sv_p) * (D / sqrt(dlvl))
pos.y = pos.y / (abs(pos.y) / D)**(
stretch) * 10 # signum modified by stretch
svn.center = pos
svn.dot_args["pos"] = "%s,%s!" % svn.center.tuple()
svn.dot_args["pin"] = "true"
if svn in nodes_remaining:
nodes_remaining.remove(svn)
node.center = pos + momentum
node.dot_args["pos"] = "%s,%s!" % node.center.tuple()
node.dot_args["pin"] = "true"
nodes_remaining.remove(node)
elif node.item.initial:
p = list(node.item.outgoing)[0]
node.center = pfunc(p.p) * 1.2 * D
node.dot_args["pos"] = "%s,%s!" % node.center.tuple()
node.dot_args["pin"] = "true"
nodes_remaining.remove(node)
sv = p.end_vertex
sv_p = [
sum(par.p[i] for par in sv.incoming) for i in (0, 1, 2)
]
svn = item_nodes[sv]
pos = pfunc(sv_p) * D
svn.center = pos
svn.dot_args["pos"] = "%s,%s!" % svn.center.tuple()
svn.dot_args["pin"] = "true"
if svn in nodes_remaining:
nodes_remaining.remove(svn)
return
while nodes_remaining:
for node in list(nodes_remaining):
pinned = [
True for edge in coming[node.item]
if item_nodes[edge.going].center
]
if pinned:
nongluons = [p for p in pinned if not p.item.gluon]
if nongluons:
edge = nongluons[0]
force_scale = None
else:
edge = pinned[0]
force_scale = 0.1
phi = edge.item.phi
px, py, pz = edge.item.p
pt = edge.item.pt
e = edge.item.e
signum = 1 if sin(phi) * px + cos(phi) * py > 0 else -1
theta = atan2(pz, pt * signum)
#phi = 0
#phi = theta
#scale = ln(edge.item.e)
scale = e / 1000.0 + pt * 10
if scale < 0.1:
scale = 0.1
elif scale > 100:
scale = 100
if force_scale:
scale = force_scale
scale *= self.options["scale"]
delta = Point2D(DR * sin(phi), DR * cos(phi)) * scale
node.center = item_nodes[edge.coming].center + delta
node.dot_args["pos"] = "%s,%s!" % node.center.tuple()
node.dot_args["pin"] = "true"
nodes_remaining.remove(node)
break
class FeynmanLayout(BaseLayout, FundamentalTool):
"""
Produces something analagous to the familiar Feynman diagram. Beware though,
event records are not really like feynman diagrams.
"""
_name = "Feynman"
_args = [Arg("gluid", Arg.bool, "label gluons")]
dummy_number = -1000
def get_subgraph(self, vertex):
if vertex.initial:
return "initial"
elif vertex.connecting:
return "connecting"
def process(self):
# TODO: do something with the edge ordering. It is currently nonsensical
def ordering(edge):
order = (1 if edge.item.gluon else 0 if edge.item.color else
2 if edge.item.anticolor else None)
return edge.going.reference, order, edge.item.reference
self.edges.sort(key=ordering)
def get_vertex(self, vertex, node_style=None):
lo = LayoutNode(vertex, width=0.1, height=0.1)
lo.label = None
lo.label_size = self.options["label_size"]
lo.subgraph = self.get_subgraph(vertex)
if node_style:
lo.dot_args.update(node_style)
# Put clusters in the same graphviz "group".
if "after_cluster" in vertex.tags:
lo.dot_args["group"] = "cluster_%i" % vertex.cluster_index
elif (all("after_cluster" in p.tags for p in vertex.outgoing)
and vertex.outgoing):
cluster_particle = (p for p in vertex.outgoing
if "after_cluster" in p.tags).next()
lo.dot_args[
"group"] = "cluster_%i" % cluster_particle.cluster_index
if vertex.initial:
# Big red initial vertices
lo.width = lo.height = 1.0
elif vertex.final:
# Don't show final particle vertices
lo.show = False
elif "cut_summary" in vertex.tags:
return None
elif "summary" in vertex.tags:
lo.width = lo.height = 1.0
elif vertex.hadronization:
# Big white hardronization vertices
lo.width = lo.height = 1.0
else:
nr_particles = len(vertex.incoming) + len(vertex.outgoing)
lo.width = lo.height = nr_particles * 0.04
return lo
def get_particle(self, particle):
if "cut_summary" in particle.tags:
dummy = ViewVertex(self.graph)
dummy.order_number = self.dummy_number
self.dummy_number -= 1
lo = LayoutEdge(particle, particle.start_vertex, dummy)
else:
lo = LayoutEdge(particle, particle.start_vertex,
particle.end_vertex)
lo.label_size = self.options["label_size"]
if "cluster" in particle.tags:
lo.label = "cluster (%.4g %seV)" % self.graph.units.pick_mag(
particle.pt)
elif (particle.gluon or particle.photon) and not self.options["gluid"]:
lo.label = None
elif "cut_summary" in particle.tags:
lo.label = None
elif "jet" in particle.tags:
jet_id = 0
for tag in particle.tags:
if tag != 'jet' and tag[:4] == 'jet_':
jet_id = int(tag[4:]) + 1
if not jet_id: print(particle.tags)
lo.label = "jet {0:d} ({1:.4g} {2:s}eV)".format(
jet_id, *self.graph.units.pick_mag(particle.pt))
else:
lo.label = particle.pdgid
#lo.dot_args["weight"] = log10(particle.e+1)*0.1 + 1
return lo
class Cut(Transform):
"""
Cut away particles from the 'outside' of the graph (default). When
'final_state' is set to force, also cut on 'inner' particles.
Instead of taking the value of the parameter of a particle
directly, it can be checked against the value of its mother or
daugther when the corresponding boolean flags are true. E.g. to
draw all particles with at least one mother with pt > 5 GeV, do:
'-tCut:cut=5:param=pt:mothers=True'
'final_state' is automatically set to 'False' if 'mothers' or
'daugthers' is True.
"""
_name = "Cut"
_args = [
Arg("cut", float, "cut value", default=5),
Arg("param", str, "parameter to cut on", default="pt"),
Arg("abs", Arg.bool, "take abs value of param", default=True),
Arg("reverse", Arg.bool, "reverse cut", default=False),
Arg("exact", Arg.bool, "exact cut", default=False),
Arg("mothers", Arg.bool, "cut on mother parameters", default=False),
Arg("daughters", Arg.bool, "cut on daughter parameters",
default=False),
Arg("final_state",
Arg.bool,
"cut on final state particles only",
default=True),
]
def __call__(self, graph_view):
cut = self.options["cut"]
param = self.options["param"]
take_abs = self.options["abs"]
reverse = self.options["reverse"]
exact = self.options["exact"]
mothers = self.options["mothers"]
daughters = self.options["daughters"]
# when cutting on daugthers don't restrict on final state particles
final_state = self.options["final_state"] if not daughters else False
if final_state:
particles = [p for p in graph_view.particles if p.final_state]
else:
particles = graph_view.particles
passed_tag = "passed_cut"
def cutter(p):
reject = True
if hasattr(p, param):
value = getattr(p, param)
if take_abs: value = abs(value)
if exact:
reject = (value != cut)
else:
reject = (value <= cut)
return (reject if not reverse else not reject)
keep = set()
def mark(item, depth):
keep.add(item)
item.tag(passed_tag)
particles_to_walk = set()
for particle in particles:
if mothers:
loop_over = particle.mothers
elif daughters:
loop_over = particle.daughters
else:
loop_over = [particle]
for p in loop_over:
if not cutter(p):
particles_to_walk.add(p)
# also display the particle whose mother or daughter passed the cut
if mothers or daughters:
particles_to_walk.add(particle)
for p in particles_to_walk:
graph_view.walk(p,
vertex_action=mark,
particle_action=mark,
ascend=True)
def pruner(item, depth):
if passed_tag in item.tags:
return ()
for vertex in graph_view.vertices:
if passed_tag in vertex.tags:
cut_daughters = [
p for p in vertex.outgoing if passed_tag not in p.tags
]
#print vertex, cut_daughters
if len(cut_daughters) > 0:
vsummary = graph_view.summarize_vertices(
set(d.end_vertex for d in cut_daughters))
vsummary.tag("cut_summary")
psummary = graph_view.summarize_particles(
set(cut_daughters))
psummary.tag("cut_summary")
for p in graph_view.particles:
if p in keep:
continue
if p.start_vertex in keep: # and not p.start_vertex.hadronization:
# Don't discard
continue
graph_view.drop(p)
#Clean out 'dangling' vertices
for vertex in graph_view.vertices:
if not vertex.incoming and not vertex.outgoing:
graph_view.drop(vertex)
@Transform.decorate("NoLoops")
def contract_loops(graph_view):
"""
Drop loops from the graph
"""
for particle in graph_view.particles:
if particle.start_vertex == particle.end_vertex:
graph_view.drop(particle)
@Transform.decorate("Pluck",
args=[
Arg("start",
float,
"start of range to keep",
default=6),
Arg("end", float, "end of range", default=0),
Arg("param",
str,
"parameter of interest",
default="pdgid"),
Arg("keep_down",
int,
"max depth to descend from vertex",
default=8),
Arg("keep_up",
int,
"max depth to ascend from vertex",
default=20)
])
class BaseLayout(Layout):
"""
Class that encapsulates the layout and styling information of the graph
"""
_args = [Arg("x", int, "width"), Arg("y", int, "height"),
Arg("ratio", float, "aspect ratio"),
Arg("label_size", float, "size of labels (1.0 is normal)", 1.0),]
_base = True
def __call__(self, graph):
self.graph = graph
self.width, self.height = self.options["x"], self.options["y"]
self.ratio = self.options["ratio"]
self.scale = 1.0
self.label_size = self.options["label_size"]
self.units = graph.units
self.subgraphs = {None: []}
self.subgraph_options = {}
self.edges = []
# create node and edge objects from the graph
self.fill_objects(graph)
# do custom processing of the layout objects here
self.process()
return self
def fill_objects(self, graph):
for vertex in graph.vertices:
self.add_object(self.get_vertex(vertex))
for particle in graph.particles:
self.add_object(self.get_particle(particle))
def process_node(self, node):
return node
def process_edge(self, edge):
return edge
def add_object(self, obj):
if obj is None:
return
elif hasattr(obj, "__iter__"):
for o in obj:
self.add_object(o)
return
obj.item.layout_objects.append(obj)
if isinstance(obj, LayoutNode):
obj = self.process_node(obj)
if obj:
self.subgraphs.setdefault(obj.subgraph, []).append(obj)
elif isinstance(obj, LayoutEdge):
obj = self.process_edge(obj)
if obj:
self.edges.append(obj)
else:
raise NotImplementedError()
def process(self):
pass
@property
def nodes(self):
return chain(*self.subgraphs.values())
@property
def subgraph_names(self):
names = sorted(self.subgraphs.keys())
# move "None" (sorted always first) to the back
return names[1:] + [None]
@property
def dot(self):
out = []
for name in self.subgraph_names:
nodelist = self.subgraphs.get(name, [])
subgraph = self.subgraph_options.get(name, [])
subgraph.extend(node.dot for node in nodelist)
if name:
subgraph.insert(0, "subgraph %s {" % name)
subgraph.append("}")
out.extend(subgraph)
for edge in self.edges:
out.append(edge.dot)
return "\n".join(out)
def update_from_plain(self, plain):
data = PlainOutput(plain)
self.width, self.height = data.width, data.height
self.scale = data.scale
for edge in self.edges:
edge.spline = data.edge_lines.get(edge.reference, None)
edge.label_center = data.edge_label.get(edge.reference, None)
for node in self.nodes:
if node.item.reference in data.nodes:
node.center, size = data.nodes[node.item.reference]
node.width, node.height = size