extract digits from selected BC ranges if any (#13625)

pillot · alcaliva · commit 253d92007c91 · 2025-03-04T14:15:11.000+01:00
diff --git a/Detectors/MUON/MCH/IO/CMakeLists.txt b/Detectors/MUON/MCH/IO/CMakeLists.txt
@@ -25,6 +25,7 @@ o2_add_library(MCHIO
            O2::DataFormatsMCH
            O2::Framework
            O2::MCHBase
+           O2::MCHDigitFiltering
            O2::SimulationDataFormat
            )
 
diff --git a/Detectors/MUON/MCH/IO/include/MCHIO/DigitReaderSpec.h b/Detectors/MUON/MCH/IO/include/MCHIO/DigitReaderSpec.h
@@ -27,7 +27,8 @@ framework::DataProcessorSpec getDigitReaderSpec(
   bool useMC,
   std::string_view specName = "mch-digit-reader",
   std::string_view outputDigitDataDescription = "DIGITS",
-  std::string_view outputDigitRofDataDescription = "DIGITROFS");
+  std::string_view outputDigitRofDataDescription = "DIGITROFS",
+  std::string_view outputDigitLabelDataDescription = "DIGITLABELS");
 }
 } // namespace o2
 
diff --git a/Detectors/MUON/MCH/IO/src/DigitReaderSpec.cxx b/Detectors/MUON/MCH/IO/src/DigitReaderSpec.cxx
@@ -17,21 +17,32 @@
 #include "MCHIO/DigitReaderSpec.h"
 
 #include <memory>
-#include <sstream>
+#include <stdexcept>
 #include <string>
-#include "DPLUtils/RootTreeReader.h"
+#include <vector>
+
+#include <fmt/format.h>
+
+#include <gsl/span>
+
+#include <TFile.h>
+#include <TTree.h>
+#include <TTreeReader.h>
+#include <TTreeReaderValue.h>
+
+#include "CommonDataFormat/IRFrame.h"
+#include "CommonUtils/IRFrameSelector.h"
+#include "CommonUtils/StringUtils.h"
+#include "DataFormatsMCH/Digit.h"
+#include "DataFormatsMCH/ROFRecord.h"
 #include "Framework/ConfigParamRegistry.h"
 #include "Framework/ControlService.h"
 #include "Framework/DataSpecUtils.h"
-#include "Framework/Logger.h"
-#include "Framework/Output.h"
 #include "Framework/Task.h"
 #include "Framework/WorkflowSpec.h"
-#include "SimulationDataFormat/MCTruthContainer.h"
+#include "MCHDigitFiltering/DigitFilterParam.h"
 #include "SimulationDataFormat/MCCompLabel.h"
-#include "DataFormatsMCH/Digit.h"
-#include "DataFormatsMCH/ROFRecord.h"
-#include "CommonUtils/StringUtils.h"
+#include "SimulationDataFormat/MCTruthContainer.h"
 
 using namespace o2::framework;
 
