Fix event selection for generated signals

Author: fgrosa

PWGHF/D2H/DataModel/ReducedDataModel.h

@@ -831,7 +831,10 @@ DECLARE_SOA_TABLE(HfMcGenRedB0s, "AOD", "HFMCGENREDB0", //! Generation-level MC
                   hf_b0_mc::EtaProng0,
                   hf_b0_mc::PtProng1,
                   hf_b0_mc::YProng1,
-                  hf_b0_mc::EtaProng1);
+                  hf_b0_mc::EtaProng1,
+                  hf_reduced_collision::HfCollisionRejectionMap,
+                  cent::CentFT0C,
+                  cent::CentFT0M);
 
 // store all configurables values used in the first part of the workflow
 // so we can use them in the B0 part
@@ -922,7 +925,10 @@ DECLARE_SOA_TABLE(HfMcGenRedBps, "AOD", "HFMCGENREDBP", //! Generation-level MC
                   hf_bplus_mc::EtaProng0,
                   hf_bplus_mc::PtProng1,
                   hf_bplus_mc::YProng1,
-                  hf_bplus_mc::EtaProng1);
+                  hf_bplus_mc::EtaProng1,
+                  hf_reduced_collision::HfCollisionRejectionMap,
+                  cent::CentFT0C,
+                  cent::CentFT0M);
 
 // store all configurables values used in the first part of the workflow
 // so we can use them in the Bplus part
@@ -1024,7 +1030,10 @@ DECLARE_SOA_TABLE(HfMcGenRedBss, "AOD", "HFMCGENREDBS", //! Generation-level MC
                   hf_bs_mc::EtaProng0,
                   hf_bs_mc::PtProng1,
                   hf_bs_mc::YProng1,
-                  hf_bs_mc::EtaProng1);
+                  hf_bs_mc::EtaProng1,
+                  hf_reduced_collision::HfCollisionRejectionMap,
+                  cent::CentFT0C,
+                  cent::CentFT0M);
 
 // store all configurables values used in the first part of the workflow
 // so we can use them in the Bs part
@@ -1108,7 +1117,10 @@ DECLARE_SOA_TABLE(HfMcGenRedLbs, "AOD", "HFMCGENREDLB", //! Generation-level MC
                   hf_lb_mc::EtaProng0,
                   hf_lb_mc::PtProng1,
                   hf_lb_mc::YProng1,
-                  hf_lb_mc::EtaProng1);
+                  hf_lb_mc::EtaProng1,
+                  hf_reduced_collision::HfCollisionRejectionMap,
+                  cent::CentFT0C,
+                  cent::CentFT0M);
 
 // store all configurables values used in the first part of the workflow
 // so we can use them in the B0 part

PWGHF/D2H/TableProducer/dataCreatorCharmHadPiReduced.cxx

@@ -51,6 +51,7 @@
 #include <Framework/O2DatabasePDGPlugin.h>
 #include <Framework/WorkflowSpec.h>
 #include <Framework/runDataProcessing.h>
+#include "Framework/RunningWorkflowInfo.h"
 #include <ReconstructionDataFormats/DCA.h>
 #include <ReconstructionDataFormats/Track.h>
 
