Detector/TPC: Added option to add MC information to track and cluster filter

- Fixed bug where unfiltered tracks were propagated instead of filtered tracks

Author: mhemmer-cern

Detectors/TPC/calibration/include/TPCCalibration/TrackDump.h

@@ -20,6 +20,7 @@
 #include <vector>
 
 #include "DataFormatsTPC/TrackTPC.h"
+#include "SimulationDataFormat/MCCompLabel.h"
 #include "CommonUtils/TreeStreamRedirector.h"
 
 /// \file TrackDump.h
@@ -76,12 +77,13 @@ class TrackDump
     ClassDefNV(TrackInfo, 1);
   };
 
-  void filter(const gsl::span<const TrackTPC> tracks, ClusterNativeAccess const& clusterIndex, const gsl::span<const o2::tpc::TPCClRefElem> clRefs);
+  void filter(const gsl::span<const TrackTPC> tracks, ClusterNativeAccess const& clusterIndex, const gsl::span<const o2::tpc::TPCClRefElem> clRefs, const gsl::span<const MCCompLabel> mcLabels);
   void finalize();
 
   std::string outputFileName{"filtered-tracks-and-clusters.root"}; ///< Name of the output file with the tree
   bool writeTracks{true};                                          ///< write global cluster information for quick drawing
   bool writeGlobal{false};                                         ///< write global cluster information for quick drawing
+  bool writeMC{false};                                             ///< write MC track information for quick drawing
  private:
   std::unique_ptr<o2::utils::TreeStreamRedirector> mTreeDump; ///< Tree writer
 };

Detectors/TPC/calibration/src/TrackDump.cxx

@@ -23,7 +23,7 @@
 using namespace o2::tpc;
 using namespace o2::tpc::constants;
 
