Writing Track MC labels to output file, adjusting default branch names

Renamed process to tpc-tracks-writer, options are now:
Data processor options: tpc-track-writer:
  --outfile arg (=tpctracks.root)       Name of the input file
  --treename arg (=o2sim)               Name of output tree
  --track-branch-name arg (=TPCTracks)  Branch name for TPC tracks
  --trackmc-branch-name arg (=TPCTracksMCTruth)
                                        Branch name for TPC track mc labels

Also catching possible error code at TPCCATracking initialization,
from AliTPCCommon v.1.4.3, error code is return on unknown option.

Author: matthiasrichter

Detectors/TPC/workflow/README.md

@@ -39,6 +39,14 @@ Important options for the `digit-reader` as initial publisher
 --tpc-sectors arg (=0-35)             TPC sector range, e.g. 5-7,8,9
 ```
 
+Options for the `tpc-track-writer` process
+```
+--outfile arg (=tpctracks.root)                Name of the input file
+--treename arg (=o2sim)                        Name of output tree
+--track-branch-name arg (=TPCTracks)           Branch name for TPC tracks
+--trackmc-branch-name arg (=TPCTracksMCTruth)  Branch name for TPC track mc labels
+```
+
 Examples:
 ```
 tpc-reco-workflow --infile tpcdigits.root --tpc-sectors 0-15
@@ -65,9 +73,16 @@ lanes, each with clusterer, converter and decoder. The tracker will fan in from
 ### Processor options
 
 #### TPC CA tracker
-The [tracker spec](src/CATrackerSpec.cxx) interfaces the [o2::TPC::TPCCATracking](reconstruction/include/TPCReconstruction/TPCCATracking.h) worker class which can be initialized using an option string. The processor spec defines the option `--tracker-option`, e.g.
+The [tracker spec](src/CATrackerSpec.cxx) interfaces the [o2::TPC::TPCCATracking](reconstruction/include/TPCReconstruction/TPCCATracking.h) worker class which can be initialized using an option string. The processor spec defines the option `--tracker-option`. Currently, the tracker should be run with options
+```
+--tracker-option "refX=83 cont"
+```
+
+The most important tracker options are:
 ```
---tracker-option ""refX=83""
+cont    activation of continuous mode
+refX=   reference x coordinate, tracker tries to propagate all track to the reference
+bz=     magnetic field
 ```
 
 ### Current limitations

Detectors/TPC/workflow/src/CATrackerSpec.cxx

@@ -28,6 +28,7 @@
 #include <memory> // for make_shared
 #include <vector>
 #include <iomanip>
+#include <stdexcept>
 
 using namespace o2::framework;
 using namespace o2::header;
@@ -67,7 +68,9 @@ DataProcessorSpec getCATrackerSpec(bool processMC, size_t fanIn)
       auto& tracker = processAttributes->tracker;
       parser = std::make_unique<ClusterGroupParser>();
       tracker = std::make_unique<o2::TPC::TPCCATracking>();
-      tracker->initialize(options.c_str());
+      if (tracker->initialize(options.c_str()) != 0) {
+        throw std::invalid_argument("TPCCATracking initialization failed");
+      }
       processAttributes->validInputs.reset();
       processAttributes->validMcInputs.reset();
     }
@@ -214,10 +217,8 @@ DataProcessorSpec getCATrackerSpec(bool processMC, size_t fanIn)
       LOG(INFO) << "found " << tracks.size() << " track(s)";
       pc.outputs().snapshot(OutputRef{ "output" }, tracks);
       if (processMC) {
-        LOG(INFO) << "have " << tracksMCTruth.getIndexedSize() << " track label(s)";
-        // have to change the writer process as well but want to convert to the RootTreeWriter tool
-        // at this occasion so we skip sending the labels for the moment
-        //pc.outputs().snapshot(OutputRef{ "mclblout" }, tracksMCTruth);
+        LOG(INFO) << "sending " << tracksMCTruth.getIndexedSize() << " track label(s)";
+        pc.outputs().snapshot(OutputRef{ "mclblout" }, tracksMCTruth);
       }
 
       validInputs.reset();
