Package iceprod :: Package modules :: Module sorting
[hide private]
[frames] | no frames]

Source Code for Module iceprod.modules.sorting

  1  from icecube import icetray, dataclasses, dataio 
  2  import sys 
  3  import logging 
  4  from I3Tray import I3Units 
  5  from icecube.dataclasses import I3Particle 
  6   
  7   
  8  primaries = [ 
  9       I3Particle.ParticleType.PPlus      , 
 10       I3Particle.ParticleType.PMinus     , 
 11       I3Particle.ParticleType.He4Nucleus , 
 12       I3Particle.ParticleType.Li7Nucleus , 
 13       I3Particle.ParticleType.Be9Nucleus , 
 14       I3Particle.ParticleType.B11Nucleus , 
 15       I3Particle.ParticleType.C12Nucleus , 
 16       I3Particle.ParticleType.N14Nucleus , 
 17       I3Particle.ParticleType.O16Nucleus , 
 18       I3Particle.ParticleType.F19Nucleus , 
 19       I3Particle.ParticleType.Ne20Nucleus, 
 20       I3Particle.ParticleType.Na23Nucleus, 
 21       I3Particle.ParticleType.Mg24Nucleus, 
 22       I3Particle.ParticleType.Al27Nucleus, 
 23       I3Particle.ParticleType.Si28Nucleus, 
 24       I3Particle.ParticleType.P31Nucleus , 
 25       I3Particle.ParticleType.S32Nucleus , 
 26       I3Particle.ParticleType.Cl35Nucleus, 
 27       I3Particle.ParticleType.Ar40Nucleus, 
 28       I3Particle.ParticleType.K39Nucleus , 
 29       I3Particle.ParticleType.Ca40Nucleus, 
 30       I3Particle.ParticleType.Sc45Nucleus, 
 31       I3Particle.ParticleType.Ti48Nucleus, 
 32       I3Particle.ParticleType.V51Nucleus , 
 33       I3Particle.ParticleType.Cr52Nucleus, 
 34       I3Particle.ParticleType.Mn55Nucleus, 
 35       I3Particle.ParticleType.Fe56Nucleus, 
 36       I3Particle.ParticleType.Gamma ] 
 37   
