new option to write digits associated to tracks

Author: pillot

Detectors/MUON/MCH/Workflow/include/MCHWorkflow/TrackWriterSpec.h

@@ -16,7 +16,8 @@
 
 namespace o2::mch
 {
-o2::framework::DataProcessorSpec getTrackWriterSpec(bool useMC, const char* specName = "mch-track-writer", const char* fileName = "mchtracks.root");
+o2::framework::DataProcessorSpec getTrackWriterSpec(bool useMC, const char* specName = "mch-track-writer",
+                                                    const char* fileName = "mchtracks.root", bool digits = false);
 }
 
 #endif

Detectors/MUON/MCH/Workflow/src/TrackFinderSpec.cxx

@@ -41,6 +41,7 @@
 #include "DataFormatsMCH/ROFRecord.h"
 #include "DataFormatsMCH/TrackMCH.h"
 #include "DataFormatsMCH/Cluster.h"
+#include "DataFormatsMCH/Digit.h"
 #include "MCHTracking/TrackParam.h"
 #include "MCHTracking/Track.h"
 #include "MCHTracking/TrackFinder.h"
@@ -57,6 +58,9 @@ using namespace o2::framework;
 class TrackFinderTask
 {
  public:
+  //_________________________________________________________________________________________________
+  TrackFinderTask(bool digits = false) : mDigits(digits) {}
+
   //_________________________________________________________________________________________________
   void init(framework::InitContext& ic)
   {
@@ -105,18 +109,22 @@ class TrackFinderTask
     auto clusterROFs = pc.inputs().get<gsl::span<ROFRecord>>("clusterrofs");
     auto clustersIn = pc.inputs().get<gsl::span<Cluster>>("clusters");
 
-    // LOG(info) << "received time frame with " << clusterROFs.size() << " interactions";
-
     // create the output messages for tracks and attached clusters
     auto& trackROFs = pc.outputs().make<std::vector<ROFRecord>>(OutputRef{"trackrofs"});
     auto& mchTracks = pc.outputs().make<std::vector<TrackMCH>>(OutputRef{"tracks"});
     auto& usedClusters = pc.outputs().make<std::vector<Cluster>>(OutputRef{"trackclusters"});
 
+    // prepare to send the associated digits if requested
+    gsl::span<const Digit> digitsIn{};
+    std::vector<Digit, o2::pmr::polymorphic_allocator<Digit>>* usedDigits(nullptr);
+    if (mDigits) {
+      digitsIn = pc.inputs().get<gsl::span<Digit>>("clusterdigits");
+      usedDigits = &pc.outputs().make<std::vector<Digit>>(OutputRef{"trackdigits"});
+    }
+
     trackROFs.reserve(clusterROFs.size());
     for (const auto& clusterROF : clusterROFs) {
 
-      // LOG(info) << "processing interaction: " << clusterROF.getBCData() << "...";
-
       // sort the input clusters of the current event per DE
       std::unordered_map<int, std::list<const Cluster*>> clusters{};
       for (const auto& cluster : clustersIn.subspan(clusterROF.getFirstIdx(), clusterROF.getNEntries())) {
@@ -131,7 +139,7 @@ class TrackFinderTask
 
       // fill the ouput messages
       int trackOffset(mchTracks.size());
-      writeTracks(tracks, mchTracks, usedClusters);
+      writeTracks(tracks, mchTracks, usedClusters, digitsIn, usedDigits);
       trackROFs.emplace_back(clusterROF.getBCData(), trackOffset, mchTracks.size() - trackOffset,
                              clusterROF.getBCWidth());
     }
@@ -144,9 +152,14 @@ class TrackFinderTask
   //_________________________________________________________________________________________________
   void writeTracks(const std::list<Track>& tracks,
                    std::vector<TrackMCH, o2::pmr::polymorphic_allocator<TrackMCH>>& mchTracks,
-                   std::vector<Cluster, o2::pmr::polymorphic_allocator<Cluster>>& usedClusters) const
+                   std::vector<Cluster, o2::pmr::polymorphic_allocator<Cluster>>& usedClusters,
+                   const gsl::span<const Digit>& digitsIn,
+                   std::vector<Digit, o2::pmr::polymorphic_allocator<Digit>>* usedDigits) const
   {
-    /// fill the output messages with tracks and attached clusters
+    /// fill the output messages with tracks and attached clusters and digits if requested
+
+    // map the location of the attached digits between the digitsIn and the usedDigits lists
+    std::unordered_map<uint32_t, uint32_t> digitLocMap{};
 
     for (const auto& track : tracks) {
 
@@ -162,26 +175,55 @@ class TrackFinderTask
                              paramAtMID.getZ(), paramAtMID.getParameters(), paramAtMID.getCovariances());
 
       for (const auto& param : track) {
+
         usedClusters.emplace_back(*param.getClusterPtr());
+
+        if (mDigits) {
+
+          // map the location of the digits associated to this cluster in the usedDigits list, if not already done
+          auto& cluster = usedClusters.back();
+          auto digitLoc = digitLocMap.emplace(cluster.firstDigit, usedDigits->size());
+
+          // add the digits associated to this cluster if not already there
+          if (digitLoc.second) {
+            auto itFirstDigit = digitsIn.begin() + cluster.firstDigit;
+            usedDigits->insert(usedDigits->end(), itFirstDigit, itFirstDigit + cluster.nDigits);
+          }
+
+          // make the cluster point to the associated digits in the usedDigits list
+          cluster.firstDigit = digitLoc.first->second;
+        }
       }
     }
   }
 
+  bool mDigits = false;                         ///< send to associated digits
   TrackFinder mTrackFinder{};                   ///< track finder
   std::chrono::duration<double> mElapsedTime{}; ///< timer
 };
 
 //_________________________________________________________________________________________________
