slhaplot

Sun, 10 Apr 2011 12:24:01 +0100

author
Andy Buckley <andy@insectnation.org>
date
Sun, 10 Apr 2011 12:24:01 +0100
changeset 141
869c16f9093b
parent 137
0aecd3e7f444
child 143
6f07611e1dc8
permissions
-rwxr-xr-x

Fix block parsing error and make 1.2.3 release

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

mercurial