pyslha.py

Sat, 21 Aug 2010 15:37:36 +0100

author
Andy Buckley <andy@insectnation.org>
date
Sat, 21 Aug 2010 15:37:36 +0100
changeset 29
764234670c09
parent 26
ffc5e155db55
child 30
2e20779d3e6e
permissions
-rw-r--r--

More SLHA-writing implementation, plus a consistency check assertion on Decay.nda == len(Decay.ids)

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

mercurial