slhaplot

Sat, 26 Feb 2011 23:21:12 +0100

author
Andy Buckley <andy@insectnation.org>
date
Sat, 26 Feb 2011 23:21:12 +0100
changeset 120
bffefe12df80
parent 116
e63729bbb044
child 121
402d6abf64b7
permissions
-rwxr-xr-x

Remove --show-gravitino: we'll always show it from now

andy@0 1 #! /usr/bin/env python
andy@10 2
andy@13 3 # from __future__ import with_statement
andy@13 4
andy@10 5 """\
andy@84 6 Usage: %prog [options] <spcfile>
andy@10 7
andy@84 8 Make a SUSY mass spectrum plot from an SLHA or ISAWIG spectrum file. If the
andy@84 9 filename ends with .isa, it will be assumed to be an ISAWIG file, otherwise
andy@84 10 it will be assumed to be an SLHA file (for which the normal extension is .spc).
andy@64 11
andy@82 12 Output is currently available as the make-plots .dat format or as LaTeX source
andy@82 13 using the PGF/TikZ graphics package. Both may be processed to make EPS or PDF
andy@82 14 images.
andy@13 15
andy@13 16 TODOs:
andy@101 17 * Allow user control over aspect ratio / geometry
andy@111 18 * PNG output (use PIL if available?)
andy@94 19 * Use proper distinction between physical, plot-logical, and plot output coords
andy@101 20 * Use scaling to allow the y coordinates to be in units of 100 GeV in TikZ output.
andy@111 21 * Distribute decay arrow start/end positions along mass lines rather than always to/from their centres?
andy@86 22 * Drop make-plots support?
andy@10 23 """
andy@10 24
andy@12 25 class XEdges(object):
andy@13 26 def __init__(self, left, offset=0.0, width=2.0):
andy@13 27 self.offset = offset
andy@13 28 self.left = left + offset
andy@13 29 self.width = width
andy@13 30 @property
andy@13 31 def right(self):
andy@13 32 return self.left + self.width
andy@12 33 @property
andy@12 34 def centre(self):
andy@12 35 return (self.left + self.right)/2.0
andy@10 36
andy@13 37
andy@13 38 class Label(object):
andy@13 39 def __init__(self, text, offset=None):
andy@13 40 self.text = text
andy@13 41 self.offset = None
andy@13 42 def __str__(self):
andy@13 43 return self.text
andy@13 44
andy@13 45
andy@101 46 ## Details classes for representing decided positions in a way independent of output format
andy@76 47
andy@76 48
andy@12 49 class ParticleDetails(object):
andy@13 50 def __init__(self, label, xnom, xoffset, color="black", labelpos="L", mass=None):
andy@10 51 self.label = label
andy@10 52 self.mass = mass
andy@13 53 self.xedges = XEdges(xnom, xoffset)
andy@11 54 self.color = color
andy@11 55 self.labelpos = labelpos
andy@10 56
andy@13 57
andy@76 58 class DecayDetails(object):
andy@86 59 def __init__(self, pidfrom, xyfrom, pidto, xyto, br, color="gray"): #, thickness=1px, label=None):
andy@76 60 self.pidfrom = pidfrom
andy@76 61 self.xyfrom = xyfrom
andy@76 62 self.pidto = pidto
andy@76 63 self.xyto = xyto
andy@86 64 self.br = br
andy@76 65 self.color = color
andy@76 66 #self.label = label
andy@76 67
andy@76 68
andy@76 69 class LabelDetails(object):
andy@76 70 def __init__(self, xy, texlabel, anchor="l", color="black"):
andy@76 71 self.xy = xy
andy@76 72 self.texlabel = texlabel
andy@76 73 ## Add non-TeX-based label rendering via this property, if needed
andy@76 74 self.textlabel = texlabel
andy@76 75 self.anchor = anchor
andy@76 76 self.color = color
andy@76 77
andy@76 78
andy@76 79
andy@10 80 XHIGGS = 0.0
andy@13 81 XSLEPTON = 5.0
andy@13 82 XGAUGINO = 10.0
andy@13 83 XSUSYQCD = 15.0
andy@10 84
andy@10 85 PDETAILS = {
andy@13 86 25 : ParticleDetails(Label(r"$h^0$"), XHIGGS, -0.2, color="blue"),
andy@13 87 35 : ParticleDetails(Label(r"$H^0$"), XHIGGS, -0.2, color="blue"),
andy@13 88 36 : ParticleDetails(Label(r"$A^0$"), XHIGGS, -0.2, color="blue"),
andy@13 89 37 : ParticleDetails(Label(r"$H^\pm$"), XHIGGS, 0.2, color="red"),
andy@13 90 1000011 : ParticleDetails(Label(r"$\tilde{\ell}_\text{L}$"), XSLEPTON, -0.2, color="blue"),
andy@13 91 2000011 : ParticleDetails(Label(r"$\tilde{\ell}_\text{R}$"), XSLEPTON, -0.2, color="blue"),
andy@13 92 1000015 : ParticleDetails(Label(r"$\tilde{\tau}_1$"), XSLEPTON, 0.2, color="red"),
andy@13 93 2000015 : ParticleDetails(Label(r"$\tilde{\tau}_2$"), XSLEPTON, 0.2, color="red"),
andy@13 94 1000012 : ParticleDetails(Label(r"$\tilde{\nu}_\text{L}$"), XSLEPTON, -0.2, color="blue"),
andy@13 95 1000016 : ParticleDetails(Label(r"$\tilde{\nu}_\tau$"), XSLEPTON, 0.2, color="red"),
andy@13 96 1000022 : ParticleDetails(Label(r"$\tilde{\chi}_1^0$"), XGAUGINO, -0.2, color="blue"),
andy@13 97 1000023 : ParticleDetails(Label(r"$\tilde{\chi}_2^0$"), XGAUGINO, -0.2, color="blue"),
andy@13 98 1000025 : ParticleDetails(Label(r"$\tilde{\chi}_3^0$"), XGAUGINO, -0.2, color="blue"),
andy@13 99 1000035 : ParticleDetails(Label(r"$\tilde{\chi}_4^0$"), XGAUGINO, -0.2, color="blue"),
andy@13 100 1000024 : ParticleDetails(Label(r"$\tilde{\chi}_1^\pm$"), XGAUGINO, 0.2, color="red"),
andy@13 101 1000037 : ParticleDetails(Label(r"$\tilde{\chi}_2^\pm$"), XGAUGINO, 0.2, color="red"),
andy@115 102 1000039 : ParticleDetails(Label(r"$\tilde{G}$"), XGAUGINO, 0.15, color="black!50!blue!30!green"),
andy@13 103 1000021 : ParticleDetails(Label(r"$\tilde{g}$"), XSUSYQCD, -0.3, color="black!50!blue!30!green"),
andy@13 104 1000001 : ParticleDetails(Label(r"$\tilde{q}_\text{L}$"), XSUSYQCD, -0.1, color="blue"),
andy@13 105 2000001 : ParticleDetails(Label(r"$\tilde{q}_\text{R}$"), XSUSYQCD, -0.1, color="blue"),
andy@13 106 1000005 : ParticleDetails(Label(r"$\tilde{b}_1$"), XSUSYQCD, 0.2, color="black!50!blue!30!green"),
andy@13 107 2000005 : ParticleDetails(Label(r"$\tilde{b}_2$"), XSUSYQCD, 0.2, color="black!50!blue!30!green"),
andy@13 108 1000006 : ParticleDetails(Label(r"$\tilde{t}_1$"), XSUSYQCD, 0.2, color="red"),
andy@13 109 2000006 : ParticleDetails(Label(r"$\tilde{t}_2$"), XSUSYQCD, 0.2, color="red")
andy@10 110 }
andy@10 111
andy@10 112
andy@73 113 import pyslha
andy@10 114 import sys, optparse
andy@73 115 parser = optparse.OptionParser(usage=__doc__, version=pyslha.__version__)
andy@13 116 parser.add_option("-o", "--out", metavar="FILE",
andy@64 117 help="write output to FILE",
andy@13 118 dest="OUTFILE", default=None)
andy@87 119 parser.add_option("-f", "--format", metavar="FORMAT",
andy@87 120 choices=["makeplots", "tikz", "tikzfrag", "tikzpdf"],
andy@82 121 help="format in which to write output. 'tikz' produces LaTeX source using the "
andy@82 122 "TikZ graphics package to render the plot, 'tikzfrag' produces the same but "
andy@82 123 "with the LaTeX preamble and document lines commented out to make it directly "
andy@87 124 "includeable as a code fragment in LaTeX document source, 'tikzpdf' produces a "
andy@87 125 "PDF file created by running pdflatex and pdfcrop on the 'tikz' output, and "
andy@101 126 "'makeplots' produces a .dat file for processing with the make-plots command. "
andy@101 127 "(default: %default)",
andy@87 128 dest="FORMAT", default="tikzpdf")
andy@86 129 parser.add_option("--preamble", metavar="FILE",
andy@86 130 help="specify a file to be inserted into LaTeX output as a special preamble",
andy@86 131 dest="PREAMBLE", default=None)
andy@62 132 parser.add_option("--minbr", "--br", metavar="BR",
andy@116 133 help="show decay lines for decays with a branching ratio of > BR, as either a "
andy@116 134 "fraction or percentage (default: show none)",
andy@116 135 dest="DECAYS_MINBR", default="1.1")
andy@111 136 parser.add_option("--decaystyle", choices=["const", "brwidth", "brcolor", "brwidth+brcolor"], metavar="STYLE",
andy@96 137 help="drawing style of decay arrows, from const/brwidth. The 'const' style draws "
andy@96 138 "all decay lines with the same width, 'brwidth' linearly scales the width of the "
andy@96 139 "decay arrow according to the decay branching ratio. Other modes such as BR-dependent "
andy@96 140 "colouring may be added later. (default: %default)",
andy@111 141 dest="DECAYS_STYLE", default="brwidth+brcolor")
andy@73 142 parser.add_option("--labels", choices=["none", "merge", "shift"], metavar="MODE",
andy@96 143 help="treatment of labels for particle IDs, from none/merge/shift. 'none' shows "
andy@96 144 "no labels at all, 'merge' combines would-be-overlapping labels into a single "
andy@96 145 "comma-separated list, and 'shift' vertically shifts the clashing labels to avoid "
andy@96 146 "collisions (default: %default)",
andy@62 147 dest="PARTICLES_LABELS", default="shift")
andy@120 148 # parser.add_option("--hide-gravitino", action="store_false",
andy@120 149 # help="hide the gravitino",
andy@120 150 # dest="SHOW_GRAVITINO", default=False)
andy@86 151
andy@86 152
andy@115 153 ## Run parser and sort out a few resulting details
andy@10 154 opts, args = parser.parse_args()
andy@62 155 opts.PARTICLES_LABELS_SHOW = (opts.PARTICLES_LABELS != "none")
andy@62 156 opts.PARTICLES_LABELS_MERGE = (opts.PARTICLES_LABELS == "merge")
andy@62 157 opts.PARTICLES_LABELS_SHIFT = (opts.PARTICLES_LABELS == "shift")
andy@116 158 #
andy@120 159 # if not opts.SHOW_GRAVITINO:
andy@120 160 # del PDETAILS[1000039]
andy@116 161 #
andy@116 162 if opts.DECAYS_MINBR.endswith("%"):
andy@116 163 opts.DECAYS_MINBR = float(opts.DECAYS_MINBR[:-1]) / 100
andy@116 164 else:
andy@116 165 opts.DECAYS_MINBR = float(opts.DECAYS_MINBR)
andy@115 166
andy@115 167 ## Check non-optional arguments
andy@10 168 if len(args) != 1:
andy@10 169 parser.print_help()
andy@10 170 sys.exit(1)
andy@76 171 opts.INFILE = args[0]
andy@10 172
andy@86 173
andy@13 174 ## Choose output file
andy@13 175 if opts.OUTFILE is None:
andy@13 176 import os
andy@76 177 o = os.path.basename(opts.INFILE)
andy@95 178 if o == "-":
andy@95 179 o = "out"
andy@13 180 if "." in o:
andy@13 181 o = o[:o.rindex(".")]
andy@79 182 ## Add format-specific suffix
andy@86 183 suffix = ".tex"
andy@87 184 if "pdf" in opts.FORMAT:
andy@87 185 suffix = ".pdf"
andy@87 186 elif opts.FORMAT == "makeplots":
andy@86 187 suffix = ".dat"
andy@86 188 opts.OUTFILE = o + suffix
andy@13 189
andy@76 190
andy@10 191 ## Read spectrum file
andy@84 192 BLOCKS, DECAYS = None, None
andy@95 193 if opts.INFILE == "-":
andy@95 194 intext = sys.stdin.read()
andy@95 195 BLOCKS, DECAYS = pyslha.readSLHA(intext)
andy@95 196 elif opts.INFILE.endswith(".isa"):
andy@84 197 BLOCKS, DECAYS = pyslha.readISAWIGFile(opts.INFILE)
andy@84 198 else:
andy@84 199 BLOCKS, DECAYS = pyslha.readSLHAFile(opts.INFILE)
andy@10 200
andy@95 201
andy@10 202 ## Set mass values in PDETAILS
andy@12 203 massblock = BLOCKS["MASS"]
andy@10 204 for pid in PDETAILS.keys():
andy@11 205 PDETAILS[pid].mass = abs(massblock.entries[pid])
andy@10 206
andy@10 207
andy@61 208 ## Decays
andy@76 209 DDETAILS = {}
andy@61 210 for pid, detail in sorted(PDETAILS.iteritems()):
andy@61 211 if DECAYS.has_key(pid):
andy@86 212 DDETAILS.setdefault(pid, {})
andy@61 213 xyfrom = (detail.xedges.centre, detail.mass)
andy@61 214 for d in DECAYS[pid].decays:
andy@61 215 if d.br > opts.DECAYS_MINBR:
andy@61 216 for pid2 in d.ids:
andy@61 217 if PDETAILS.has_key(pid2):
andy@61 218 xyto = (PDETAILS[pid2].xedges.centre, PDETAILS[pid2].mass)
andy@86 219 DDETAILS[pid][pid2] = DecayDetails(pid, xyfrom, pid2, xyto, d.br)
andy@120 220 if DDETAILS.has_key(pid) and not DDETAILS[pid]:
andy@86 221 del DDETAILS[pid]
andy@61 222
andy@61 223
andy@13 224 ## Labels
andy@76 225 PLABELS = []
andy@55 226 if opts.PARTICLES_LABELS_SHOW:
andy@60 227 class MultiLabel(object):
andy@55 228 def __init__(self, label=None, x=None, y=None, anchor=None):
andy@60 229 self.labels = [(label, x, y)] or []
andy@55 230 self.anchor = anchor or "l"
andy@60 231
andy@60 232 def __len__(self):
andy@60 233 return len(self.labels)
andy@60 234
andy@55 235 @property
andy@55 236 def joinedlabel(self):
andy@60 237 return r",\,".join(l[0] for l in self.labels)
andy@60 238
andy@55 239 @property
andy@55 240 def avgx(self):
andy@60 241 return sum(l[1] for l in self.labels)/float(len(self))
andy@60 242 @property
andy@60 243 def minx(self):
andy@60 244 return min(l[1] for l in self.labels)/float(len(self))
andy@60 245 @property
andy@60 246 def maxx(self):
andy@60 247 return max(l[1] for l in self.labels)/float(len(self))
andy@60 248
andy@55 249 @property
andy@55 250 def avgy(self):
andy@60 251 return sum(l[2] for l in self.labels)/float(len(self))
andy@60 252 @property
andy@60 253 def miny(self):
andy@60 254 return min(l[2] for l in self.labels)/float(len(self))
andy@60 255 @property
andy@60 256 def maxy(self):
andy@60 257 return max(l[2] for l in self.labels)/float(len(self))
andy@60 258
andy@55 259 def add(self, label, x, y):
andy@60 260 self.labels.append((label, x, y))
andy@60 261 self.labels = sorted(self.labels, key=lambda l : l[2])
andy@55 262 return self
andy@60 263 def get(self):
andy@60 264 for i in self.labels:
andy@60 265 yield i
andy@55 266
andy@55 267 def rel_err(a, b):
andy@55 268 return abs((a-b)/(a+b)/2.0)
andy@55 269
andy@61 270 ## Use max mass to work out the height of a text line in mass units
andy@61 271 maxmass = None
andy@61 272 for pid, pdetail in sorted(PDETAILS.iteritems()):
andy@61 273 maxmass = max(pdetail.mass, maxmass)
andy@61 274 text_height_in_mass_units = maxmass/22.0
andy@61 275 ##
andy@61 276 ## Merge colliding labels
andy@55 277 reallabels = []
andy@13 278 for pid, pdetail in sorted(PDETAILS.iteritems()):
andy@13 279 labelx = None
andy@13 280 offset = pdetail.label.offset or 0.2
andy@13 281 anchor = None
andy@13 282 if pdetail.xedges.offset <= 0:
andy@13 283 labelx = pdetail.xedges.left - offset
andy@13 284 anchor = "r"
andy@13 285 else:
andy@13 286 labelx = pdetail.xedges.right + offset
andy@13 287 anchor = "l"
andy@55 288 labely = pdetail.mass
andy@115 289 ## Avoid hitting the 0 mass line/border
andy@115 290 if labely < 0.6*text_height_in_mass_units:
andy@115 291 labely = 0.6*text_height_in_mass_units
andy@115 292
andy@55 293 text = pdetail.label.text
andy@55 294 ## Check for collisions
andy@55 295 collision = False
andy@62 296 if opts.PARTICLES_LABELS_SHIFT or opts.PARTICLES_LABELS_MERGE:
andy@62 297 for i, rl in enumerate(reallabels):
andy@62 298 if anchor == rl.anchor and abs(labelx - rl.avgx) < 0.5:
andy@62 299 import math
andy@75 300 if labely > rl.miny - text_height_in_mass_units and labely < rl.maxy + text_height_in_mass_units:
andy@62 301 reallabels[i] = rl.add(text, labelx, labely)
andy@62 302 collision = True
andy@62 303 break
andy@55 304 if not collision:
andy@60 305 reallabels.append(MultiLabel(text, labelx, labely, anchor))
andy@76 306 ## Calculate position shifts and fill PLABELS
andy@55 307 for rl in reallabels:
andy@62 308 if len(rl) == 1 or opts.PARTICLES_LABELS_MERGE:
andy@76 309 PLABELS.append(LabelDetails((rl.avgx, rl.avgy), rl.joinedlabel, anchor=rl.anchor))
andy@60 310 else:
andy@61 311 num_gaps = len(rl)-1
andy@61 312 yrange_old = rl.maxy - rl.miny
andy@61 313 yrange_nom = num_gaps * text_height_in_mass_units
andy@61 314 yrange = max(yrange_old, yrange_nom)
andy@61 315 ydiff = yrange - yrange_old
andy@60 316 for i, (t, x, y) in enumerate(rl.get()):
andy@61 317 ydiff_per_line = ydiff/num_gaps
andy@61 318 # TODO: Further improvement using relative or average positions?
andy@61 319 newy = y + (i - num_gaps/2.0) * ydiff_per_line
andy@76 320 PLABELS.append(LabelDetails((x, newy), t, anchor=rl.anchor))
andy@76 321
andy@76 322
andy@87 323 ## Function for writing out the generated source
andy@87 324 def writeout(out, outfile=opts.OUTFILE):
andy@87 325 f = sys.stdout
andy@87 326 if outfile != "-":
andy@87 327 f = open(outfile, "w")
andy@87 328 f.write(out)
andy@87 329 if f is not sys.stdout:
andy@87 330 f.close()
andy@76 331
andy@76 332 out = ""
andy@79 333 ## MAKE-PLOTS FORMAT
andy@79 334 if opts.FORMAT == "makeplots":
andy@76 335
andy@76 336 ## Write plot header
andy@76 337 out += "# SUSY mass/decay spectrum plot, created by pyslha/slhaplot from %s\n" % opts.INFILE
andy@76 338 out += "# http://pypi.python.org/pypi/pyslha\n"
andy@76 339 out += "\n"
andy@76 340 out += "# BEGIN PLOT\n"
andy@76 341 if opts.PARTICLES_LABELS_MERGE:
andy@76 342 ## Need more space if labels are to be merged horizontally
andy@76 343 out += "XMin=-4\n"
andy@76 344 out += "XMax=20\n"
andy@76 345 else:
andy@76 346 out += "XMin=-3\n"
andy@76 347 out += "XMax=19\n"
andy@76 348 # if opts.LOGSCALE:
andy@76 349 # out += "LogY=1\n"
andy@76 350 # else:
andy@76 351 # out += "LogY=0\n"
andy@76 352 out += "YMin=0\n"
andy@76 353 out += "#XCustomTicks=1.0 Higgs 6.0 Sleptons 11.0 Gauginos 16.0 Squarks\n"
andy@76 354 out += "XCustomTicks=-10.0 ~\n"
andy@76 355 out += "YLabel=Mass / GeV\n"
andy@76 356 out += "DrawSpecialFirst=1\n"
andy@76 357 out += "# END PLOT\n\n"
andy@76 358
andy@76 359
andy@76 360 ## Mass lines
andy@76 361 for pid, pdetail in sorted(PDETAILS.iteritems()):
andy@76 362 out += """
andy@76 363 # BEGIN HISTOGRAM %s
andy@76 364 ErrorBars=1
andy@76 365 LineWidth=1pt
andy@76 366 LineColor=%s
andy@76 367 %f %f %e 0
andy@76 368 # END HISTOGRAM\n""" % ("pid"+str(pid), pdetail.color,
andy@76 369 pdetail.xedges.left, pdetail.xedges.right,
andy@76 370 pdetail.mass)
andy@76 371 out += "\n"
andy@76 372
andy@76 373
andy@76 374 ## Decay arrows
andy@76 375 for pidfrom, todict in sorted(DDETAILS.iteritems()):
andy@76 376 for pidto, dd in sorted(todict.iteritems()):
andy@76 377 out += r"""
andy@76 378 # BEGIN SPECIAL decay_%d_%d
andy@76 379 \psset{arrowsize=0.1}
andy@76 380 \psline[linestyle=dashed,dash=3px 2px,linecolor=%s]{->}\physicscoor(%f,%f)\physicscoor(%f,%f)
andy@76 381 # END SPECIAL
andy@76 382 """ % (dd.pidfrom, dd.pidto, dd.color, dd.xyfrom[0], dd.xyfrom[1], dd.xyto[0], dd.xyto[1])
andy@76 383
andy@76 384
andy@76 385 ## Particle labels
andy@76 386 out += "\n\n"
andy@76 387 out += "# BEGIN SPECIAL labels\n"
andy@76 388 for ld in PLABELS:
andy@76 389 out += r"\rput[%s]\physicscoor(%f,%f){\small %s}" % (ld.anchor, ld.xy[0], ld.xy[1], ld.texlabel) + "\n"
andy@55 390 out += "# END SPECIAL\n"
andy@10 391
andy@87 392 ## Write it out
andy@87 393 writeout(out)
andy@79 394
andy@79 395
andy@79 396 ## TIKZ FORMAT
andy@82 397 if "tikz" in opts.FORMAT:
andy@82 398
andy@86 399 ## Comment out the preamble etc. if only the TikZ fragment is wanted
andy@82 400 c = ""
andy@82 401 if opts.FORMAT == "tikzfrag":
andy@82 402 c = "%"
andy@79 403
andy@79 404 ## Write LaTeX header
andy@86 405 out += "%% http://pypi.python.org/pypi/pyslha\n\n"
andy@82 406 out += c + "\\documentclass[11pt]{article}\n"
andy@82 407 out += c + "\\usepackage{amsmath,amssymb}\n"
andy@109 408 out += c + "\\usepackage[margin=0cm,paperwidth=15.2cm,paperheight=9.8cm]{geometry}\n"
andy@82 409 out += c + "\\usepackage{tikz}\n"
andy@82 410 out += c + "\\pagestyle{empty}\n"
andy@82 411 out += c + "\n"
andy@86 412 ## Insert user-specified preamble file
andy@86 413 if opts.PREAMBLE is not None:
andy@86 414 out += c + "%% User-supplied preamble\n"
andy@86 415 try:
andy@86 416 fpre = open(opts.PREAMBLE, "r")
andy@86 417 for line in fpre:
andy@86 418 out += c + line
andy@86 419 except:
andy@86 420 sys.stderr.write("Could not read preamble file %s -- fallback to using \\input\n" % opts.PREAMBLE)
andy@86 421 out += c + "\\input{%s}\n" % opts.PREAMBLE.replace(".tex", "")
andy@86 422 else:
andy@86 423 out += c + "%% Default preamble\n"
andy@86 424 out += c + "\\usepackage[osf]{mathpazo}\n"
andy@86 425 #
andy@86 426 out += c + "\n"
andy@82 427 out += c + "\\begin{document}\n"
andy@82 428 out += c + "\\thispagestyle{empty}\n\n"
andy@79 429
andy@80 430 ## Get coord space size: horizontal range is fixed by make-plots
andy@80 431 xmin = -3.0
andy@80 432 xmax = 19.0
andy@80 433 if opts.PARTICLES_LABELS_MERGE:
andy@80 434 ## Need more space if labels are to be merged horizontally
andy@80 435 xmin -= 1.0
andy@80 436 xmax += 1.0
andy@80 437 xdiff = xmax - xmin
andy@93 438 XWIDTH = 22.0
andy@93 439 def scalex(x):
andy@93 440 return x * XWIDTH/xdiff
andy@80 441
andy@93 442 ASPECTRATIO = 0.7 #0.618
andy@93 443 ydiff = ASPECTRATIO * XWIDTH
andy@80 444 ymin = 0.0
andy@80 445 ymax = ymin + ydiff
andy@79 446 ## Get range of masses needed
andy@79 447 maxmass = max(pd.mass for pid, pd in PDETAILS.iteritems())
andy@79 448 maxdisplaymass = maxmass * 1.1
andy@79 449 if maxdisplaymass % 100 != 0:
andy@79 450 maxdisplaymass = ((maxdisplaymass + 100) // 100) * 100
andy@80 451 yscale = (ymax-ymin)/maxdisplaymass
andy@79 452
andy@79 453 ## Write TikZ header
andy@87 454 out += "\\centering\n"
andy@80 455 out += "\\begin{tikzpicture}[scale=0.6]\n"
andy@79 456
andy@101 457 out += " %% y-scalefactor (GeV -> coords) = %e\n\n" % yscale
andy@101 458
andy@79 459 ## Draw the plot boundary and y-ticks
andy@80 460 out += " %% Frame\n"
andy@93 461 out += " \\draw (%f,%f) rectangle (%f,%f);\n" % (scalex(xmin), ymin, scalex(xmax), ymax)
andy@80 462 out += " %% y-ticks\n"
andy@80 463 for mtick in xrange(0, int(maxdisplaymass) + 100, 100):
andy@80 464 ytick = mtick * yscale
andy@93 465 out += " \\draw (%f,%f) node[left] {%d};\n" % (scalex(xmin), ytick, mtick)
andy@80 466 if mtick > 0 and mtick < maxdisplaymass:
andy@80 467 ## The 0.3 needs to be in the plot coords
andy@93 468 out += " \\draw (%f,%f) -- (%f,%f);\n" % (scalex(xmin+0.3), ytick, scalex(xmin), ytick)
andy@93 469 out += " \\draw (%f,%f) node[left,rotate=90] {Mass / GeV};\n" % (scalex(xmin-2.0), ymax)
andy@79 470
andy@81 471 ## Decay arrows
andy@86 472 if DDETAILS:
andy@86 473 out += "\n %% Decay arrows\n"
andy@86 474 for pidfrom, todict in sorted(DDETAILS.iteritems()):
andy@86 475 for pidto, dd in sorted(todict.iteritems()):
andy@96 476 out += " %% decay_%d_%d, BR=%0.1f%%\n" % (dd.pidfrom, dd.pidto, dd.br*100)
andy@111 477
andy@96 478 def scalethickness(br):
andy@96 479 if opts.DECAYS_STYLE == "const":
andy@96 480 return 0.8
andy@111 481 elif "brwidth" in opts.DECAYS_STYLE:
andy@96 482 return 1.0 * br
andy@96 483 else:
andy@111 484 raise Exception("Unexpected problem with unknown decay line style option: please contact the PySLHA authors!")
andy@111 485
andy@111 486 def scalecolor(br):
andy@111 487 if opts.DECAYS_STYLE == "const":
andy@111 488 return None
andy@111 489 elif "brcolor" in opts.DECAYS_STYLE:
andy@115 490 return "black!"+str(60*dd.br + 10)
andy@111 491 else:
andy@111 492 raise Exception("Unexpected problem with unknown decay line style option: please contact the PySLHA authors!")
andy@111 493
andy@96 494 out += " \\draw[-stealth,line width=%0.2fpt,dashed,color=%s] (%f,%f) -- (%f,%f);\n" % \
andy@111 495 (scalethickness(dd.br), scalecolor(dd.br) or dd.color,
andy@111 496 scalex(dd.xyfrom[0]), yscale*dd.xyfrom[1], scalex(dd.xyto[0]), yscale*dd.xyto[1])
andy@111 497
andy@79 498
andy@81 499 ## Draw mass lines
andy@86 500 if PDETAILS:
andy@86 501 out += "\n %% Particle lines\n"
andy@86 502 for pid, pdetail in sorted(PDETAILS.iteritems()):
andy@86 503 y = pdetail.mass*yscale
andy@86 504 out += " %% pid%s\n" % str(pid)
andy@86 505 out += " \\draw[color=%s,thick] (%f,%f) -- (%f,%f);\n" % \
andy@93 506 (pdetail.color, scalex(pdetail.xedges.left), y, scalex(pdetail.xedges.right), y)
andy@79 507
andy@82 508 ## Particle labels
andy@86 509 if PLABELS:
andy@86 510 out += "\n %% Particle labels\n"
andy@86 511 for ld in PLABELS:
andy@86 512 anchors_pstricks_tikz = { "r" : "left", "l" : "right" }
andy@86 513 out += " \\draw (%f,%f) node[%s] {\small %s};\n" % \
andy@93 514 (scalex(ld.xy[0]), yscale*ld.xy[1], anchors_pstricks_tikz[ld.anchor], ld.texlabel)
andy@79 515
andy@79 516 ## Write TikZ footer
andy@82 517 out += "\end{tikzpicture}\n\n"
andy@79 518
andy@79 519 ## Write LaTeX footer
andy@82 520 out += c + "\end{document}\n"
andy@79 521
andy@79 522
andy@87 523 ## Write output
andy@87 524 if opts.FORMAT != "tikzpdf":
andy@87 525 writeout(out)
andy@87 526 else:
andy@87 527 ## Run LaTeX and pdfcrop
andy@87 528 import tempfile, shutil, subprocess
andy@87 529 tmpdir = tempfile.mkdtemp()
andy@87 530 writeout(out, os.path.join(tmpdir, "mytmp.tex"))
andy@87 531 ok = True
andy@90 532 ## Test for pdflatex
andy@90 533 if ok:
andy@90 534 p = subprocess.Popen(["which", "pdflatex"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@90 535 rtn = p.wait()
andy@90 536 if rtn != 0:
andy@90 537 sys.stderr.write("pdflatex could not be found: tikzpdf format mode cannot work\n")
andy@90 538 ok = False
andy@90 539 ## Test for tikz package
andy@90 540 if ok:
andy@90 541 p = subprocess.Popen(["which", "kpsewhich"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@90 542 rtn = p.wait()
andy@90 543 if rtn != 0:
andy@90 544 sys.stderr.write("WARNING: kpsewhich could not be found: check for tikz package cannot be run\n")
andy@90 545 else:
andy@90 546 p = subprocess.Popen(["kpsewhich", "tikz.sty"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@90 547 rtn = p.wait()
andy@90 548 if rtn != 0:
andy@90 549 sys.stderr.write("tikz.sty could not be found: tikzpdf format mode cannot work\n")
andy@90 550 ok = False
andy@87 551 try:
andy@88 552 p = subprocess.Popen(["pdflatex", "\scrollmode\input", "mytmp.tex"],
andy@88 553 stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=tmpdir)
andy@87 554 p.wait()
andy@89 555 shutil.copyfile(os.path.join(tmpdir, "mytmp.pdf"), opts.OUTFILE)
andy@87 556 except Exception, e:
andy@87 557 sys.stderr.write("pdflatex could not be run: tikzpdf format mode cannot work\n")
andy@87 558 ok = False
andy@87 559 shutil.rmtree(tmpdir)
andy@87 560 if not ok:
andy@87 561 sys.exit(3)
andy@86 562
andy@79 563
andy@79 564 ## UNRECOGNISED FORMAT!
andy@76 565 else:
andy@76 566 print "Other formats not currently supported! How did we even get here?!"
andy@86 567 sys.exit(2)

mercurial