slhaplot

Mon, 07 Mar 2011 10:48:37 +0000

author
Andy Buckley <andy@insectnation.org>
date
Mon, 07 Mar 2011 10:48:37 +0000
changeset 134
323754f1d261
parent 133
5e27f4121fd1
child 136
1218f20d4cf5
permissions
-rwxr-xr-x

Use more semantic handling of the format string specifier, and try to avoid using constructs not available in SLC5's native Python

andy@0 1 #! /usr/bin/env python
andy@10 2
andy@10 3 """\
andy@123 4 Usage: %prog [options] <spcfile> [<spcfile2> ...]
andy@10 5
andy@84 6 Make a SUSY mass spectrum plot from an SLHA or ISAWIG spectrum file. If the
andy@84 7 filename ends with .isa, it will be assumed to be an ISAWIG file, otherwise
andy@84 8 it will be assumed to be an SLHA file (for which the normal extension is .spc).
andy@64 9
andy@121 10 Output is currently rendered via the LaTeX PGF/TikZ graphics package: this may
andy@131 11 be obtained as PDF (by default), EPS, PNG or as LaTeX source which can be edited or
andy@131 12 compiled into any LaTeX-supported form. By default the output file name(s) are
andy@131 13 the same as the input names but with the file extension replaced with one appropriate
andy@131 14 to the output file format.
andy@13 15
andy@13 16 TODOs:
andy@131 17 * Merge labels if shifting fails (cf. "poi" test spectrum file).
andy@132 18 * Avoid overlap/too-tight clustering of y-axis ticks. Looks fine for *most*
andy@132 19 spectra I've seen with 100 GeV intervals, but it should be handled properly.
andy@121 20 * Allow user to provide a file which defines the particle line x-positions, labels, etc.
andy@131 21 * Allow use of --outname to specify a list of output base names for multiple inputs.
andy@131 22 * Use proper distinction between physical, plot-logical, and plot output coords.
andy@131 23 * Allow user control over aspect ratio / geometry.
andy@101 24 * Use scaling to allow the y coordinates to be in units of 100 GeV in TikZ output.
andy@132 25 * Distribute decay arrow start/end positions along mass lines rather than always
andy@132 26 to/from their centres?
andy@10 27 """
andy@10 28
andy@12 29 class XEdges(object):
andy@13 30 def __init__(self, left, offset=0.0, width=2.0):
andy@13 31 self.offset = offset
andy@13 32 self.left = left + offset
andy@13 33 self.width = width
andy@13 34 @property
andy@13 35 def right(self):
andy@13 36 return self.left + self.width
andy@12 37 @property
andy@12 38 def centre(self):
andy@12 39 return (self.left + self.right)/2.0
andy@10 40
andy@13 41
andy@13 42 class Label(object):
andy@13 43 def __init__(self, text, offset=None):
andy@13 44 self.text = text
andy@13 45 self.offset = None
andy@13 46 def __str__(self):
andy@13 47 return self.text
andy@13 48
andy@13 49
andy@101 50 ## Details classes for representing decided positions in a way independent of output format
andy@76 51
andy@76 52
andy@12 53 class ParticleDetails(object):
andy@13 54 def __init__(self, label, xnom, xoffset, color="black", labelpos="L", mass=None):
andy@10 55 self.label = label
andy@10 56 self.mass = mass
andy@13 57 self.xedges = XEdges(xnom, xoffset)
andy@11 58 self.color = color
andy@11 59 self.labelpos = labelpos
andy@10 60
andy@13 61
andy@76 62 class DecayDetails(object):
andy@86 63 def __init__(self, pidfrom, xyfrom, pidto, xyto, br, color="gray"): #, thickness=1px, label=None):
andy@76 64 self.pidfrom = pidfrom
andy@76 65 self.xyfrom = xyfrom
andy@76 66 self.pidto = pidto
andy@76 67 self.xyto = xyto
andy@86 68 self.br = br
andy@76 69 self.color = color
andy@76 70 #self.label = label
andy@76 71
andy@76 72
andy@76 73 class LabelDetails(object):
andy@76 74 def __init__(self, xy, texlabel, anchor="l", color="black"):
andy@76 75 self.xy = xy
andy@76 76 self.texlabel = texlabel
andy@76 77 ## Add non-TeX-based label rendering via this property, if needed
andy@76 78 self.textlabel = texlabel
andy@76 79 self.anchor = anchor
andy@76 80 self.color = color
andy@76 81
andy@76 82
andy@134 83 # ## Python version workaround
andy@134 84 # if not "any" in dir():
andy@134 85 # def any(*args):
andy@134 86 # for i in args:
andy@134 87 # if i: return True
andy@134 88 # return False
andy@134 89
andy@134 90
andy@133 91 class OutputFormatSpec(object):
andy@133 92 """Object to abstract translation of semi-arbitrary format strings into
andy@133 93 something more semantically queryable."""
andy@134 94
andy@134 95 def __init__(self, fmtstr):
andy@134 96 self.format_string = fmtstr.lower()
andy@134 97 if "tikz" not in self.format_string:
andy@134 98 self.format_string = "tikz" + self.format_string
andy@134 99 if self.format_string == "tikz":
andy@134 100 self.format_string = "tikztex"
andy@134 101 elif self.format_string == "tikzfrag":
andy@134 102 self.format_string = "tikztexfrag"
andy@134 103 if "frag" in self.format_string and any(f in self.format_string for f in ["pdf", "eps", "png"]):
andy@134 104 logging.error("Oops! You can't currently use LaTeX fragment output together with graphics "
andy@134 105 "formats, since the graphics can't be built from the incomplete LaTeX "
andy@134 106 "file. We'll fix this, but for now you will have to run slhaplot twice: "
andy@134 107 "once for the LaTeX fragment, and another time for the graphical output "
andy@134 108 "formats. Exiting...")
andy@134 109 sys.exit(1)
andy@134 110
andy@134 111 def make_tex(self):
andy@134 112 return ("tex" in self.format_string and not "frag" in self.format_string)
andy@134 113
andy@134 114 def make_texfrag(self):
andy@134 115 return ("texfrag" in self.format_string)
andy@134 116
andy@134 117 def make_pdf(self):
andy@134 118 return ("pdf" in self.format_string)
andy@134 119
andy@134 120 def make_eps(self):
andy@134 121 return ("eps" in self.format_string)
andy@134 122
andy@134 123 def make_png(self):
andy@134 124 return ("png" in self.format_string)
andy@134 125
andy@134 126
andy@134 127 def need_tikz(self):
andy@134 128 for f in ["pdf", "eps", "png"]:
andy@134 129 if f in self.format_string:
andy@134 130 return True
andy@134 131 return False
andy@134 132
andy@134 133 def need_epslatex(self):
andy@134 134 for f in ["eps"]:
andy@134 135 if f in self.format_string:
andy@134 136 return True
andy@134 137 return False
andy@134 138
andy@134 139 def need_pdflatex(self):
andy@134 140 for f in ["pdf", "png"]:
andy@134 141 if f in self.format_string:
andy@134 142 return True
andy@134 143 return False
andy@134 144
andy@134 145 def need_convert(self):
andy@134 146 for f in ["png"]:
andy@134 147 if f in self.format_string:
andy@134 148 return True
andy@134 149 return False
andy@134 150
andy@134 151 def need_compilation(self):
andy@134 152 return self.need_epslatex() or self.need_pdflatex()
andy@134 153
andy@134 154 def file_extensions(self):
andy@134 155 return [f for f in ["tex", "pdf", "eps", "png"] if f in self.format_string]
andy@134 156
andy@133 157
andy@133 158
andy@73 159 import pyslha
andy@129 160 import sys, optparse, logging
andy@73 161 parser = optparse.OptionParser(usage=__doc__, version=pyslha.__version__)
andy@131 162 parser.add_option("-o", "--outname", metavar="NAME",
andy@131 163 help="write output to NAME.suffix, i.e. the suffix will be automatically "
andy@131 164 "generated and the argument to this command now just specifies the base "
andy@131 165 "of the output to which the extension is appended: this allows multiple "
andy@131 166 "formats to be written simultaneously. If you provide a file extension "
andy@131 167 "as part of the NAME argument, it will be treated as part of the base "
andy@131 168 "name and another extension will be automatically added. Note that this "
andy@131 169 "option can only be used if only processing a single input file, as you "
andy@131 170 "presumably don't want all the input spectra to overwrite each other!",
andy@131 171 dest="OUTNAME", default=None)
andy@87 172 parser.add_option("-f", "--format", metavar="FORMAT",
andy@131 173 help="format in which to write output. 'tex' produces LaTeX source using the "
andy@131 174 "TikZ graphics package to render the plot, 'texfrag' produces the same but "
andy@82 175 "with the LaTeX preamble and document lines commented out to make it directly "
andy@131 176 "includeable as a code fragment in LaTeX document source, and 'pdf' produces "
andy@131 177 "a PDF file created by running pdflatex and pdfcrop on the 'tex' output. You "
andy@131 178 "may also combine multiple formats just by listing them all in the format "
andy@131 179 "string, e.g. 'png,pdf,tex' (default: %default)",
andy@131 180 dest="FORMAT", default="pdf")
andy@86 181 parser.add_option("--preamble", metavar="FILE",
andy@86 182 help="specify a file to be inserted into LaTeX output as a special preamble",
andy@86 183 dest="PREAMBLE", default=None)
andy@62 184 parser.add_option("--minbr", "--br", metavar="BR",
andy@116 185 help="show decay lines for decays with a branching ratio of > BR, as either a "
andy@116 186 "fraction or percentage (default: show none)",
andy@116 187 dest="DECAYS_MINBR", default="1.1")
andy@111 188 parser.add_option("--decaystyle", choices=["const", "brwidth", "brcolor", "brwidth+brcolor"], metavar="STYLE",
andy@96 189 help="drawing style of decay arrows, from const/brwidth. The 'const' style draws "
andy@96 190 "all decay lines with the same width, 'brwidth' linearly scales the width of the "
andy@96 191 "decay arrow according to the decay branching ratio. Other modes such as BR-dependent "
andy@96 192 "colouring may be added later. (default: %default)",
andy@111 193 dest="DECAYS_STYLE", default="brwidth+brcolor")
andy@73 194 parser.add_option("--labels", choices=["none", "merge", "shift"], metavar="MODE",
andy@96 195 help="treatment of labels for particle IDs, from none/merge/shift. 'none' shows "
andy@96 196 "no labels at all, 'merge' combines would-be-overlapping labels into a single "
andy@96 197 "comma-separated list, and 'shift' vertically shifts the clashing labels to avoid "
andy@96 198 "collisions (default: %default)",
andy@62 199 dest="PARTICLES_LABELS", default="shift")
andy@129 200 verbgroup = optparse.OptionGroup(parser, "Verbosity control")
andy@129 201 parser.add_option("-l", dest="NATIVE_LOG_STRS", action="append",
andy@129 202 default=[], help="set a log level in the Rivet library")
andy@129 203 verbgroup.add_option("-v", "--verbose", action="store_const", const=logging.DEBUG, dest="LOGLEVEL",
andy@129 204 default=logging.INFO, help="print debug (very verbose) messages")
andy@129 205 verbgroup.add_option("-q", "--quiet", action="store_const", const=logging.WARNING, dest="LOGLEVEL",
andy@129 206 default=logging.INFO, help="be very quiet")
andy@129 207 parser.add_option_group(verbgroup)
andy@86 208
andy@86 209
andy@131 210 ## Run parser and configure the logging level
andy@10 211 opts, args = parser.parse_args()
andy@131 212 logging.basicConfig(level=opts.LOGLEVEL, format="%(message)s")
andy@131 213
andy@131 214 ## Create some boolean flags from the chosen particle label clash-avoidance scheme
andy@62 215 opts.PARTICLES_LABELS_SHOW = (opts.PARTICLES_LABELS != "none")
andy@62 216 opts.PARTICLES_LABELS_MERGE = (opts.PARTICLES_LABELS == "merge")
andy@62 217 opts.PARTICLES_LABELS_SHIFT = (opts.PARTICLES_LABELS == "shift")
andy@131 218
andy@131 219 ## Parsing the branching ratio string
andy@116 220 if opts.DECAYS_MINBR.endswith("%"):
andy@123 221 opts.DECAYS_MINBR = float(opts.DECAYS_MINBR[:-1]) / 100.0
andy@116 222 else:
andy@116 223 opts.DECAYS_MINBR = float(opts.DECAYS_MINBR)
andy@131 224
andy@134 225 ## Output format handling: convert string arg to a more semantically queryable type
andy@134 226 opts.FORMAT = OutputFormatSpec(opts.FORMAT)
andy@123 227
andy@115 228
andy@115 229 ## Check non-optional arguments
andy@123 230 INFILES = args
andy@123 231 if len(INFILES) == 0:
andy@10 232 parser.print_help()
andy@10 233 sys.exit(1)
andy@131 234 if len(INFILES) > 1 and opts.OUTNAME is not None:
andy@131 235 logging.error("Multiple input files specified with --outname... not a good plan! Exiting for your own good...")
andy@123 236 sys.exit(1)
andy@10 237
andy@86 238
andy@123 239 ## Test for external packages
andy@123 240 import subprocess
andy@131 241
andy@131 242 ## Test for tikz package if rendering the LaTeX source
andy@134 243 if opts.FORMAT.need_tikz():
andy@123 244 p = subprocess.Popen(["which", "kpsewhich"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@123 245 rtn = p.wait()
andy@123 246 if rtn != 0:
andy@129 247 logging.warning("WARNING: kpsewhich could not be found: check for tikz package cannot be run")
andy@86 248 else:
andy@123 249 p = subprocess.Popen(["kpsewhich", "tikz.sty"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@123 250 rtn = p.wait()
andy@123 251 if rtn != 0:
andy@131 252 logging.error("LaTeX tikz.sty could not be found: graphical format modes cannot work")
andy@87 253 sys.exit(3)
andy@86 254
andy@131 255 ## Test for pdflatex if we need to make a PDF
andy@134 256 if opts.FORMAT.need_pdflatex():
andy@131 257 p = subprocess.Popen(["which", "pdflatex"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@131 258 rtn = p.wait()
andy@131 259 if rtn != 0:
andy@131 260 logging.error("pdflatex could not be found: PDF format mode (and dependent ones like PNG) cannot work")
andy@131 261 sys.exit(3)
andy@131 262
andy@131 263 ## Test for convert if we need to make a bitmap format
andy@134 264 if opts.FORMAT.need_convert():
andy@131 265 p = subprocess.Popen(["which", "convert"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@131 266 rtn = p.wait()
andy@131 267 if rtn != 0:
andy@131 268 logging.error("convert could not be found: PNG format mode cannot work")
andy@131 269 sys.exit(3)
andy@131 270
andy@131 271 ## Test for latex, dvips and ps2eps if making an EPS
andy@134 272 if opts.FORMAT.need_epslatex():
andy@131 273 p = subprocess.Popen(["which", "latex"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@131 274 rtn = p.wait()
andy@131 275 if rtn != 0:
andy@131 276 logging.error("latex could not be found: EPS format mode cannot work")
andy@131 277 sys.exit(3)
andy@131 278 p = subprocess.Popen(["which", "dvips"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@131 279 rtn = p.wait()
andy@131 280 if rtn != 0:
andy@131 281 logging.error("dvips could not be found: EPS format mode cannot work")
andy@131 282 sys.exit(3)
andy@131 283 p = subprocess.Popen(["which", "ps2eps"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@131 284 rtn = p.wait()
andy@131 285 if rtn != 0:
andy@131 286 logging.error("ps2eps could not be found: EPS format mode cannot work")
andy@131 287 sys.exit(3)
andy@131 288
andy@79 289
andy@123 290 ## Loop over input spectrum files
andy@123 291 for infile in INFILES:
andy@123 292
andy@123 293 ## Choose output file
andy@131 294 outname = opts.OUTNAME
andy@131 295 if outname is None:
andy@123 296 import os
andy@123 297 o = os.path.basename(infile)
andy@123 298 if o == "-":
andy@123 299 o = "out"
andy@131 300 elif "." in o:
andy@123 301 o = o[:o.rindex(".")]
andy@131 302 outname = o
andy@123 303
andy@123 304 ## Info for the user
andy@134 305 extlist = opts.FORMAT.file_extensions()
andy@131 306 extstr = ",".join(extlist)
andy@131 307 if len(extlist) > 1:
andy@131 308 extstr = "{" + extstr + "}"
andy@131 309 logging.info("Plotting %s -> %s.%s" % (infile, outname, extstr))
andy@123 310
andy@123 311
andy@123 312 ## Read spectrum file
andy@123 313 BLOCKS, DECAYS = None, None
andy@123 314 # print BLOCKS
andy@123 315 if infile == "-":
andy@123 316 intext = sys.stdin.read()
andy@123 317 BLOCKS, DECAYS = pyslha.readSLHA(intext)
andy@123 318 elif infile.endswith(".isa"):
andy@123 319 BLOCKS, DECAYS = pyslha.readISAWIGFile(infile)
andy@123 320 else:
andy@123 321 BLOCKS, DECAYS = pyslha.readSLHAFile(infile)
andy@123 322 # print BLOCKS
andy@123 323
andy@123 324
andy@123 325 ## Define particle rendering details (may be adapted based on input file, so it *really*
andy@123 326 ## does need to be redefined in each loop over spectrum files!)
andy@123 327 XHIGGS = 0.0
andy@123 328 XSLEPTON = 5.0
andy@123 329 XGAUGINO = 10.0
andy@123 330 XSUSYQCD = 15.0
andy@123 331 PDETAILS = {
andy@123 332 25 : ParticleDetails(Label(r"$h^0$"), XHIGGS, -0.2, color="blue"),
andy@123 333 35 : ParticleDetails(Label(r"$H^0$"), XHIGGS, -0.2, color="blue"),
andy@123 334 36 : ParticleDetails(Label(r"$A^0$"), XHIGGS, -0.2, color="blue"),
andy@123 335 37 : ParticleDetails(Label(r"$H^\pm$"), XHIGGS, 0.2, color="red"),
andy@123 336 1000011 : ParticleDetails(Label(r"$\tilde{\ell}_\text{L}$"), XSLEPTON, -0.2, color="blue"),
andy@123 337 2000011 : ParticleDetails(Label(r"$\tilde{\ell}_\text{R}$"), XSLEPTON, -0.2, color="blue"),
andy@123 338 1000015 : ParticleDetails(Label(r"$\tilde{\tau}_1$"), XSLEPTON, 0.2, color="red"),
andy@123 339 2000015 : ParticleDetails(Label(r"$\tilde{\tau}_2$"), XSLEPTON, 0.2, color="red"),
andy@123 340 1000012 : ParticleDetails(Label(r"$\tilde{\nu}_\text{L}$"), XSLEPTON, -0.2, color="blue"),
andy@123 341 1000016 : ParticleDetails(Label(r"$\tilde{\nu}_\tau$"), XSLEPTON, 0.2, color="red"),
andy@123 342 1000022 : ParticleDetails(Label(r"$\tilde{\chi}_1^0$"), XGAUGINO, -0.2, color="blue"),
andy@123 343 1000023 : ParticleDetails(Label(r"$\tilde{\chi}_2^0$"), XGAUGINO, -0.2, color="blue"),
andy@123 344 1000025 : ParticleDetails(Label(r"$\tilde{\chi}_3^0$"), XGAUGINO, -0.2, color="blue"),
andy@123 345 1000035 : ParticleDetails(Label(r"$\tilde{\chi}_4^0$"), XGAUGINO, -0.2, color="blue"),
andy@123 346 1000024 : ParticleDetails(Label(r"$\tilde{\chi}_1^\pm$"), XGAUGINO, 0.2, color="red"),
andy@123 347 1000037 : ParticleDetails(Label(r"$\tilde{\chi}_2^\pm$"), XGAUGINO, 0.2, color="red"),
andy@123 348 1000039 : ParticleDetails(Label(r"$\tilde{G}$"), XGAUGINO, 0.15, color="black!50!blue!30!green"),
andy@123 349 1000021 : ParticleDetails(Label(r"$\tilde{g}$"), XSUSYQCD, -0.3, color="black!50!blue!30!green"),
andy@123 350 1000001 : ParticleDetails(Label(r"$\tilde{q}_\text{L}$"), XSUSYQCD, -0.1, color="blue"),
andy@123 351 2000001 : ParticleDetails(Label(r"$\tilde{q}_\text{R}$"), XSUSYQCD, -0.1, color="blue"),
andy@123 352 1000005 : ParticleDetails(Label(r"$\tilde{b}_1$"), XSUSYQCD, 0.2, color="black!50!blue!30!green"),
andy@123 353 2000005 : ParticleDetails(Label(r"$\tilde{b}_2$"), XSUSYQCD, 0.2, color="black!50!blue!30!green"),
andy@123 354 1000006 : ParticleDetails(Label(r"$\tilde{t}_1$"), XSUSYQCD, 0.2, color="red"),
andy@123 355 2000006 : ParticleDetails(Label(r"$\tilde{t}_2$"), XSUSYQCD, 0.2, color="red")
andy@123 356 }
andy@123 357
andy@123 358
andy@123 359 ## Set mass values in PDETAILS
andy@123 360 massblock = BLOCKS["MASS"]
andy@123 361 for pid in PDETAILS.keys():
andy@123 362 if massblock.entries.has_key(pid):
andy@123 363 PDETAILS[pid].mass = abs(massblock.entries[pid])
andy@123 364 else:
andy@123 365 del PDETAILS[pid]
andy@123 366
andy@123 367
andy@123 368 ## Decays
andy@123 369 DDETAILS = {}
andy@123 370 for pid, detail in sorted(PDETAILS.iteritems()):
andy@123 371 if DECAYS.has_key(pid):
andy@123 372 DDETAILS.setdefault(pid, {})
andy@123 373 xyfrom = (detail.xedges.centre, detail.mass)
andy@123 374 for d in DECAYS[pid].decays:
andy@123 375 if d.br > opts.DECAYS_MINBR:
andy@123 376 for pid2 in d.ids:
andy@123 377 if PDETAILS.has_key(pid2):
andy@123 378 xyto = (PDETAILS[pid2].xedges.centre, PDETAILS[pid2].mass)
andy@123 379 DDETAILS[pid][pid2] = DecayDetails(pid, xyfrom, pid2, xyto, d.br)
andy@123 380 if DDETAILS.has_key(pid) and not DDETAILS[pid]:
andy@123 381 del DDETAILS[pid]
andy@123 382
andy@123 383
andy@123 384 ## Labels
andy@123 385 PLABELS = []
andy@123 386 if opts.PARTICLES_LABELS_SHOW:
andy@123 387 class MultiLabel(object):
andy@123 388 def __init__(self, label=None, x=None, y=None, anchor=None):
andy@123 389 self.labels = [(label, x, y)] or []
andy@123 390 self.anchor = anchor or "l"
andy@123 391
andy@123 392 def __len__(self):
andy@123 393 return len(self.labels)
andy@123 394
andy@123 395 @property
andy@123 396 def joinedlabel(self):
andy@123 397 return r",\,".join(l[0] for l in self.labels)
andy@123 398
andy@123 399 @property
andy@123 400 def avgx(self):
andy@123 401 return sum(l[1] for l in self.labels)/float(len(self))
andy@123 402 @property
andy@123 403 def minx(self):
andy@123 404 return min(l[1] for l in self.labels)/float(len(self))
andy@123 405 @property
andy@123 406 def maxx(self):
andy@123 407 return max(l[1] for l in self.labels)/float(len(self))
andy@123 408
andy@123 409 @property
andy@123 410 def avgy(self):
andy@123 411 return sum(l[2] for l in self.labels)/float(len(self))
andy@123 412 @property
andy@123 413 def miny(self):
andy@123 414 return min(l[2] for l in self.labels)/float(len(self))
andy@123 415 @property
andy@123 416 def maxy(self):
andy@123 417 return max(l[2] for l in self.labels)/float(len(self))
andy@123 418
andy@123 419 def add(self, label, x, y):
andy@123 420 self.labels.append((label, x, y))
andy@123 421 self.labels = sorted(self.labels, key=lambda l : l[2])
andy@123 422 return self
andy@123 423 def get(self):
andy@123 424 for i in self.labels:
andy@123 425 yield i
andy@123 426
andy@123 427 def rel_err(a, b):
andy@123 428 return abs((a-b)/(a+b)/2.0)
andy@123 429
andy@123 430 ## Use max mass to work out the height of a text line in mass units
andy@123 431 maxmass = None
andy@123 432 for pid, pdetail in sorted(PDETAILS.iteritems()):
andy@123 433 maxmass = max(pdetail.mass, maxmass)
andy@123 434 text_height_in_mass_units = maxmass/22.0
andy@123 435 ##
andy@123 436 ## Merge colliding labels
andy@123 437 reallabels = []
andy@123 438 for pid, pdetail in sorted(PDETAILS.iteritems()):
andy@123 439 labelx = None
andy@123 440 offset = pdetail.label.offset or 0.2
andy@123 441 anchor = None
andy@123 442 if pdetail.xedges.offset <= 0:
andy@123 443 labelx = pdetail.xedges.left - offset
andy@123 444 anchor = "r"
andy@123 445 else:
andy@123 446 labelx = pdetail.xedges.right + offset
andy@123 447 anchor = "l"
andy@123 448 labely = pdetail.mass
andy@123 449 ## Avoid hitting the 0 mass line/border
andy@123 450 if labely < 0.6*text_height_in_mass_units:
andy@123 451 labely = 0.6*text_height_in_mass_units
andy@123 452
andy@123 453 text = pdetail.label.text
andy@123 454 ## Check for collisions
andy@123 455 collision = False
andy@123 456 if opts.PARTICLES_LABELS_SHIFT or opts.PARTICLES_LABELS_MERGE:
andy@123 457 for i, rl in enumerate(reallabels):
andy@123 458 if anchor == rl.anchor and abs(labelx - rl.avgx) < 0.5:
andy@123 459 import math
andy@123 460 if labely > rl.miny - text_height_in_mass_units and labely < rl.maxy + text_height_in_mass_units:
andy@123 461 reallabels[i] = rl.add(text, labelx, labely)
andy@123 462 collision = True
andy@123 463 break
andy@123 464 if not collision:
andy@123 465 reallabels.append(MultiLabel(text, labelx, labely, anchor))
andy@123 466 ## Calculate position shifts and fill PLABELS
andy@123 467 for rl in reallabels:
andy@123 468 if len(rl) == 1 or opts.PARTICLES_LABELS_MERGE:
andy@123 469 PLABELS.append(LabelDetails((rl.avgx, rl.avgy), rl.joinedlabel, anchor=rl.anchor))
andy@123 470 else:
andy@123 471 num_gaps = len(rl)-1
andy@123 472 yrange_old = rl.maxy - rl.miny
andy@123 473 yrange_nom = num_gaps * text_height_in_mass_units
andy@123 474 yrange = max(yrange_old, yrange_nom)
andy@123 475 ydiff = yrange - yrange_old
andy@123 476 for i, (t, x, y) in enumerate(rl.get()):
andy@123 477 ydiff_per_line = ydiff/num_gaps
andy@123 478 # TODO: Further improvement using relative or average positions?
andy@123 479 newy = y + (i - num_gaps/2.0) * ydiff_per_line
andy@123 480 PLABELS.append(LabelDetails((x, newy), t, anchor=rl.anchor))
andy@123 481
andy@123 482
andy@123 483 ## Function for writing out the generated source
andy@123 484 def writeout(out, outfile):
andy@123 485 f = sys.stdout
andy@123 486 if outfile != "-":
andy@123 487 f = open(outfile, "w")
andy@123 488 f.write(out)
andy@123 489 if f is not sys.stdout:
andy@123 490 f.close()
andy@123 491
andy@123 492 out = ""
andy@123 493
andy@123 494 ## TIKZ FORMAT
andy@134 495 # TODO: Remove this test?
andy@134 496 if "tikz" in opts.FORMAT.format_string:
andy@123 497
andy@123 498 ## Comment out the preamble etc. if only the TikZ fragment is wanted
andy@123 499 c = ""
andy@134 500 if opts.FORMAT.make_texfrag():
andy@123 501 c = "%"
andy@123 502
andy@123 503 ## Write LaTeX header
andy@123 504 out += "%% http://pypi.python.org/pypi/pyslha\n\n"
andy@123 505 out += c + "\\documentclass[11pt]{article}\n"
andy@123 506 out += c + "\\usepackage{amsmath,amssymb}\n"
andy@123 507 out += c + "\\usepackage[margin=0cm,paperwidth=15.2cm,paperheight=9.8cm]{geometry}\n"
andy@123 508 out += c + "\\usepackage{tikz}\n"
andy@123 509 out += c + "\\pagestyle{empty}\n"
andy@123 510 out += c + "\n"
andy@123 511 ## Insert user-specified preamble file
andy@123 512 if opts.PREAMBLE is not None:
andy@123 513 out += c + "%% User-supplied preamble\n"
andy@123 514 try:
andy@123 515 fpre = open(opts.PREAMBLE, "r")
andy@123 516 for line in fpre:
andy@123 517 out += c + line
andy@123 518 except:
andy@129 519 logging.warning("Could not read preamble file %s -- fallback to using \\input" % opts.PREAMBLE)
andy@123 520 out += c + "\\input{%s}\n" % opts.PREAMBLE.replace(".tex", "")
andy@123 521 else:
andy@123 522 out += c + "%% Default preamble\n"
andy@123 523 out += c + "\\usepackage[osf]{mathpazo}\n"
andy@123 524 #
andy@123 525 out += c + "\n"
andy@123 526 out += c + "\\begin{document}\n"
andy@123 527 out += c + "\\thispagestyle{empty}\n\n"
andy@123 528
andy@123 529 ## Get coord space size: horizontal range is fixed by make-plots
andy@123 530 xmin = -3.0
andy@123 531 xmax = 19.0
andy@123 532 if opts.PARTICLES_LABELS_MERGE:
andy@123 533 ## Need more space if labels are to be merged horizontally
andy@123 534 xmin -= 1.0
andy@123 535 xmax += 1.0
andy@123 536 xdiff = xmax - xmin
andy@123 537 XWIDTH = 22.0
andy@123 538 def scalex(x):
andy@123 539 return x * XWIDTH/xdiff
andy@123 540
andy@123 541 ASPECTRATIO = 0.7 #0.618
andy@123 542 ydiff = ASPECTRATIO * XWIDTH
andy@123 543 ymin = 0.0
andy@123 544 ymax = ymin + ydiff
andy@123 545 ## Get range of masses needed
andy@123 546 maxmass = max(pd.mass for pid, pd in PDETAILS.iteritems())
andy@123 547 maxdisplaymass = maxmass * 1.1
andy@123 548 if maxdisplaymass % 100 != 0:
andy@123 549 maxdisplaymass = ((maxdisplaymass + 100) // 100) * 100
andy@123 550 yscale = (ymax-ymin)/maxdisplaymass
andy@123 551
andy@123 552 ## Write TikZ header
andy@123 553 out += "\\centering\n"
andy@123 554 out += "\\begin{tikzpicture}[scale=0.6]\n"
andy@123 555
andy@123 556 out += " %% y-scalefactor (GeV -> coords) = %e\n\n" % yscale
andy@123 557
andy@123 558 ## Draw the plot boundary and y-ticks
andy@123 559 out += " %% Frame\n"
andy@123 560 out += " \\draw (%f,%f) rectangle (%f,%f);\n" % (scalex(xmin), ymin, scalex(xmax), ymax)
andy@123 561 out += " %% y-ticks\n"
andy@123 562 for mtick in xrange(0, int(maxdisplaymass) + 100, 100):
andy@123 563 ytick = mtick * yscale
andy@123 564 out += " \\draw (%f,%f) node[left] {%d};\n" % (scalex(xmin), ytick, mtick)
andy@123 565 if mtick > 0 and mtick < maxdisplaymass:
andy@123 566 ## The 0.3 needs to be in the plot coords
andy@123 567 out += " \\draw (%f,%f) -- (%f,%f);\n" % (scalex(xmin+0.3), ytick, scalex(xmin), ytick)
andy@123 568 out += " \\draw (%f,%f) node[left,rotate=90] {Mass / GeV};\n" % (scalex(xmin-2.0), ymax)
andy@123 569
andy@123 570 ## Decay arrows
andy@123 571 if DDETAILS:
andy@123 572 out += "\n %% Decay arrows\n"
andy@123 573 for pidfrom, todict in sorted(DDETAILS.iteritems()):
andy@123 574 for pidto, dd in sorted(todict.iteritems()):
andy@123 575 out += " %% decay_%d_%d, BR=%0.1f%%\n" % (dd.pidfrom, dd.pidto, dd.br*100)
andy@123 576
andy@123 577 def scalethickness(br):
andy@123 578 if opts.DECAYS_STYLE == "const":
andy@123 579 return 0.8
andy@123 580 elif "brwidth" in opts.DECAYS_STYLE:
andy@123 581 return 1.0 * br
andy@123 582 else:
andy@123 583 raise Exception("Unexpected problem with unknown decay line style option: please contact the PySLHA authors!")
andy@123 584
andy@123 585 def scalecolor(br):
andy@123 586 if opts.DECAYS_STYLE == "const":
andy@123 587 return None
andy@123 588 elif "brcolor" in opts.DECAYS_STYLE:
andy@123 589 return "black!"+str(60*dd.br + 10)
andy@123 590 else:
andy@123 591 raise Exception("Unexpected problem with unknown decay line style option: please contact the PySLHA authors!")
andy@123 592
andy@123 593 out += " \\draw[-stealth,line width=%0.2fpt,dashed,color=%s] (%f,%f) -- (%f,%f);\n" % \
andy@123 594 (scalethickness(dd.br), scalecolor(dd.br) or dd.color,
andy@123 595 scalex(dd.xyfrom[0]), yscale*dd.xyfrom[1], scalex(dd.xyto[0]), yscale*dd.xyto[1])
andy@123 596
andy@123 597 ## Draw mass lines
andy@123 598 if PDETAILS:
andy@123 599 out += "\n %% Particle lines\n"
andy@123 600 for pid, pdetail in sorted(PDETAILS.iteritems()):
andy@123 601 y = pdetail.mass*yscale
andy@123 602 out += " %% pid%s\n" % str(pid)
andy@123 603 out += " \\draw[color=%s,thick] (%f,%f) -- (%f,%f);\n" % \
andy@123 604 (pdetail.color, scalex(pdetail.xedges.left), y, scalex(pdetail.xedges.right), y)
andy@123 605
andy@123 606 ## Particle labels
andy@123 607 if PLABELS:
andy@123 608 out += "\n %% Particle labels\n"
andy@123 609 for ld in PLABELS:
andy@123 610 anchors_pstricks_tikz = { "r" : "left", "l" : "right" }
andy@123 611 out += " \\draw (%f,%f) node[%s] {\small %s};\n" % \
andy@123 612 (scalex(ld.xy[0]), yscale*ld.xy[1], anchors_pstricks_tikz[ld.anchor], ld.texlabel)
andy@123 613
andy@123 614 ## Write TikZ footer
andy@123 615 out += "\end{tikzpicture}\n\n"
andy@123 616
andy@123 617 ## Write LaTeX footer
andy@123 618 out += c + "\end{document}\n"
andy@123 619
andy@123 620
andy@123 621 ## Write output
andy@134 622 if opts.FORMAT.make_tex():
andy@131 623 writeout(out, outname+".tex")
andy@131 624
andy@134 625 if opts.FORMAT.need_compilation():
andy@131 626 ## Run LaTeX
andy@123 627 import tempfile, shutil, subprocess
andy@123 628 tmpdir = tempfile.mkdtemp()
andy@123 629 writeout(out, os.path.join(tmpdir, "mytmp.tex"))
andy@131 630
andy@131 631 def mktmpprocess(cmdlist):
andy@131 632 "Convenience method to reduce repeated subprocess.Popen call noise in what follows."
andy@131 633 return subprocess.Popen(cmdlist, cwd=tmpdir,
andy@131 634 stdout=subprocess.PIPE, stderr=subprocess.PIPE)
andy@131 635
andy@131 636 ## Processing via PDF
andy@134 637 if opts.FORMAT.need_pdflatex():
andy@131 638 try:
andy@131 639 p = mktmpprocess(["pdflatex", r"\scrollmode\input", "mytmp.tex"])
andy@131 640 p.wait()
andy@134 641 if opts.FORMAT.make_pdf():
andy@131 642 shutil.copyfile(os.path.join(tmpdir, "mytmp.pdf"), outname+".pdf")
andy@131 643 except Exception, e:
andy@131 644 logging.error("pdflatex could not be run: PDF and/or PNG format mode cannot work")
andy@131 645 shutil.rmtree(tmpdir)
andy@131 646 sys.exit(3)
andy@131 647 # TODO: Can we use a finally block to make sure that the tmpdir is cleared up, despite the 'sys.exit's?
andy@131 648
andy@131 649 ## Turning the PDF into a PNG if required
andy@134 650 if opts.FORMAT.make_png():
andy@131 651 try:
andy@131 652 p = mktmpprocess(["convert", "-density", "150", "-flatten", "mytmp.pdf", "mytmp.png"])
andy@131 653 p.wait()
andy@131 654 shutil.copyfile(os.path.join(tmpdir, "mytmp.png"), outname+".png")
andy@131 655 except Exception, e:
andy@131 656 logging.error("convert could not be run: PNG format mode cannot work")
andy@131 657 shutil.rmtree(tmpdir)
andy@131 658 sys.exit(3)
andy@131 659
andy@131 660 ## Making a PS or EPS
andy@134 661 if opts.FORMAT.make_eps():
andy@131 662 try:
andy@131 663 p = mktmpprocess(["latex", r"\scrollmode\input", "mytmp.tex"])
andy@131 664 p.wait()
andy@131 665 except Exception, e:
andy@131 666 logging.error("latex could not be run: EPS format mode cannot work")
andy@131 667 shutil.rmtree(tmpdir)
andy@131 668 sys.exit(3)
andy@131 669 try:
andy@131 670 p = mktmpprocess(["dvips", "mytmp.dvi", "-o", "mytmp.ps"])
andy@131 671 p.wait()
andy@131 672 except Exception, e:
andy@131 673 logging.error("dvips could not be run: EPS format mode cannot work")
andy@131 674 shutil.rmtree(tmpdir)
andy@131 675 sys.exit(3)
andy@131 676 try:
andy@131 677 p = mktmpprocess(["ps2eps", "mytmp.ps"])
andy@131 678 p.wait()
andy@131 679 shutil.copyfile(os.path.join(tmpdir, "mytmp.eps"), outname+".eps")
andy@131 680 except Exception, e:
andy@131 681 logging.error("ps2eps could not be run: EPS format mode cannot work")
andy@131 682 shutil.rmtree(tmpdir)
andy@131 683 sys.exit(3)
andy@131 684
andy@123 685 shutil.rmtree(tmpdir)
andy@123 686
andy@123 687
andy@123 688 ## UNRECOGNISED FORMAT!
andy@123 689 else:
andy@129 690 logging.error("Other formats not currently supported! How did we even get here?!")
andy@123 691 sys.exit(2)

mercurial