TPC SCD: write and read the MC truth of the residual tracks via the ROOT files (#15943)

With --enable-mc the interpolation workflow writes the TrackDataMC vector as branch tracksMC of o2residuals_tpc.root
(with --send-track-data), and the residual aggregator reads it back from there, so the MC truth no longer needs the
aggregator to run directly behind the interpolation workflow (disable-root-input). This is the setup of the O2DPG MC
workflow, where both run as separate tasks.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: matthias-kleiner

Detectors/GlobalTrackingWorkflow/tpcinterpolationworkflow/include/TPCInterpolationWorkflow/TPCResidualWriterSpec.h

@@ -22,7 +22,7 @@ namespace tpc
 {
 
 /// create a processor spec
-framework::DataProcessorSpec getTPCResidualWriterSpec(bool writeTrackData, bool debugOutput);
+framework::DataProcessorSpec getTPCResidualWriterSpec(bool writeTrackData, bool debugOutput, bool writeTrackDataMC);
 
 } // namespace tpc
 } // namespace o2

Detectors/GlobalTrackingWorkflow/tpcinterpolationworkflow/include/TPCInterpolationWorkflow/TPCUnbinnedResidualReaderSpec.h

@@ -30,26 +30,28 @@ namespace tpc
 class TPCUnbinnedResidualReader : public o2::framework::Task
 {
  public:
-  TPCUnbinnedResidualReader(bool trkInput) : mTrackInput(trkInput){};
+  TPCUnbinnedResidualReader(bool trkInput, bool mcInput) : mTrackInput(trkInput), mMCInput(mcInput) {};
   ~TPCUnbinnedResidualReader() override = default;
   void init(o2::framework::InitContext& ic) final;
   void run(o2::framework::ProcessingContext& pc) final;
 
  private:
   void connectTree();
   bool mTrackInput = false;
+  bool mMCInput = false;
   std::unique_ptr<TFile> mFileIn;
   std::unique_ptr<TTree> mTreeIn;
   std::string mInFileName;
   std::string mInTreeName;
   std::vector<UnbinnedResid> mUnbinnedResid, *mUnbinnedResidPtr = &mUnbinnedResid;
   std::vector<DetInfoResid> mDetInfoUnbRes, *mDetInfoUnbResPtr = &mDetInfoUnbRes;
   std::vector<TrackData> mTrackData, *mTrackDataPtr = &mTrackData;
+  std::vector<TrackDataMC> mTrackDataMC, *mTrackDataMCPtr = &mTrackDataMC;
   std::vector<TrackDataCompact> mTrackDataCompact, *mTrackDataCompactPtr = &mTrackDataCompact;
 };
 
-/// read unbinned TPC residuals and reference tracks from a root file
-framework::DataProcessorSpec getUnbinnedTPCResidualsReaderSpec(bool trkInput);
+/// read unbinned TPC residuals and reference tracks (with mcInput also their MC truth) from a root file
+framework::DataProcessorSpec getUnbinnedTPCResidualsReaderSpec(bool trkInput, bool mcInput);
 
 } // namespace tpc
 } // namespace o2

Detectors/GlobalTrackingWorkflow/tpcinterpolationworkflow/src/TPCResidualWriterSpec.cxx

@@ -29,7 +29,7 @@ namespace tpc
 template <typename T>
 using BranchDefinition = MakeRootTreeWriterSpec::BranchDefinition<T>;
 
-DataProcessorSpec getTPCResidualWriterSpec(bool writeTrackData, bool debugOutput)
+DataProcessorSpec getTPCResidualWriterSpec(bool writeTrackData, bool debugOutput, bool writeTrackDataMC)
 {
   bool writeUnfiltered = SpacePointsCalibConfParam::Instance().writeUnfiltered;
   return MakeRootTreeWriterSpec("tpc-residuals-writer",
@@ -39,6 +39,7 @@ DataProcessorSpec getTPCResidualWriterSpec(bool writeTrackData, bool debugOutput
                                 BranchDefinition<std::vector<DetInfoResid>>{InputSpec{"detInfo", "GLO", "DETINFORES"}, "detInfo"},
                                 BranchDefinition<std::vector<TrackDataCompact>>{InputSpec{"trackRefs", "GLO", "TRKREFS"}, "trackRefs"},
                                 BranchDefinition<std::vector<TrackData>>{InputSpec{"tracks", "GLO", "TRKDATA"}, "tracks", (writeTrackData ? 1 : 0)},
+                                BranchDefinition<std::vector<TrackDataMC>>{InputSpec{"tracksMC", "GLO", "TRKDATAMC"}, "tracksMC", (writeTrackDataMC ? 1 : 0)},
                                 BranchDefinition<std::vector<TrackDataExtended>>{InputSpec{"trackExt", "GLO", "TRKDATAEXT"}, "trackExt", (debugOutput ? 1 : 0)})();
 }
 

Detectors/GlobalTrackingWorkflow/tpcinterpolationworkflow/src/TPCUnbinnedResidualReaderSpec.cxx

@@ -52,6 +52,12 @@ void TPCUnbinnedResidualReader::connectTree()
   if (mTrackInput) {
     mTreeIn->SetBranchAddress("tracks", &mTrackDataPtr);
   }
+  if (mMCInput) {
+    if (!mTreeIn->GetBranch("tracksMC")) {
+      LOG(fatal) << "No tracksMC branch found in the unbinned residuals tree, the interpolation workflow must run with --enable-mc";
+    }
+    mTreeIn->SetBranchAddress("tracksMC", &mTrackDataMCPtr);
+  }
   LOG(info) << "Loaded tree from " << mInFileName << " with " << mTreeIn->GetEntries() << " entries";
 }
 
