slhaplot

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

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

Add credit to ChangeLog

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

mercurial