Merge pull request #7497 from chiarazampolli/AODproducer

Use mc flag to be able to run with data

Author: chiarazampolli

Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h

@@ -191,13 +191,15 @@ typedef boost::unordered_map<Triplet_t, int, TripletHash, TripletEqualTo> Triple
 class AODProducerWorkflowDPL : public Task
 {
  public:
-  AODProducerWorkflowDPL(GID::mask_t src, std::shared_ptr<DataRequest> dataRequest) : mInputSources(src), mDataRequest(dataRequest) {}
+  AODProducerWorkflowDPL(GID::mask_t src, std::shared_ptr<DataRequest> dataRequest, bool useMC = true) : mInputSources(src), mDataRequest(dataRequest), mUseMC(useMC) {}
   ~AODProducerWorkflowDPL() override = default;
   void init(InitContext& ic) final;
   void run(ProcessingContext& pc) final;
   void endOfStream(framework::EndOfStreamContext& ec) final;
 
  private:
+  bool mUseMC = true;
+
   const float cSpeed = 0.029979246f; // speed of light in TOF units
 
   GID::mask_t mInputSources;

Detectors/AOD/src/AODProducerWorkflowSpec.cxx

@@ -977,17 +977,15 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   auto v0sCursor = v0sBuilder.cursor<o2::aod::StoredV0s>();
   auto zdcCursor = zdcBuilder.cursor<o2::aod::Zdcs>();
 
-  o2::steer::MCKinematicsReader mcReader("collisioncontext.root");
-  const auto mcContext = mcReader.getDigitizationContext();
-  const auto& mcRecords = mcContext->getEventRecords();
-  const auto& mcParts = mcContext->getEventParts();
-
-  LOG(DEBUG) << "FOUND " << mcRecords.size() << " records";
-  LOG(DEBUG) << "FOUND " << mcParts.size() << " parts";
+  std::unique_ptr<o2::steer::MCKinematicsReader> mcReader;
+  if (mUseMC) {
+    mcReader = std::make_unique<o2::steer::MCKinematicsReader>("collisioncontext.root");
+    LOG(DEBUG) << "FOUND " << mcReader->getDigitizationContext()->getEventRecords().size()
+               << " records" << mcReader->getDigitizationContext()->getEventParts().size() << " parts";
+  }
 
   std::map<uint64_t, int> bcsMap;
-  collectBCs(fddRecPoints, ft0RecPoints, fv0RecPoints, primVertices, mcRecords, bcsMap);
-
+  collectBCs(fddRecPoints, ft0RecPoints, fv0RecPoints, primVertices, mUseMC ? mcReader->getDigitizationContext()->getEventRecords() : std::vector<o2::InteractionTimeRecord>{}, bcsMap);
   const auto* dh = o2::header::get<o2::header::DataHeader*>(pc.inputs().getFirstValid(true).header);
   o2::InteractionRecord startIR = {0, dh->firstTForbit};
 
@@ -1070,40 +1068,44 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   // correct ordering when iterating over container elements
   std::map<std::pair<int, int>, int> mcColToEvSrc;
 
-  // TODO: figure out collision weight
-  float mcColWeight = 1.;
-  // filling mcCollision table
-  int nMCCollisions = mcContext->getNCollisions();
-  for (int iCol = 0; iCol < nMCCollisions; iCol++) {
-    auto time = mcRecords[iCol].getTimeNS();
-    auto globalBC = mcRecords[iCol].toLong();
-    auto item = bcsMap.find(globalBC);
-    int bcID = -1;
-    if (item != bcsMap.end()) {
-      bcID = item->second;
-    } else {
-      LOG(FATAL) << "Error: could not find a corresponding BC ID for MC collision; BC = " << globalBC << ", mc collision = " << iCol;
-    }
-    auto& colParts = mcParts[iCol];
-    for (auto colPart : colParts) {
-      auto eventID = colPart.entryID;
-      auto sourceID = colPart.sourceID;
-      if (sourceID == 0) { // embedding: using background event info
-        // FIXME:
-        // use generators' names for generatorIDs (?)
-        short generatorID = sourceID;
-        auto& header = mcReader.getMCEventHeader(sourceID, eventID);
-        mcCollisionsCursor(0,
-                           bcID,
-                           generatorID,
-                           truncateFloatFraction(header.GetX(), mCollisionPosition),
-                           truncateFloatFraction(header.GetY(), mCollisionPosition),
-                           truncateFloatFraction(header.GetZ(), mCollisionPosition),
-                           truncateFloatFraction(time, mCollisionPosition),
-                           truncateFloatFraction(mcColWeight, mCollisionPosition),
-                           header.GetB());
+  if (mUseMC) {
+    // TODO: figure out collision weight
+    float mcColWeight = 1.;
+    // filling mcCollision table
+    int nMCCollisions = mcReader->getDigitizationContext()->getNCollisions();
+    const auto& mcRecords = mcReader->getDigitizationContext()->getEventRecords();
+    const auto& mcParts = mcReader->getDigitizationContext()->getEventParts();
+    for (int iCol = 0; iCol < nMCCollisions; iCol++) {
+      auto time = mcRecords[iCol].getTimeNS();
+      auto globalBC = mcRecords[iCol].toLong();
+      auto item = bcsMap.find(globalBC);
+      int bcID = -1;
+      if (item != bcsMap.end()) {
+        bcID = item->second;
+      } else {
+        LOG(FATAL) << "Error: could not find a corresponding BC ID for MC collision; BC = " << globalBC << ", mc collision = " << iCol;
+      }
+      auto& colParts = mcParts[iCol];
+      for (auto colPart : colParts) {
+        auto eventID = colPart.entryID;
+        auto sourceID = colPart.sourceID;
+        if (sourceID == 0) { // embedding: using background event info
+          // FIXME:
+          // use generators' names for generatorIDs (?)
+          short generatorID = sourceID;
+          auto& header = mcReader->getMCEventHeader(sourceID, eventID);
+          mcCollisionsCursor(0,
+                             bcID,
+                             generatorID,
+                             truncateFloatFraction(header.GetX(), mCollisionPosition),
+                             truncateFloatFraction(header.GetY(), mCollisionPosition),
+                             truncateFloatFraction(header.GetZ(), mCollisionPosition),
+                             truncateFloatFraction(time, mCollisionPosition),
+                             truncateFloatFraction(mcColWeight, mCollisionPosition),
+                             header.GetB());
+        }
+        mcColToEvSrc.emplace(std::pair<int, int>(eventID, sourceID), iCol); // point background and injected signal events to one collision
       }
-      mcColToEvSrc.emplace(std::pair<int, int>(eventID, sourceID), iCol); // point background and injected signal events to one collision
     }
   }
 
@@ -1182,12 +1184,14 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
               ft0RecPoint.getTrigger().triggersignals);
   }
 
-  // filling MC collision labels
-  for (auto& label : primVerLabels) {
-    auto it = mcColToEvSrc.find(std::pair<int, int>(label.getEventID(), label.getSourceID()));
-    int32_t mcCollisionID = it != mcColToEvSrc.end() ? it->second : -1;
-    uint16_t mcMask = 0; // todo: set mask using normalized weights?
-    mcColLabelsCursor(0, mcCollisionID, mcMask);
+  if (mUseMC) {
+    // filling MC collision labels
+    for (auto& label : primVerLabels) {
+      auto it = mcColToEvSrc.find(std::pair<int, int>(label.getEventID(), label.getSourceID()));
+      int32_t mcCollisionID = it != mcColToEvSrc.end() ? it->second : -1;
+      uint16_t mcMask = 0; // todo: set mask using normalized weights?
+      mcColLabelsCursor(0, mcCollisionID, mcMask);
+    }
   }
 
   // hash map for track indices of secondary vertices
@@ -1291,26 +1295,27 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
 
   bcsMap.clear();
 
-  // filling mc particles table
-  fillMCParticlesTable(mcReader,
-                       mcParticlesCursor,
-                       primVer2TRefs,
-                       primVerGIs,
-                       recoData,
-                       mcColToEvSrc);
-
-  mcColToEvSrc.clear();
-
-  // ------------------------------------------------------
-  // filling track labels
-
-  // need to go through labels in the same order as for tracks
-  fillMCTrackLabelsTable(mcTrackLabelCursor, mcMFTTrackLabelCursor, mcFwdTrackLabelCursor, primVer2TRefs.back(), primVerGIs, recoData);
-  for (int iref = 0; iref < primVer2TRefs.size() - 1; iref++) {
-    auto& trackRef = primVer2TRefs[iref];
-    fillMCTrackLabelsTable(mcTrackLabelCursor, mcMFTTrackLabelCursor, mcFwdTrackLabelCursor, trackRef, primVerGIs, recoData);
+  if (mUseMC) {
+    // filling mc particles table
+    fillMCParticlesTable(*mcReader,
+                         mcParticlesCursor,
+                         primVer2TRefs,
+                         primVerGIs,
+                         recoData,
+                         mcColToEvSrc);
+
+    mcColToEvSrc.clear();
+
+    // ------------------------------------------------------
+    // filling track labels
+
+    // need to go through labels in the same order as for tracks
+    fillMCTrackLabelsTable(mcTrackLabelCursor, mcMFTTrackLabelCursor, mcFwdTrackLabelCursor, primVer2TRefs.back(), primVerGIs, recoData);
+    for (int iref = 0; iref < primVer2TRefs.size() - 1; iref++) {
+      auto& trackRef = primVer2TRefs[iref];
+      fillMCTrackLabelsTable(mcTrackLabelCursor, mcMFTTrackLabelCursor, mcFwdTrackLabelCursor, trackRef, primVerGIs, recoData);
+    }
   }
-
   mToStore.clear();
 
   pc.outputs().snapshot(Output{"TFN", "TFNumber", 0, Lifetime::Timeframe}, tfNumber);
@@ -1363,7 +1368,7 @@ DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool useMC)
     "aod-producer-workflow",
     dataRequest->inputs,
     outputs,
-    AlgorithmSpec{adaptFromTask<AODProducerWorkflowDPL>(src, dataRequest)},
+    AlgorithmSpec{adaptFromTask<AODProducerWorkflowDPL>(src, dataRequest, useMC)},
     Options{
       ConfigParamSpec{"aod-timeframe-id", VariantType::Int64, -1L, {"Set timeframe number"}},
       ConfigParamSpec{"enable-truncation", VariantType::Int, 1, {"Truncation parameter: 1 -- on, != 1 -- off"}},