slhaplot

Mon, 15 Jul 2013 16:50:38 +0200

author
Andy Buckley <andy@insectnation.org>
date
Mon, 15 Jul 2013 16:50:38 +0200
changeset 244
269bc4611cdd
parent 236
d65f2e971a6a
child 245
ca52fe2c0d71
permissions
-rwxr-xr-x

Add credit to ChangeLog

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

mercurial