@@ -123,6 +124,7 @@ struct HfDataCreatorJpsiHadReduced {
   Produces<aod::HfCfgBpToJpsi> rowCandidateConfigBplus;
   Produces<aod::HfCfgBsToJpsis> rowCandidateConfigBs;
 
+  Configurable<bool> skipRejectedCollisions{"skipRejectedCollisions", true, "skips collisions rejected by the event selection, instead of flagging only"};
   Configurable<bool> propagateToPCA{"propagateToPCA", true, "create tracks version propagated to PCA"};
   Configurable<bool> useAbsDCA{"useAbsDCA", false, "Minimise abs. distance rather than chi2"};
   Configurable<bool> useWeightedFinalPCA{"useWeightedFinalPCA", false, "Recalculate vertex position using track covariances, effective only if useAbsDCA is true"};
@@ -163,10 +165,12 @@ struct HfDataCreatorJpsiHadReduced {
   using TracksSel = soa::Join<aod::TracksWDcaExtra, aod::TracksPidPi, aod::PidTpcTofFullPi, aod::TracksPidKa, aod::PidTpcTofFullKa>;
   using TracksPidWithSelAndMc = soa::Join<TracksPidWithSel, aod::McTrackLabels>;
   using CollisionsWCMcLabels = soa::Join<aod::Collisions, aod::McCollisionLabels, aod::EvSels>;
+  using BCsInfo = soa::Join<aod::BCsWithTimestamps, aod::BcSels>;
 
   Preslice<aod::HfCand2ProngWPid> candsJpsiPerCollision = aod::track_association::collisionId;
   Preslice<aod::TrackAssoc> trackIndicesPerCollision = aod::track_association::collisionId;
   PresliceUnsorted<CollisionsWCMcLabels> colPerMcCollision = aod::mccollisionlabel::mcCollisionId;
+  Preslice<aod::McParticles> mcParticlesPerMcCollision = aod::mcparticle::mcCollisionId;
 
   o2::base::Propagator::MatCorrType noMatCorr = o2::base::Propagator::MatCorrType::USEMatCorrNONE;
   int runNumber;
@@ -175,13 +179,15 @@ struct HfDataCreatorJpsiHadReduced {
   bool isHfCandBhadConfigFilled = false;
 
   o2::hf_evsel::HfEventSelection hfEvSel;
+  o2::hf_evsel::HfEventSelectionMc hfEvSelMc;
+
   o2::vertexing::DCAFitterN<2> df2;
   o2::vertexing::DCAFitterN<3> df3;
   o2::vertexing::DCAFitterN<4> df4;
 
   HistogramRegistry registry{"registry"};
 
-  void init(InitContext const&)
+  void init(InitContext& initContext)
   {
     std::array<int, 4> doProcess = {doprocessJpsiKData, doprocessJpsiKMc, doprocessJpsiPhiData, doprocessJpsiPhiMc};
     if (std::accumulate(doProcess.begin(), doProcess.end(), 0) != 1) {
@@ -271,6 +277,19 @@ struct HfDataCreatorJpsiHadReduced {
       invMass2JpsiHadMin = (MassBS - invMassWindowJpsiHad) * (MassBS - invMassWindowJpsiHad);
       invMass2JpsiHadMax = (MassBS + invMassWindowJpsiHad) * (MassBS + invMassWindowJpsiHad);
     }
+
+    // init HF event selection helper
+    hfEvSel.init(registry);
+    if (doprocessJpsiKMc || doprocessJpsiPhiMc) {
+      const auto& workflows = initContext.services().get<RunningWorkflowInfo const>();
+      for (const DeviceSpec& device : workflows.devices) {
+        if (device.name.compare("hf-data-creator-jpsi-had-reduced") == 0) {
+          // init HF event selection helper
+          hfEvSelMc.init(device, registry);
+          break;
+        }
+      }
+    }
   }
 
   /// Topological cuts
@@ -513,10 +532,23 @@ struct HfDataCreatorJpsiHadReduced {
   }
 
   template <uint8_t decChannel>
-  void runMcGen(aod::McParticles const& particlesMc)
+  void runMcGen(aod::McCollision const& mcCollision,
+                aod::McParticles const& particlesMc,
+                CollisionsWCMcLabels const& collisions,
+                BCsInfo const&)
   {
+    // Check event selection
+    float centDummy{-1.f}, centFT0C{-1.f}, centFT0M{-1.f};
+    const auto collSlice = collisions.sliceBy(colPerMcCollision, mcCollision.globalIndex());
+    auto hfRejMap = hfEvSelMc.getHfMcCollisionRejectionMask<BCsInfo, o2::hf_centrality::CentralityEstimator::None>(mcCollision, collSlice, centDummy);
+    if (skipRejectedCollisions && hfRejMap!= 0) {
+      return;
+    }
+
+    const auto mcParticlesPerMcColl = particlesMc.sliceBy(mcParticlesPerMcCollision, mcCollision.globalIndex());
+
     // Match generated particles.
-    for (const auto& particle : particlesMc) {
+    for (const auto& particle : mcParticlesPerMcColl) {
       int8_t sign{0}, flag{0}, channel{0};
       if constexpr (decChannel == DecayChannel::BplusToJpsiK) {
         // B+ → J/Psi K+ → (µ+µ-) K+
@@ -557,7 +589,7 @@ struct HfDataCreatorJpsiHadReduced {
         }
         rowHfBpMcGenReduced(flag, channel, ptParticle, yParticle, etaParticle,
                             ptProngs[0], yProngs[0], etaProngs[0],
-                            ptProngs[1], yProngs[1], etaProngs[1]);
+                            ptProngs[1], yProngs[1], etaProngs[1], hfRejMap, centFT0C, centFT0M);
       } else if constexpr (decChannel == DecayChannel::BsToJpsiPhi) {
         // Bs → J/Psi phi → (µ+µ-) (K+K-)
         if (RecoDecay::isMatchedMCGen<true>(particlesMc, particle, Pdg::kBS, std::array{static_cast<int>(Pdg::kJPsi), +kKPlus, -kKPlus}, true, &sign, 2)) {
@@ -597,22 +629,29 @@ struct HfDataCreatorJpsiHadReduced {
         }
         rowHfBsMcGenReduced(flag, channel, ptParticle, yParticle, etaParticle,
                             ptProngs[0], yProngs[0], etaProngs[0],
-                            ptProngs[1], yProngs[1], etaProngs[1]);
+                            ptProngs[1], yProngs[1], etaProngs[1], hfRejMap, centFT0C, centFT0M);
       }
     } // gen
   }
 
   // Jpsi candidate selection
-  template <bool doMc, uint8_t decChannel, typename Coll, typename JpsiCands, typename TTracks, typename PParticles>
+  template <bool doMc, uint8_t decChannel, typename Coll, typename JpsiCands, typename TTracks, typename PParticles, typename BBCs>
   void runDataCreation(Coll const& collision,
                        JpsiCands const& candsJpsi,
                        aod::TrackAssoc const& trackIndices,
                        TTracks const&,
                        PParticles const& particlesMc,
                        uint64_t const& indexCollisionMaxNumContrib,
-                       aod::BCsWithTimestamps const&)
+                       BBCs const&)
   {
 
+    registry.fill(HIST("hEvents"), 1 + Event::Processed);
+    float centrality = -1.f;
+    auto hfRejMap = hfEvSel.getHfCollisionRejectionMask<true, o2::hf_centrality::CentralityEstimator::None, aod::BCsWithTimestamps>(collision, centrality, ccdb, registry);
+    if (skipRejectedCollisions && hfRejMap != 0) {
+      return;
+    }
+
     // helpers for ReducedTables filling
     int indexHfReducedCollision = hfReducedCollision.lastIndex() + 1;
     // std::map where the key is the track.globalIndex() and
@@ -960,14 +999,11 @@ struct HfDataCreatorJpsiHadReduced {
       }
     } // candsJpsi loop
 
-    registry.fill(HIST("hEvents"), 1 + Event::Processed);
     if (!fillHfReducedCollision) {
       registry.fill(HIST("hEvents"), 1 + Event::NoCharmHadPiSelected);
       return;
     }
     registry.fill(HIST("hEvents"), 1 + Event::CharmHadPiSelected);
-    float centrality = -1.f;
-    uint16_t hfRejMap = hfEvSel.getHfCollisionRejectionMask<true, o2::hf_centrality::CentralityEstimator::None, aod::BCsWithTimestamps>(collision, centrality, ccdb, registry);
     // fill collision table if it contains a J/Psi K pair at minimum
     hfReducedCollision(collision.posX(), collision.posY(), collision.posZ(), collision.numContrib(), hfRejMap, bz);
     hfReducedCollExtra(collision.covXX(), collision.covXY(), collision.covYY(),
@@ -1048,8 +1084,8 @@ struct HfDataCreatorJpsiHadReduced {
                       aod::TrackAssoc const& trackIndices,
                       TracksPidWithSelAndMc const& tracks,
                       aod::McParticles const& particlesMc,
-                      aod::BCsWithTimestamps const& bcs,
-                      McCollisions const&)
+                      BCsInfo const& bcs,
+                      McCollisions const& mcCollisions)
   {
     // store configurables needed for B+ workflow
     if (!isHfCandBhadConfigFilled) {
@@ -1074,7 +1110,9 @@ struct HfDataCreatorJpsiHadReduced {
     }
     // handle normalization by the right number of collisions
     hfCollisionCounter(collisions.tableSize(), zvtxColl, sel8Coll, zvtxAndSel8Coll, zvtxAndSel8CollAndSoftTrig, allSelColl);
-    runMcGen<DecayChannel::BplusToJpsiK>(particlesMc);
+    for (const auto& mcCollision: mcCollisions) {
+      runMcGen<DecayChannel::BplusToJpsiK>(mcCollision, particlesMc, collisions, bcs);
+    }
   }
   PROCESS_SWITCH(HfDataCreatorJpsiHadReduced, processJpsiKMc, "Process J/Psi K with MC info", false);
 
@@ -1083,8 +1121,8 @@ struct HfDataCreatorJpsiHadReduced {
                         aod::TrackAssoc const& trackIndices,
                         TracksPidWithSelAndMc const& tracks,
                         aod::McParticles const& particlesMc,
-                        aod::BCsWithTimestamps const& bcs,
-                        McCollisions const&)
+                        BCsInfo const& bcs,
+                        McCollisions const& mcCollisions)
   {
     // store configurables needed for B+ workflow
     if (!isHfCandBhadConfigFilled) {
@@ -1109,7 +1147,9 @@ struct HfDataCreatorJpsiHadReduced {
     }
     // handle normalization by the right number of collisions
     hfCollisionCounter(collisions.tableSize(), zvtxColl, sel8Coll, zvtxAndSel8Coll, zvtxAndSel8CollAndSoftTrig, allSelColl);
-    runMcGen<DecayChannel::BsToJpsiPhi>(particlesMc);
+    for (const auto& mcCollision: mcCollisions) {
+      runMcGen<DecayChannel::BsToJpsiPhi>(mcCollision, particlesMc, collisions, bcs);
+    }
   }
   PROCESS_SWITCH(HfDataCreatorJpsiHadReduced, processJpsiPhiMc, "Process J/Psi phi with MC info", false);
 };

PWGHF/D2H/TableProducer/dataCreatorJpsiHadReduced.cxx

@@ -995,7 +995,7 @@ struct HfCandidateCreator2ProngExpressions {
       const auto mcParticlesPerMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, mcCollision.globalIndex());
       // Slice the collisions table to get the collision info for the current MC collision
       float centrality{-1.f};
-      uint16_t rejectionMask{0};
+      uint32_t rejectionMask{0u};
       int nSplitColl = 0;
       if constexpr (centEstimator == CentralityEstimator::FT0C) {
         const auto collSlice = collInfos.sliceBy(colPerMcCollisionFT0C, mcCollision.globalIndex());

PWGHF/TableProducer/candidateCreator2Prong.cxx

@@ -1284,7 +1284,7 @@ struct HfCandidateCreator3ProngExpressions {
       const auto mcParticlesPerMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, mcCollision.globalIndex());
       // Slice the collisions table to get the collision info for the current MC collision
       float centrality{-1.f};
-      uint16_t rejectionMask{0};
+      uint32_t rejectionMask{0u};
       int nSplitColl = 0;
       if constexpr (centEstimator == CentralityEstimator::FT0C) {
         const auto collSlice = collInfos.sliceBy(colPerMcCollisionFT0C, mcCollision.globalIndex());

PWGHF/TableProducer/candidateCreator3Prong.cxx

@@ -549,7 +549,7 @@ struct HfCandidateCreatorCascadeMc {
       const auto mcParticlesPerMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, mcCollision.globalIndex());
       // Slice the collisions table to get the collision info for the current MC collision
       float centrality{-1.f};
-      uint16_t rejectionMask{0};
+      uint32_t rejectionMask{0u};
       int nSplitColl = 0;
       if constexpr (centEstimator == CentralityEstimator::FT0C) {
         const auto collSlice = collInfos.sliceBy(colPerMcCollisionFT0C, mcCollision.globalIndex());

PWGHF/TableProducer/candidateCreatorCascade.cxx

@@ -714,7 +714,7 @@ struct HfCandidateCreatorDstarExpressions {
       const auto mcParticlesPerMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, mcCollision.globalIndex());
       // Slice the collisions table to get the collision info for the current MC collision
       float centrality{-1.f};
-      uint16_t rejectionMask{0};
+      uint32_t rejectionMask{0u};
       int nSplitColl = 0;
       if constexpr (centEstimator == CentralityEstimator::FT0C) {
         const auto collSlice = collInfos.sliceBy(colPerMcCollisionFT0C, mcCollision.globalIndex());

PWGHF/TableProducer/candidateCreatorDstar.cxx

@@ -566,9 +566,8 @@ struct HfCandidateSigmac0plusplusMc {
       /// In case of need, readapt the code templetizing the function
       auto mcCollision = particle.mcCollision();
       float centrality{-1.f};
-      uint16_t rejectionMask{0};
       const auto collSlice = collInfos.sliceBy(colPerMcCollision, mcCollision.globalIndex());
-      rejectionMask = hfEvSelMc.getHfMcCollisionRejectionMask<BCsInfo, o2::hf_centrality::CentralityEstimator::None>(mcCollision, collSlice, centrality);
+      auto rejectionMask = hfEvSelMc.getHfMcCollisionRejectionMask<BCsInfo, o2::hf_centrality::CentralityEstimator::None>(mcCollision, collSlice, centrality);
       hfEvSelMc.fillHistograms<o2::hf_centrality::CentralityEstimator::None>(mcCollision, rejectionMask, 0);
       if (rejectionMask != 0) {
         // at least one event selection not satisfied --> reject gen particles from this collision

PWGHF/TableProducer/candidateCreatorSigmac0plusplus.cxx

@@ -1070,7 +1070,7 @@ struct HfCandidateCreatorXicToXiPiPiExpressions {
       const auto mcParticlesPerMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, mcCollision.globalIndex());
       // Slice the collisions table to get the collision info for the current MC collision
       float centrality{-1.f};
-      uint16_t rejectionMask{0};
+      uint32_t rejectionMask{0u};
       int nSplitColl = 0;
       if constexpr (centEstimator == o2::hf_centrality::CentralityEstimator::FT0C) {
         const auto collSlice = collInfos.sliceBy(colPerMcCollisionFT0C, mcCollision.globalIndex());

PWGHF/TableProducer/candidateCreatorXicToXiPiPi.cxx

@@ -261,7 +261,7 @@ struct HfTaskMcValidationGen {
     if (storeOccupancy) {
       occupancy = getOccupancyGenColl(recoCollisions, OccupancyEstimator::Its);
     }
-    uint16_t rejectionMask{0};
+    uint32_t rejectionMask{0u};
     if constexpr (centEstimator == CentralityEstimator::FT0C) {
       rejectionMask = hfEvSelMc.getHfMcCollisionRejectionMask<BCsInfo, centEstimator>(mcCollision, recoCollisions, centrality);
     } else if constexpr (centEstimator == CentralityEstimator::FT0M) {