From a6aa329cc9b33ab8447ffe36a860276cdb07410d Mon Sep 17 00:00:00 2001
From: pillot <pillot@subatech.in2p3.fr>
Date: Fri, 25 Oct 2024 15:49:47 +0200
Subject: [PATCH] add option to change digit labels data description

---
 .../MUON/MCH/IO/include/MCHIO/DigitWriterSpec.h   |  3 ++-
 Detectors/MUON/MCH/IO/src/DigitWriterSpec.cxx     | 15 +++++++--------
 .../MUON/MCH/IO/src/digits-writer-workflow.cxx    |  4 +++-
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/Detectors/MUON/MCH/IO/include/MCHIO/DigitWriterSpec.h b/Detectors/MUON/MCH/IO/include/MCHIO/DigitWriterSpec.h
index 66909ec5972b8..c615774c1add9 100644
--- a/Detectors/MUON/MCH/IO/include/MCHIO/DigitWriterSpec.h
+++ b/Detectors/MUON/MCH/IO/include/MCHIO/DigitWriterSpec.h
@@ -25,7 +25,8 @@ o2::framework::DataProcessorSpec getDigitWriterSpec(
   std::string_view specName = "mch-digit-writer",
   std::string_view outfile = "mchdigits.root",
   std::string_view inputDigitDataDescription = "DIGITS",
-  std::string_view inputDigitRofDataDescription = "DIGITROFS");
+  std::string_view inputDigitRofDataDescription = "DIGITROFS",
+  std::string_view inputDigitLabelDataDescription = "DIGITLABELS");
 
 } // end namespace mch
 } // end namespace o2
diff --git a/Detectors/MUON/MCH/IO/src/DigitWriterSpec.cxx b/Detectors/MUON/MCH/IO/src/DigitWriterSpec.cxx
index 9d5755bd7911c..84a20d855dab7 100644
--- a/Detectors/MUON/MCH/IO/src/DigitWriterSpec.cxx
+++ b/Detectors/MUON/MCH/IO/src/DigitWriterSpec.cxx
@@ -38,9 +38,7 @@ o2::framework::DataProcessorSpec getMCHDigitWriterSpec(bool mctruth)
                                 1, // default number of events
                                 BranchDefinition<std::vector<o2::mch::Digit>>{InputSpec{"mchdigits", "MCH", "DIGITS"}, "MCHDigit"},
                                 BranchDefinition<std::vector<o2::mch::ROFRecord>>{InputSpec{"mchrofrecords", "MCH", "DIGITROFS"}, "MCHROFRecords"},
-                                BranchDefinition<o2::dataformats::MCTruthContainer<o2::MCCompLabel>>{InputSpec{"mchdigitlabels", "MCH", "DIGITSLABELS"}, "MCHMCLabels", mctruth ? 1 : 0}
-                                // add more branch definitions (for example Monte Carlo labels here)
-                                )();
+                                BranchDefinition<o2::dataformats::MCTruthContainer<o2::MCCompLabel>>{InputSpec{"mchdigitlabels", "MCH", "DIGITSLABELS"}, "MCHMCLabels", mctruth ? 1 : 0})();
 }
 
 o2::framework::DataProcessorSpec getDigitWriterSpec(
@@ -48,23 +46,24 @@ o2::framework::DataProcessorSpec getDigitWriterSpec(
   std::string_view specName,
   std::string_view outfile,
   std::string_view inputDigitDataDescription,
-  std::string_view inputDigitRofDataDescription)
+  std::string_view inputDigitRofDataDescription,
+  std::string_view inputDigitLabelDataDescription)
 {
   std::string input =
-    fmt::format("digits:MCH/{};rofs:MCH/{}",
-                inputDigitDataDescription, inputDigitRofDataDescription);
+    fmt::format("digits:MCH/{};rofs:MCH/{};labels:MCH/{}",
+                inputDigitDataDescription, inputDigitRofDataDescription, inputDigitLabelDataDescription);
 
   framework::Inputs inputs{framework::select(input.c_str())};
   auto rofs = std::find_if(inputs.begin(), inputs.end(), [](const framework::InputSpec& is) { return is.binding == "rofs"; });
   auto digits = std::find_if(inputs.begin(), inputs.end(), [](const framework::InputSpec& is) { return is.binding == "digits"; });
+  auto labels = std::find_if(inputs.begin(), inputs.end(), [](const framework::InputSpec& is) { return is.binding == "labels"; });
   return framework::MakeRootTreeWriterSpec(
     std::string(specName).c_str(),
     std::string(outfile).c_str(),
     framework::MakeRootTreeWriterSpec::TreeAttributes{"o2sim", "Tree MCH Digits"},
     BranchDefinition<std::vector<ROFRecord>>{framework::InputSpec{*rofs}, "MCHROFRecords"},
     BranchDefinition<std::vector<Digit>>{framework::InputSpec{*digits}, "MCHDigit"},
-    BranchDefinition<o2::dataformats::MCTruthContainer<o2::MCCompLabel>>{
-      framework::InputSpec{"mchdigitlabels", "MCH", "DIGITSLABELS"}, "MCHMCLabels", useMC ? 1 : 0})();
+    BranchDefinition<o2::dataformats::MCTruthContainer<o2::MCCompLabel>>{framework::InputSpec{*labels}, "MCHMCLabels", useMC ? 1 : 0})();
 }
 
 } // end namespace mch
diff --git a/Detectors/MUON/MCH/IO/src/digits-writer-workflow.cxx b/Detectors/MUON/MCH/IO/src/digits-writer-workflow.cxx
index b806039ac9062..4915a9de483f2 100644
--- a/Detectors/MUON/MCH/IO/src/digits-writer-workflow.cxx
+++ b/Detectors/MUON/MCH/IO/src/digits-writer-workflow.cxx
@@ -47,6 +47,7 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
 {
   workflowOptions.emplace_back(ConfigParamSpec{"input-digits-data-description", VariantType::String, "DIGITS", {"description string for the input digits data"}});
   workflowOptions.emplace_back(ConfigParamSpec{"input-digitrofs-data-description", VariantType::String, "DIGITROFS", {"description string for the input digit rofs data"}});
+  workflowOptions.emplace_back(ConfigParamSpec{"input-digitlabels-data-description", VariantType::String, "DIGITLABELS", {"description string for the input digit labels data"}});
   workflowOptions.emplace_back(ConfigParamSpec{"enable-mc", VariantType::Bool, false, {" output MC labels if available "}});
   workflowOptions.emplace_back(ConfigParamSpec{"mch-digit-outfile", VariantType::String, "mchdigits.root", {"name of digit root file"}});
 }
@@ -65,5 +66,6 @@ WorkflowSpec defineDataProcessing(const ConfigContext& cc)
     name,
     cc.options().get<std::string>("mch-digit-outfile"),
     cc.options().get<std::string>("input-digits-data-description"),
-    cc.options().get<std::string>("input-digitrofs-data-description"))};
+    cc.options().get<std::string>("input-digitrofs-data-description"),
+    cc.options().get<std::string>("input-digitlabels-data-description"))};
 }