slhaplot

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

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

Better newline handling in final SLHA output formatting

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

mercurial