slhaplot

Sun, 28 Apr 2013 21:13:25 +0200

author
Andy Buckley <andy@insectnation.org>
date
Sun, 28 Apr 2013 21:13:25 +0200
changeset 199
80d4b063103a
parent 180
6e4c9fde492e
child 206
6e091122c7d5
permissions
-rwxr-xr-x

Various fixes to ISAWIG output (and conversion to use new tuple-indexing on mixing matrices and single-value accessing on ALPHA).

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

mercurial