[PWGUD] Added configurable generatorIds in Mc[D,S]GCandProducer (#10427)

Author: pbuehler

PWGUD/TableProducer/DGCandProducer.cxx

@@ -473,6 +473,7 @@ struct McDGCandProducer {
   Produces<aod::UDMcTrackLabels> outputMcTrackLabels;
 
   // save all McTruth, even if the collisions is not reconstructed
+  Configurable<std::vector<int>> generatorIds{"generatorIds", std::vector<int>{-1}, "MC generatorIds to process"};
   Configurable<bool> saveAllMcCollisions{"saveAllMcCollisions", true, "save all McCollisions"};
 
   using CCs = soa::Join<aod::Collisions, aod::EvSels, aod::McCollisionLabels>;
@@ -686,6 +687,7 @@ struct McDGCandProducer {
 
     // loop over McCollisions and UDCCs simultaneously
     auto mccol = mccols.iteratorAt(0);
+    auto mcOfInterest = std::find(generatorIds->begin(), generatorIds->end(), mccol.getGeneratorId()) != generatorIds->end();
     auto lastmccol = mccols.iteratorAt(mccols.size() - 1);
     auto mccolAtEnd = false;
 
@@ -729,7 +731,9 @@ struct McDGCandProducer {
 
         // If the dgcand has an associated McCollision then the McCollision and all associated
         // McParticles are saved
-        if (mcdgId >= 0) {
+        // but only consider generated events of interest
+        if (mcdgId >= 0 && mcOfInterest) {
+
           if (mcColIsSaved.find(mcdgId) == mcColIsSaved.end()) {
             // update UDMcCollisions
             LOGF(debug, "  writing mcCollision %d to UDMcCollisions", mcdgId);
@@ -789,7 +793,8 @@ struct McDGCandProducer {
         // this is case 2.
 
         // update UDMcCollisions and UDMcParticles
-        if (mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
+        // but only consider generated events of interest
+        if (mcOfInterest && mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
 
           // update UDMcCollisions
           LOGF(debug, "  writing mcCollision %d to UDMcCollisions", mccolId);
@@ -804,6 +809,7 @@ struct McDGCandProducer {
         // advance mccol
         if (mccol != lastmccol) {
           mccol++;
+          mcOfInterest = std::find(generatorIds->begin(), generatorIds->end(), mccol.getGeneratorId()) != generatorIds->end();
           mccolId = mccol.globalIndex();
         } else {
           mccolAtEnd = true;
@@ -826,8 +832,11 @@ struct McDGCandProducer {
 
     // loop over McCollisions
     for (auto const& mccol : mccols) {
-      int64_t mccolId = mccol.globalIndex();
+      // only consider generated events of interest
+      if (std::find(generatorIds->begin(), generatorIds->end(), mccol.getGeneratorId()) == generatorIds->end())
+        continue;
 
+      int64_t mccolId = mccol.globalIndex();
       // update UDMcCollisions and UDMcParticles
       if (mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
 

PWGUD/TableProducer/SGCandProducer.cxx

@@ -366,6 +366,7 @@ struct McSGCandProducer {
   Produces<aod::UDMcTrackLabels> outputMcTrackLabels;
 
   // save all McTruth, even if the collisions is not reconstructed
+  Configurable<std::vector<int>> generatorIds{"generatorIds", std::vector<int>{-1}, "MC generatorIds to process"};
   Configurable<bool> saveAllMcCollisions{"saveAllMcCollisions", true, "save all McCollisions"};
 
   using CCs = soa::Join<aod::Collisions, aod::EvSels, aod::McCollisionLabels>;
@@ -555,6 +556,7 @@ struct McSGCandProducer {
 
     // loop over McCollisions and UDCCs simultaneously
     auto mccol = mccols.iteratorAt(0);
+    auto mcOfInterest = std::find(generatorIds->begin(), generatorIds->end(), mccol.getGeneratorId()) != generatorIds->end();
     auto lastmccol = mccols.iteratorAt(mccols.size() - 1);
     auto mccolAtEnd = false;
 
@@ -565,7 +567,6 @@ struct McSGCandProducer {
     // advance dgcand and mccol until both are AtEnd
     int64_t mccolId = mccol.globalIndex();
     int64_t mcsgId = -1;
-
     bool goon = true;
     while (goon) {
       auto globBC = mccol.bc_as<BCs>().globalBC();
@@ -599,7 +600,8 @@ struct McSGCandProducer {
 
         // If the sgcand has an associated McCollision then the McCollision and all associated
         // McParticles are saved
-        if (mcsgId >= 0) {
+        // but only consider generated events of interest
+        if (mcsgId >= 0 && mcOfInterest) {
           if (mcColIsSaved.find(mcsgId) == mcColIsSaved.end()) {
             if (verboseInfoMC)
               LOGF(info, "  Saving McCollision %d", mcsgId);
@@ -662,10 +664,10 @@ struct McSGCandProducer {
           LOGF(info, "Doing case 2");
 
         // update UDMcCollisions and UDMcParticles
-        if (mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
+        // but only consider generated events of interest
+        if (mcOfInterest && mcColIsSaved.find(mccolId) == mcColIsSaved.end()) {
           if (verboseInfoMC)
             LOGF(info, "  Saving McCollision %d", mccolId);
-
           // update UDMcCollisions
           updateUDMcCollisions(mccol, globBC);
           mcColIsSaved[mccolId] = outputMcCollisions.lastIndex();
@@ -678,6 +680,7 @@ struct McSGCandProducer {
         // advance mccol
         if (mccol != lastmccol) {
           mccol++;
+          mcOfInterest = std::find(generatorIds->begin(), generatorIds->end(), mccol.getGeneratorId()) != generatorIds->end();
           mccolId = mccol.globalIndex();
         } else {
           mccolAtEnd = true;