@@ -261,7 +262,7 @@ DataProcessorSpec getCATrackerSpec(bool processMC, size_t fanIn)
 
   return DataProcessorSpec{ "tpc-tracker", // process id
                             { createInputSpecs(processMC) },
-                            { createOutputSpecs(false /*create onece writer process has been changed*/) },
+                            { createOutputSpecs(processMC) },
                             AlgorithmSpec(initFunction),
                             Options{
                               { "tracker-options", VariantType::String, "", { "Option string passed to tracker" } },

Detectors/TPC/workflow/src/RecoWorkflow.cxx

@@ -144,7 +144,7 @@ framework::WorkflowSpec getWorkflow(bool propagateMC, int nLanes, std::string cf
 
   if (outputType == OutputType::Tracks) {
     specs.emplace_back(o2::TPC::getCATrackerSpec(propagateMC, nLanes));
-    specs.emplace_back(o2::TPC::getRootFileWriterSpec());
+    specs.emplace_back(o2::TPC::getRootFileWriterSpec(propagateMC));
   } else if (outputType == OutputType::DecodedClusters) {
     specs.emplace_back(DataProcessorSpec{ "writer",
                                           { createInputSpec() },

Detectors/TPC/workflow/src/RootFileWriterSpec.cxx

@@ -17,31 +17,41 @@
 #include "Framework/ControlService.h"
 #include "Framework/CallbackService.h"
 #include "DataFormatsTPC/TrackTPC.h"
+#include "SimulationDataFormat/MCTruthContainer.h"
+#include "SimulationDataFormat/MCCompLabel.h"
 #include <TFile.h>
 #include <TTree.h>
 #include <memory> // for make_shared, make_unique, unique_ptr
 #include <stdexcept>
 
 using namespace o2::framework;
 
+using MCLabelContainer = o2::dataformats::MCTruthContainer<o2::MCCompLabel>;
+
 namespace o2
 {
 namespace TPC
 {
 /// create a processor spec
 /// read simulated TPC digits from file and publish
-DataProcessorSpec getRootFileWriterSpec()
+DataProcessorSpec getRootFileWriterSpec(bool writeMC)
 {
-  auto initFunction = [](InitContext& ic) {
+  auto initFunction = [writeMC](InitContext& ic) {
     // get the option from the init context
     auto filename = ic.options().get<std::string>("outfile");
     auto treename = ic.options().get<std::string>("treename");
+    auto trackBranchName = ic.options().get<std::string>("track-branch-name");
+    auto trackMcBranchName = ic.options().get<std::string>("trackmc-branch-name");
     auto nevents = ic.options().get<int>("nevents");
 
     auto outputfile = std::make_shared<TFile>(filename.c_str(), "RECREATE");
     auto outputtree = std::make_shared<TTree>(treename.c_str(), treename.c_str());
     auto tracks = std::make_shared<std::vector<o2::TPC::TrackTPC>>();
-    auto trackbranch = outputtree->Branch("Tracks", tracks.get());
+    auto mclabels = std::make_shared<MCLabelContainer>();
+    outputtree->Branch(trackBranchName.c_str(), tracks.get());
+    if (writeMC) {
+      outputtree->Branch(trackMcBranchName.c_str(), mclabels.get());
+    }
 
     // the callback to be set as hook at stop of processing for the framework
     auto finishWriting = [outputfile, outputtree]() {
@@ -54,11 +64,15 @@ DataProcessorSpec getRootFileWriterSpec()
     // using by-copy capture of the worker instance shared pointer
     // the shared pointer makes sure to clean up the instance when the processing
     // function gets out of scope
-    auto processingFct = [outputfile, outputtree, tracks, nevents](ProcessingContext& pc) {
+    auto processingFct = [outputfile, outputtree, tracks, mclabels, nevents, writeMC](ProcessingContext& pc) {
       static int eventCount = 0;
       auto indata = pc.inputs().get<std::vector<o2::TPC::TrackTPC>>("input");
       LOG(INFO) << "RootFileWriter: get " << indata.size() << " track(s)";
       *tracks.get() = std::move(indata);
+      if (writeMC) {
+        auto mcdata = pc.inputs().get<MCLabelContainer*>("mcinput");
+        *mclabels.get() = *mcdata;
+      }
       LOG(INFO) << "RootFileWriter: write " << tracks->size() << " track(s)";
       outputtree->Fill();
 
@@ -72,13 +86,26 @@ DataProcessorSpec getRootFileWriterSpec()
     return processingFct;
   };
 
-  return DataProcessorSpec{ "writer",
-                            { InputSpec{ "input", "TPC", "TRACKS", 0, Lifetime::Timeframe } }, // track input
-                            {},                                                                // no output
+  auto createInputSpecs = [](bool makeMcInput) {
+    std::vector<InputSpec> inputSpecs{
+      InputSpec{ { "input" }, "TPC", "TRACKS", 0, Lifetime::Timeframe },
+    };
+    if (makeMcInput) {
+      constexpr o2::header::DataDescription datadesc("TRACKMCLBL");
+      inputSpecs.emplace_back(InputSpec{ "mcinput", "TPC", datadesc, 0, Lifetime::Timeframe });
+    }
+    return std::move(inputSpecs);
+  };
+
+  return DataProcessorSpec{ "tpc-track-writer",
+                            { createInputSpecs(writeMC) },
+                            {},
                             AlgorithmSpec(initFunction),
                             Options{
                               { "outfile", VariantType::String, "tpctracks.root", { "Name of the input file" } },
-                              { "treename", VariantType::String, "Tracks", { "Name of tree for tracks" } },
+                              { "treename", VariantType::String, "o2sim", { "Name of output tree" } },
+                              { "track-branch-name", VariantType::String, "TPCTracks", { "Branch name for TPC tracks" } },
+                              { "trackmc-branch-name", VariantType::String, "TPCTracksMCTruth", { "Branch name for TPC track mc labels" } },
                               { "nevents", VariantType::Int, 1, { "terminate after n events" } },
                             } };
 }

Detectors/TPC/workflow/src/RootFileWriterSpec.h

@@ -20,7 +20,7 @@ namespace o2
 namespace TPC
 {
 
-framework::DataProcessorSpec getRootFileWriterSpec();
+framework::DataProcessorSpec getRootFileWriterSpec(bool writeMC = false);
 
 } // end namespace TPC
 } // end namespace o2