python/tools/trackselectionRefitting.py

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