TPC SCD: store the MC truth of the residual tracks (--enable-mc)

The interpolation workflow had MC switched off ("not yet implemented"). With --enable-mc (opt-in, since the
workflow also runs on data without --disable-mc) and --send-track-data, a TrackDataMC vector aligned 1:1 with
the TrackData output is sent as GLO/TRKDATAMC: MC labels of the ITS-TPC part of the seeding track and of its
ITS and TPC parts (flag for a fake ITS-TPC match), the truth at the ITS outer parameters (nearest ITS track
reference, propagated with the workflow's material correction and the true mass to the x and alpha of
TrackData::par, with the distance to that reference), the truth at the TPC entrance (first TPC track reference
in time, sector frame) with the distance between it, propagated to the innermost TPC cluster of the track, and
that cluster (flags loopers, wrong legs and fakes), and the PDG code. Only the ITS, TPC and ITS-TPC track
labels are read; the kinematics are looked up event by event and released after each event. The residual
aggregator writes the truth as branch trkMC of the trackData tree with --enable-mc.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: matthias-kleiner

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

@@ -23,12 +23,14 @@
 #include "DetectorsBase/GRPGeomHelper.h"
 #include "TPCCalibration/VDriftHelper.h"
 #include "DataFormatsITSMFT/TopologyDictionary.h"
+#include "Steer/MCKinematicsReader.h"
 
 using namespace o2::framework;
 
 namespace o2::globaltracking
 {
 struct DataRequest;
+struct RecoContainer;
 } // namespace o2::globaltracking
 
 namespace o2
@@ -40,14 +42,15 @@ class TPCInterpolationDPL : public Task
  public:
   TPCInterpolationDPL(std::shared_ptr<o2::globaltracking::DataRequest> dr, o2::dataformats::GlobalTrackID::mask_t src, o2::dataformats::GlobalTrackID::mask_t srcMap, std::shared_ptr<o2::base::GRPGeomRequest> gr, bool useMC,
                       bool processITSTPConly, bool sendTrackData, bool debugOutput, bool extDetResid) : mDataRequest(dr), mSources(src), mSourcesMap(srcMap), mGGCCDBRequest(gr), mUseMC(useMC), mProcessITSTPConly(processITSTPConly), mSendTrackData(sendTrackData), mDebugOutput(debugOutput), mExtDetResid(extDetResid) {}
-  ~TPCInterpolationDPL() override = default;
+  ~TPCInterpolationDPL() override;
   void init(InitContext& ic) final;
   void run(ProcessingContext& pc) final;
   void endOfStream(EndOfStreamContext& ec) final;
   void finaliseCCDB(ConcreteDataMatcher& matcher, void* obj) final;
 
  private:
   void updateTimeDependentParams(ProcessingContext& pc);
+  void fillMCTruth(const o2::globaltracking::RecoContainer& recoData);
   o2::tpc::TrackInterpolation mInterpolation;                    ///< track interpolation engine
   std::shared_ptr<o2::globaltracking::DataRequest> mDataRequest; ///< steers the input
   std::shared_ptr<o2::base::GRPGeomRequest> mGGCCDBRequest;
@@ -56,6 +59,8 @@ class TPCInterpolationDPL : public Task
   o2::dataformats::GlobalTrackID::mask_t mSources{};        ///< which input sources are configured
   o2::dataformats::GlobalTrackID::mask_t mSourcesMap{};     ///< possible subset of mSources specifically for map creation
   bool mUseMC{false}; ///< MC flag
+  std::unique_ptr<o2::steer::MCKinematicsReader> mMCReader; ///< MC kinematics and track references (MC only)
+  std::vector<TrackDataMC> mTrackDataMC;                    ///< MC truth aligned with the TrackData output (MC only)
   bool mProcessITSTPConly{false}; ///< should also tracks without outer point (ITS-TPC only) be processed?
   bool mProcessSeeds{false};      ///< process not only most complete track, but also its shorter parts
   bool mDebugOutput{false};       ///< add more information to the output (track points of ITS, TRD and TOF)

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

