[Common] Extend option to reassociate PV contributors to different collisions (#15929)

Author: fgrosa

Common/Core/CollisionAssociation.h

@@ -46,6 +46,12 @@ enum TrackSelection {
   QualityTracksITS = 2
 };
 
+enum PVContrReassocOpt {
+  Disabled = 0,
+  OnlySameBc = 1,
+  SameBcAndLowMult = 2
+};
+
 } // namespace track_association
 } // namespace o2::aod
 
@@ -59,10 +65,11 @@ class CollisionAssociation
   void setNumSigmaForTimeCompat(float nSigma) { mNumSigmaForTimeCompat = nSigma; }
   void setTimeMargin(float timeMargin) { mTimeMargin = timeMargin; }
   void setTrackSelectionOptionForStdAssoc(int option) { mTrackSelection = option; }
-  void setUsePvAssociation(bool enable = true) { mUsePvAssociation = enable; }
+  void setUsePvAssociation(int option = o2::aod::track_association::PVContrReassocOpt::OnlySameBc) { mUsePvAssociation = option; }
   void setIncludeUnassigned(bool enable = true) { mIncludeUnassigned = enable; }
   void setFillTableOfCollIdsPerTrack(bool fill = true) { mFillTableOfCollIdsPerTrack = fill; }
   void setBcWindow(int bcWindow = 115) { mBcWindowForOneSigma = bcWindow; }
