python/tools/trackselectionRefitting.py

0001 import sys
0002 import FWCore.ParameterSet.Config as cms
0003
0004 def customlog(s):
0005     print("# MSG-i trackselectionRefitting:  %s" % s)
0006
0007 def getSequence(process, collection,
0008                 saveCPU = False,
0009                 TTRHBuilder = "WithAngleAndTemplate",
0010                 usePixelQualityFlag = None,
0011                 openMassWindow = False,
0012                 cosmicsDecoMode = False,
0013                 cosmicsZeroTesla = True,
0014                 momentumConstraint = None,
0015                 cosmicTrackSplitting = False,
0016                 isPVValidation = False,
0017                 use_d0cut = True,
0018                 g4Refitting = False):
0019     """This function returns a cms.Sequence containing as last element the
0020     module 'FinalTrackRefitter', which can be used as cms.InputTag for
0021     subsequent processing steps.
0022     The modules in the sequence are already attached to the given `process`
0023     object using the given track collection `collection` and the given
0024     optional arguments.
0025
0026     Arguments:
0027     - `process`: 'cms.Process' object to which the modules of the sequence will
0028                  be attached.
0029     - `collection`: String indicating the input track collection.
0030     - `saveCPU`: If set to 'True', some steps are merged to reduce CPU time.
0031                  Reduces a little the accuracy of the results.
0032                  This option is currently not recommended.
0033     - `TTRHBuilder`: Option used for the Track(Re)Fitter modules.
0034     - `usePixelQualityFlag`: Option used for the TrackHitFilter module.
0035                              Defaults to 'True' but is automatically set to
0036                              'False' if a `TTRHBuilder` without templates is
0037                              used.
0038                              If this is still wanted for some reason, one can
0039                              explicitely specify it as 'True'.
0040     - `openMassWindow`: Used to configure the TwoBodyDecaySelector for ZMuMu.
0041     - `cosmicsDecoMode`: If set to 'True' a lower Signal/Noise cut is used.
0042     - `cosmicsZeroTesla`: If set to 'True' a 0T-specific selection is used.
0043     - `momentumConstraint`: If you want to apply a momentum constraint for the
0044                             track refitting, e.g. for CRUZET data, you need
0045                             to provide here the name of the constraint module.
0046     - `cosmicTrackSplitting`: If set to 'True' cosmic tracks are split before the
0047                               second track refitter.
0048     - `isPVValidation`: If set to 'True' most of the selection cuts are overridden
0049                         to allow unbiased selection of tracks for vertex refitting
0050     - `use_d0cut`: If 'True' (default), apply a cut |d0| < 50.
0051     """
0052
0053     ###################################################
0054     # resolve default values incl. consistency checks #
0055     ###################################################
0056
0057     customlog("g4Refitting=%s" % g4Refitting)
0058
0059     if usePixelQualityFlag is None:
0060         if "Template" not in TTRHBuilder:
0061             usePixelQualityFlag = False # not defined without templates
0062             customlog("Using 'TTRHBuilder' without templates %s" % TTRHBuilder)
0063             customlog(" --> Turning off pixel quality flag in hit filter.")
0064         else:
0065             usePixelQualityFlag = True # default for usage with templates
0066
0067
0068     #############################
0069     ## setting general options ##
0070     #############################
0071
0072     options = {"TrackHitFilter": {},
0073                "TrackFitter": {},
0074                "TrackRefitter": {},
0075                "TrackSelector": {},
0076                "geopro": {} }
0077
0078     options["TrackSelector"]["HighPurity"] = {
0079         "trackQualities": ["highPurity"],
0080         "filter": True,
0081         "etaMin": -3.0,
0082         "etaMax": 3.0,
0083         "pMin": 8.0
0084         }
0085     options["TrackSelector"]["Alignment"] = {
0086         "filter": True,
0087         "pMin": 3.0,
0088         "nHitMin2D": 2,
0089         "d0Min": -50.0,
0090         "d0Max": 50.0,
0091         "etaMin": -3.0,
0092         "etaMax": 3.0,
0093         "nHitMin": 8,
0094         "chi2nMax": 9999.0
0095         }
0096     options["TrackRefitter"]["First"] = {
0097         "NavigationSchool": "",
0098         "TTRHBuilder": TTRHBuilder,
0099         }
0100     options["TrackRefitter"]["Second"] = {
0101         "NavigationSchool": "",
0102         "TTRHBuilder": TTRHBuilder,
0103         }
0104     options["TrackHitFilter"]["Tracker"] = {
0105         "useTrajectories": True,
0106         "minimumHits": 8,
0107         "commands": cms.vstring("keep PXB", "keep PXE", "keep TIB", "keep TID",
0108                                 "keep TOB", "keep TEC"),
0109         "replaceWithInactiveHits": True,
0110         "rejectBadStoNHits": True,
0111         "rejectLowAngleHits": True,
0112         "usePixelQualityFlag": usePixelQualityFlag,
0113         "StoNcommands": cms.vstring("ALL 12.0"),
0114         "TrackAngleCut": 0.087,
0115         }
0116     options["TrackFitter"]["HitFilteredTracks"] = {
0117         "NavigationSchool": "",
0118         "TTRHBuilder": TTRHBuilder,
0119         }
0120     options["geopro"][""] = {
0121         }
0122
0123     if g4Refitting:
0124         options["TrackRefitter"]["Second"] = {
0125             "AlgorithmName" : cms.string('undefAlgorithm'),
0126             "Fitter" : cms.string('G4eFitterSmoother'),
0127             "GeometricInnerState" : cms.bool(False),
0128             "MeasurementTracker" : cms.string(''),
0129             "MeasurementTrackerEvent" : cms.InputTag("MeasurementTrackerEvent"),
0130             "NavigationSchool" : cms.string('SimpleNavigationSchool'),  # Correct?
0131             "Propagator" : cms.string('Geant4ePropagator'),
0132             "TTRHBuilder" : cms.string('WithAngleAndTemplate'),
0133             "TrajectoryInEvent" : cms.bool(True),
0134             "beamSpot" : cms.InputTag("offlineBeamSpot"),
0135             "constraint" : cms.string(''),
0136             "src" : cms.InputTag("AlignmentTrackSelector"),
0137             "srcConstr" : cms.InputTag(""),
0138             "useHitsSplitting" : cms.bool(False),
0139             "usePropagatorForPCA" : cms.bool(True)   # not sure whether it is needed
0140         }
0141
0142     #########################################
0143     ## setting collection specific options ##
0144     #########################################
0145     isCosmics = False
0146
0147     if collection in ("ALCARECOTkAlMinBias", "generalTracks",
0148                       "ALCARECOTkAlMinBiasHI", "hiGeneralTracks",
0149                       "ALCARECOTkAlJetHT", "ALCARECOTkAlDiMuonVertexTracks",
0150                       "hltMergedTracks"):
0151         options["TrackSelector"]["Alignment"].update({
0152                 "ptMin": 1.0,
0153                 "pMin": 8.,
0154                 })
0155         options["TrackHitFilter"]["Tracker"].update({
0156                 "minimumHits": 10,
0157                 })
0158     elif collection in ("ALCARECOTkAlCosmicsCTF0T",
0159                         "ALCARECOTkAlCosmicsCosmicTF0T",
0160                         "ALCARECOTkAlCosmicsInCollisions"):
0161         isCosmics = True
0162         options["TrackSelector"]["HighPurity"] = {} # drop high purity cut
0163         if not cosmicsDecoMode:
0164             options["TrackHitFilter"]["Tracker"].update({
0165                     "StoNcommands": cms.vstring("ALL 18.0")
0166                     })
0167         if cosmicsZeroTesla:
0168             options["TrackHitFilter"]["Tracker"].update({
0169                     "TrackAngleCut": 0.1 # Run-I: 0.087 for 0T
0170                     })
0171         else:
0172             options["TrackHitFilter"]["Tracker"].update({
0173                     "TrackAngleCut": 0.1 # Run-I: 0.35 for 3.8T
0174                     })
0175         options["TrackSelector"]["Alignment"].update({
0176                 "pMin": 4.0,
0177                 "etaMin": -99.0,
0178                 "etaMax": 99.0,
0179                 "applyMultiplicityFilter": True,
0180                 "maxMultiplicity": 1
0181                 })
0182         if cosmicTrackSplitting:
0183             options["TrackSplitting"] = {}
0184             options["TrackSplitting"]["TrackSplitting"] = {}
0185         if not use_d0cut:
0186             options["TrackSelector"]["Alignment"].update({
0187                     "d0Min": -99999.0,
0188                     "d0Max": 99999.0,
0189                     })
0190     elif collection in ("ALCARECOTkAlMuonIsolated",
0191                         "ALCARECOTkAlMuonIsolatedHI",
0192                         "ALCARECOTkAlMuonIsolatedPA"):
0193         options["TrackSelector"]["Alignment"].update({
0194                 ("minHitsPerSubDet", "inPIXEL"): 1,
0195                 "ptMin": 5.0,
0196                 "nHitMin": 10,
0197                 "applyMultiplicityFilter": True,
0198                 "maxMultiplicity": 1,
0199                 })
0200     elif collection in ("ALCARECOTkAlZMuMu",
0201                         "ALCARECOTkAlZMuMuHI",
0202                         "ALCARECOTkAlZMuMuPA",
0203                         "ALCARECOTkAlDiMuon"):
0204         options["TrackSelector"]["Alignment"].update({
0205                 "ptMin": 15.0,
0206                 "etaMin": -3.0,
0207                 "etaMax": 3.0,
0208                 "nHitMin": 10,
0209                 "applyMultiplicityFilter": True,
0210                 "minMultiplicity": 2,
0211                 "maxMultiplicity": 2,
0212                 ("minHitsPerSubDet", "inPIXEL"): 1,
0213                 ("TwoBodyDecaySelector", "applyChargeFilter"): True,
0214                 ("TwoBodyDecaySelector", "charge"): 0,
0215                 ("TwoBodyDecaySelector",
0216                  "applyMassrangeFilter"): not openMassWindow,
0217                 ("TwoBodyDecaySelector", "minXMass"): 85.8,
0218                 ("TwoBodyDecaySelector", "maxXMass"): 95.8,
0219                 ("TwoBodyDecaySelector", "daughterMass"): 0.105
0220                 })
0221         options["TrackHitFilter"]["Tracker"].update({
0222                 "minimumHits": 10,
0223                 })
0224     elif collection == "ALCARECOTkAlUpsilonMuMu":
0225         options["TrackSelector"]["Alignment"].update({
0226                 "ptMin": 3.0,
0227                 "etaMin": -2.4,
0228                 "etaMax": 2.4,
0229                 "nHitMin": 10,
0230                 "applyMultiplicityFilter": True,
0231                 "minMultiplicity": 2,
0232                 "maxMultiplicity": 2,
0233                 ("minHitsPerSubDet", "inPIXEL"): 1,
0234                 ("TwoBodyDecaySelector", "applyChargeFilter"): True,
0235                 ("TwoBodyDecaySelector", "charge"): 0,
0236                 ("TwoBodyDecaySelector",
0237                  "applyMassrangeFilter"): not openMassWindow,
0238                 ("TwoBodyDecaySelector", "minXMass"): 9.2,
0239                 ("TwoBodyDecaySelector", "maxXMass"): 9.7,
0240                 ("TwoBodyDecaySelector", "daughterMass"): 0.105
0241                 })
0242         options["TrackHitFilter"]["Tracker"].update({
0243                 "minimumHits": 10,
0244                 })
0245     elif collection == "ALCARECOTkAlJpsiMuMu":
0246         options["TrackSelector"]["Alignment"].update({
0247                 "ptMin": 1.0,
0248                 "etaMin": -2.4,
0249                 "etaMax": 2.4,
0250                 "nHitMin": 10,
0251                 "applyMultiplicityFilter": True,
0252                 "minMultiplicity": 2,
0253                 "maxMultiplicity": 2,
0254                 ("minHitsPerSubDet", "inPIXEL"): 1,
0255                 ("TwoBodyDecaySelector", "applyChargeFilter"): True,
0256                 ("TwoBodyDecaySelector", "charge"): 0,
0257                 ("TwoBodyDecaySelector",
0258                  "applyMassrangeFilter"): not openMassWindow,
0259                 ("TwoBodyDecaySelector", "minXMass"): 2.7,
0260                 ("TwoBodyDecaySelector", "maxXMass"): 3.4,
0261                 ("TwoBodyDecaySelector", "daughterMass"): 0.105
0262                 })
0263         options["TrackHitFilter"]["Tracker"].update({
0264                 "minimumHits": 10,
0265                 })
0266     else:
0267         raise ValueError("Unknown input track collection: {}".format(collection))
0268
0269     if cosmicTrackSplitting and not isCosmics:
0270         raise ValueError("Can only do cosmic track splitting for cosmics.")
0271
0272
0273
0274     ####################
0275     ## save CPU time? ##
0276     ####################
0277
0278     if saveCPU:
0279         if cosmicTrackSplitting:
0280             raise ValueError("Can't turn on both saveCPU and cosmicTrackSplitting at the same time")
0281         mods = [("TrackSelector", "Alignment", {"method": "load"}),
0282                 ("TrackRefitter", "First", {"method": "load",
0283                                             "clone": True}),
0284                 ("TrackHitFilter", "Tracker", {"method": "load"}),
0285                 ("TrackFitter", "HitFilteredTracks", {"method": "import"})]
0286         options["TrackSelector"]["Alignment"].update(
0287             options["TrackSelector"]["HighPurity"])
0288     elif cosmicTrackSplitting:
0289         mods = [("TrackRefitter", "First", {"method": "load",
0290                                             "clone": True}),
0291                 ("TrackSelector", "Alignment", {"method": "load"}),
0292                 ("TrackSplitting", "TrackSplitting", {"method": "load"}),
0293                 ("TrackFitter", "HitFilteredTracks", {"method": "import"}),
0294                 ("TrackRefitter", "Second", {"method": "load",
0295                                              "clone": True})]
0296     elif g4Refitting:
0297         mods = [("TrackSelector", "HighPurity", {"method": "import"}),
0298                 ("TrackRefitter", "First", {"method": "load",
0299                                             "clone": True}),
0300                 ("TrackHitFilter", "Tracker", {"method": "load"}),
0301                 ("TrackFitter", "HitFilteredTracks", {"method": "import"}),
0302                 ("TrackSelector", "Alignment", {"method": "load"}),
0303                 #("geopro","", {"method": "load"}),
0304                 ("TrackRefitter", "Second", {"method": "load",
0305                                              "clone": True})]
0306         if isCosmics: mods = mods[1:] # skip high purity selector for cosmics
0307     else:
0308         mods = [("TrackSelector", "HighPurity", {"method": "import"}),
0309                 ("TrackRefitter", "First", {"method": "load",
0310                                             "clone": True}),
0311                 ("TrackHitFilter", "Tracker", {"method": "load"}),
0312                 ("TrackFitter", "HitFilteredTracks", {"method": "import"}),
0313                 ("TrackSelector", "Alignment", {"method": "load"}),
0314                 ("TrackRefitter", "Second", {"method": "load",
0315                                              "clone": True})]
0316         if isCosmics: mods = mods[1:] # skip high purity selector for cosmics
0317
0318     #############################
0319     ## PV Validation cuts tune ##
0320     #############################
0321
0322     if isPVValidation:
0323         options["TrackSelector"]["HighPurity"].update({
0324                 "trackQualities": [],
0325                 "pMin": 0.
0326                 })
0327         options["TrackSelector"]["Alignment"].update({
0328                 "pMin" :      0.,
0329                 "ptMin" :     0.,
0330                 "nHitMin2D" : 0,
0331                 "nHitMin"   : 0,
0332                 "d0Min" : -999999.0,
0333                 "d0Max" :  999999.0,
0334                 "dzMin" : -999999.0,
0335                 "dzMax" :  999999.0
0336                 })
0337
0338     ################################
0339     ## apply momentum constraint? ##
0340     ################################
0341
0342     if momentumConstraint is not None:
0343         for mod in options["TrackRefitter"]:
0344             momconstrspecs = momentumConstraint.split(',')
0345             if len(momconstrspecs)==1:
0346                 options["TrackRefitter"][mod].update({
0347                     "constraint": "momentum",
0348                     "srcConstr": momconstrspecs[0]
0349                     })
0350             else:
0351                 options["TrackRefitter"][mod].update({
0352                     "constraint": momconstrspecs[1],
0353                     "srcConstr": momconstrspecs[0]
0354                     })
0355
0356
0357
0358     #######################################################
0359     # load offline beam spot module required by the refit #
0360     #######################################################
0361     process.load("RecoVertex.BeamSpotProducer.BeamSpot_cff")
0362
0363     ###############################
0364     ## put the sequence together ##
0365     ###############################
0366
0367     modules = []
0368     src = collection
0369     prevsrc = None
0370     for mod in mods[:-1]:
0371         src, prevsrc = _getModule(process, src, mod[0], "".join(reversed(mod[:-1])),
0372                                   options[mod[0]][mod[1]], isCosmics = isCosmics, prevsrc = prevsrc,
0373                                   **(mod[2])), src
0374         modules.append(getattr(process, src))
0375     else:
0376         if mods[-1][-1]["method"] == "load" and \
0377                 not mods[-1][-1].get("clone", False):
0378             customlog("Name of the last module needs to be modifiable.")
0379             sys.exit(1)
0380
0381         if g4Refitting:
0382             customlog("Here we must include geopro first")
0383             process.load('Configuration.StandardSequences.GeometryDB_cff')
0384             process.load("TrackPropagation.Geant4e.geantRefit_cff")
0385             modules.append(getattr(process,"geopro"))
0386
0387         src = _getModule(process, src, mods[-1][0], "FinalTrackRefitter",
0388                          options[mods[-1][0]][mods[-1][1]],
0389                          isCosmics = isCosmics, **(mods[-1][2]))
0390         modules.append(getattr(process, src))
0391
0392     moduleSum = process.offlineBeamSpot        # first element of the sequence
0393     if g4Refitting:
0394         # g4Refitter needs measurements
0395         moduleSum += getattr(process,"MeasurementTrackerEvent")
0396
0397     for module in modules:
0398         # Spply srcConstr fix here
0399         if hasattr(module,"srcConstr"):
0400            strSrcConstr = module.srcConstr.getModuleLabel()
0401            if strSrcConstr:
0402                procsrcconstr = getattr(process,strSrcConstr)
0403                if hasattr(procsrcconstr,"src"): # Momentum or track parameter constraints
0404                   if procsrcconstr.src != module.src:
0405                      module.srcConstr=''
0406                      module.constraint=''
0407                   else:
0408                      moduleSum += procsrcconstr # Add constraint
0409                elif hasattr(procsrcconstr,"srcTrk"): # Vertex constraint
0410                   if procsrcconstr.srcTrk != module.src:
0411                      module.srcConstr=''
0412                      module.constraint=''
0413                   else:
0414                      procsrcconstrsrcvtx = getattr(process,procsrcconstr.srcVtx.getModuleLabel())
0415                      if type(procsrcconstrsrcvtx) is cms.EDFilter: # If source of vertices is itself a filter (e.g. good PVs)
0416                         procsrcconstrsrcvtxprefilter = getattr(process,procsrcconstrsrcvtx.src.getModuleLabel())
0417                         moduleSum += procsrcconstrsrcvtxprefilter # Add vertex source to constraint before filter
0418                      moduleSum += procsrcconstrsrcvtx # Add vertex source to constraint
0419                      moduleSum += procsrcconstr # Add constraint
0420
0421         moduleSum += module # append the other modules
0422
0423     return cms.Sequence(moduleSum)
0424
0425 ###############################
0426 ###############################
0427 ###                         ###
0428 ###   Auxiliary functions   ###
0429 ###                         ###
0430 ###############################
0431 ###############################
0432
0433
0434 def _getModule(process, src, modType, moduleName, options, **kwargs):
0435     """General function for attaching the module of type `modType` to the
0436     cms.Process `process` using `options` for customization and `moduleName` as
0437     the name of the new attribute of `process`.
0438
0439     Arguments:
0440     - `process`: 'cms.Process' object to which the module is attached.
0441     - `src`: cms.InputTag for this module.
0442     - `modType`: Type of the requested module.
0443     - `options`: Dictionary with customized values for the module's options.
0444     - `**kwargs`: Used to supply options at construction time of the module.
0445     """
0446
0447     objTuple = globals()["_"+modType](kwargs)
0448     method = kwargs.get("method")
0449     if method == "import":
0450         __import__(objTuple[0])
0451         obj = getattr(sys.modules[objTuple[0]], objTuple[1]).clone()
0452     elif method == "load":
0453         process.load(objTuple[0])
0454         if kwargs.get("clone", False):
0455             obj = getattr(process, objTuple[1]).clone(src=src)
0456         else:
0457             obj = getattr(process, objTuple[1])
0458             moduleName = objTuple[1]
0459     else:
0460         customlog("Unknown method: %s" % method)
0461         sys.exit(1)
0462
0463     if modType == "TrackSplitting":
0464         #track splitting takes the TrackSelector as tracks
0465         # and the first TrackRefitter as tjTkAssociationMapTag
0466         _customSetattr(obj, "tracks", src)
0467         _customSetattr(obj, "tjTkAssociationMapTag", kwargs["prevsrc"])
0468     else:
0469         obj.src = src
0470
0471     for option in options:
0472         _customSetattr(obj, option, options[option])
0473
0474     if moduleName is not objTuple[1]:
0475         setattr(process, moduleName, obj)
0476     return moduleName
0477
0478
0479 def _TrackHitFilter(kwargs):
0480     """Returns TrackHitFilter module name.
0481
0482     Arguments:
0483     - `kwargs`: Not used in this function.
0484     """
0485
0486     return ("RecoTracker.FinalTrackSelectors.TrackerTrackHitFilter_cff",
0487             "TrackerTrackHitFilter")
0488
0489
0490 def _TrackSelector(kwargs):
0491     """Returns TrackSelector module name.
0492
0493     Arguments:
0494     - `kwargs`: Not used in this function.
0495     """
0496
0497     return ("Alignment.CommonAlignmentProducer.AlignmentTrackSelector_cfi",
0498             "AlignmentTrackSelector")
0499
0500
0501 def _TrackFitter(kwargs):
0502     """Returns TrackFitter module name.
0503
0504     Arguments:
0505     - `kwargs`: Used to supply options at construction time of the object.
0506     """
0507
0508     isCosmics = kwargs.get("isCosmics", False)
0509     if isCosmics:
0510         return ("RecoTracker.TrackProducer.CTFFinalFitWithMaterialP5_cff",
0511                 "ctfWithMaterialTracksCosmics")
0512     else:
0513         return ("RecoTracker.TrackProducer.CTFFinalFitWithMaterial_cff",
0514                 "ctfWithMaterialTracks")
0515
0516
0517 def _TrackRefitter(kwargs):
0518     """Returns TrackRefitter module name.
0519
0520     Arguments:
0521     - `kwargs`: Used to supply options at construction time of the object.
0522     """
0523
0524     isCosmics = kwargs.get("isCosmics", False)
0525     if isCosmics:
0526         return ("RecoTracker.TrackProducer.TrackRefitters_cff",
0527                 "TrackRefitterP5")
0528     else:
0529         return ("RecoTracker.TrackProducer.TrackRefitters_cff",
0530                 "TrackRefitter")
0531
0532 def _TrackSplitting(kwargs):
0533     return ("RecoTracker.FinalTrackSelectors.cosmicTrackSplitter_cfi",
0534             "cosmicTrackSplitter")
0535
0536 def _geopro(kwargs):
0537     return ("TrackPropagation.Geant4e.geantRefit_cff","geopro")
0538
0539
0540 def _customSetattr(obj, attr, val):
0541     """Sets the attribute `attr` of the object `obj` using the value `val`.
0542     `attr` can be a string or a tuple of strings, if one wants to set an
0543     attribute of an attribute, etc.
0544
0545     Arguments:
0546     - `obj`: Object, which must have a '__dict__' attribute.
0547     - `attr`: String or tuple of strings describing the attribute's name.
0548     - `val`: value of the attribute.
0549     """
0550
0551     if isinstance(attr, tuple) and len(attr) > 1:
0552         _customSetattr(getattr(obj, attr[0]), attr[1:], val)
0553     else:
0554         if isinstance(attr, tuple): attr = attr[0]
0555         setattr(obj, attr, val)
0556