@@ -68,14 +74,17 @@ void TPCUnbinnedResidualReader::run(ProcessingContext& pc)
     LOG(info) << "Pushing " << mTrackData.size() << " reference tracks for these residuals";
     pc.outputs().snapshot(Output{"GLO", "TRKDATA", 0}, mTrackData);
   }
+  if (mMCInput) {
+    pc.outputs().snapshot(Output{"GLO", "TRKDATAMC", 0}, mTrackDataMC);
+  }
 
   if (mTreeIn->GetReadEntry() + 1 >= mTreeIn->GetEntries()) {
     pc.services().get<ControlService>().endOfStream();
     pc.services().get<ControlService>().readyToQuit(QuitRequest::Me);
   }
 }
 
-DataProcessorSpec getUnbinnedTPCResidualsReaderSpec(bool trkInput)
+DataProcessorSpec getUnbinnedTPCResidualsReaderSpec(bool trkInput, bool mcInput)
 {
   std::vector<OutputSpec> outputs;
   outputs.emplace_back("GLO", "UNBINNEDRES", 0, Lifetime::Timeframe);
@@ -84,12 +93,15 @@ DataProcessorSpec getUnbinnedTPCResidualsReaderSpec(bool trkInput)
   if (trkInput) {
     outputs.emplace_back("GLO", "TRKDATA", 0, Lifetime::Timeframe);
   }
+  if (mcInput) {
+    outputs.emplace_back("GLO", "TRKDATAMC", 0, Lifetime::Timeframe);
+  }
 
   return DataProcessorSpec{
     "TPCUnbinnedResidualReader",
     Inputs{},
     outputs,
-    AlgorithmSpec{adaptFromTask<TPCUnbinnedResidualReader>(trkInput)},
+    AlgorithmSpec{adaptFromTask<TPCUnbinnedResidualReader>(trkInput, mcInput)},
     Options{
       {"input-filename", VariantType::String, "o2residuals_tpc.root", {"Name of the input file"}},
       {"input-dir", VariantType::String, "none", {"Input directory"}},

Detectors/GlobalTrackingWorkflow/tpcinterpolationworkflow/src/tpc-interpolation-workflow.cxx

@@ -113,7 +113,7 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
 
   specs.emplace_back(o2::tpc::getTPCInterpolationSpec(srcClusters, srcVtx, srcTracks, srcTracksMap, useMC, processITSTPConly, sendTrackData, debugOutput, extDetResid, doStag));
   if (!configcontext.options().get<bool>("disable-root-output")) {
-    specs.emplace_back(o2::tpc::getTPCResidualWriterSpec(sendTrackData, debugOutput));
+    specs.emplace_back(o2::tpc::getTPCResidualWriterSpec(sendTrackData, debugOutput, useMC && sendTrackData));
   }
 
   // MC labels only for the ITS-TPC tracks and their ITS and TPC parts (see getTPCInterpolationSpec)

Detectors/GlobalTrackingWorkflow/tpcinterpolationworkflow/src/tpc-residual-aggregator.cxx

@@ -53,9 +53,6 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
     LOG(error) << "MC truth input requires the track input (enable-track-input), will be ignored";
     mcInput = false;
   }
-  if (mcInput && !configcontext.options().get<bool>("disable-root-input")) {
-    LOG(fatal) << "MC truth input is only supported directly from the interpolation workflow (disable-root-input)";
-  }
 
   bool writeUnbinnedResiduals = false;
   bool writeBinnedResiduals = false;
@@ -85,7 +82,7 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
 
   WorkflowSpec specs;
   if (!configcontext.options().get<bool>("disable-root-input")) {
-    specs.emplace_back(o2::tpc::getUnbinnedTPCResidualsReaderSpec(trkInput));
+    specs.emplace_back(o2::tpc::getUnbinnedTPCResidualsReaderSpec(trkInput, mcInput));
   }
   specs.emplace_back(getTPCResidualAggregatorSpec(trkInput, ctpInput, writeUnbinnedResiduals, writeBinnedResiduals, writeTrackData, mcInput));