slhaplot

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

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

Fix block parsing error and make 1.2.3 release

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

mercurial