+  void setMaxPvContributorsForLowMultReassoc(int pvContributorsMax) { mMaxPvContributorsForLowMultReassoc = pvContributorsMax; }
 
   template <typename TTracks, typename Slice, typename Assoc, typename RevIndices>
   void runStandardAssoc(o2::aod::Collisions const& collisions,
@@ -234,7 +241,7 @@ class CollisionAssociation
           float trackTime = 0;
           float trackTimeRes = 0;
           if constexpr (isCentralBarrel) {
-            if (mUsePvAssociation && trackInWindow.isPVContributor()) {
+            if ((mUsePvAssociation == o2::aod::track_association::PVContrReassocOpt::OnlySameBc && trackInWindow.isPVContributor()) || (mUsePvAssociation == o2::aod::track_association::PVContrReassocOpt::SameBcAndLowMult && trackInWindow.isPVContributor() && trackInWindow.collision().numContrib() > mMaxPvContributorsForLowMultReassoc)) {
               trackTime = trackInWindow.collision().collisionTime(); // if PV contributor, we assume the time to be the one of the collision
               trackTimeRes = o2::constants::lhc::LHCBunchSpacingNS;  // 1 BC
             } else {
@@ -252,7 +259,7 @@ class CollisionAssociation
 
           float thresholdTime = 0.;
           if constexpr (isCentralBarrel) {
-            if (mUsePvAssociation && trackInWindow.isPVContributor()) {
+            if ((mUsePvAssociation == o2::aod::track_association::PVContrReassocOpt::OnlySameBc && trackInWindow.isPVContributor()) || (mUsePvAssociation == o2::aod::track_association::PVContrReassocOpt::SameBcAndLowMult && trackInWindow.isPVContributor() && trackInWindow.collision().numContrib() > mMaxPvContributorsForLowMultReassoc)) {
               thresholdTime = trackTimeRes;
             } else if (TESTBIT(trackInWindow.flags(), o2::aod::track::TrackTimeResIsRange)) {
               // the track time resolution is a range, not a gaussian resolution
@@ -306,7 +313,8 @@ class CollisionAssociation
   float mNumSigmaForTimeCompat{4.};                                                  // number of sigma for time compatibility
   float mTimeMargin{500.};                                                           // additional time margin in ns
   int mTrackSelection{o2::aod::track_association::TrackSelection::GlobalTrackWoDCA}; // track selection for central barrel tracks (standard association only)
-  bool mUsePvAssociation{true};                                                      // use the information of PV contributors
+  int mUsePvAssociation{1};                                                          // use the information of PV contributors (0: off, 1: reassociate only with collisions in the same BC, 2: reassociate with collisions in the same BC, or within time resolution in case of low mult PVs)
+  int mMaxPvContributorsForLowMultReassoc{10};                                       // maximum value of PV contributors to reassociate tracks that are PV contributors to other vertices
   bool mIncludeUnassigned{true};                                                     // include tracks that were originally not assigned to any collision
   bool mFillTableOfCollIdsPerTrack{false};                                           // fill additional table with vectors of compatible collisions per track
   int mBcWindowForOneSigma{115};                                                     // BC window to be multiplied by the number of sigmas to define maximum window to be considered

Common/TableProducer/fwdtrackToCollisionAssociator.cxx

@@ -63,7 +63,7 @@ struct FwdTrackToCollisionAssociation {
     collisionAssociator.setNumSigmaForTimeCompat(nSigmaForTimeCompat);
     collisionAssociator.setTimeMargin(timeMargin);
     collisionAssociator.setTrackSelectionOptionForStdAssoc(track_association::TrackSelection::None);
-    collisionAssociator.setUsePvAssociation(false);
+    collisionAssociator.setUsePvAssociation(track_association::PVContrReassocOpt::Disabled);
     collisionAssociator.setIncludeUnassigned(includeUnassigned);
     collisionAssociator.setFillTableOfCollIdsPerTrack(fillTableOfCollIdsPerTrack);
     collisionAssociator.setBcWindow(bcWindowForOneSigma);

Common/TableProducer/trackToCollisionAssociator.cxx

@@ -39,10 +39,11 @@ struct TrackToCollisionAssociation {
   Configurable<float> nSigmaForTimeCompat{"nSigmaForTimeCompat", 4.f, "number of sigmas for time compatibility"};
   Configurable<float> timeMargin{"timeMargin", 0.f, "time margin in ns added to uncertainty because of uncalibrated TPC"};
   Configurable<int> setTrackSelections{"setTrackSelections", 1, "flag to apply track selections: -1=minimal track selection for Run 2 (standard association); 0=none; 1=global track w/o DCA selection; 2=only ITS quality"};
-  Configurable<bool> usePVAssociation{"usePVAssociation", true, "if the track is a PV contributor, use the collision time for it"};
+  Configurable<int> usePVAssociation{"usePVAssociation", 1, "use information of PV contribution in reassociation; 0 -> off, 1 -> reassociate only with collisions in the same BC, 2 -> as 1 only for tracks associated to high multiplicity vertices"};
   Configurable<bool> includeUnassigned{"includeUnassigned", false, "consider also tracks which are not assigned to any collision"};
   Configurable<bool> fillTableOfCollIdsPerTrack{"fillTableOfCollIdsPerTrack", false, "fill additional table with vector of collision ids per track"};
   Configurable<int> bcWindowForOneSigma{"bcWindowForOneSigma", 60, "BC window to be multiplied by the number of sigmas to define maximum window to be considered"};
+  Configurable<int> maxPvContributorsForLowMultReassoc{"maxPvContributorsForLowMultReassoc", 10, "Maximum number of PV contributors to consider a collision at low multiplicity and reassociate tracks even if PV contributors if enabled"};
 
   CollisionAssociation<true> collisionAssociator;
 
@@ -67,6 +68,7 @@ struct TrackToCollisionAssociation {
     collisionAssociator.setIncludeUnassigned(includeUnassigned);
     collisionAssociator.setFillTableOfCollIdsPerTrack(fillTableOfCollIdsPerTrack);
     collisionAssociator.setBcWindow(bcWindowForOneSigma);
+    collisionAssociator.setMaxPvContributorsForLowMultReassoc(maxPvContributorsForLowMultReassoc);
   }
 
   void processAssocWithTime(Collisions const& collisions, TracksWithSel const& tracksUnfiltered, TracksWithSelFilter const& tracks, AmbiguousTracks const& ambiguousTracks, BCs const& bcs)