-o2::framework::DataProcessorSpec getTrackFinderSpec(const char* specName)
+o2::framework::DataProcessorSpec getTrackFinderSpec(const char* specName, bool digits)
 {
+  std::vector<InputSpec> inputSpecs{};
+  inputSpecs.emplace_back(InputSpec{"clusterrofs", "MCH", "CLUSTERROFS", 0, Lifetime::Timeframe});
+  inputSpecs.emplace_back(InputSpec{"clusters", "MCH", "GLOBALCLUSTERS", 0, Lifetime::Timeframe});
+
+  std::vector<OutputSpec> outputSpecs{};
+  outputSpecs.emplace_back(OutputSpec{{"trackrofs"}, "MCH", "TRACKROFS", 0, Lifetime::Timeframe});
+  outputSpecs.emplace_back(OutputSpec{{"tracks"}, "MCH", "TRACKS", 0, Lifetime::Timeframe});
+  outputSpecs.emplace_back(OutputSpec{{"trackclusters"}, "MCH", "TRACKCLUSTERS", 0, Lifetime::Timeframe});
+
+  if (digits) {
+    inputSpecs.emplace_back(InputSpec{"clusterdigits", "MCH", "CLUSTERDIGITS", 0, Lifetime::Timeframe});
+    outputSpecs.emplace_back(OutputSpec{{"trackdigits"}, "MCH", "TRACKDIGITS", 0, Lifetime::Timeframe});
+  }
+
   return DataProcessorSpec{
     specName,
-    Inputs{InputSpec{"clusterrofs", "MCH", "CLUSTERROFS", 0, Lifetime::Timeframe},
-           InputSpec{"clusters", "MCH", "GLOBALCLUSTERS", 0, Lifetime::Timeframe}},
-    Outputs{OutputSpec{{"trackrofs"}, "MCH", "TRACKROFS", 0, Lifetime::Timeframe},
-            OutputSpec{{"tracks"}, "MCH", "TRACKS", 0, Lifetime::Timeframe},
-            OutputSpec{{"trackclusters"}, "MCH", "TRACKCLUSTERS", 0, Lifetime::Timeframe}},
-    AlgorithmSpec{adaptFromTask<TrackFinderTask>()},
+    inputSpecs,
+    outputSpecs,
+    AlgorithmSpec{adaptFromTask<TrackFinderTask>(digits)},
     Options{{"l3Current", VariantType::Float, -30000.0f, {"L3 current"}},
             {"dipoleCurrent", VariantType::Float, -6000.0f, {"Dipole current"}},
             {"grp-file", VariantType::String, o2::base::NameConf::getGRPFileName(), {"Name of the grp file"}},

Detectors/MUON/MCH/Workflow/src/TrackFinderSpec.h

@@ -24,7 +24,7 @@ namespace o2
 namespace mch
 {
 
-o2::framework::DataProcessorSpec getTrackFinderSpec(const char* specName = "mch-track-finder");
+o2::framework::DataProcessorSpec getTrackFinderSpec(const char* specName = "mch-track-finder", bool digits = false);
 
 } // end namespace mch
 } // end namespace o2

Detectors/MUON/MCH/Workflow/src/TrackWriterSpec.cxx

@@ -17,6 +17,7 @@
 #include "DataFormatsMCH/TrackMCH.h"
 #include "Framework/Logger.h"
 #include "DataFormatsMCH/Cluster.h"
+#include "DataFormatsMCH/Digit.h"
 #include <vector>
 
 using namespace o2::framework;
@@ -27,14 +28,15 @@ namespace o2::mch
 template <typename T>
 using BranchDefinition = MakeRootTreeWriterSpec::BranchDefinition<T>;
 
-DataProcessorSpec getTrackWriterSpec(bool useMC, const char* specName, const char* fileName)
+DataProcessorSpec getTrackWriterSpec(bool useMC, const char* specName, const char* fileName, bool digits)
 {
   return MakeRootTreeWriterSpec(specName,
                                 fileName,
                                 MakeRootTreeWriterSpec::TreeAttributes{"o2sim", "Tree MCH Standalone Tracks"},
                                 BranchDefinition<std::vector<TrackMCH>>{InputSpec{"tracks", "MCH", "TRACKS"}, "tracks"},
                                 BranchDefinition<std::vector<ROFRecord>>{InputSpec{"trackrofs", "MCH", "TRACKROFS"}, "trackrofs"},
                                 BranchDefinition<std::vector<Cluster>>{InputSpec{"trackclusters", "MCH", "TRACKCLUSTERS"}, "trackclusters"},
+                                BranchDefinition<std::vector<Digit>>{InputSpec{"trackdigits", "MCH", "TRACKDIGITS"}, "trackdigits", digits ? 1 : 0},
                                 BranchDefinition<std::vector<o2::MCCompLabel>>{InputSpec{"tracklabels", "MCH", "TRACKLABELS"}, "tracklabels", useMC ? 1 : 0})();
 }
 

Detectors/MUON/MCH/Workflow/src/clusters-to-tracks-workflow.cxx

@@ -24,12 +24,15 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
 {
   workflowOptions.emplace_back("configKeyValues", VariantType::String, "",
                                ConfigParamSpec::HelpString{"Semicolon separated key=value strings"});
+  workflowOptions.emplace_back("digits", VariantType::Bool, false,
+                               ConfigParamSpec::HelpString{"Send associated digits"});
 }
 
 #include "Framework/runDataProcessing.h"
 
 WorkflowSpec defineDataProcessing(const ConfigContext& configcontext)
 {
   o2::conf::ConfigurableParam::updateFromString(configcontext.options().get<std::string>("configKeyValues"));
-  return WorkflowSpec{o2::mch::getTrackFinderSpec()};
+  bool digits = configcontext.options().get<bool>("digits");
+  return WorkflowSpec{o2::mch::getTrackFinderSpec("mch-track-finder", digits)};
 }

Detectors/MUON/MCH/Workflow/src/reco-workflow.cxx

@@ -43,6 +43,7 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
     {"disable-root-input", o2::framework::VariantType::Bool, false, {"disable root-files input reader"}},
     {"disable-root-output", o2::framework::VariantType::Bool, false, {"do not write output root files"}},
     {"disable-mc", o2::framework::VariantType::Bool, false, {"disable MC propagation even if available"}},
+    {"digits", VariantType::Bool, false, {"Write digits associated to tracks"}},
     {"configKeyValues", VariantType::String, "", {"Semicolon separated key=value strings"}}};
   o2::raw::HBFUtilsInitializer::addConfigOption(options);
   std::swap(workflowOptions, options);
@@ -56,6 +57,7 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
 
   auto disableRootOutput = configcontext.options().get<bool>("disable-root-output");
   auto disableRootInput = configcontext.options().get<bool>("disable-root-input");
+  auto digits = configcontext.options().get<bool>("digits");
   auto useMC = !configcontext.options().get<bool>("disable-mc");
 
   o2::conf::ConfigurableParam::updateFromString(configcontext.options().get<std::string>("configKeyValues"));
@@ -79,15 +81,15 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
                                                       "TC-F-DIGITROFS"));
   specs.emplace_back(o2::mch::getClusterFinderOriginalSpec("mch-cluster-finder"));
   specs.emplace_back(o2::mch::getClusterTransformerSpec());
