Add possibility to apply signal filtering for MC with embedding (#14698)

* Add possibility to apply signal filtering for MC with embedding
* Add protection for signal filtering to be enabled only with emdedding

Author: fgrosa

Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h

@@ -256,6 +256,7 @@ class AODProducerWorkflowDPL : public Task
 
   int mNThreads = 1;
   bool mUseMC = true;
+  bool mUseSigFiltMC = false; // enable signal filtering for MC with embedding
   bool mEnableSV = true; // enable secondary vertices
   bool mEnableFITextra = false;
   bool mFieldON = false;

Detectors/AOD/src/AODProducerWorkflowSpec.cxx

@@ -947,13 +947,17 @@ void clearMCKeepStore(std::vector<std::vector<std::unordered_map<int, int>>>& st
 }
 
 // helper function to add a particle/track to the MC keep store
-void keepMCParticle(std::vector<std::vector<std::unordered_map<int, int>>>& store, int source, int event, int track, int value = 1)
+void keepMCParticle(std::vector<std::vector<std::unordered_map<int, int>>>& store, int source, int event, int track, int value = 1, bool useSigFilt = false)
 {
   if (track < 0) {
     LOG(warn) << "trackID is smaller than 0. Neglecting";
     return;
   }
-  store[source][event][track] = value;
+  if (useSigFilt && source == 0) {
+    store[source][event][track] = -1;
+  } else {
+    store[source][event][track] = value;
+  }
 }
 
 void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader& mcReader,
@@ -982,7 +986,7 @@ void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader&
     if (!mcLabel.isValid()) {
       return;
     }
-    keepMCParticle(mToStore, mcLabel.getSourceID(), mcLabel.getEventID(), mcLabel.getTrackID());
+    keepMCParticle(mToStore, mcLabel.getSourceID(), mcLabel.getEventID(), mcLabel.getTrackID(), 1, mUseSigFiltMC);
   };
 
   // mark reconstructed MC particles to store them into the table
@@ -997,7 +1001,7 @@ void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader&
           if (!mcTruth.isValid()) {
             continue;
           }
-          keepMCParticle(mToStore, mcTruth.getSourceID(), mcTruth.getEventID(), mcTruth.getTrackID());
+          keepMCParticle(mToStore, mcTruth.getSourceID(), mcTruth.getEventID(), mcTruth.getTrackID(), 1, mUseSigFiltMC);
           // treating contributors of global tracks
           auto contributorsGID = data.getSingleDetectorRefs(trackIndex);
           if (contributorsGID[GIndex::Source::TPC].isIndexSet()) {
@@ -1012,7 +1016,7 @@ void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader&
               if (!mcLabel.isValid()) {
                 continue;
               }
-              keepMCParticle(mToStore, mcLabel.getSourceID(), mcLabel.getEventID(), mcLabel.getTrackID());
+              keepMCParticle(mToStore, mcLabel.getSourceID(), mcLabel.getEventID(), mcLabel.getTrackID(), 1, mUseSigFiltMC);
             }
           }
         }
@@ -1026,7 +1030,7 @@ void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader&
       if (!mcTruth.isValid()) {
         continue;
       }
-      keepMCParticle(mToStore, mcTruth.getSourceID(), mcTruth.getEventID(), mcTruth.getTrackID());
+      keepMCParticle(mToStore, mcTruth.getSourceID(), mcTruth.getEventID(), mcTruth.getTrackID(), 1, mUseSigFiltMC);
     }
   }
   if (mInputSources[GIndex::PHS]) {
@@ -1035,7 +1039,7 @@ void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader&
       if (!mcTruth.isValid()) {
         continue;
       }
-      keepMCParticle(mToStore, mcTruth.getSourceID(), mcTruth.getEventID(), mcTruth.getTrackID());
+      keepMCParticle(mToStore, mcTruth.getSourceID(), mcTruth.getEventID(), mcTruth.getTrackID(), 1, mUseSigFiltMC);
     }
   }
   using namespace aodmchelpers;
@@ -1743,6 +1747,8 @@ void AODProducerWorkflowDPL::init(InitContext& ic)
     LOG(info) << "The Run number will be obtained from DPL headers";
   }
 
+  mUseSigFiltMC = ic.options().get<bool>("mc-signal-filt");
+
   // set no truncation if selected by user
   if (mTruncate != 1) {
     LOG(info) << "Truncation is not used!";
@@ -2061,6 +2067,24 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
     int totalNParts = 0;
     for (int iCol = 0; iCol < nMCCollisions; iCol++) {
       totalNParts += mcParts[iCol].size();
+
+      // if signal filtering enabled, let's check if there are more than one source; otherwise fatalise
+      if (mUseSigFiltMC) {
+        std::vector<int> sourceIDs{};
+        auto& colParts = mcParts[iCol];
+        for (auto colPart : colParts) {
+          int sourceID = colPart.sourceID;
+          if (std::find(sourceIDs.begin(), sourceIDs.end(), sourceID) == sourceIDs.end()) {
+            sourceIDs.push_back(sourceID);
+          }
+          if (sourceIDs.size() > 1) { // we found more than one, exit
+            break;
+          }
+        }
+        if (sourceIDs.size() <= 1) {
+          LOGP(fatal, "Signal filtering cannot be enabled without embedding. Please fix the configuration either enabling the embedding, or turning off the signal filtering.");
+        }
+      }
     }
     mcCollisionsCursor.reserve(totalNParts);
 
@@ -2098,7 +2122,9 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
                          0,
                          sourceID);
         }
-        mcColToEvSrc.emplace_back(std::vector<int>{iCol, sourceID, eventID}); // point background and injected signal events to one collision
+        if (sourceID != 0 || !mUseSigFiltMC) {
+          mcColToEvSrc.emplace_back(std::vector<int>{iCol, sourceID, eventID}); // point background and injected signal events to one collision
+        }
       }
     }
   }
@@ -3326,7 +3352,7 @@ DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, boo
       ConfigParamSpec{"trackqc-tpc-pt", VariantType::Float, 0.2f, {"Keep TPC standalone track with this pt"}},
       ConfigParamSpec{"with-streamers", VariantType::String, "", {"Bit-mask to steer writing of intermediate streamer files"}},
       ConfigParamSpec{"seed", VariantType::Int, 0, {"Set seed for random generator used for sampling (0 (default) means using a random_device)"}},
-    }};
+      ConfigParamSpec{"mc-signal-filt", VariantType::Bool, false, {"Enable usage of signal filtering (only for MC with embedding)"}}}};
 }
 
 } // namespace o2::aodproducer