@@ -41,7 +41,7 @@ namespace calibration
 class ResidualAggregatorDevice : public o2::framework::Task
 {
  public:
-  ResidualAggregatorDevice(std::shared_ptr<o2::base::GRPGeomRequest> req, bool trackInput, bool ctpInput, bool writeUnbinnedResiduals, bool writeBinnedResiduals, bool writeTrackData, std::shared_ptr<o2::globaltracking::DataRequest> dataRequest) : mCCDBRequest(req), mTrackInput(trackInput), mCTPInput(ctpInput), mWriteUnbinnedResiduals(writeUnbinnedResiduals), mWriteBinnedResiduals(writeBinnedResiduals), mWriteTrackData(writeTrackData), mDataRequest(dataRequest) {}
+  ResidualAggregatorDevice(std::shared_ptr<o2::base::GRPGeomRequest> req, bool trackInput, bool ctpInput, bool writeUnbinnedResiduals, bool writeBinnedResiduals, bool writeTrackData, bool mcInput, std::shared_ptr<o2::globaltracking::DataRequest> dataRequest) : mCCDBRequest(req), mTrackInput(trackInput), mCTPInput(ctpInput), mWriteUnbinnedResiduals(writeUnbinnedResiduals), mWriteBinnedResiduals(writeBinnedResiduals), mWriteTrackData(writeTrackData), mMCInput(mcInput), mDataRequest(dataRequest) {}
 
   void init(o2::framework::InitContext& ic) final
   {
@@ -97,6 +97,7 @@ class ResidualAggregatorDevice : public o2::framework::Task
     mAggregator->setWriteBinnedResiduals(mWriteBinnedResiduals);
     mAggregator->setWriteUnbinnedResiduals(mWriteUnbinnedResiduals);
     mAggregator->setWriteTrackData(mWriteTrackData);
+    mAggregator->setWriteTrackDataMC(mMCInput);
     mAggregator->setCompression(ic.options().get<int>("compression"));
   }
 
@@ -141,6 +142,14 @@ class ResidualAggregatorDevice : public o2::framework::Task
       trkData.emplace(pc.inputs().get<gsl::span<o2::tpc::TrackData>>("trkData"));
       trkDataPtr = &trkData.value();
     }
+    // MC truth of the track data (optional, MC only)
+    const gsl::span<const o2::tpc::TrackDataMC>* trkDataMCPtr = nullptr;
+    using trkDataMCType = std::decay_t<decltype(pc.inputs().get<gsl::span<o2::tpc::TrackDataMC>>(""))>;
+    std::optional<trkDataMCType> trkDataMC;
+    if (mMCInput) {
+      trkDataMC.emplace(pc.inputs().get<gsl::span<o2::tpc::TrackDataMC>>("trkDataMC"));
+      trkDataMCPtr = &trkDataMC.value();
+    }
     // CTP lumi input (optional)
     const o2::ctp::LumiInfo* lumi = nullptr;
     using lumiDataType = std::decay_t<decltype(pc.inputs().get<o2::ctp::LumiInfo>(""))>;
@@ -152,7 +161,7 @@ class ResidualAggregatorDevice : public o2::framework::Task
 
     o2::base::TFIDInfoHelper::fillTFIDInfo(pc, mAggregator->getCurrentTFInfo());
     LOG(detail) << "Processing TF " << mAggregator->getCurrentTFInfo().tfCounter << " with " << trkData->size() << " tracks and " << residualsData.size() << " unbinned residuals associated to them";
-    mAggregator->process(residualsData, residualsDataDet, trackRefs, trkDataPtr, lumi);
+    mAggregator->process(residualsData, residualsDataDet, trackRefs, trkDataPtr, trkDataMCPtr, lumi);
     std::chrono::duration<double, std::milli> runDuration = std::chrono::high_resolution_clock::now() - runStartTime;
     LOGP(debug, "Duration for run method: {} ms. From this taken for time dependent param update: {} ms",
          std::chrono::duration_cast<std::chrono::milliseconds>(runDuration).count(),
@@ -205,6 +214,7 @@ class ResidualAggregatorDevice : public o2::framework::Task
   bool mWriteBinnedResiduals{false};   ///< flag, whether to write binned residuals to output file
   bool mWriteUnbinnedResiduals{false}; ///< flag, whether to write unbinned residuals to output file
   bool mWriteTrackData{false};         ///< flag, whether to write track data to output file
+  bool mMCInput{false};                ///< flag whether to expect the MC truth of the track data as input
   bool mRunStopRequested{false};       ///< flag in case the run was stopped
   bool mInitDone{false};               ///< flag whether initialization was done for current run
 };