-  specs.emplace_back(o2::mch::getTrackFinderSpec("mch-track-finder"));
+  specs.emplace_back(o2::mch::getTrackFinderSpec("mch-track-finder", digits));
   if (useMC) {
     specs.emplace_back(o2::mch::getTrackMCLabelFinderSpec("mch-track-mc-label-finder",
                                                           "TC-F-DIGITROFS",
                                                           "F-DIGITLABELS"));
   }
 
   if (!disableRootOutput) {
-    specs.emplace_back(o2::mch::getTrackWriterSpec(useMC, "mch-track-writer", "mchtracks.root"));
+    specs.emplace_back(o2::mch::getTrackWriterSpec(useMC, "mch-track-writer", "mchtracks.root", digits));
   }
 
   // configure dpl timer to inject correct firstTFOrbit: start from the 1st orbit of TF containing 1st sampled orbit

Detectors/MUON/MCH/Workflow/src/tracks-writer-workflow.cxx

@@ -19,12 +19,14 @@ using namespace o2::framework;
 void customize(std::vector<ConfigParamSpec>& workflowOptions)
 {
   workflowOptions.emplace_back("enable-mc", VariantType::Bool, false, ConfigParamSpec::HelpString{"Propagate MC info"});
+  workflowOptions.emplace_back("digits", VariantType::Bool, false, ConfigParamSpec::HelpString{"Write associated digits"});
 }
 
 #include "Framework/runDataProcessing.h"
 
 WorkflowSpec defineDataProcessing(const ConfigContext& config)
 {
   bool useMC = config.options().get<bool>("enable-mc");
-  return WorkflowSpec{o2::mch::getTrackWriterSpec(useMC)};
+  bool digits = config.options().get<bool>("digits");
+  return WorkflowSpec{o2::mch::getTrackWriterSpec(useMC, "mch-track-writer", "mchtracks.root", digits)};
 }