@@ -43,72 +54,171 @@ namespace mch
 class DigitsReaderDeviceDPL
 {
  public:
-  DigitsReaderDeviceDPL(bool useMC, const std::vector<header::DataDescription>& descriptions)
-    : mUseMC(useMC), mDescriptions(descriptions) {}
+  DigitsReaderDeviceDPL(bool useMC) : mUseMC(useMC)
+  {
+    if (mUseMC) {
+      mLabels = std::make_unique<TTreeReaderValue<dataformats::MCTruthContainer<MCCompLabel>>>(mTreeReader, "MCHMCLabels");
+    }
+  }
 
   void init(InitContext& ic)
   {
-    auto filename = o2::utils::Str::concat_string(o2::utils::Str::rectifyDirectory(ic.options().get<std::string>("input-dir")), ic.options().get<std::string>("mch-digit-infile"));
+    auto fileName = utils::Str::concat_string(utils::Str::rectifyDirectory(ic.options().get<std::string>("input-dir")),
+                                              ic.options().get<std::string>("mch-digit-infile"));
+    connectTree(fileName);
 
-    if (mUseMC) {
-      mReader = std::make_unique<RootTreeReader>("o2sim", filename.c_str(), -1,
-                                                 RootTreeReader::PublishingMode::Single,
-                                                 RootTreeReader::BranchDefinition<std::vector<Digit>>{
-                                                   Output{header::gDataOriginMCH, mDescriptions[0], 0, Lifetime::Timeframe}, "MCHDigit"},
-                                                 RootTreeReader::BranchDefinition<std::vector<ROFRecord>>{
-                                                   Output{header::gDataOriginMCH, mDescriptions[1], 0, Lifetime::Timeframe}, "MCHROFRecords"},
-                                                 RootTreeReader::BranchDefinition<dataformats::MCTruthContainer<MCCompLabel>>{
-                                                   Output{header::gDataOriginMCH, mDescriptions[2], 0, Lifetime::Timeframe}, "MCHMCLabels"});
-    } else {
-      mReader = std::make_unique<RootTreeReader>("o2sim", filename.c_str(), -1,
-                                                 RootTreeReader::PublishingMode::Single,
-                                                 RootTreeReader::BranchDefinition<std::vector<Digit>>{
-                                                   Output{header::gDataOriginMCH, mDescriptions[0], 0, Lifetime::Timeframe}, "MCHDigit"},
-                                                 RootTreeReader::BranchDefinition<std::vector<ROFRecord>>{
-                                                   Output{header::gDataOriginMCH, mDescriptions[1], 0, Lifetime::Timeframe}, "MCHROFRecords"});
+    if (ic.options().hasOption("ignore-irframes") && !ic.options().get<bool>("ignore-irframes")) {
+      mUseIRFrames = true;
     }
+
+    mTimeOffset = ic.options().get<bool>("no-time-offset") ? 0 : DigitFilterParam::Instance().timeOffset;
   }
 
   void run(ProcessingContext& pc)
   {
-    if ((++(*mReader))(pc) == false) {
-      pc.services().get<ControlService>().endOfStream();
+    if (mUseIRFrames) {
+      sendNextIRFrames(pc);
+    } else {
+      sendNextTF(pc);
     }
   }
 
  private:
-  std::unique_ptr<RootTreeReader> mReader{};
-  std::vector<header::DataDescription> mDescriptions{};
+  TTreeReader mTreeReader{};
+  TTreeReaderValue<std::vector<ROFRecord>> mRofs = {mTreeReader, "MCHROFRecords"};
+  TTreeReaderValue<std::vector<Digit>> mDigits = {mTreeReader, "MCHDigit"};
+  std::unique_ptr<TTreeReaderValue<dataformats::MCTruthContainer<MCCompLabel>>> mLabels{};
   bool mUseMC = true;
+  bool mUseIRFrames = false;
+  int mTimeOffset = 0;
+
+  void connectTree(std::string fileName)
+  {
+    auto file = TFile::Open(fileName.c_str());
+    if (!file || file->IsZombie()) {
+      throw std::invalid_argument(fmt::format("Opening file {} failed", fileName));
+    }
+
+    auto tree = file->Get<TTree>("o2sim");
+    if (!tree) {
+      throw std::invalid_argument(fmt::format("Tree o2sim not found in {}", fileName));
+    }
+    mTreeReader.SetTree(tree);
+    mTreeReader.Restart();
+  }
+
+  void sendNextTF(ProcessingContext& pc)
+  {
+    // load the next TF and check its validity (missing branch, ...)
+    if (!mTreeReader.Next()) {
+      throw std::invalid_argument(mTreeReader.fgEntryStatusText[mTreeReader.GetEntryStatus()]);
+    }
+
+    // send the whole TF
+    pc.outputs().snapshot(OutputRef{"rofs"}, *mRofs);
+    pc.outputs().snapshot(OutputRef{"digits"}, *mDigits);
+    if (mUseMC) {
+      pc.outputs().snapshot(OutputRef{"labels"}, **mLabels);
+    }
+
+    // stop here if it was the last one
+    if (mTreeReader.GetCurrentEntry() + 1 >= mTreeReader.GetEntries()) {
+      pc.services().get<ControlService>().endOfStream();
+      pc.services().get<ControlService>().readyToQuit(QuitRequest::Me);
+    }
+  }
+
+  void sendNextIRFrames(ProcessingContext& pc)
+  {
+    std::vector<ROFRecord> rofs{};
+    std::vector<Digit> digits{};
+    dataformats::MCTruthContainer<MCCompLabel> labels{};
+
+    // get the IR frames to select
+    auto irFrames = pc.inputs().get<gsl::span<dataformats::IRFrame>>("driverInfo");
+
+    if (!irFrames.empty()) {
+      utils::IRFrameSelector irfSel{};
+      irfSel.setSelectedIRFrames(irFrames, 0, 0, -mTimeOffset, true);
+      const auto irMin = irfSel.getIRFrames().front().getMin();
+      const auto irMax = irfSel.getIRFrames().back().getMax();
+
+      // load the first TF if not already done
+      bool loadNextTF = mTreeReader.GetCurrentEntry() < 0;
+
+      while (true) {
+        // load the next TF if requested
+        if (loadNextTF && !mTreeReader.Next()) {
+          throw std::invalid_argument(mTreeReader.fgEntryStatusText[mTreeReader.GetEntryStatus()]);
+        }
+
+        // look for selected ROFs in this TF and copy them
+        if (!mRofs->empty() && mRofs->front().getBCData() <= irMax &&
+            mRofs->back().getBCData() + mRofs->back().getBCWidth() - 1 >= irMin) {
+          for (const auto& rof : *mRofs) {
+            if (irfSel.check({rof.getBCData(), rof.getBCData() + rof.getBCWidth() - 1}) != -1) {
+              rofs.emplace_back(rof);
+              rofs.back().setDataRef(digits.size(), rof.getNEntries());
+              digits.insert(digits.end(), mDigits->begin() + rof.getFirstIdx(), mDigits->begin() + rof.getFirstIdx() + rof.getNEntries());
+              if (mUseMC) {
+                for (auto i = 0; i < rof.getNEntries(); ++i) {
+                  labels.addElements(labels.getIndexedSize(), (*mLabels)->getLabels(rof.getFirstIdx() + i));
+                }
+              }
+            }
+          }
+        }
+
+        // move to the next TF if needed and if any
+        if ((mRofs->empty() || mRofs->back().getBCData() + mRofs->back().getBCWidth() - 1 < irMax) &&
+            mTreeReader.GetCurrentEntry() + 1 < mTreeReader.GetEntries()) {
+          loadNextTF = true;
+          continue;
+        }
+
+        break;
+      }
+    }
+
+    // send the selected data
+    pc.outputs().snapshot(OutputRef{"rofs"}, rofs);
+    pc.outputs().snapshot(OutputRef{"digits"}, digits);
+    if (mUseMC) {
+      pc.outputs().snapshot(OutputRef{"labels"}, labels);
+    }
+
+    // stop here if they were the last IR frames to select
+    if (irFrames.empty() || irFrames.back().isLast()) {
+      pc.services().get<ControlService>().endOfStream();
+      pc.services().get<ControlService>().readyToQuit(QuitRequest::Me);
+    }
+  }
 };
 
-framework::DataProcessorSpec getDigitReaderSpec(
-  bool useMC,
-  std::string_view specName,
-  std::string_view outputDigitDataDescription,
-  std::string_view outputDigitRofDataDescription)
+DataProcessorSpec getDigitReaderSpec(bool useMC, std::string_view specName,
+                                     std::string_view outputDigitDataDescription,
+                                     std::string_view outputDigitRofDataDescription,
+                                     std::string_view outputDigitLabelDataDescription)
 {
-  std::vector<OutputSpec> outputs;
-  std::vector<header::DataDescription> descriptions;
-  std::stringstream ss;
-  ss << "A:" << header::gDataOriginMCH.as<std::string>() << "/" << outputDigitDataDescription << "/0";
-  ss << ";B:" << header::gDataOriginMCH.as<std::string>() << "/" << outputDigitRofDataDescription << "/0";
+  std::string output = fmt::format("digits:MCH/{}/0;rofs:MCH/{}/0", outputDigitDataDescription, outputDigitRofDataDescription);
   if (useMC) {
-    ss << ";C:" << header::gDataOriginMCH.as<std::string>() << "/DIGITLABELS/0";
+    output += fmt::format(";labels:MCH/{}/0", outputDigitLabelDataDescription);
   }
-  auto matchers = select(ss.str().c_str());
+
+  std::vector<OutputSpec> outputs;
+  auto matchers = select(output.c_str());
   for (auto& matcher : matchers) {
     outputs.emplace_back(DataSpecUtils::asOutputSpec(matcher));
-    descriptions.emplace_back(DataSpecUtils::asConcreteDataDescription(matcher));
   }
 
   return DataProcessorSpec{
     std::string(specName),
     Inputs{},
     outputs,
-    AlgorithmSpec{adaptFromTask<DigitsReaderDeviceDPL>(useMC, descriptions)},
+    AlgorithmSpec{adaptFromTask<DigitsReaderDeviceDPL>(useMC)},
     Options{{"mch-digit-infile", VariantType::String, "mchdigits.root", {"Name of the input file"}},
-            {"input-dir", VariantType::String, "none", {"Input directory"}}}};
+            {"input-dir", VariantType::String, "none", {"Input directory"}},
+            {"no-time-offset", VariantType::Bool, false, {"no time offset between IRFrames and digits"}}}};
 }
 } // namespace mch
 } // namespace o2
