slhaplot

Mon, 29 Apr 2013 15:04:31 +0200

author
Andy Buckley <andy@insectnation.org>
date
Mon, 29 Apr 2013 15:04:31 +0200
changeset 214
fa07ed634b18
parent 209
42fe24d30c22
child 236
d65f2e971a6a
permissions
-rwxr-xr-x

Better newline handling in final SLHA output formatting

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

mercurial