38 -def IsTrack(particle):
39 track = False 40 track = track or particle.GetShape() in [ 41 I3Particle.ParticleShape.InfiniteTrack, 42 I3Particle.ParticleShape.StartingTrack, 43 I3Particle.ParticleShape.StoppingTrack, 44 I3Particle.ParticleShape.ContainedTrack ] 45 46 track = track or particle.GetType() in [ 47 I3Particle.ParticleType.MuPlus, 48 I3Particle.ParticleType.MuMinus, 49 I3Particle.ParticleType.TauPlus, 50 I3Particle.ParticleType.TauMinus, 51 I3Particle.ParticleType.STauPlus, 52 I3Particle.ParticleType.STauMinus, 53 I3Particle.ParticleType.Monopole ] 54 55 track = track or (particle.GetShape() == I3Particle.ParticleShape.Primary and particle.GetType() in primaries) 56 return track 57
58 -def IsCascade(particle):
59 60 csc = False 61 csc = csc or particle.GetShape() == I3Particle.ParticleShape.Cascade 62 csc = csc or particle.GetType() in [ 63 I3Particle.ParticleType.EPlus, 64 I3Particle.ParticleType.EMinus, 65 I3Particle.ParticleType.Brems, 66 I3Particle.ParticleType.DeltaE, 67 I3Particle.ParticleType.PairProd, 68 I3Particle.ParticleType.NuclInt, 69 I3Particle.ParticleType.Hadrons, 70 I3Particle.ParticleType.Pi0, 71 I3Particle.ParticleType.PiPlus, 72 I3Particle.ParticleType.PiMinus] 73 74 csc = csc or (particle.GetShape() == Primary and particle.GetType() in primaries) 75 return csc 76 77
78 -class ZenithWriter(icetray.I3Module):
79
80 - def __init__(self, context):
81 icetray.I3Module.__init__(self, context) 82 self.context = context 83 self.logger = logging.getLogger("ZenithSorter") 84 85 self.configWritten = False 86 self.framecounter = 0 87 self.path = "physics.%02u.i3" 88 self.mcTreeName = "I3MCTree" 89 self.writeconfig = True 90 self.nbin = 1 91 self.zenmin = 0.0*I3Units.deg 92 self.zenmax = 89.0*I3Units.deg 93 self.ofs = [] 94 self.bins = [] 95 96 self.AddParameter("Filename","The file we'll write to.", self.path); 97 self.AddParameter("MCTreeName","Name of I3MCTree in frame", self.mcTreeName); 98 self.AddParameter("Bins","Number of zenith bins",self.nbin); 99 self.AddParameter("ZenithMin","Zenith angle lower limit",self.zenmin); 100 self.AddParameter("ZenithMax","Zenith angle upper limit",self.zenmax); 101 self.AddOutBox("OutBox");
102 103
104 - def Configure(self):
105 self.path = self.GetParameter("Filename") 106 self.mcTreeName = self.GetParameter("MCTreeName") 107 self.nbin = self.GetParameter("Bins") 108 self.zenmin = self.GetParameter("ZenithMin") 109 self.zenmax = self.GetParameter("ZenithMax") 110 111 # Compute binning of zenith angular distribution for events 112 dzen = (self.zenmax - self.zenmin)/float(self.nbin); 113 self.bins = map(lambda i: self.zenmin+i*dzen, range(self.nbin+1)); 114 115 print "Sorting in %u zenith bins with zenmin = %f, zenmax = %f" % (self.nbin, self.zenmin , self.zenmax); 116 for i in range(self.nbin): 117 filename = self.path % i; 118 self.ofs.append(dataio.I3File(filename, dataio.I3File.Mode.Writing))
119 120
121 - def Physics(self, frame):
122 123 print frame 124 isTrack = False; 125 zen = -1; 126 zenmin = 999; 127 zenmax = -999; 128 129 mcList = []; 130 particletype = dataclasses.I3Particle.ParticleType.unknown; 131 mcTree = frame.Get(self.mcTreeName); 132 133 self.framecounter += 1 134 frame.Put("FrameIndex", icetray.I3Int(self.framecounter)); 135 136 for cmctrack in mcTree.GetInIce(): 137 zen = cmctrack.GetDir().GetZenith(); 138 if cmctrack.GetType() > 0: 139 type = cmctrack.GetType(); 140 if zen == zen and ( IsTrack(cmctrack) or IsCascade(cmctrack) ) : 141 zenmin = min(zenmin,zen); 142 zenmax = max(zenmax,zen); 143 isTrack = IsTrack(cmctrack) or IsCascade(cmctrack); 144 145 if not isTrack and type != I3Particle.ParticleType.unknown : 146 self.logger.error("Unable to find Track or Cascade in particle list! ") 147 self.logger.error("particle type is %d" % type ) 148 149 for i in range(self.nbin): 150 if zenmin >= self.bins[i] and zenmin < self.bins[i+1]: 151 self.ofs[i].push(frame); 152 elif zenmax >= self.bins[i] and zenmax < self.bins[i+1]: 153 self.ofs[i].push(frame); 154 break; 155 156 if zenmin < self.zenmin or zenmax > self.zenmax : 157 self.logger.error("zenith %f outside range (%2.1f,%2.1f)" % (zen,self.bins[0],self.bins[self.nbin] )) 158 159 self.PushFrame(frame,"OutBox"); 160 self.logger.debug("physics... done")
161
162 - def Finish(self):
163 for file in self.ofs: file.close() 164 self.logger.info("FileMerger... done")
165 166
167 -class PhysicsFile(dataio.I3File):
168 current = None 169
170 - def pop_physics(self):
171 tmp = self.current 172 self.current = dataio.I3File.pop_physics(self) 173 return tmp
174
175 - def peak(self):
176 if not self.current: 177 self.current = dataio.I3File.pop_physics(self) 178 return self.current
179
180 -class FileMerger(icetray.I3Module):
181
182 - def __init__(self, context):
183 icetray.I3Module.__init__(self, context) 184 self.logger = logging.getLogger("FileMerger") 185 self.counter = 0 186 self.pattern = "physics.[0-9]*.i3" 187 188 self.AddParameter("FilePattern","The file we'll write to.", self.pattern); 189 self.AddParameter("FileNameList","The file we'll write to.", []); 190 self.AddParameter("IndexName","I3Int for merging","FrameIndex"); 191 self.AddParameter("HitMapName","Name of HitSeriesMap","MCHitSeriesMap"); 192 self.AddOutBox("OutBox");
193
194 - def Configure(self):
195 import glob 196 self.pattern = self.GetParameter("FilePattern") 197 filelist = self.GetParameter("FileNameList") 198 self.index = self.GetParameter("IndexName") 199 self.hitmap = self.GetParameter("HitMapName") 200 if filelist: 201 self.logger.warn("FileNameList was given. FilePattern will be ignored") 202 else: 203 filelist = glob.glob(self.pattern) 204 print "opening", filelist 205 self.ofs = map(lambda x: PhysicsFile(x, dataio.I3File.Mode.Reading), filelist)
206
207 - def Process(self):
208 209 indices = [ f.peak().Get(self.index).value for f in self.ofs if f.peak() ] 210 if not indices: 211 self.RequestSuspension() 212 return 213 214 next = min(indices) 215 frames_to_merge = [ f.pop_physics() for f in self.ofs if f.peak() and f.peak().Get(self.index).value == next ] 216 217 # grab the first frame in case they are degenerate 218 frame = frames_to_merge[0] 219 220 # Merge hits from degenerate frames 221 if frame.Has(self.hitmap): 222 hitmap = reduce(self.MergeHits, map(lambda x: x.Get(self.hitmap), frames_to_merge)) 223 frame.Delete(self.hitmap) 224 frame.Put(self.hitmap,hitmap) 225 frame.Delete(self.index) 226 227 if len(frames_to_merge) > 1: print "index",next, ", duplicate frames",len(frames_to_merge) 228 self.PushFrame(frame)
229 230
231 - def MergeHits(self,map1, map2):
232 233 if not map2: return map1 234 # Add the coinc hits 235 for omkey,hits in map2.items(): 236 if not map1.has_key(omkey): 237 map1[omkey] = hits; 238 else: 239 newSeries = hits; 240 for hit in hits: 241 map1[omkey].append(hit); 242 243 for omkey,hits in map1.items(): 244 hits = list(hits) 245 246 def compare_times(a,b): 247 if a > b:return 1 248 elif a == b:return 0 249 return -1
250 251 hits.sort(compare_times); 252 newhits = dataclasses.vector_I3MCHit() 253 for hit in hits: newhits.append(hit) 254 map1[omkey] = newhits 255 256 return map1
257
258 - def Finish(self):
259 for file in self.ofs: file.close()
260 261
262 -class HitFilter(icetray.I3Module):
263
264 - def __init__(self, context):
265 icetray.I3Module.__init__(self, context) 266 self.context = context 267 self.logger = logging.getLogger("sorting.HitFilter") 268 269 self.AddParameter("MCHitSeriesName","Name of HitSeriesMap","MCHitSeriesMap"); 270 self.AddParameter("Threshold","Minimun number of hits",1); 271 self.AddOutBox("OutBox");
272
273 - def Configure(self):
274 self.hitmap = self.GetParameter("MCHitSeriesName") 275 self.threshold = self.GetParameter("Threshold")
276
277 - def Physics(self, frame):
278 279 if len(frame.Get(self.hitmap).items()) < self.threshold: 280 self.logger.debug("Event does not contain enough hits") 281 return 282 283 self.PushFrame(frame,"OutBox"); 284 return
285