diff --git a/Detectors/MUON/MCH/IO/src/digits-reader-workflow.cxx b/Detectors/MUON/MCH/IO/src/digits-reader-workflow.cxx
@@ -35,7 +35,7 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
     {"configKeyValues", VariantType::String, "", {"Semicolon separated key=value strings"}},
     {"mch-output-digits-data-description", VariantType::String, "DIGITS", {"description string for the output digits message"}},
     {"mch-output-digitrofs-data-description", VariantType::String, "DIGITROFS", {"description string for the output digit rofs message"}},
-  };
+    {"mch-output-digitlabels-data-description", VariantType::String, "DIGITLABELS", {"description string for the output digit labels message"}}};
   workflowOptions.insert(workflowOptions.end(), options.begin(), options.end());
   o2::raw::HBFUtilsInitializer::addConfigOption(workflowOptions);
 }
@@ -55,7 +55,8 @@ WorkflowSpec defineDataProcessing(ConfigContext const& cc)
     useMC,
     name,
     cc.options().get<std::string>("mch-output-digits-data-description"),
-    cc.options().get<std::string>("mch-output-digitrofs-data-description")));
+    cc.options().get<std::string>("mch-output-digitrofs-data-description"),
+    cc.options().get<std::string>("mch-output-digitlabels-data-description")));
 
   o2::raw::HBFUtilsInitializer hbfIni(cc, specs);
 

Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,7 @@ o2_add_library(MCHIO`
`25`	`25`	`O2::DataFormatsMCH`
`26`	`26`	`O2::Framework`
`27`	`27`	`O2::MCHBase`
	`28`	`+ O2::MCHDigitFiltering`
`28`	`29`	`O2::SimulationDataFormat`
`29`	`30`	`)`
`30`	`31`
Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,8 @@ framework::DataProcessorSpec getDigitReaderSpec(`
`27`	`27`	`bool useMC,`
`28`	`28`	`std::string_view specName = "mch-digit-reader",`
`29`	`29`	`std::string_view outputDigitDataDescription = "DIGITS",`
`30`		`- std::string_view outputDigitRofDataDescription = "DIGITROFS");`
	`30`	`+ std::string_view outputDigitRofDataDescription = "DIGITROFS",`
	`31`	`+ std::string_view outputDigitLabelDataDescription = "DIGITLABELS");`
`31`	`32`	`}`
`32`	`33`	`} // namespace o2`
`33`	`34`