[PWGUD] enable skimmers to save all McCollisions, even if they are not reconstructed (AliceO2Group#10393)

Author: pbuehler

PWGUD/Core/UDHelpers.h

@@ -122,7 +122,7 @@ T compatibleBCs(B const& bc, uint64_t const& meanBC, int const& deltaBC, T const
 
   // check [min,max]BC to overlap with [bcs.iteratorAt([0,bcs.size() - 1])
   if (maxBC < bcs.iteratorAt(0).globalBC() || minBC > bcs.iteratorAt(bcs.size() - 1).globalBC()) {
-    LOGF(info, "<compatibleBCs> No overlap of [%d, %d] and [%d, %d]", minBC, maxBC, bcs.iteratorAt(0).globalBC(), bcs.iteratorAt(bcs.size() - 1).globalBC());
+    LOGF(debug, "<compatibleBCs> No overlap of [%d, %d] and [%d, %d]", minBC, maxBC, bcs.iteratorAt(0).globalBC(), bcs.iteratorAt(bcs.size() - 1).globalBC());
     return T{{bcs.asArrowTable()->Slice(0, 0)}, static_cast<uint64_t>(0)};
   }
 

PWGUD/TableProducer/DGCandProducer.cxx

@@ -472,6 +472,9 @@ struct McDGCandProducer {
   Produces<aod::UDMcCollsLabels> outputMcCollsLabels;
   Produces<aod::UDMcTrackLabels> outputMcTrackLabels;
 
+  // save all McTruth, even if the collisions is not reconstructed
+  Configurable<bool> saveAllMcCollisions{"saveAllMcCollisions", true, "save all McCollisions"};
+
   using CCs = soa::Join<aod::Collisions, aod::EvSels, aod::McCollisionLabels>;
   using BCs = soa::Join<aod::BCsWithTimestamps, aod::BcSels, aod::Run3MatchedToBCSparse>;
   using TCs = soa::Join<aod::Tracks, aod::McTrackLabels>;
@@ -670,30 +673,10 @@ struct McDGCandProducer {
     }
   }
 
-  void init(InitContext& context)
-  {
-    // add histograms for the different process functions
-    if (context.mOptions.get<bool>("processMCTruth")) {
-      registry.add("mcTruth/collisions", "Number of associated collisions", {HistType::kTH1F, {{11, -0.5, 10.5}}});
-      registry.add("mcTruth/collType", "Collision type", {HistType::kTH1F, {{5, -0.5, 4.5}}});
-      registry.add("mcTruth/IVMpt", "Invariant mass versus p_{T}", {HistType::kTH2F, {{150, 0.0, 3.0}, {150, 0.0, 3.0}}});
-    }
-  }
-
-  // process function for MC data
-  // save the MC truth of all events of interest and of the DG events
-  void processMCTruth(aod::McCollisions const& mccols, aod::McParticles const& mcparts,
-                      UDCCs const& dgcands, UDTCs const& udtracks,
-                      CCs const& /*collisions*/, BCs const& /*bcs*/, TCs const& /*tracks*/)
+  // updating McTruth data and links to reconstructed data
+  void procWithDgCand(aod::McCollisions const& mccols, aod::McParticles const& mcparts,
+                      UDCCs const& dgcands, UDTCs const& udtracks)
   {
-    LOGF(debug, "Number of McCollisions %d", mccols.size());
-    LOGF(debug, "Number of DG candidates %d", dgcands.size());
-    LOGF(debug, "Number of UD tracks %d", udtracks.size());
-    if (dgcands.size() <= 0) {
-      LOGF(info, "No DG candidates to save!");
-      return;
-    }
-
     // use a hash table to keep track of the McCollisions which have been added to the UDMcCollision table
     // {McCollisionId : udMcCollisionId}
     // similar for the McParticles which have been added to the UDMcParticle table
@@ -703,18 +686,19 @@ struct McDGCandProducer {
 
     // loop over McCollisions and UDCCs simultaneously
     auto mccol = mccols.iteratorAt(0);
-    auto dgcand = dgcands.iteratorAt(0);
     auto lastmccol = mccols.iteratorAt(mccols.size() - 1);
+    auto mccolAtEnd = false;
+
+    auto dgcand = dgcands.iteratorAt(0);
     auto lastdgcand = dgcands.iteratorAt(dgcands.size() - 1);
+    auto dgcandAtEnd = false;
 
     // advance dgcand and mccol until both are AtEnd
     int64_t mccolId = mccol.globalIndex();
     int64_t mcdgId = -1;
     int64_t colId = -1;
-    auto dgcandAtEnd = dgcand == lastdgcand;
-    auto mccolAtEnd = mccol == lastmccol;
-    bool goon = !dgcandAtEnd || !mccolAtEnd;
 
+    bool goon = true;
     while (goon) {
       // check if dgcand has an associated Collision and McCollision
       if (dgcand.has_collision()) {
@@ -825,10 +809,71 @@ struct McDGCandProducer {
           mccolAtEnd = true;
         }
       }
-      LOGF(debug, "  UDMcCollsLabels %d (of %d) UDMcCollisions %d", outputMcCollsLabels.lastIndex(), dgcands.size() - 1, outputMcCollisions.lastIndex());
+      LOGF(info, "  UDMcCollsLabels %d (of %d) UDMcCollisions %d", outputMcCollsLabels.lastIndex(), dgcands.size() - 1, outputMcCollisions.lastIndex());
       goon = !dgcandAtEnd || !mccolAtEnd;
     }
   }
+
+  // updating McTruth data only
+  void procWithoutDgCand(aod::McCollisions const& mccols, aod::McParticles const& mcparts)
+  {
+    // use a hash table to keep track of the McCollisions which have been added to the UDMcCollision table
+    // {McCollisionId : udMcCollisionId}
+    // similar for the McParticles which have been added to the UDMcParticle table
+    // {McParticleId : udMcParticleId}
+    std::map<int64_t, int64_t> mcColIsSaved;
+    std::map<int64_t, int64_t> mcPartIsSaved;
+
+    // loop over McCollisions
+    for (auto const& mccol : mccols) {
+      int64_t mccolId = mccol.globalIndex();
+
+      // update UDMcCollisions and UDMcParticles
+      if (mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
+
+        // update UDMcCollisions
+        LOGF(debug, "  writing mcCollision %d to UDMcCollisions", mccolId);
+        updateUDMcCollisions(mccol);
+        mcColIsSaved[mccolId] = outputMcCollisions.lastIndex();
+
+        // update UDMcParticles
+        auto mcPartsSlice = mcparts.sliceBy(mcPartsPerMcCollision, mccolId);
+        updateUDMcParticles(mcPartsSlice, mcColIsSaved[mccolId], mcPartIsSaved);
+      }
+    }
+  }
+
+  void init(InitContext& context)
+  {
+    // add histograms for the different process functions
+    if (context.mOptions.get<bool>("processMCTruth")) {
+      LOGF(info, "Preparing histograms for processMCTruth.");
+      registry.add("mcTruth/collisions", "Number of associated collisions", {HistType::kTH1F, {{11, -0.5, 10.5}}});
+      registry.add("mcTruth/collType", "Collision type", {HistType::kTH1F, {{5, -0.5, 4.5}}});
+      registry.add("mcTruth/IVMpt", "Invariant mass versus p_{T}", {HistType::kTH2F, {{150, 0.0, 3.0}, {150, 0.0, 3.0}}});
+    }
+  }
+
+  // process function for MC data
+  // save the MC truth of all events of interest and of the DG events
+  void processMCTruth(aod::McCollisions const& mccols, aod::McParticles const& mcparts,
+                      UDCCs const& dgcands, UDTCs const& udtracks,
+                      CCs const& /*collisions*/, BCs const& /*bcs*/, TCs const& /*tracks*/)
+  {
+    LOGF(info, "Number of McCollisions %d", mccols.size());
+    LOGF(info, "Number of DG candidates %d", dgcands.size());
+    LOGF(info, "Number of UD tracks %d", udtracks.size());
+
+    if (mccols.size() > 0) {
+      if (dgcands.size() > 0) {
+        procWithDgCand(mccols, mcparts, dgcands, udtracks);
+      } else {
+        if (saveAllMcCollisions) {
+          procWithoutDgCand(mccols, mcparts);
+        }
+      }
+    }
+  }
   PROCESS_SWITCH(McDGCandProducer, processMCTruth, "Produce MC tables", false);
 
   void processDummy(aod::Collisions const& /*collisions*/)

PWGUD/TableProducer/SGCandProducer.cxx

@@ -365,6 +365,9 @@ struct McSGCandProducer {
   Produces<aod::UDMcCollsLabels> outputMcCollsLabels;
   Produces<aod::UDMcTrackLabels> outputMcTrackLabels;
 
+  // save all McTruth, even if the collisions is not reconstructed
+  Configurable<bool> saveAllMcCollisions{"saveAllMcCollisions", true, "save all McCollisions"};
+
   using CCs = soa::Join<aod::Collisions, aod::EvSels, aod::McCollisionLabels>;
   using BCs = soa::Join<aod::BCsWithTimestamps, aod::BcSels, aod::Run3MatchedToBCSparse>;
   using TCs = soa::Join<aod::Tracks, aod::McTrackLabels>;
@@ -539,33 +542,10 @@ struct McSGCandProducer {
     }
   }
 
-  void init(InitContext& context)
+  // updating McTruth data and links to reconstructed data
+  void procWithSgCand(aod::McCollisions const& mccols, aod::McParticles const& mcparts,
+                      UDCCs const& sgcands, UDTCs const& udtracks)
   {
-    // add histograms for the different process functions
-    if (context.mOptions.get<bool>("processMC")) {
-      registry.add("mcTruth/collisions", "Number of associated collisions", {HistType::kTH1F, {{11, -0.5, 10.5}}});
-      registry.add("mcTruth/collType", "Collision type", {HistType::kTH1F, {{5, -0.5, 4.5}}});
-      registry.add("mcTruth/IVMpt", "Invariant mass versus p_{T}", {HistType::kTH2F, {{150, 0.0, 3.0}, {150, 0.0, 3.0}}});
-    }
-  }
-
-  // process function for MC data
-  // save the MC truth of all events of interest and of the DG events
-  void processMC(aod::McCollisions const& mccols, aod::McParticles const& mcparts,
-                 UDCCs const& sgcands, UDTCs const& udtracks,
-                 CCs const& /*collisions*/, BCs const& /*bcs*/, TCs const& /*tracks*/)
-  {
-    if (verboseInfoMC) {
-      LOGF(info, "Number of McCollisions %d", mccols.size());
-      LOGF(info, "Number of SG candidates %d", sgcands.size());
-      LOGF(info, "Number of UD tracks %d", udtracks.size());
-    }
-    if (sgcands.size() <= 0) {
-      if (verboseInfoMC)
-        LOGF(info, "No DG candidates to save!");
-      return;
-    }
-
     // use a hash table to keep track of the McCollisions which have been added to the UDMcCollision table
     // {McCollisionId : udMcCollisionId}
     // similar for the McParticles which have been added to the UDMcParticle table
@@ -575,21 +555,20 @@ struct McSGCandProducer {
 
     // loop over McCollisions and UDCCs simultaneously
     auto mccol = mccols.iteratorAt(0);
-    auto sgcand = sgcands.iteratorAt(0);
     auto lastmccol = mccols.iteratorAt(mccols.size() - 1);
+    auto mccolAtEnd = false;
+
+    auto sgcand = sgcands.iteratorAt(0);
     auto lastsgcand = sgcands.iteratorAt(sgcands.size() - 1);
+    auto sgcandAtEnd = false;
 
     // advance dgcand and mccol until both are AtEnd
     int64_t mccolId = mccol.globalIndex();
     int64_t mcsgId = -1;
-    // int64_t colId = -1;
-    auto sgcandAtEnd = sgcand == lastsgcand;
-    auto mccolAtEnd = mccol == lastmccol;
-    bool goon = !sgcandAtEnd || !mccolAtEnd;
+
+    bool goon = true;
     while (goon) {
-      auto bcIter = mccol.bc_as<BCs>();
-      uint64_t globBC = bcIter.globalBC();
-      // uint64_t globBC = 0;
+      auto globBC = mccol.bc_as<BCs>().globalBC();
       // check if dgcand has an associated McCollision
       if (sgcand.has_collision()) {
         auto sgcandCol = sgcand.collision_as<CCs>();
@@ -600,7 +579,6 @@ struct McSGCandProducer {
           mcsgId = -1;
         }
       } else {
-        //  colId = -1;
         mcsgId = -1;
       }
       if (verboseInfoMC)
@@ -659,7 +637,7 @@ struct McSGCandProducer {
                 auto mcPart = track.mcParticle();
                 auto mcCol = mcPart.mcCollision();
                 if (mcColIsSaved.find(mcCol.globalIndex()) == mcColIsSaved.end()) {
-                  updateUDMcCollisions(mcCol, mcCol.bc_as<BCs>().globalBC());
+                  updateUDMcCollisions(mcCol, globBC);
                   mcColIsSaved[mcCol.globalIndex()] = outputMcCollisions.lastIndex();
                 }
                 updateUDMcParticle(mcPart, mcColIsSaved[mcCol.globalIndex()], mcPartIsSaved);
@@ -687,6 +665,7 @@ struct McSGCandProducer {
         if (mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
           if (verboseInfoMC)
             LOGF(info, "  Saving McCollision %d", mccolId);
+
           // update UDMcCollisions
           updateUDMcCollisions(mccol, globBC);
           mcColIsSaved[mccolId] = outputMcCollisions.lastIndex();
@@ -710,7 +689,71 @@ struct McSGCandProducer {
         LOGF(info, "End of loop mcsgId %d mccolId %d", mcsgId, mccolId);
     }
   }
+
+  // updating McTruth data only
+  void procWithoutSgCand(aod::McCollisions const& mccols, aod::McParticles const& mcparts)
+  {
+    // use a hash table to keep track of the McCollisions which have been added to the UDMcCollision table
+    // {McCollisionId : udMcCollisionId}
+    // similar for the McParticles which have been added to the UDMcParticle table
+    // {McParticleId : udMcParticleId}
+    std::map<int64_t, int64_t> mcColIsSaved;
+    std::map<int64_t, int64_t> mcPartIsSaved;
+
+    // loop over McCollisions
+    for (auto const& mccol : mccols) {
+      int64_t mccolId = mccol.globalIndex();
+      uint64_t globBC = mccol.bc_as<BCs>().globalBC();
+
+      // update UDMcCollisions and UDMcParticles
+      if (mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
+        if (verboseInfoMC)
+          LOGF(info, "  Saving McCollision %d", mccolId);
+
+        // update UDMcCollisions
+        updateUDMcCollisions(mccol, globBC);
+        mcColIsSaved[mccolId] = outputMcCollisions.lastIndex();
+
+        // update UDMcParticles
+        auto mcPartsSlice = mcparts.sliceBy(mcPartsPerMcCollision, mccolId);
+        updateUDMcParticles(mcPartsSlice, mcColIsSaved[mccolId], mcPartIsSaved);
+      }
+    }
+  }
+
+  void init(InitContext& context)
+  {
+    // add histograms for the different process functions
+    if (context.mOptions.get<bool>("processMC")) {
+      registry.add("mcTruth/collisions", "Number of associated collisions", {HistType::kTH1F, {{11, -0.5, 10.5}}});
+      registry.add("mcTruth/collType", "Collision type", {HistType::kTH1F, {{5, -0.5, 4.5}}});
+      registry.add("mcTruth/IVMpt", "Invariant mass versus p_{T}", {HistType::kTH2F, {{150, 0.0, 3.0}, {150, 0.0, 3.0}}});
+    }
+  }
+
+  // process function for MC data
+  // save the MC truth of all events of interest and of the DG events
+  void processMC(aod::McCollisions const& mccols, aod::McParticles const& mcparts,
+                 UDCCs const& sgcands, UDTCs const& udtracks,
+                 CCs const& /*collisions*/, BCs const& /*bcs*/, TCs const& /*tracks*/)
+  {
+    if (verboseInfoMC) {
+      LOGF(info, "Number of McCollisions %d", mccols.size());
+      LOGF(info, "Number of SG candidates %d", sgcands.size());
+      LOGF(info, "Number of UD tracks %d", udtracks.size());
+    }
+    if (mccols.size() > 0) {
+      if (sgcands.size() > 0) {
+        procWithSgCand(mccols, mcparts, sgcands, udtracks);
+      } else {
+        if (saveAllMcCollisions) {
+          procWithoutSgCand(mccols, mcparts);
+        }
+      }
+    }
+  }
   PROCESS_SWITCH(McSGCandProducer, processMC, "Produce MC tables", false);
+
   void processDummy(aod::Collisions const& /*collisions*/)
   {
     // do nothing