Skip to content

Commit

Permalink
extract digits from selected BC ranges if any
Browse files Browse the repository at this point in the history
  • Loading branch information
@pillot
pillot committed Oct 25, 2024
1 parent 260d9da commit e454f13
Show file tree
Hide file tree
Showing 4 changed files with 162 additions and 49 deletions.
1 change: 1 addition & 0 deletions Detectors/MUON/MCH/IO/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ o2_add_library(MCHIO
O2::DataFormatsMCH
O2::Framework
O2::MCHBase
O2::MCHDigitFiltering
O2::SimulationDataFormat
)

Expand Down
3 changes: 2 additions & 1 deletion Detectors/MUON/MCH/IO/include/MCHIO/DigitReaderSpec.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
202 changes: 156 additions & 46 deletions Detectors/MUON/MCH/IO/src/DigitReaderSpec.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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}, "MCHDigit"},
RootTreeReader::BranchDefinition<std::vector<ROFRecord>>{
Output{header::gDataOriginMCH, mDescriptions[1], 0}, "MCHROFRecords"},
RootTreeReader::BranchDefinition<dataformats::MCTruthContainer<MCCompLabel>>{
Output{header::gDataOriginMCH, mDescriptions[2], 0}, "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}, "MCHDigit"},
RootTreeReader::BranchDefinition<std::vector<ROFRecord>>{
Output{header::gDataOriginMCH, mDescriptions[1], 0}, "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
5 changes: 3 additions & 2 deletions Detectors/MUON/MCH/IO/src/digits-reader-workflow.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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);

Expand Down

0 comments on commit e454f13

Please sign in to comment.