@@ -214,7 +224,7 @@ class ResidualAggregatorDevice : public o2::framework::Task
 namespace framework
 {
 
-DataProcessorSpec getTPCResidualAggregatorSpec(bool trackInput, bool ctpInput, bool writeUnbinnedResiduals, bool writeBinnedResiduals, bool writeTrackData)
+DataProcessorSpec getTPCResidualAggregatorSpec(bool trackInput, bool ctpInput, bool writeUnbinnedResiduals, bool writeBinnedResiduals, bool writeTrackData, bool mcInput = false)
 {
   std::shared_ptr<o2::globaltracking::DataRequest> dataRequest = std::make_shared<o2::globaltracking::DataRequest>();
   if (ctpInput) {
@@ -227,6 +237,9 @@ DataProcessorSpec getTPCResidualAggregatorSpec(bool trackInput, bool ctpInput, b
   inputs.emplace_back("trackRefs", "GLO", "TRKREFS");
   if (trackInput) {
     inputs.emplace_back("trkData", "GLO", "TRKDATA");
+    if (mcInput) {
+      inputs.emplace_back("trkDataMC", "GLO", "TRKDATAMC");
+    }
   }
   auto ccdbRequest = std::make_shared<o2::base::GRPGeomRequest>(true,                           // orbitResetTime
                                                                 true,                           // GRPECS=true
@@ -240,7 +253,7 @@ DataProcessorSpec getTPCResidualAggregatorSpec(bool trackInput, bool ctpInput, b
     "residual-aggregator",
     inputs,
     Outputs{},
-    AlgorithmSpec{adaptFromTask<o2::calibration::ResidualAggregatorDevice>(ccdbRequest, trackInput, ctpInput, writeUnbinnedResiduals, writeBinnedResiduals, writeTrackData, dataRequest)},
+    AlgorithmSpec{adaptFromTask<o2::calibration::ResidualAggregatorDevice>(ccdbRequest, trackInput, ctpInput, writeUnbinnedResiduals, writeBinnedResiduals, writeTrackData, trackInput && mcInput, dataRequest)},
     Options{
       {"sec-per-slot", VariantType::UInt32, 600u, {"number of seconds per calibration time slot (put 0 for infinite slot length)"}},
       {"updateInterval", VariantType::UInt32, 6'000u, {"update interval in number of TFs (only used in case slot length is infinite)"}},

Detectors/GlobalTrackingWorkflow/tpcinterpolationworkflow/src/TPCInterpolationSpec.cxx

@@ -13,6 +13,8 @@
 
 #include <vector>
 #include <unordered_map>
+#include <algorithm>
+#include <cmath>
 
 #include "DataFormatsITS/TrackITS.h"
 #include "ITSBase/GeometryTGeo.h"
@@ -33,6 +35,9 @@
 #include "Framework/ConfigParamRegistry.h"
 #include "Framework/ControlService.h"
 #include "Framework/DeviceSpec.h"
+#include "Steer/MCKinematicsReader.h"
+#include "SimulationDataFormat/TrackReference.h"
+#include "SimulationDataFormat/O2DatabasePDG.h"
 
 using namespace o2::framework;
 using namespace o2::globaltracking;
@@ -44,6 +49,8 @@ namespace o2
 namespace tpc
 {
 
+TPCInterpolationDPL::~TPCInterpolationDPL() = default;
+
 void TPCInterpolationDPL::init(InitContext& ic)
 {
   //-------- init geometry and field --------//
@@ -59,6 +66,16 @@ void TPCInterpolationDPL::init(InitContext& ic)
   int lane = ic.services().get<const o2::framework::DeviceSpec>().inputTimesliceId;
   int maxLanes = ic.services().get<const o2::framework::DeviceSpec>().maxInputTimeslices;
   mInterpolation.setLane(lane, maxLanes);
+  if (mUseMC) {
+    if (!mSendTrackData) {
+      LOG(warning) << "MC truth is stored aligned with the track data, but send-track-data is not set: no MC truth will be sent";
+    }
+    mMCReader = std::make_unique<o2::steer::MCKinematicsReader>();
+    auto mcContext = ic.options().get<std::string>("mc-collision-context");
+    if (!mMCReader->initFromDigitContext(mcContext)) {
+      LOG(fatal) << "Could not initialize the MC kinematics reader from " << mcContext;
+    }
+  }
 }
 
 void TPCInterpolationDPL::updateTimeDependentParams(ProcessingContext& pc)
@@ -149,9 +166,168 @@ void TPCInterpolationDPL::run(ProcessingContext& pc)
   if (mDebugOutput) {
     pc.outputs().snapshot(Output{"GLO", "TRKDATAEXT", 0}, mInterpolation.getTrackDataExtended());
   }
+  if (mUseMC && mSendTrackData) {
+    fillMCTruth(recoData);
+    pc.outputs().snapshot(Output{"GLO", "TRKDATAMC", 0}, mTrackDataMC);
+  }
   mInterpolation.reset();
 }
 
+void TPCInterpolationDPL::fillMCTruth(const RecoContainer& recoData)
+{
+  // MC truth for every stored TrackData: labels of the ITS-TPC part of the seed and of its ITS and TPC parts, the truth at
+  // the ITS outer parameters (the ITS track reference nearest to TrackData::par, propagated to its x with the material
+  // correction of the workflow and the mass of the true particle) and the truth at the TPC entrance (the first TPC track
+  // reference in time, in the sector frame)
+  const auto& trkData = mInterpolation.getReferenceTracks();
+  mTrackDataMC.clear();
+  mTrackDataMC.resize(trkData.size());
+  struct Lookup {
+    o2::MCCompLabel lbl;
+    uint32_t idx;
+    bool its; // ITS outer (true) or TPC entrance (false)
+  };
+  std::vector<Lookup> lookups;
+  lookups.reserve(2 * trkData.size());
+  for (size_t i = 0; i < trkData.size(); ++i) {
+    auto& mc = mTrackDataMC[i];
+    auto gidSet = recoData.getSingleDetectorRefs(trkData[i].gid);
+    auto gidITS = gidSet[GTrackID::ITS].isIndexSet() ? gidSet[GTrackID::ITS] : gidSet[GTrackID::ITSAB];
+    if (gidSet[GTrackID::ITSTPC].isIndexSet()) {
+      mc.label = recoData.getTrackMCLabel(gidSet[GTrackID::ITSTPC]);
+    }
+    if (gidITS.isIndexSet()) {
+      mc.labelITS = recoData.getTrackMCLabel(gidITS);
+    }
+    if (gidSet[GTrackID::TPC].isIndexSet()) {
+      mc.labelTPC = recoData.getTrackMCLabel(gidSet[GTrackID::TPC]);
+    }
+    if (mc.labelITS.isValid() && mc.labelTPC.isValid() && mc.labelITS.getTrackEventSourceID() != mc.labelTPC.getTrackEventSourceID()) {
+      mc.flags |= TrackDataMC::FakeITSTPC;
+    }
+    const auto& lblITS = mc.labelITS.isValid() ? mc.labelITS : mc.label;
+    const auto& lblTPC = mc.labelTPC.isValid() ? mc.labelTPC : mc.label;
+    if (lblITS.isValid()) {
+      lookups.push_back({lblITS, uint32_t(i), true});
+    }
+    if (lblTPC.isValid()) {
+      lookups.push_back({lblTPC, uint32_t(i), false});
+    }
+  }
+  // the reader loads the kinematics of a whole event (can be >100 MB): process event by event and release it right after
+  std::sort(lookups.begin(), lookups.end(), [](const Lookup& a, const Lookup& b) {
+    return a.lbl.getSourceID() != b.lbl.getSourceID() ? a.lbl.getSourceID() < b.lbl.getSourceID() : a.lbl.getEventID() < b.lbl.getEventID();
+  });
+  auto pdgToPID = [](int pdg) {
+    switch (std::abs(pdg)) {
+      case 11:
+        return o2::track::PID(o2::track::PID::Electron);
+      case 13:
+        return o2::track::PID(o2::track::PID::Muon);
+      case 321:
+        return o2::track::PID(o2::track::PID::Kaon);
+      case 2212:
+        return o2::track::PID(o2::track::PID::Proton);
+      case 1000010020:
+        return o2::track::PID(o2::track::PID::Deuteron);
+      case 1000010030:
+        return o2::track::PID(o2::track::PID::Triton);
+      case 1000020030:
+        return o2::track::PID(o2::track::PID::Helium3);
+      case 1000020040:
+        return o2::track::PID(o2::track::PID::Alpha);
+      default:
+        return o2::track::PID(o2::track::PID::Pion);
+    }
+  };
+  auto refToPar = [&pdgToPID](const o2::TrackReference& ref, int charge, int pdg, bool sectorAlpha) {
+    std::array<float, 3> xyz{ref.X(), ref.Y(), ref.Z()};
+    std::array<float, 3> pxyz{ref.Px(), ref.Py(), ref.Pz()};
+    return o2::track::TrackPar(xyz, pxyz, charge, sectorAlpha, pdgToPID(pdg));
+  };
+  const auto matCorr = static_cast<o2::base::Propagator::MatCorrType>(mMatCorr);
+  auto prop = o2::base::Propagator::Instance();
+  int curSrc = -1;
+  int curEv = -1;
+  for (const auto& lk : lookups) {
+    const auto& lbl = lk.lbl;
+    if (lbl.getSourceID() != curSrc || lbl.getEventID() != curEv) {
+      if (curSrc >= 0) {
+        mMCReader->releaseTracksForSourceAndEvent(curSrc, curEv);
+      }
+      curSrc = lbl.getSourceID();
+      curEv = lbl.getEventID();
+    }
+    const auto& trk = trkData[lk.idx];
+    auto& mc = mTrackDataMC[lk.idx];
+    const auto* mcTrk = mMCReader->getTrack(lbl);
+    int pdg = mcTrk ? mcTrk->GetPdgCode() : 0;
+    const auto* pPDG = mcTrk ? O2DatabasePDG::Instance()->GetParticle(pdg) : nullptr;
+    int charge = pPDG ? int(std::lround(pPDG->Charge() / 3.)) : 0; // TParticlePDG charge is in units of |e|/3
+    auto refs = mMCReader->getTrackRefs(lbl.getSourceID(), lbl.getEventID(), lbl.getTrackID());
+    if (lk.its) { // ITS outer: track reference of the ITS part nearest to TrackData::par
+      mc.pdg = pdg;
+      const o2::TrackReference* best = nullptr;
+      float bestD2 = 1e30f;
+      auto xyzReco = trk.par.getXYZGlo();
+      for (const auto& ref : refs) {
+        if (ref.getDetectorId() != DetID::ITS) {
+          continue;
+        }
+        float dx = ref.X() - xyzReco.X();
+        float dy = ref.Y() - xyzReco.Y();
+        float dz = ref.Z() - xyzReco.Z();
+        float d2 = dx * dx + dy * dy + dz * dz;
+        if (d2 < bestD2) {
+          bestD2 = d2;
+          best = &ref;
+        }
+      }
+      if (best && charge) {
+        auto par = refToPar(*best, charge, pdg, false);
+        if (par.rotateParam(trk.par.getAlpha()) && prop->PropagateToXBxByBz(par, trk.par.getX(), 0.999f, o2::base::Propagator::MAX_STEP, matCorr)) {
+          mc.parITSOut = par;
+          mc.distITSRef = std::sqrt(bestD2);
+          mc.flags |= TrackDataMC::HasITSOut;
+        }
+      }
+    } else { // TPC entrance: first TPC track reference in time of the TPC part
+      if (!mc.labelITS.isValid() && !mc.label.isValid()) {
+        mc.pdg = pdg; // no ITS lookup for this track
+      }
+      const o2::TrackReference* first = nullptr;
+      for (const auto& ref : refs) {
+        if (ref.getDetectorId() == DetID::TPC && (!first || ref.getTime() < first->getTime())) {
+          first = &ref;
+        }
+      }
+      if (first && charge) {
+        mc.parTPCIn = refToPar(*first, charge, pdg, true);
+        mc.flags |= TrackDataMC::HasTPCIn;
+        // association check (loopers, wrong leg, fake): truth at the innermost TPC cluster of the track vs that cluster
+        const auto& clRes = mInterpolation.getClusterResiduals();
+        const UnbinnedResid* inner = nullptr;
+        for (int ic = trk.clIdx.getFirstEntry(); ic < trk.clIdx.getFirstEntry() + trk.clIdx.getEntries(); ++ic) {
+          if (clRes[ic].row < constants::MAXGLOBALPADROW && (!inner || clRes[ic].row < inner->row)) {
+            inner = &clRes[ic];
+          }
+        }
+        if (inner) {
+          auto par = mc.parTPCIn;
+          float yCl = inner->y * param::MaxY / 0x7fff + inner->dy * param::MaxResid / 0x7fff;
+          float zCl = inner->z * param::MaxZ / 0x7fff + inner->dz * param::MaxResid / 0x7fff;
+          if (par.rotateParam(o2::math_utils::sector2Angle(inner->sec)) && prop->PropagateToXBxByBz(par, param::RowX[inner->row], 0.999f, o2::base::Propagator::MAX_STEP, matCorr)) {
+            mc.distTPCRef = std::hypot(par.getY() - yCl, par.getZ() - zCl);
+          }
+        }
+      }
+    }
+  }
+  if (curSrc >= 0) {
+    mMCReader->releaseTracksForSourceAndEvent(curSrc, curEv);
+  }
+}
+
 void TPCInterpolationDPL::endOfStream(EndOfStreamContext& ec)
 {
   mInterpolation.finalize();
@@ -165,14 +341,13 @@ DataProcessorSpec getTPCInterpolationSpec(GTrackID::mask_t srcCls, GTrackID::mas
   dataRequest->setITSPerLayer(itsStag);
   std::vector<OutputSpec> outputs;
 
-  if (useMC) {
-    LOG(fatal) << "MC usage must be disabled for this workflow, since it is not yet implemented";
+  dataRequest->requestTracks(srcVtx, false);
+  dataRequest->requestClusters(srcCls, false);
+  dataRequest->requestPrimaryVertices(false);
+  if (useMC) { // the MC truth needs only the labels of the ITS-TPC tracks and of their ITS and TPC parts
+    dataRequest->requestTracks(GTrackID::getSourcesMask("ITS,TPC,ITS-TPC"), true);
   }
 
-  dataRequest->requestTracks(srcVtx, useMC);
-  dataRequest->requestClusters(srcCls, useMC);
-  dataRequest->requestPrimaryVertices(useMC);
-
   auto ggRequest = std::make_shared<o2::base::GRPGeomRequest>(false,                             // orbitResetTime
                                                               true,                              // GRPECS=true
                                                               true,                              // GRPLHCIF
@@ -191,6 +366,9 @@ DataProcessorSpec getTPCInterpolationSpec(GTrackID::mask_t srcCls, GTrackID::mas
   if (debugOutput) {
     outputs.emplace_back("GLO", "TRKDATAEXT", 0, Lifetime::Timeframe);
   }
+  if (useMC && sendTrackData) {
+    outputs.emplace_back("GLO", "TRKDATAMC", 0, Lifetime::Timeframe);
+  }
 
   return DataProcessorSpec{
     "tpc-track-interpolation",
@@ -200,7 +378,8 @@ DataProcessorSpec getTPCInterpolationSpec(GTrackID::mask_t srcCls, GTrackID::mas
     Options{
       {"matCorrType", VariantType::Int, 2, {"material correction type (definition in Propagator.h)"}},
       {"sec-per-slot", VariantType::UInt32, 300u, {"number of seconds per calibration time slot (put 0 for infinite slot length)"}},
-      {"process-seeds", VariantType::Bool, false, {"do not remove duplicates, e.g. for ITS-TPC-TRD track also process its seeding ITS-TPC part"}}}};
+      {"process-seeds", VariantType::Bool, false, {"do not remove duplicates, e.g. for ITS-TPC-TRD track also process its seeding ITS-TPC part"}},
+      {"mc-collision-context", VariantType::String, "collisioncontext.root", {"collision context used to access the MC kinematics and track references (MC only)"}}}};
 }
 
 } // namespace tpc