-void TrackDump::filter(const gsl::span<const TrackTPC> tracks, ClusterNativeAccess const& clusterIndex, const gsl::span<const o2::tpc::TPCClRefElem> clRefs)
+void TrackDump::filter(const gsl::span<const TrackTPC> tracks, ClusterNativeAccess const& clusterIndex, const gsl::span<const o2::tpc::TPCClRefElem> clRefs, const gsl::span<const o2::MCCompLabel> mcLabels)
 {
   if (!mTreeDump && outputFileName.size()) {
     mTreeDump = std::make_unique<utils::TreeStreamRedirector>(outputFileName.data(), "recreate");
@@ -32,6 +32,7 @@ void TrackDump::filter(const gsl::span<const TrackTPC> tracks, ClusterNativeAcce
   std::vector<TrackInfo> tpcTracks;
   std::vector<std::vector<ClusterGlobal>> clustersGlobalEvent;
   std::vector<ClusterGlobal>* clustersGlobal{};
+  std::vector<o2::MCCompLabel> tpcTracksMCTruth;
 
   GPUCA_NAMESPACE::gpu::GPUTPCGeometry gpuGeom;
 
@@ -57,6 +58,11 @@ void TrackDump::filter(const gsl::span<const TrackTPC> tracks, ClusterNativeAcce
       }
     }
   }
+  if (writeMC) {
+    for (const auto& mcLabel : mcLabels) {
+      tpcTracksMCTruth.emplace_back(mcLabel);
+    }
+  }
 
   if (mTreeDump) {
     auto& tree = (*mTreeDump) << "tpcrec";
@@ -67,6 +73,9 @@ void TrackDump::filter(const gsl::span<const TrackTPC> tracks, ClusterNativeAcce
     if (writeGlobal) {
       tree << "cls" << clustersGlobalEvent;
     }
+    if (writeMC) {
+      tree << "TPCTracksMCTruth=" << tpcTracksMCTruth;
+    }
     tree << "\n";
     //  << "clusters=" << clInfoVec
   }

Detectors/TPC/workflow/include/TPCWorkflow/TrackAndClusterFilterSpec.h

@@ -25,7 +25,7 @@ namespace o2::tpc
 
 /// create a processor spec
 /// read simulated TPC clusters from file and publish
-o2::framework::DataProcessorSpec getTrackAndClusterFilterSpec(const std::string dataDescriptionStr = "TRACKS");
+o2::framework::DataProcessorSpec getTrackAndClusterFilterSpec(const std::string dataDescriptionStr = "TRACKS", const bool writeMC = false);
 
 } // namespace o2::tpc
 

Detectors/TPC/workflow/src/TrackAndClusterFilterSpec.cxx

@@ -38,6 +38,7 @@ class TrackAndClusterFilterDevice : public o2::framework::Task
 {
  public:
   TrackAndClusterFilterDevice() = default;
+  TrackAndClusterFilterDevice(bool writeMC) { mTrackDump.writeMC = writeMC; };
 
   void init(o2::framework::InitContext& ic) final
   {
@@ -65,15 +66,23 @@ class TrackAndClusterFilterDevice : public o2::framework::Task
     // }
     const auto tracks = pc.inputs().get<gsl::span<o2::tpc::TrackTPC>>("trackTPC");
     const auto clRefs = pc.inputs().get<gsl::span<o2::tpc::TPCClRefElem>>("trackTPCClRefs");
+    const auto mcLabel = (mTrackDump.writeMC) ? pc.inputs().get<gsl::span<o2::MCCompLabel>>("trackTPCMCTruth") : gsl::span<o2::MCCompLabel>();
     const auto& clustersInputs = getWorkflowTPCInput(pc);
 
     std::vector<TrackTPC> filteredTracks;
-    std::copy_if(tracks.begin(), tracks.end(), std::back_inserter(filteredTracks),
-                 [this](const auto& track) { return mCuts.goodTrack(track); });
-
-    LOGP(info, "Filtered {} good tracks out of {} total tpc tracks", filteredTracks.size(), tracks.size());
-
-    mTrackDump.filter(tracks, clustersInputs->clusterIndex, clRefs);
+    std::vector<o2::MCCompLabel> filteredMCLabels;
+    for (size_t iTrack = 0; iTrack < tracks.size(); iTrack++) {
+      if (mCuts.goodTrack(tracks[iTrack])) {
+        filteredTracks.emplace_back(tracks[iTrack]);
+        if (mTrackDump.writeMC) {
+          filteredMCLabels.emplace_back(mcLabel[iTrack]);
+        }
+      }
+    }
+
+    LOGP(info, "Filtered {} good tracks with {} MC labels out of {} total tpc tracks", filteredTracks.size(), filteredMCLabels.size(), tracks.size());
+
+    mTrackDump.filter(filteredTracks, clustersInputs->clusterIndex, clRefs, filteredMCLabels);
   }
 
   void endOfStream(o2::framework::EndOfStreamContext& /*ec*/) final
@@ -91,7 +100,7 @@ class TrackAndClusterFilterDevice : public o2::framework::Task
   TrackCuts mCuts{};
 };
 
-DataProcessorSpec getTrackAndClusterFilterSpec(const std::string dataDescriptionStr)
+DataProcessorSpec getTrackAndClusterFilterSpec(const std::string dataDescriptionStr, const bool writeMC)
 {
   std::vector<InputSpec> inputs;
   std::vector<OutputSpec> outputs;
@@ -105,12 +114,15 @@ DataProcessorSpec getTrackAndClusterFilterSpec(const std::string dataDescription
   inputs.emplace_back("trackTPC", "TPC", dataDescription, 0, Lifetime::Timeframe);
   inputs.emplace_back("trackTPCClRefs", "TPC", "CLUSREFS", 0, Lifetime::Timeframe);
   inputs.emplace_back("clusTPC", ConcreteDataTypeMatcher{"TPC", "CLUSTERNATIVE"}, Lifetime::Timeframe);
+  if (writeMC) {
+    inputs.emplace_back("trackTPCMCTruth", "TPC", "TRACKSMCLBL", 0, Lifetime::Timeframe);
+  }
 
   return DataProcessorSpec{
     "tpc-track-and-cluster-filter",
     inputs,
     outputs,
-    AlgorithmSpec{adaptFromTask<TrackAndClusterFilterDevice>()},
+    AlgorithmSpec{adaptFromTask<TrackAndClusterFilterDevice>(writeMC)},
     Options{
       {"output-file", VariantType::String, "filtered-track-and-clusters.root", {"output file name"}},
       {"write-tracks", VariantType::Bool, true, {"dump filtered tracks and clusters"}},

Detectors/TPC/workflow/src/track-and-cluster-filter.cxx

@@ -41,6 +41,7 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
   std::vector<ConfigParamSpec> options{
     // {"use-digit-input", VariantType::Bool, false, {"use TPC digits as input, instead of raw data"}},
     {"data-description", VariantType::String, "TRACKS", {"Can be used to select filterd tracks, e.g. 'LASERTRACKS', 'MIPS'"}},
+    {"write-mc", VariantType::Bool, false, {"Can be used to write mc labels, e.g. 'TPCTracksMCTruth'"}},
   };
 
   std::swap(workflowOptions, options);
@@ -52,10 +53,11 @@ WorkflowSpec defineDataProcessing(ConfigContext const& cfg)
 {
   // const bool useDigitsAsInput = cfg.options().get<bool>("use-digit-input");
   const auto dataDescription = cfg.options().get<std::string>("data-description");
+  const auto writeMC = cfg.options().get<bool>("write-mc");
 
   WorkflowSpec specs;
 
-  specs.emplace_back(o2::tpc::getTrackAndClusterFilterSpec(dataDescription));
+  specs.emplace_back(o2::tpc::getTrackAndClusterFilterSpec(dataDescription, writeMC));
 
   return specs;
 }