pyslha.py

changeset 189
60c73b489420
parent 188
e706e3b2a647
child 192
453a523cba25
     1.1 --- a/pyslha.py	Thu Apr 25 11:43:16 2013 +0200
     1.2 +++ b/pyslha.py	Fri Apr 26 01:00:15 2013 +0200
     1.3 @@ -18,7 +18,7 @@
     1.4  user-tweaking of the generated output.
     1.5  
     1.6  TODOs:
     1.7 - * Opportunistically use ordereddict to preserve write order
     1.8 + * Store dict entries as tuple keys, e.g. myblock.entries[1,2] rather than recursive dicts => v1.6.0
     1.9   * Preserve comments from read -> write (needs full-line/inline comment separation?)
    1.10   * Split writeSLHA into writeSLHA{Blocks,Decays}
    1.11   * Identify HERWIG decay matrix element to use in ISAWIG
    1.12 @@ -26,9 +26,21 @@
    1.13  """
    1.14  
    1.15  __author__ = "Andy Buckley <andy.buckley@cern.ch"
    1.16 -__version__ = "1.4.4"
    1.17 +__version__ = "1.5.0"
    1.18  
    1.19  
    1.20 +def _mkdict():
    1.21 +    """Try to return an empty ordered dict, but fall back to normal dict if necessary"""
    1.22 +    try:
    1.23 +        from collections import OrderedDict
    1.24 +        return OrderedDict()
    1.25 +    except:
    1.26 +        try:
    1.27 +            from ordereddict import OrderedDict
    1.28 +            return OrderedDict()
    1.29 +        except:
    1.30 +            return dict()
    1.31 +
    1.32  def _autotype(var):
    1.33      """Automatically convert strings to numerical types if possible."""
    1.34      if type(var) is not str:
    1.35 @@ -67,20 +79,22 @@
    1.36      """
    1.37      def __init__(self, name, q=None):
    1.38          self.name = name
    1.39 -        self.entries = {} # TODO: Use ordereddict?
    1.40 +        self.entries = _mkdict()
    1.41          self.q = _autotype(q)
    1.42  
    1.43      def add_entry(self, entry):
    1.44          #print entry
    1.45          nextparent = self.entries
    1.46 +        if type(entry) is str:
    1.47 +            raise Exception("Block entries must be tuples or lists")
    1.48          if len(entry) < 2:
    1.49 -            raise Exception("Block entries must be at least a 2-tuple")
    1.50 +            raise Exception("Block entry tuples must have at least two entries")
    1.51          #print "in", entry
    1.52          entry = map(_autotype, entry)
    1.53          #print "out", entry
    1.54          for e in entry[:-2]:
    1.55              if e is not entry[-1]:
    1.56 -                nextparent = nextparent.setdefault(e, {}) # TODO: Use ordereddict?
    1.57 +                nextparent = nextparent.setdefault(e, _mkdict())
    1.58          nextparent[entry[-2]] = entry[-1]
    1.59          #print self.entries
    1.60  
    1.61 @@ -181,8 +195,8 @@
    1.62      If the ignorenobr parameter is True, do not store decay entries with a
    1.63      branching ratio of zero.
    1.64      """
    1.65 -    blocks = {} # TODO: Use ordereddict?
    1.66 -    decays = {} # TODO: Use ordereddict?
    1.67 +    blocks = _mkdict()
    1.68 +    decays = _mkdict()
    1.69      #
    1.70      import re
    1.71      currentblock = None
    1.72 @@ -262,14 +276,14 @@
    1.73      sep = "   "
    1.74      out = ""
    1.75      def dict_hier_strs(d, s=""):
    1.76 -        if type(d) is dict: # TODO: Use ordereddict?
    1.77 -            for k, v in sorted(d.iteritems()):
    1.78 +        if isinstance(d, dict):
    1.79 +            for k, v in d.iteritems():
    1.80                  for s2 in dict_hier_strs(v, s + sep + _autostr(k)):
    1.81                      yield s2
    1.82          else:
    1.83              yield s + sep + _autostr(d)
    1.84      ## Blocks
    1.85 -    for bname, b in sorted(blocks.iteritems()):
    1.86 +    for bname, b in blocks.iteritems():
    1.87          namestr = b.name
    1.88          if b.q is not None:
    1.89              namestr += (" Q= " + fmte) % float(b.q)
    1.90 @@ -278,9 +292,9 @@
    1.91              out += sep + s + "\n"
    1.92          out += "\n"
    1.93      ## Decays
    1.94 -    for pid, particle in sorted(decays.iteritems()):
    1.95 +    for pid, particle in decays.iteritems():
    1.96          out += ("DECAY %d " + fmte + "\n") % (particle.pid, particle.totalwidth or -1)
    1.97 -        for d in sorted(particle.decays):
    1.98 +        for d in particle.decays:
    1.99              if d.br > 0.0 or not ignorenobr:
   1.100                  products_str = "   ".join(map(str, d.ids))
   1.101                  out += sep + fmte % d.br + sep + "%d" % len(d.ids) + sep + products_str + "\n"
   1.102 @@ -510,8 +524,8 @@
   1.103      branching ratio of zero.
   1.104      """
   1.105  
   1.106 -    blocks = {} # TODO: Use ordereddict?
   1.107 -    decays = {} # TODO: Use ordereddict?
   1.108 +    blocks = _mkdict()
   1.109 +    decays = _mkdict()
   1.110      LINES = isastr.splitlines()
   1.111  
   1.112      def getnextvalidline():

mercurial