pyslha.py

Sat, 21 Aug 2010 16:30:21 +0100

author
Andy Buckley <andy@insectnation.org>
date
Sat, 21 Aug 2010 16:30:21 +0100
changeset 31
7d41d901fdd0
parent 30
2e20779d3e6e
child 32
34fa1801ce34
permissions
-rw-r--r--

Adding ignorenobr and reformatting a bit

andy@1 1 #! /usr/bin/env python
andy@1 2
andy@8 3 """\
andy@8 4 A simple but flexible parser of SUSY Les Houches Accord (SLHA) model and decay files.
andy@25 5
andy@30 6 TODO: Fix treatment of single-entry blocks, e.g. ALPHA
andy@8 7 """
andy@8 8
andy@8 9 from __future__ import with_statement
andy@26 10
andy@8 11 __author__ = "Andy Buckley <andy.buckley@cern.ch"
andy@30 12 __version__ = "0.2.0"
andy@8 13
andy@1 14 import re
andy@1 15
andy@4 16 def _autotype(var):
andy@30 17 """Automatically convert strings to numerical types if possible."""
andy@30 18 if type(var) is not str:
andy@4 19 return var
andy@4 20 if var.isdigit():
andy@4 21 return int(var)
andy@4 22 try:
andy@4 23 f = float(var)
andy@4 24 return f
andy@4 25 except ValueError:
andy@4 26 return var
andy@4 27
andy@30 28 def _autostr(var):
andy@30 29 """Automatically numerical types to the right sort of string."""
andy@30 30 if type(var) is float:
andy@30 31 return "%e" % var
andy@30 32 return str(var)
andy@30 33
andy@30 34
andy@4 35
andy@12 36 class Block(object):
andy@8 37 """
andy@8 38 Object representation of any BLOCK elements read from the SLHA file. Blocks
andy@8 39 have a name, may have an associated Q value, and then a collection of data
andy@8 40 entries, stored as a recursive dictionary. Types in the dictionary are
andy@8 41 numeric (int or float) when a cast from the string in the file has been
andy@8 42 possible.
andy@8 43 """
andy@5 44 def __init__(self, name, q=None):
andy@1 45 self.name = name
andy@1 46 self.entries = {}
andy@5 47 self.q = _autotype(q)
andy@1 48
andy@1 49 def add_entry(self, entry):
andy@1 50 #print entry
andy@1 51 nextparent = self.entries
andy@1 52 if len(entry) < 2:
andy@1 53 raise Exception("Block entries must be at least a 2-tuple")
andy@4 54 #print "in", entry
andy@4 55 entry = map(_autotype, entry)
andy@4 56 #print "out", entry
andy@1 57 for e in entry[:-2]:
andy@1 58 if e is not entry[-1]:
andy@1 59 nextparent = nextparent.setdefault(e, {})
andy@1 60 nextparent[entry[-2]] = entry[-1]
andy@1 61 #print self.entries
andy@1 62
andy@1 63 def __cmp__(self, other):
andy@31 64 return cmp(self.name, other.name)
andy@1 65
andy@1 66 def __str__(self):
andy@1 67 s = self.name
andy@5 68 if self.q is not None:
andy@5 69 s += " (Q=%s)" % self.q
andy@1 70 s += "\n"
andy@1 71 s += str(self.entries)
andy@1 72 return s
andy@1 73
andy@1 74
andy@12 75 class Decay(object):
andy@8 76 """
andy@8 77 Object representing a decay entry on a particle decribed by the SLHA file.
andy@8 78 'Decay' objects are not a direct representation of a DECAY block in an SLHA
andy@8 79 file... that role, somewhat confusingly, is taken by the Particle class.
andy@8 80
andy@8 81 Decay objects have three properties: a branching ratio, br, an nda number
andy@12 82 (number of daughters == len(ids)), and a tuple of PDG PIDs to which the
andy@12 83 decay occurs. The PDG ID of the particle whose decay this represents may
andy@12 84 also be stored, but this is normally known via the Particle in which the
andy@12 85 decay is stored.
andy@8 86 """
andy@8 87 def __init__(self, br, nda, ids, parentid=None):
andy@8 88 self.parentid = parentid
andy@6 89 self.br = br
andy@6 90 self.nda = nda
andy@6 91 self.ids = ids
andy@29 92 assert(self.nda == len(self.ids))
andy@6 93
andy@6 94 def __cmp__(self, other):
andy@31 95 return cmp(other.br, self.br)
andy@6 96
andy@6 97 def __str__(self):
andy@12 98 return "%e %s" % (self.br, self.ids)
andy@6 99
andy@6 100
andy@12 101 class Particle(object):
andy@8 102 """
andy@8 103 Representation of a single, specific particle, decay block from an SLHA
andy@8 104 file. These objects are not themselves called 'Decay', since that concept
andy@8 105 applies more naturally to the various decays found inside this
andy@8 106 object. Particle classes store the PDG ID (pid) of the particle being
andy@8 107 represented, and optionally the mass (mass) and total decay width
andy@8 108 (totalwidth) of that particle in the SLHA scenario. Masses may also be found
andy@8 109 via the MASS block, from which the Particle.mass property is filled, if at
andy@8 110 all. They also store a list of Decay objects (decays) which are probably the
andy@8 111 item of most interest.
andy@8 112 """
andy@6 113 def __init__(self, pid, totalwidth=None, mass=None):
andy@6 114 self.pid = pid
andy@6 115 self.totalwidth = totalwidth
andy@6 116 self.mass = mass
andy@6 117 self.decays = []
andy@6 118
andy@6 119 def add_decay(self, br, nda, ids):
andy@6 120 self.decays.append(Decay(br, nda, ids))
andy@6 121 self.decays.sort()
andy@6 122
andy@6 123 def __cmp__(self, other):
andy@6 124 if abs(self.pid) == abs(other.pid):
andy@31 125 return cmp(self.pid, other.pid)
andy@31 126 return cmp(abs(self.pid), abs(other.pid))
andy@6 127
andy@6 128 def __str__(self):
andy@6 129 s = str(self.pid)
andy@7 130 if self.mass is not None:
andy@7 131 s += " : mass = %e GeV" % self.mass
andy@6 132 if self.totalwidth is not None:
andy@7 133 s += " : total width = %e GeV" % self.totalwidth
andy@6 134 for d in self.decays:
andy@12 135 if d.br > 0.0:
andy@12 136 s += "\n %s" % d
andy@6 137 return s
andy@1 138
andy@1 139
andy@31 140 def readSLHAFile(spcfilename, **kwargs):
andy@21 141 """
andy@21 142 Read an SLHA file, returning dictionaries of blocks and decays.
andy@31 143
andy@31 144 Other keyword parameters are passed to readSLHA.
andy@21 145 """
andy@21 146 with open(spcfilename, "r") as f:
andy@31 147 return readSLHA(f.read(), kwargs)
andy@21 148
andy@21 149
andy@31 150 def readSLHA(spcstr, ignorenobr=False):
andy@21 151 """
andy@31 152 Read an SLHA definition from a string, returning dictionaries of blocks and
andy@31 153 decays.
andy@31 154
andy@31 155 If the ignorenobr parameter is True, do not store decay entries with a
andy@31 156 branching ratio of zero.
andy@21 157 """
andy@1 158 blocks = {}
andy@1 159 decays = {}
andy@21 160 #
andy@21 161 currentblock = None
andy@21 162 currentdecay = None
andy@21 163 for line in spcstr.splitlines():
andy@21 164 ## Handle (ignore) comment lines
andy@21 165 if line.startswith("#"):
andy@21 166 continue
andy@21 167 if "#" in line:
andy@21 168 line = line[:line.index("#")]
andy@21 169
andy@21 170 ## Handle BLOCK/DECAY start lines
andy@21 171 if line.upper().startswith("BLOCK"):
andy@21 172 match = re.match(r"BLOCK\s+(\w+)\s+(Q=\s*.+)?", line.upper())
andy@21 173 if not match:
andy@8 174 continue
andy@21 175 blockname = match.group(1)
andy@21 176 qstr = match.group(2)
andy@21 177 if qstr is not None:
andy@21 178 qstr = qstr[2:].strip()
andy@21 179 currentblock = blockname
andy@21 180 currentdecay = None
andy@21 181 blocks[blockname] = Block(blockname, q=qstr)
andy@21 182 elif line.upper().startswith("DECAY"):
andy@21 183 match = re.match(r"DECAY\s+(\d+)\s+([\d\.E+-]+).*", line.upper())
andy@21 184 if not match:
andy@21 185 continue
andy@21 186 pdgid = int(match.group(1))
andy@21 187 width = float(match.group(2))
andy@21 188 currentblock = "DECAY"
andy@21 189 currentdecay = pdgid
andy@21 190 decays[pdgid] = Particle(pdgid, width)
andy@21 191 else:
andy@21 192 ## In-block line
andy@21 193 if currentblock is not None:
andy@21 194 items = line.split()
andy@21 195 # TODO: Sort out tuple item types: autoconvert integers and floats
andy@21 196 if len(items) < 1:
andy@6 197 continue
andy@21 198 if currentblock != "DECAY":
andy@21 199 #print currentblock
andy@21 200 if len(items) < 2:
andy@21 201 ## Treat the ALPHA block differently
andy@21 202 blocks[currentblock].value = _autotype(items[0])
andy@21 203 blocks[currentblock].add_entry((0, items[0]))
andy@8 204 else:
andy@21 205 blocks[currentblock].add_entry(items)
andy@21 206 else:
andy@21 207 br = float(items[0])
andy@21 208 nda = int(items[1])
andy@21 209 ids = map(int, items[2:])
andy@31 210 if br > 0.0 or not ignorenobr:
andy@31 211 decays[currentdecay].add_decay(br, nda, ids)
andy@1 212
andy@8 213 ## Try to populate Particle masses from the MASS block
andy@8 214 for pid in blocks["MASS"].entries.keys():
andy@8 215 if decays.has_key(pid):
andy@8 216 decays[pid].mass = blocks["MASS"].entries[pid]
andy@8 217
andy@1 218 return blocks, decays
andy@1 219
andy@1 220
andy@21 221
andy@31 222 def writeSLHAFile(spcfilename, blocks, decays, **kwargs):
andy@29 223 """
andy@29 224 Write an SLHA file from the supplied blocks and decays dicts.
andy@31 225
andy@31 226 Other keyword parameters are passed to writeSLHA.
andy@29 227 """
andy@29 228 with open(spcfilename, "w") as f:
andy@31 229 f.write(writeSLHA(blocks, decays, kwargs))
andy@29 230
andy@29 231
andy@31 232 def writeSLHA(blocks, decays, ignorenobr=False):
andy@29 233 """
andy@29 234 Return an SLHA definition as a string, from the supplied blocks and decays dicts.
andy@29 235 """
andy@31 236 sep = " "
andy@29 237 out = ""
andy@30 238 def dict_hier_strs(d, s=""):
andy@30 239 if type(d) is dict:
andy@30 240 for k, v in sorted(d.iteritems()):
andy@31 241 for s2 in dict_hier_strs(v, s + sep + _autostr(k)):
andy@30 242 yield s2
andy@30 243 else:
andy@31 244 yield s + sep + _autostr(d)
andy@30 245 ## Blocks
andy@30 246 for bname, b in sorted(blocks.iteritems()):
andy@29 247 namestr = b.name
andy@29 248 if b.q is not None:
andy@29 249 namestr += " Q= %e" % b.q
andy@29 250 out += "BLOCK %s\n" % namestr
andy@30 251 for s in dict_hier_strs(b.entries):
andy@31 252 out += sep + s + "\n"
andy@29 253 out += "\n"
andy@30 254 ## Decays
andy@30 255 for pid, particle in sorted(decays.iteritems()):
andy@29 256 out += "DECAY %d %e\n" % (particle.pid, particle.totalwidth)
andy@30 257 for d in sorted(particle.decays):
andy@31 258 if d.br > 0.0 or not ignorenobr:
andy@31 259 products_str = " ".join(map(str, d.ids))
andy@31 260 out += sep + "%e" % d.br + sep + "%d" % len(d.ids) + sep + products_str + "\n"
andy@29 261 out += "\n"
andy@29 262 return out
andy@29 263
andy@29 264
andy@29 265
andy@1 266 if __name__ == "__main__":
andy@1 267 import sys
andy@1 268 for a in sys.argv[1:]:
andy@29 269 blocks, decays = readSLHAFile(a)
andy@3 270
andy@5 271 for bname, b in sorted(blocks.iteritems()):
andy@5 272 print b
andy@5 273 print
andy@3 274
andy@4 275 print blocks["MASS"].entries[25]
andy@6 276 print
andy@6 277
andy@6 278 for p in sorted(decays.values()):
andy@6 279 print p
andy@6 280 print
andy@29 281
andy@31 282 print writeSLHA(blocks, decays, ignorenobr=True)

mercurial