python/tools/trackselectionRefitting.py

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