File indexing completed on 2024-04-06 12:33:24
0001 from builtins import range
0002 import math
0003 import collections
0004
0005 import ROOT
0006
0007 from Validation.RecoTrack.plotting.ntupleEnum import *
0008 from Validation.RecoTrack.plotting.ntupleEnum import _Enum
0009
0010 class _Collection(object):
0011 """Adaptor class representing a collection of objects.
0012
0013 Concrete collection classes should inherit from this class.
0014
0015 """
0016 def __init__(self, tree, sizeBranch, objclass):
0017 """Constructor.
0018
0019 Arguments:
0020 tree -- TTree object
0021 sizeBranch -- Name of the branch to be used in size()
0022 objclass -- Class to be used for the objects in __getitem__()
0023 """
0024 super(_Collection, self).__init__()
0025 self._tree = tree
0026 self._sizeBranch = sizeBranch
0027 self._objclass = objclass
0028
0029 def size(self):
0030 """Number of objects in the collection."""
0031 return int(getattr(self._tree, self._sizeBranch).size())
0032
0033 def __len__(self):
0034 """Number of objects in the collection."""
0035 return self.size()
0036
0037 def __getitem__(self, index):
0038 """Get object 'index' in the collection."""
0039 return self._objclass(self._tree, index)
0040
0041 def __iter__(self):
0042 """Returns generator for the objects."""
0043 for index in range(self.size()):
0044 yield self._objclass(self._tree, index)
0045
0046 class _Object(object):
0047 """Adaptor class representing a single object in a collection.
0048
0049 The member variables of the object are obtained from the branches
0050 with common prefix and a given index.
0051
0052 Concrete object classes should inherit from this class.
0053 """
0054 def __init__(self, tree, index, prefix):
0055 """Constructor.
0056
0057 Arguments:
0058 tree -- TTree object
0059 index -- Index for this object
0060 prefix -- Prefix of the branchs
0061 """
0062 super(_Object, self).__init__()
0063 self._tree = tree
0064 self._index = index
0065 self._prefix = prefix
0066
0067 def __getattr__(self, attr):
0068 """Return object member variable.
0069
0070 'attr' is translated as a branch in the TTree (<prefix>_<attr>).
0071 """
0072 self._checkIsValid()
0073 val = getattr(self._tree, self._prefix+"_"+attr)[self._index]
0074 return lambda: val
0075
0076 def _checkIsValid(self):
0077 """Raise an exception if the object index is not valid."""
0078 if not self.isValid():
0079 raise Exception("%s is not valid" % self.__class__.__name__)
0080
0081 def isValid(self):
0082 """Check if object index is valid."""
0083 return self._index != -1
0084
0085 def index(self):
0086 """Return object index."""
0087 return self._index
0088
0089 class _DetIdStrAdaptor(object):
0090 """Adaptor class for objects containgin DetId (hits)."""
0091 def __init__(self):
0092 super(_DetIdStrAdaptor, self).__init__()
0093
0094 def layerStr(self):
0095 """Returns a string describing the layer of the hit."""
0096 self._checkIsValid()
0097 get = lambda name: getattr(self._tree, self._prefix+"_"+name)[self._index]
0098 subdet = get("subdet")
0099 side = ""
0100 isPhase2OTBarrel = (subdet == SubDet.TOB and hasattr(self._tree, self._prefix+"_isLower"))
0101 if subdet in [SubDet.FPix, SubDet.TID, SubDet.TEC] or isPhase2OTBarrel:
0102 sideNum = get("side")
0103 if sideNum == 1:
0104 side = "-"
0105 elif sideNum == 2:
0106 side = "+"
0107 elif isPhase2OTBarrel and sideNum == 3:
0108 side = ""
0109 else:
0110 side = "?"
0111 return "%s%d%s" % (SubDet.toString(subdet),
0112 getattr(self._tree, self._prefix+"_layer")[self._index],
0113 side)
0114
0115 def detIdStr(self):
0116 """Returns a string describing the DetId fields."""
0117 self._checkIsValid
0118 get = lambda name: getattr(self._tree, self._prefix+"_"+name)[self._index]
0119 isPhase2 = hasattr(self._tree, self._prefix+"_isLower")
0120 def stereo():
0121 if isPhase2:
0122 if get("isLower"):
0123 return " isLower"
0124 if get("isUpper"):
0125 return " isUpper"
0126 if get("isStack"):
0127 return " isStack"
0128 else:
0129 if get("isStereo"):
0130 return " isStereo"
0131 if get("isRPhi"):
0132 return " isRPhi"
0133 if get("isGlued"):
0134 return " isGlued"
0135 return ""
0136
0137 subdet = get("subdet")
0138 if subdet == SubDet.BPix:
0139 return "ladder {} module {}".format(get("ladder"), get("module"))
0140 if subdet == SubDet.FPix:
0141 return "blade {} panel {} module {}".format(get("blade"), get("panel"), get("module"))
0142 if subdet == SubDet.TIB:
0143 return "side {} order {} string {} module {}{}".format(get("side"), get("order"), get("string"), get("module"), stereo())
0144 if subdet == SubDet.TID:
0145 return "ring {} order {} module {}{}".format(get("ring"), get("order"), get("module"), stereo())
0146 if subdet == SubDet.TOB:
0147 if isPhase2:
0148 return "rod {} module {}{}".format(get("rod"), get("module"), stereo())
0149 else:
0150 return "side {} rod {} module {}{}".format(get("side"), get("rod"), get("module"), stereo())
0151 if subdet == SubDet.TEC:
0152 return "order {} petal {} ring {} module {}{}".format(get("order"), get("petalNumber"), get("ring"), get("module"), stereo())
0153 raise Exception("Unknown subdet %d" % subdet)
0154
0155 class _HitObject(_Object, _DetIdStrAdaptor):
0156 """Adaptor class for pixel/strip hit objects."""
0157 def __init__(self, tree, index, prefix):
0158 """Constructor.
0159
0160 Arguments:
0161 tree -- TTree object
0162 index -- Index for this object
0163 prefix -- Prefix of the branchs
0164 """
0165 """Constructor
0166 """
0167 super(_HitObject, self).__init__(tree, index, prefix)
0168
0169 def ntracks(self):
0170 """Returns number of tracks containing this hit."""
0171 self._checkIsValid()
0172 return getattr(self._tree, self._prefix+"_trkIdx")[self._index].size()
0173
0174 def tracks(self):
0175 """Returns generator for tracks containing this hit.
0176
0177 The generator returns Track objects
0178 """
0179 self._checkIsValid()
0180 for itrack in getattr(self._tree, self._prefix+"_trkIdx")[self._index]:
0181 yield Track(self._tree, itrack)
0182
0183 def nseeds(self):
0184 """Returns number of seeds containing this hit."""
0185 self._checkIsValid()
0186 return getattr(self._tree, self._prefix+"_seeIdx")[self._index].size()
0187
0188 def seeds(self):
0189 """Returns generator for tracks containing this hit.
0190
0191 The generator returns Seed objects
0192 """
0193 self._checkIsValid()
0194 for iseed in getattr(self._tree, self._prefix+"_seeIdx")[self._index]:
0195 yield Seed(self._tree, iseed)
0196
0197 def r(self):
0198 return math.sqrt(self.x()**2 + self.y()**2)
0199
0200 def r3D(self):
0201 return math.sqrt(self.x()**2 + self.y()**2 + self.z()**2)
0202
0203 class _RecoHitAdaptor(object):
0204 """Adaptor class for objects containing hits (e.g. tracks)"""
0205 def __init__(self):
0206 super(_RecoHitAdaptor, self).__init__()
0207
0208 def _hits(self):
0209 """Internal method to generate pairs of hit index and type."""
0210 for ihit, hitType in zip(self.hitIdx(), self.hitType()):
0211 yield (ihit, hitType)
0212
0213 def hits(self):
0214 """Returns generator for hits.
0215
0216 Generator returns PixelHit/StripHit/GluedHit/Phase2OT depending on the
0217 hit type.
0218
0219 """
0220 for ihit, hitType in self._hits():
0221 if hitType == 0:
0222 yield PixelHit(self._tree, ihit)
0223 elif hitType == 1:
0224 yield StripHit(self._tree, ihit)
0225 elif hitType == 2:
0226 yield GluedHit(self._tree, ihit)
0227 elif hitType == 3:
0228 yield InvalidHit(self._tree, ihit)
0229 elif hitType == 4:
0230 yield Phase2OTHit(self._tree, ihit)
0231 else:
0232 raise Exception("Unknown hit type %d" % hitType)
0233
0234 def pixelHits(self):
0235 """Returns generator for pixel hits."""
0236 self._checkIsValid()
0237 for ihit, hitType in self._hits():
0238 if hitType != 0:
0239 continue
0240 yield PixelHit(self._tree, ihit)
0241
0242 def stripHits(self):
0243 """Returns generator for strip hits."""
0244 self._checkIsValid()
0245 for ihit, hitType in self._hits():
0246 if hitType != 1:
0247 continue
0248 yield StripHit(self._tree, ihit)
0249
0250 def gluedHits(self):
0251 """Returns generator for matched strip hits."""
0252 self._checkIsValid()
0253 for ihit, hitType in self._hits():
0254 if hitType != 2:
0255 continue
0256 yield GluedHit(self._tree, ihit)
0257
0258 def invalidHits(self):
0259 """Returns generator for invalid hits."""
0260 self._checkIsValid()
0261 for ihit, hitType in self._hits():
0262 if hitType != 3:
0263 continue
0264 yield InvalidHit(self._tree, ihit)
0265
0266 def phase2OTHits(self):
0267 """Returns generator for phase2 outer tracker hits."""
0268 self._checkIsValid()
0269 for ihit, hitType in self._hits():
0270 if hitType != 4:
0271 continue
0272 yield Phase2OTHit(self._tree, ihit)
0273
0274 class _SimHitMatchAdaptor(object):
0275 """Adaptor class for objects containing or matched to SimHits (e.g. reco hits)."""
0276 def __init__(self):
0277 super(_SimHitMatchAdaptor, self).__init__()
0278
0279 def _nMatchedSimHits(self):
0280 """Internal method to get the number of matched SimHits."""
0281 return getattr(self._tree, self._prefix+"_simHitIdx")[self._index].size()
0282
0283 def nMatchedSimHits(self):
0284 """Returns the number of matched SimHits."""
0285 self._checkIsValid()
0286 return self._nMatchedSimHits()
0287
0288 def matchedSimHitInfos(self):
0289 """Returns a generator for matched SimHits.
0290
0291 The generator returns SimHitMatchInfo objects.
0292 """
0293 self._checkIsValid()
0294 for imatch in range(self._nMatchedSimHits()):
0295 yield SimHitMatchInfo(self._tree, self._index, imatch, self._prefix)
0296
0297 class _TrackingParticleMatchAdaptor(object):
0298 """Adaptor class for objects matched to TrackingParticles."""
0299 def __init__(self):
0300 super(_TrackingParticleMatchAdaptor, self).__init__()
0301
0302 def _nMatchedTrackingParticles(self):
0303 """Internal method to get the number of matched TrackingParticles."""
0304 return getattr(self._tree, self._prefix+"_simTrkIdx")[self._index].size()
0305
0306 def nMatchedTrackingParticles(self):
0307 """Returns the number of matched TrackingParticles."""
0308 self._checkIsValid()
0309 return self._nMatchedTrackingParticles()
0310
0311 def matchedTrackingParticleInfos(self):
0312 """Returns a generator for matched TrackingParticles.
0313
0314 The generator returns TrackingParticleMatchInfo objects.
0315
0316 """
0317 self._checkIsValid()
0318 for imatch in range(self._nMatchedTrackingParticles()):
0319 yield TrackingParticleMatchInfo(self._tree, self._index, imatch, self._prefix)
0320
0321 def bestMatchingTrackingParticle(self):
0322 """Returns best-matching TrackingParticle, even for fake tracks, or None if there is no best-matching TrackingParticle.
0323
0324 Best-matching is defined as the one with largest number of
0325 hits matched to the hits of a track (>= 3). If there are many
0326 fulfilling the same number of hits, the one inducing the
0327 innermost hit of the track is chosen.
0328 """
0329 idx = self.bestSimTrkIdx()
0330 if idx < 0:
0331 return None
0332 return TrackingParticle(self._tree, idx)
0333
0334 def bestMatchingTrackingParticleShareFrac(self):
0335 """Fraction of shared hits with reco hits as denominator for best-matching TrackingParticle."""
0336 return self.bestSimTrkShareFrac()
0337
0338 def bestMatchingTrackingParticleShareFracSimDenom(self):
0339 """Fraction of shared hits with TrackingParticle::numberOfTrackerHits() as denominator for best-matching TrackingParticle."""
0340 return self.bestSimTrkShareFracSimDenom()
0341
0342 def bestMatchingTrackingParticleShareFracSimClusterDenom(self):
0343 """Fraction of shared hits with number of reco clusters associated to a TrackingParticle as denominator for best-matching TrackingParticle."""
0344 return self.bestSimTrkShareFracSimClusterDenom()
0345
0346 def bestMatchingTrackingParticleNormalizedChi2(self):
0347 """Normalized chi2 calculated from track parameters+covariance matrix and TrackingParticle parameters for best-matching TrackingParticle."""
0348 return self.bestSimTrkNChi2()
0349
0350 def bestMatchingTrackingParticleFromFirstHit(self):
0351 """Returns best-matching TrackingParticle, even for fake tracks, or None if there is no best-matching TrackingParticle.
0352
0353 Best-matching is defined as the one with largest number of
0354 hits matched to the hits of a track (>= 3) starting from the
0355 beginning of the track. If there are many fulfilling the same
0356 number of hits, "a first TP" is chosen (a bit arbitrary, but
0357 should be rare".
0358 """
0359 idx = self.bestFromFirstHitSimTrkIdx()
0360 if idx < 0:
0361 return None
0362 return TrackingParticle(self._tree, idx)
0363
0364 def bestMatchingTrackingParticleFromFirstHitShareFrac(self):
0365 """Fraction of shared hits with reco hits as denominator for best-matching TrackingParticle starting from the first hit of a track."""
0366 return self.bestFromFirstHitSimTrkShareFrac()
0367
0368 def bestMatchingTrackingParticleFromFirstHitShareFracSimDenom(self):
0369 """Fraction of shared hits with TrackingParticle::numberOfTrackerHits() as denominator for best-matching TrackingParticle starting from the first hit of a track."""
0370 return self.bestFromFirstHitSimTrkShareFracSimDenom()
0371
0372 def bestMatchingTrackingParticleFromFirstHitShareFracSimClusterDenom(self):
0373 """Fraction of shared hits with number of reco clusters associated to a TrackingParticle as denominator for best-matching TrackingParticle starting from the first hit of a track."""
0374 return self.bestFromFirstHitSimTrkShareFracSimClusterDenom()
0375
0376 def bestMatchingTrackingParticleFromFirstHitNormalizedChi2(self):
0377 """Normalized chi2 calculated from track parameters+covariance matrix and TrackingParticle parameters for best-matching TrackingParticle starting from the first hit of a track."""
0378 return self.bestFromFirstHitSimTrkNChi2()
0379
0380
0381 class TrackingNtuple(object):
0382 """Class abstracting the whole ntuple/TTree.
0383
0384 Main benefit is to provide nice interface for
0385 - iterating over events
0386 - querying whether hit/seed information exists
0387
0388 Note that to iteratate over the evets with zip(), you should use
0389 itertools.izip() instead.
0390 """
0391 def __init__(self, fileName, tree="trackingNtuple/tree"):
0392 """Constructor.
0393
0394 Arguments:
0395 fileName -- String for path to the ROOT file
0396 tree -- Name of the TTree object inside the ROOT file (default: 'trackingNtuple/tree')
0397 """
0398 super(TrackingNtuple, self).__init__()
0399 self._file = ROOT.TFile.Open(fileName)
0400 self._tree = self._file.Get(tree)
0401 self._entries = self._tree.GetEntriesFast()
0402
0403 def file(self):
0404 return self._file
0405
0406 def tree(self):
0407 return self._tree
0408
0409 def nevents(self):
0410 return self._entries
0411
0412 def hasHits(self):
0413 """Returns true if the ntuple has hit information."""
0414 return hasattr(self._tree, "pix_isBarrel")
0415
0416 def hasSeeds(self):
0417 """Returns true if the ntuple has seed information."""
0418 return hasattr(self._tree, "see_fitok")
0419
0420 def __iter__(self):
0421 """Returns generator for iterating over TTree entries (events)
0422
0423 Generator returns Event objects.
0424
0425 """
0426 for jentry in range(self._entries):
0427
0428 ientry = self._tree.LoadTree( jentry )
0429 if ientry < 0: break
0430
0431 nb = self._tree.GetEntry( jentry )
0432 if nb <= 0: continue
0433
0434 yield Event(self._tree, jentry)
0435
0436 def getEvent(self, index):
0437 """Returns Event for a given index"""
0438 ientry = self._tree.LoadTree(index)
0439 if ientry < 0: return None
0440 nb = self._tree.GetEntry(ientry)
0441 if nb <= 0: None
0442
0443 return Event(self._tree, ientry)
0444
0445
0446 class Event(object):
0447 """Class abstracting a single event.
0448
0449 Main benefit is to provide nice interface to get various objects
0450 or collections of objects.
0451 """
0452 def __init__(self, tree, entry):
0453 """Constructor.
0454
0455 Arguments:
0456 tree -- TTree object
0457 entry -- Entry number in the tree
0458 """
0459 super(Event, self).__init__()
0460 self._tree = tree
0461 self._entry = entry
0462
0463 def entry(self):
0464 return self._entry
0465
0466 def event(self):
0467 """Returns event number."""
0468 return self._tree.event
0469
0470 def lumi(self):
0471 """Returns lumisection number."""
0472 return self._tree.lumi
0473
0474 def run(self):
0475 """Returns run number."""
0476 return self._tree.run
0477
0478 def eventId(self):
0479 """Returns (run, lumi, event) tuple."""
0480 return (self._tree.run, self._tree.lumi, self._tree.event)
0481
0482 def eventIdStr(self):
0483 """Returns 'run:lumi:event' string."""
0484 return "%d:%d:%d" % self.eventId()
0485
0486 def beamspot(self):
0487 """Returns BeamSpot object."""
0488 return BeamSpot(self._tree)
0489
0490 def tracks(self):
0491 """Returns Tracks object."""
0492 return Tracks(self._tree)
0493
0494 def pixelHits(self):
0495 """Returns PixelHits object."""
0496 return PixelHits(self._tree)
0497
0498 def stripHits(self):
0499 """Returns StripHits object."""
0500 return StripHits(self._tree)
0501
0502 def gluedHits(self):
0503 """Returns GluedHits object."""
0504 return GluedHits(self._tree)
0505
0506 def phase2OTHits(self):
0507 """Returns Phase2OTHits object."""
0508 return Phase2OTHits(self._tree)
0509
0510 def seeds(self):
0511 """Returns Seeds object."""
0512 return Seeds(self._tree)
0513
0514 def trackingParticles(self):
0515 """Returns TrackingParticles object."""
0516 return TrackingParticles(self._tree)
0517
0518 def vertices(self):
0519 """Returns Vertices object."""
0520 return Vertices(self._tree)
0521
0522 def trackingVertices(self):
0523 """Returns TrackingVertices object."""
0524 return TrackingVertices(self._tree)
0525
0526
0527 class BeamSpot(object):
0528 """Class representing the beam spot."""
0529 def __init__(self, tree):
0530 """Constructor.
0531
0532 Arguments:
0533 tree -- TTree object
0534 """
0535 super(BeamSpot, self).__init__()
0536 self._tree = tree
0537 self._prefix = "bsp"
0538
0539 def __getattr__(self, attr):
0540 """Return object member variable.
0541
0542 'attr' is translated as a branch in the TTree (bsp_<attr>).
0543 """
0544 val = getattr(self._tree, self._prefix+"_"+attr)
0545 return lambda: val
0546
0547
0548 class SimHitMatchInfo(_Object):
0549 """Class representing a match to a SimHit.
0550
0551 The point of this class is to provide, in addition to the matched
0552 SimHit, also other information about the match (e.g. fraction of
0553 charge from this SimHit).
0554 """
0555 def __init__(self, tree, index, shindex, prefix):
0556 """Constructor.
0557
0558 Arguments:
0559 tree -- TTree object
0560 index -- Index of the hit matched to SimHit
0561 shindex -- Index of the SimHit match (second index in _simHitIdx branch)
0562 prefix -- String for prefix of the object (track/seed/hit) matched to TrackingParticle
0563 """
0564 super(SimHitMatchInfo, self).__init__(tree, index, prefix)
0565 self._shindex = shindex
0566
0567 def __getattr__(self, attr):
0568 """Custom __getattr__ because of the second index needed to access the branch."""
0569 val = super(SimHitMatchInfo, self).__getattr__(attr)()[self._shindex]
0570 return lambda: val
0571
0572 def simHit(self):
0573 """Returns matched SimHit."""
0574 self._checkIsValid()
0575 return SimHit(self._tree, getattr(self._tree, self._prefix+"_simHitIdx")[self._index][self._shindex])
0576
0577 class TrackingParticleMatchInfo(_Object):
0578
0579 """Class representing a match to a TrackingParticle.
0580
0581 The point of this class is to provide, in addition to the matched
0582 TrackingParticle, also other information about the match (e.g.
0583 shared hit fraction for tracks/seeds).
0584 """
0585 def __init__(self, tree, index, tpindex, prefix):
0586 """Constructor.
0587
0588 Arguments:
0589 tree -- TTree object
0590 index -- Index of the object (track/seed) matched to TrackingParticle
0591 tpindex -- Index of the TrackingParticle match (second index in _simTrkIdx branch)
0592 prefix -- String for prefix of the object (track/seed) matched to TrackingParticle
0593 """
0594 super(TrackingParticleMatchInfo, self).__init__(tree, index, prefix)
0595 self._tpindex = tpindex
0596
0597 def __getattr__(self, attr):
0598 """Custom __getattr__ because of the second index needed to access the branch.
0599
0600 Note that when mapping the 'attr' to a branch, a 'simTrk' is
0601 prepended and the first letter of 'attr' is turned to upper
0602 case.
0603 """
0604 val = super(TrackingParticleMatchInfo, self).__getattr__("simTrk"+attr[0].upper()+attr[1:])()[self._tpindex]
0605 return lambda: val
0606
0607 def trackingParticle(self):
0608 """Returns matched TrackingParticle."""
0609 return TrackingParticle(self._tree, self.idx())
0610
0611 class TrackMatchInfo(_Object):
0612 """Class representing a match to a Track.
0613
0614 The point of this class is to provide, in addition to the matched
0615 Track, also other information about the match (e.g. shared hit fraction.
0616 """
0617 def __init__(self, tree, index, trkindex, prefix):
0618 """Constructor.
0619
0620 Arguments:
0621 tree -- TTree object
0622 index -- Index of the object (TrackingParticle) matched to track
0623 trkindex -- Index of the track match (second index in _trkIdx branch)
0624 prefix -- String for prefix of the object (TrackingParticle) matched to track
0625 """
0626 super(TrackMatchInfo, self).__init__(tree, index, prefix)
0627 self._trkindex = trkindex
0628
0629 def __getattr__(self, attr):
0630 """Custom __getattr__ because of the second index needed to access the branch.
0631
0632 Note that when mapping the 'attr' to a branch, a 'trk' is
0633 prepended and the first letter of 'attr' is turned to upper
0634 case.
0635 """
0636 val = super(TrackMatchInfo, self).__getattr__("trk"+attr[0].upper()+attr[1:])()[self._trkindex]
0637 return lambda: val
0638
0639 def track(self):
0640 """Returns matched Track."""
0641 return Track(self._tree, self.idx())
0642
0643 class SeedMatchInfo(_Object):
0644 """Class representing a match to a Seed.
0645
0646 The point of this class is to provide an interface compatible with
0647 all other "MatchInfo" classes
0648
0649 """
0650 def __init__(self, tree, index, seedindex, prefix):
0651 """Constructor.
0652
0653 Arguments:
0654 tree -- TTree object
0655 index -- Index of the object (TrackingParticle) matched to seed
0656 seedindex -- Index of the seed match (second index in _trkIdx branch)
0657 prefix -- String for prefix of the object (TrackingParticle) matched to seed
0658 """
0659 super(SeedMatchInfo, self).__init__(tree, index, prefix)
0660 self._seedindex = seedindex
0661
0662 def seed(self):
0663 """Returns matched Seed."""
0664 self._checkIsValid()
0665 return Seed(self._tree, getattr(self._tree, self._prefix+"_seedIdx")[self._index][self._seedindex])
0666
0667
0668 class Track(_Object, _RecoHitAdaptor, _TrackingParticleMatchAdaptor):
0669 """Class presenting a track."""
0670 def __init__(self, tree, index):
0671 """Constructor.
0672
0673 Arguments:
0674 tree -- TTree object
0675 index -- Index of the track
0676 """
0677 super(Track, self).__init__(tree, index, "trk")
0678
0679 def seed(self):
0680 """Returns Seed of the track."""
0681 self._checkIsValid()
0682 return Seed(self._tree, self._tree.trk_seedIdx[self._index])
0683
0684 def vertex(self):
0685 """Returns Vertex that used this track in its fit."""
0686 self._checkIsValid()
0687 return Vertex(self._tree, self._tree.trk_vtxIdx[self._index])
0688
0689 def ptPull(self):
0690 tp = self.bestMatchingTrackingParticle()
0691 if tp is None:
0692 return None
0693 return (self.pt() - tp.pca_pt())/self.ptErr()
0694
0695 def thetaPull(self):
0696 tp = self.bestMatchingTrackingParticle()
0697 if tp is None:
0698 return None
0699 return (getattr(self, "lambda")() - tp.pca_lambda())/self.lambdaErr()
0700
0701 def phiPull(self):
0702 tp = self.bestMatchingTrackingParticle()
0703 if tp is None:
0704 return None
0705 return (self.phi() - tp.pca_phi())/self.phiErr()
0706
0707 def dxyPull(self):
0708 tp = self.bestMatchingTrackingParticle()
0709 if tp is None:
0710 return None
0711 return (self.dxy() - tp.pca_dxy())/self.dxyErr()
0712
0713 def dzPull(self):
0714 tp = self.bestMatchingTrackingParticle()
0715 if tp is None:
0716 return None
0717 return (self.dz() - tp.pca_dz())/self.dzErr()
0718
0719 class Tracks(_Collection):
0720 """Class presenting a collection of tracks."""
0721 def __init__(self, tree):
0722 """Constructor.
0723
0724 Arguments:
0725 tree -- TTree object
0726 """
0727 super(Tracks, self).__init__(tree, "trk_pt", Track)
0728
0729
0730 class PixelHit(_HitObject, _SimHitMatchAdaptor):
0731 """Class representing a pixel hit."""
0732 def __init__(self, tree, index):
0733 """Constructor.
0734
0735 Arguments:
0736 tree -- TTree object
0737 index -- Index of the hit
0738 """
0739 super(PixelHit, self).__init__(tree, index, "pix")
0740
0741 def isValidHit(self):
0742 return True
0743
0744 class PixelHits(_Collection):
0745 """Class presenting a collection of pixel hits."""
0746 def __init__(self, tree):
0747 """Constructor.
0748
0749 Arguments:
0750 tree -- TTree object
0751 """
0752 super(PixelHits, self).__init__(tree, "pix_isBarrel", PixelHit)
0753
0754
0755 class StripHit(_HitObject, _SimHitMatchAdaptor):
0756 """Class representing a strip hit."""
0757 def __init__(self, tree, index):
0758 """Constructor.
0759
0760 Arguments:
0761 tree -- TTree object
0762 index -- Index of the hit
0763 """
0764 super(StripHit, self).__init__(tree, index, "str")
0765
0766 def isValidHit(self):
0767 return True
0768
0769 class StripHits(_Collection):
0770 """Class presenting a collection of strip hits."""
0771 def __init__(self, tree):
0772 """Constructor.
0773
0774 Arguments:
0775 tree -- TTree object
0776 """
0777 super(StripHits, self).__init__(tree, "str_isBarrel", StripHit)
0778
0779
0780 class GluedHit(_Object, _DetIdStrAdaptor):
0781 """Class representing a matched strip hit."""
0782 def __init__(self, tree, index):
0783 """Constructor.
0784
0785 Arguments:
0786 tree -- TTree object
0787 index -- Index of the hit
0788 """
0789 super(GluedHit, self).__init__(tree, index, "glu")
0790
0791 def isValidHit(self):
0792 return True
0793
0794 def monoHit(self):
0795 """Returns a StripHit for the mono hit."""
0796 self._checkIsValid()
0797 return StripHit(self._tree, self._tree.glu_monoIdx[self._index])
0798
0799 def stereoHit(self):
0800 """Returns a StripHit for the stereo hit."""
0801 self._checkIsValid()
0802 return StripHit(self._tree, self._tree.glu_stereoIdx[self._index])
0803
0804 def nseeds(self):
0805 """Returns the number of seeds containing this hit."""
0806 self._checkIsValid()
0807 return self._tree.glu_seeIdx[self._index].size()
0808
0809 def seeds(self):
0810 """Returns generator for seeds containing this hit.
0811
0812 The generator returns Seed objects
0813 """
0814 self._checkIsValid()
0815 for iseed in self._tree.glu_seeIdx[self._index]:
0816 yield Seed(self._tree, iseed)
0817
0818 class GluedHits(_Collection):
0819 """Class presenting a collection of matched strip hits."""
0820 def __init__(self, tree):
0821 """Constructor.
0822
0823 Arguments:
0824 tree -- TTree object
0825 """
0826 super(GluedHits, self).__init__(tree, "glu_isBarrel", GluedHit)
0827
0828
0829 class InvalidHit(_Object, _DetIdStrAdaptor):
0830
0831 Type = _Enum(
0832 missing = 1,
0833 inactive = 2,
0834 bad = 3,
0835 missing_inner = 4,
0836 missing_outer = 5
0837 )
0838
0839 """Class representing an invalid hit."""
0840 def __init__(self, tree, index):
0841 """Constructor.
0842
0843 Arguments:
0844 tree -- TTree object
0845 index -- Index of the hit
0846 """
0847 super(InvalidHit, self).__init__(tree, index, "inv")
0848
0849 def isValidHit(self):
0850 return False
0851
0852 def layerStr(self):
0853 """Returns a string describing the layer of the hit."""
0854 invalid_type = self._tree.inv_type[self._index]
0855 return super(InvalidHit, self).layerStr() + " (%s)"%InvalidHit.Type.toString(invalid_type)
0856
0857
0858 class Phase2OTHit(_HitObject, _SimHitMatchAdaptor):
0859 """Class representing a phase2 OT hit."""
0860 def __init__(self, tree, index):
0861 """Constructor.
0862
0863 Arguments:
0864 tree -- TTree object
0865 index -- Index of the hit
0866 """
0867 super(Phase2OTHit, self).__init__(tree, index, "ph2")
0868
0869 def isValidHit(self):
0870 return True
0871
0872 class Phase2OTHits(_Collection):
0873 """Class presenting a collection of phase2 OT hits."""
0874 def __init__(self, tree):
0875 """Constructor.
0876
0877 Arguments:
0878 tree -- TTree object
0879 """
0880 super(Phase2OTHits, self).__init__(tree, "ph2_isBarrel", Phase2OTHit)
0881
0882
0883 def _seedOffsetForAlgo(tree, algo):
0884 """Internal function for returning a pair of indices for the beginning of seeds of a given 'algo', and the one-beyond-last index of the seeds."""
0885 for ioffset, offset in enumerate(tree.see_offset):
0886 if tree.see_algo[offset] == algo:
0887 next_offset = tree.see_offset[ioffset+1] if ioffset < tree.see_offset.size()-1 else tree.see_algo.size()
0888 return (offset, next_offset)
0889 return (-1, -1)
0890
0891 class Seed(_Object, _RecoHitAdaptor, _TrackingParticleMatchAdaptor):
0892 """Class presenting a seed."""
0893 def __init__(self, tree, index):
0894 """Constructor.
0895
0896 Arguments:
0897 tree -- TTree object
0898 index -- Index of the seed
0899 """
0900 super(Seed, self).__init__(tree, index, "see")
0901
0902 def indexWithinAlgo(self):
0903 """Returns the seed index within the seeds of the same algo.
0904
0905 In case of errors, -1 is returned.
0906 """
0907 self._checkIsValid()
0908 algo = self._tree.see_algo[self._index]
0909 (offset, next_offset) = _seedOffsetForAlgo(self._tree, algo)
0910 if offset == -1:
0911 return -1
0912 return self._index - offset
0913
0914 def track(self):
0915 """Returns Track that was made from this seed."""
0916 self._checkIsValid()
0917 return Track(self._tree, self._tree.see_trkIdx[self._index])
0918
0919 class Seeds(_Collection):
0920 """Class presenting a collection of seeds."""
0921 def __init__(self, tree):
0922 """Constructor.
0923
0924 Arguments:
0925 tree -- TTree object
0926 """
0927 super(Seeds, self).__init__(tree, "see_pt", Seed)
0928
0929 def nSeedsForAlgo(self, algo):
0930 """Returns the number of seeds for a given 'algo'."""
0931 (offset, next_offset) = _seedOffsetForAlgo(self._tree, algo)
0932 return next_offset - offset
0933
0934 def seedsForAlgo(self, algo):
0935 """Returns generator iterating over the seeds of a given 'algo'.
0936
0937 Generator returns Seed object.
0938 """
0939 (offset, next_offset) = _seedOffsetForAlgo(self._tree, algo)
0940 for isee in range(offset, next_offset):
0941 yield Seed(self._tree, isee)
0942
0943 def seedForAlgo(self, algo, iseed):
0944 """Returns Seed of index 'iseed' for 'algo'."""
0945 (offset, next_offset) = _seedOffsetForAlgo(self._tree, algo)
0946 if iseed >= (next_offset-offset):
0947 raise Exception("Seed index %d is larger than the number of seeds %d for algo %d (%s)" % (iseed, next_offset-offset, algo, Algo.toString(algo)))
0948 return Seed(self._tree, offset+iseed)
0949
0950
0951 class SimHit(_Object, _DetIdStrAdaptor, _RecoHitAdaptor):
0952 """Class representing a SimHit which has not induced a RecHit."""
0953 def __init__(self, tree, index):
0954 """Constructor.
0955
0956 Arguments:
0957 tree -- TTree object
0958 index -- Index of the SimHit
0959 """
0960 super(SimHit, self).__init__(tree, index, "simhit")
0961
0962 def nRecHits(self):
0963 self._checkIsValid()
0964 return self._tree.simhit_hitIdx[self._index].size()
0965
0966 def trackingParticle(self):
0967 self._checkIsValid()
0968 return TrackingParticle(self._tree, getattr(self._tree, self._prefix+"_simTrkIdx")[self._index])
0969
0970
0971 class TrackingParticle(_Object):
0972 """Class representing a TrackingParticle."""
0973 def __init__(self, tree, index):
0974 """Constructor.
0975
0976 Arguments:
0977 tree -- TTree object
0978 index -- Index of the TrackingParticle
0979 """
0980 super(TrackingParticle, self).__init__(tree, index, "sim")
0981
0982 def _nMatchedTracks(self):
0983 """Internal function to get the number of matched tracks."""
0984 return self._tree.sim_trkIdx[self._index].size()
0985
0986 def nMatchedTracks(self):
0987 """Returns the number of matched tracks."""
0988 self._checkIsValid()
0989 return self._nMatchedTracks()
0990
0991 def matchedTrackInfos(self):
0992 """Returns a generator for matched tracks.
0993
0994 The generator returns TrackMatchInfo objects.
0995 """
0996 self._checkIsValid()
0997 for imatch in range(self._nMatchedTracks()):
0998 yield TrackMatchInfo(self._tree, self._index, imatch, self._prefix)
0999
1000 def bestMatchingTrack(self):
1001 """Returns best-matching track, even for non-reconstructed TrackingParticles, or None, if there is no best-matching track.
1002
1003 Best-matching is defined as the one with largest number of
1004 hits matched to the hits of a TrackingParticle (>= 3). If
1005 there are many fulfilling the same number of hits, the one
1006 inducing the innermost hit of the TrackingParticle is chosen.
1007 """
1008 self._checkIsValid()
1009 if self._nMatchedTracks() == 1:
1010 return next(self.matchedTrackInfos()).track()
1011
1012 tracks = collections.OrderedDict()
1013 for hit in self.simHits():
1014 for recHit in hit.hits():
1015 for track in recHit.tracks():
1016 if track.index() in tracks:
1017 tracks[track.index()] += 1
1018 else:
1019 tracks[track.index()] = 1
1020
1021 best = (None, 2)
1022 for trackIndex, nhits in tracks.items():
1023 if nhits > best[1]:
1024 best = (trackIndex, nhits)
1025 if best[0] is None:
1026 return None
1027 return Tracks(self._tree)[best[0]]
1028
1029 def _nMatchedSeeds(self):
1030 """Internal function to get the number of matched seeds."""
1031 return self._tree.sim_seedIdx[self._index].size()
1032
1033 def nMatchedSeeds(self):
1034 """Returns the number of matched seeds."""
1035 self._checkIsValid()
1036 return self._nMatchedSeeds()
1037
1038 def matchedSeedInfos(self):
1039 """Returns a generator for matched tracks.
1040
1041 The generator returns SeedMatchInfo objects.
1042 """
1043 self._checkIsValid()
1044 for imatch in range(self._nMatchedSeeds()):
1045 yield SeedMatchInfo(self._tree, self._index, imatch, self._prefix)
1046
1047 def nSimHits(self):
1048 self._checkIsValid()
1049 return self.simHitIdx().size()
1050
1051 def simHits(self):
1052 """Returns generator for SimHits."""
1053 self._checkIsValid()
1054 for ihit in self.simHitIdx():
1055 yield SimHit(self._tree, ihit)
1056
1057 def parentVertex(self):
1058 """Returns the parent TrackingVertex."""
1059 self._checkIsValid()
1060 return TrackingVertex(self._tree, self._tree.sim_parentVtxIdx[self._index])
1061
1062 def decayVertices(self):
1063 """Returns a generator for decay vertices.
1064
1065 The generator returns TrackingVertex objects.
1066 """
1067 self._checkIsValid()
1068 for ivtx in self._tree.sim_decayVtxIdx[self._index]:
1069 yield TrackingVertex(self._tree, ivtx)
1070
1071 def isLooper(self):
1072 """Returns True if this TrackingParticle is a looper.
1073
1074 Note that the check involves looping over the SimHits, so it is not too cheap."""
1075 self._checkIsValid()
1076 prevr = 0
1077 for ihit in self.simHitIdx():
1078 hit = SimHit(self._tree, ihit)
1079 r = hit.x()**2 + hit.y()**2
1080 if r < prevr:
1081 return True
1082 prevr = r
1083 return False
1084
1085
1086 class TrackingParticles(_Collection):
1087 """Class presenting a collection of TrackingParticles."""
1088 def __init__(self, tree):
1089 """Constructor.
1090
1091 Arguments:
1092 tree -- TTree object
1093 """
1094 super(TrackingParticles, self).__init__(tree, "sim_pt", TrackingParticle)
1095
1096
1097 class Vertex(_Object):
1098 """Class presenting a primary vertex."""
1099 def __init__(self, tree, index):
1100 """Constructor.
1101
1102 Arguments:
1103 tree -- TTree object
1104 index -- Index of the vertex
1105 """
1106 super(Vertex, self).__init__(tree, index, "vtx")
1107
1108 def nTracks(self):
1109 """Returns the number of tracks used in the vertex fit."""
1110 self._checkIsValid()
1111 return self._tree.vtx_trkIdx[self._index].size()
1112
1113 def tracks(self):
1114 """Returns a generator for the tracks used in the vertex fit.
1115
1116 The generator returns Track object.
1117 """
1118 self._checkIsValid()
1119 for itrk in self._tree.vtx_trkIdx[self._index]:
1120 yield Track(self._tree, itrk)
1121
1122 class Vertices(_Collection):
1123 """Class presenting a collection of vertices."""
1124 def __init__(self, tree):
1125 """Constructor.
1126
1127 Arguments:
1128 tree -- TTree object
1129 """
1130 super(Vertices, self).__init__(tree, "vtx_valid", Vertex)
1131
1132
1133 class TrackingVertex(_Object):
1134 """Class representing a TrackingVertex."""
1135 def __init__(self, tree, index):
1136 """Constructor.
1137
1138 Arguments:
1139 tree -- TTree object
1140 index -- Index of the TrackingVertex
1141 """
1142 super(TrackingVertex, self).__init__(tree, index, "simvtx")
1143
1144 def nSourceTrackingParticles(self):
1145 """Returns the number of source TrackingParticles."""
1146 self._checkIsValid()
1147 return self._tree.simvtx_sourceSimIdx[self._index].size()
1148
1149 def nDaughterTrackingParticles(self):
1150 """Returns the number of daughter TrackingParticles."""
1151 self._checkIsValid()
1152 return self._tree.simvtx_daughterSimIdx[self._index].size()
1153
1154 def sourceTrackingParticles(self):
1155 """Returns a generator for the source TrackingParticles."""
1156 self._checkIsValid()
1157 for isim in self._tree.simvtx_sourceSimIdx[self._index]:
1158 yield TrackingParticle(self._tree, isim)
1159
1160 def daughterTrackingParticles(self):
1161 """Returns a generator for the daughter TrackingParticles."""
1162 self._checkIsValid()
1163 for isim in self._tree.simvtx_daughterSimIdx[self._index]:
1164 yield TrackingParticle(self._tree, isim)
1165
1166 class TrackingVertices(_Collection, TrackingVertex):
1167 """Class presenting a collection of TrackingVertices."""
1168 def __init__(self, tree):
1169 """Constructor.
1170
1171 Arguments:
1172 tree -- TTree object
1173 """
1174 super(TrackingVertex, self).__init__(tree, "simvtx_x")