From e454f13b1415bd0de6e36d125247909d2d98feac Mon Sep 17 00:00:00 2001 From: pillot Date: Fri, 18 Oct 2024 18:47:21 +0200 Subject: [PATCH] extract digits from selected BC ranges if any --- Detectors/MUON/MCH/IO/CMakeLists.txt | 1 + .../MCH/IO/include/MCHIO/DigitReaderSpec.h | 3 +- Detectors/MUON/MCH/IO/src/DigitReaderSpec.cxx | 202 ++++++++++++++---- .../MCH/IO/src/digits-reader-workflow.cxx | 5 +- 4 files changed, 162 insertions(+), 49 deletions(-) diff --git a/Detectors/MUON/MCH/IO/CMakeLists.txt b/Detectors/MUON/MCH/IO/CMakeLists.txt index 18a8630a7fed0..ed329b44922d8 100644 --- 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 index a7fd59b61a739..c061a260cf6f8 100644 --- 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 index 541e6c5fa26c0..8fc7595e918e6 100644 --- 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 -#include +#include #include -#include "DPLUtils/RootTreeReader.h" +#include + +#include + +#include + +#include +#include +#include +#include + +#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& descriptions) - : mUseMC(useMC), mDescriptions(descriptions) {} + DigitsReaderDeviceDPL(bool useMC) : mUseMC(useMC) + { + if (mUseMC) { + mLabels = std::make_unique>>(mTreeReader, "MCHMCLabels"); + } + } void init(InitContext& ic) { - auto filename = o2::utils::Str::concat_string(o2::utils::Str::rectifyDirectory(ic.options().get("input-dir")), ic.options().get("mch-digit-infile")); + auto fileName = utils::Str::concat_string(utils::Str::rectifyDirectory(ic.options().get("input-dir")), + ic.options().get("mch-digit-infile")); + connectTree(fileName); - if (mUseMC) { - mReader = std::make_unique("o2sim", filename.c_str(), -1, - RootTreeReader::PublishingMode::Single, - RootTreeReader::BranchDefinition>{ - Output{header::gDataOriginMCH, mDescriptions[0], 0}, "MCHDigit"}, - RootTreeReader::BranchDefinition>{ - Output{header::gDataOriginMCH, mDescriptions[1], 0}, "MCHROFRecords"}, - RootTreeReader::BranchDefinition>{ - Output{header::gDataOriginMCH, mDescriptions[2], 0}, "MCHMCLabels"}); - } else { - mReader = std::make_unique("o2sim", filename.c_str(), -1, - RootTreeReader::PublishingMode::Single, - RootTreeReader::BranchDefinition>{ - Output{header::gDataOriginMCH, mDescriptions[0], 0}, "MCHDigit"}, - RootTreeReader::BranchDefinition>{ - Output{header::gDataOriginMCH, mDescriptions[1], 0}, "MCHROFRecords"}); + if (ic.options().hasOption("ignore-irframes") && !ic.options().get("ignore-irframes")) { + mUseIRFrames = true; } + + mTimeOffset = ic.options().get("no-time-offset") ? 0 : DigitFilterParam::Instance().timeOffset; } void run(ProcessingContext& pc) { - if ((++(*mReader))(pc) == false) { - pc.services().get().endOfStream(); + if (mUseIRFrames) { + sendNextIRFrames(pc); + } else { + sendNextTF(pc); } } private: - std::unique_ptr mReader{}; - std::vector mDescriptions{}; + TTreeReader mTreeReader{}; + TTreeReaderValue> mRofs = {mTreeReader, "MCHROFRecords"}; + TTreeReaderValue> mDigits = {mTreeReader, "MCHDigit"}; + std::unique_ptr>> 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("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().endOfStream(); + pc.services().get().readyToQuit(QuitRequest::Me); + } + } + + void sendNextIRFrames(ProcessingContext& pc) + { + std::vector rofs{}; + std::vector digits{}; + dataformats::MCTruthContainer labels{}; + + // get the IR frames to select + auto irFrames = pc.inputs().get>("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().endOfStream(); + pc.services().get().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 outputs; - std::vector descriptions; - std::stringstream ss; - ss << "A:" << header::gDataOriginMCH.as() << "/" << outputDigitDataDescription << "/0"; - ss << ";B:" << header::gDataOriginMCH.as() << "/" << outputDigitRofDataDescription << "/0"; + std::string output = fmt::format("digits:MCH/{}/0;rofs:MCH/{}/0", outputDigitDataDescription, outputDigitRofDataDescription); if (useMC) { - ss << ";C:" << header::gDataOriginMCH.as() << "/DIGITLABELS/0"; + output += fmt::format(";labels:MCH/{}/0", outputDigitLabelDataDescription); } - auto matchers = select(ss.str().c_str()); + + std::vector 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(useMC, descriptions)}, + AlgorithmSpec{adaptFromTask(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 index 0ec46f2fc4d3b..d6c91023de53a 100644 --- 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& 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("mch-output-digits-data-description"), - cc.options().get("mch-output-digitrofs-data-description"))); + cc.options().get("mch-output-digitrofs-data-description"), + cc.options().get("mch-output-digitlabels-data-description"))); o2::raw::HBFUtilsInitializer hbfIni(cc, specs);