00001
00002 '''
00003 Produce some user data based on GenHeader and SimHeaders.
00004 '''
00005
00006 from UserTagging.UserTaggingAlg import UserTaggingAlg, SUCCESS, FAILURE
00007 from UserTagging.Models import TagSet, Tag, Data, Int, Float, IntArray, FloatArray
00008
00009 class GenSimData(UserTaggingAlg):
00010
00011 def initTagList(self):
00012
00013 self.tag = self.tagSet.GenSimTag = Tag()
00014 self.data = self.tag.dataSet.GenSimData = Data('/Event/UserData/SimHistsExample/GenSim')
00015 self.gen = self.data.GenSimData.Gen
00016 self.sim = self.data.GenSimData.Sim
00017
00018 self.gen.number = Int()
00019 self.gen.pid = IntArray()
00020 self.gen.mass = FloatArray()
00021 self.gen.energy = FloatArray()
00022 self.gen.x, self.gen.y, self.gen.z, self.gen.dx, self.gen.dy, self.gen.dz = \
00023 map (lambda x: FloatArray(), range(6))
00024
00025 self.sim.nhits = Int()
00026 self.sim.weight = FloatArray()
00027 self.sim.detid = IntArray()
00028 self.sim.time = FloatArray()
00029
00030 return
00031
00032 def check(self,evt):
00033
00034 gh = evt['/Event/Gen/GenHeader']
00035 if not gh:
00036 self.error('No GenHeader found')
00037 return FAILURE
00038 self._fill_gen(gh)
00039
00040 sh = evt['/Event/Sim/SimHeader']
00041 if not sh:
00042 self.error('No SimHeader found')
00043 return FAILURE
00044 self._fill_sim(sh)
00045
00046 self.tag.tagIt()
00047
00048 return SUCCESS
00049
00050 def _fill_gen(self,gh):
00051 ge = gh.event()
00052 if not ge:
00053 self.error('Failed to get a GenEvent')
00054 return FAILURE
00055
00056 self.gen.number = ge.particles_size()
00057 for part in irange(ge.particles_begin(),ge.particles_end()):
00058 self.gen.pid.append(part.pdg_id())
00059 self.gen.mass.append(part.generated_mass())
00060 mom = part.momentum()
00061 mag = mom.mag()
00062 self.gen.energy.append(mom.e())
00063 vtx = part.production_vertex()
00064 self.gen.x.append(vtx.x())
00065 self.gen.y.append(vtx.y())
00066 self.gen.z.append(vtx.z())
00067 self.gen.dx.append(mom.x()/mag)
00068 self.gen.dy.append(mom.y()/mag)
00069 self.gen.dz.append(mom.z()/mag)
00070 return SUCCESS
00071
00072 def _fill_sim(self,sh):
00073 shh = sh.hits()
00074 hcmap = shh.hitCollection()
00075 for shcoll in irange(hcmap.begin(),hcmap.end()):
00076 detid = shcoll.first
00077 shcont = shcoll.second.collection()
00078
00079 for hit in irange(shcont.begin(),shcont.end()):
00080
00081 self.sim.detid.append(detid)
00082 self.sim.weight.append(hit.weight())
00083 self.sim.time.append(hit.time())
00084
00085 continue
00086 continue
00087 self.sim.nhits = len(self.sim.detid)
00088 return SUCCESS
00089
