pyslha.py

Tue, 03 Aug 2010 17:29:41 +0200

author
Andy Buckley <andy@insectnation.org>
date
Tue, 03 Aug 2010 17:29:41 +0200
changeset 8
c97f2f581685
parent 7
a28827356061
child 12
558b3b392500
permissions
-rw-r--r--

Adding docstrings and making the Particle masses get populated via the MASS block contents

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

mercurial