slhaplot

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

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

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

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

mercurial