TRD tracker can use transformed trklts directly

- calibration can be added to TrackletTransformer
- GPUTRDSpacePointInternal does not need fields for labels (is stored in separate array), covariance matrix (is calculated on-the-fly anyway since it depends on the track inclincation) and ID (tracklets and space points are kept in order now. for the sorting by HCId we fill an index array)

Author: martenole

Detectors/TRD/workflow/include/TRDWorkflow/TRDGlobalTrackingSpec.h

@@ -28,7 +28,7 @@ namespace trd
 class TRDGlobalTracking : public o2::framework::Task
 {
  public:
-  TRDGlobalTracking(bool useMC) : mUseMC(useMC) {}
+  TRDGlobalTracking(bool useMC, bool useTrkltTransf) : mUseMC(useMC), mUseTrackletTransform(useTrkltTransf) {}
   ~TRDGlobalTracking() override = default;
   void init(o2::framework::InitContext& ic) final;
   void run(o2::framework::ProcessingContext& pc) final;
@@ -40,11 +40,12 @@ class TRDGlobalTracking : public o2::framework::Task
   o2::gpu::GPUChainTracking* mChainTracking{nullptr}; ///< TRD tracker is run in the tracking chain
   std::unique_ptr<GeometryFlat> mFlatGeo{nullptr};    ///< flat TRD geometry
   bool mUseMC{false};                                 ///< MC flag
+  bool mUseTrackletTransform{false};                  ///< if true, output from TrackletTransformer is used instead of uncalibrated Tracklet64 directly
   TStopwatch mTimer;
 };
 
 /// create a processor spec
-framework::DataProcessorSpec getTRDGlobalTrackingSpec(bool useMC);
+framework::DataProcessorSpec getTRDGlobalTrackingSpec(bool useMC, bool useTrkltTransf);
 
 } // namespace trd
 } // namespace o2

Detectors/TRD/workflow/include/TRDWorkflow/TRDTrackingWorkflow.h

@@ -20,7 +20,7 @@ namespace o2
 namespace trd
 {
 
-framework::WorkflowSpec getTRDTrackingWorkflow(bool disableRootInp, bool disableRootOut);
+framework::WorkflowSpec getTRDTrackingWorkflow(bool disableRootInp, bool disableRootOut, bool useTrackletTransformer);
 
 } // namespace trd
 } // namespace o2

Detectors/TRD/workflow/include/TRDWorkflow/TRDTrackletReaderSpec.h

@@ -19,6 +19,7 @@
 #include "Framework/DataProcessorSpec.h"
 #include "Framework/Task.h"
 #include "DataFormatsTRD/Tracklet64.h"
+#include "DataFormatsTRD/CalibratedTracklet.h"
 #include "DataFormatsTRD/TriggerRecord.h"
 #include "SimulationDataFormat/MCCompLabel.h"
 
@@ -30,26 +31,31 @@ namespace trd
 class TRDTrackletReader : public o2::framework::Task
 {
  public:
-  TRDTrackletReader(bool useMC) : mUseMC(useMC) {}
+  TRDTrackletReader(bool useMC, bool useTrkltTransf) : mUseMC(useMC), mUseTrackletTransform(useTrkltTransf) {}
   ~TRDTrackletReader() override = default;
   void init(o2::framework::InitContext& ic) final;
   void run(o2::framework::ProcessingContext& pc) final;
 
  private:
-  void connectTree(const std::string& filename);
+  void connectTree();
+  void connectTreeCTracklet();
   bool mUseMC{false};
-  std::unique_ptr<TFile> mFile;
-  std::unique_ptr<TTree> mTree;
-  std::string mInFileName{"trdtracklets.root"};
-  std::string mInTreeName{"o2sim"};
+  bool mUseTrackletTransform{false};
+  std::unique_ptr<TFile> mFileTrklt;
+  std::unique_ptr<TTree> mTreeTrklt;
+  std::unique_ptr<TFile> mFileCTrklt;
+  std::unique_ptr<TTree> mTreeCTrklt;
+  std::string mInFileNameTrklt{"trdtracklets.root"};
+  std::string mInTreeNameTrklt{"o2sim"};
+  std::vector<o2::trd::CalibratedTracklet> mTrackletsCal, *mTrackletsCalPtr = &mTrackletsCal;
   std::vector<o2::trd::Tracklet64> mTracklets, *mTrackletsPtr = &mTracklets;
   std::vector<o2::trd::TriggerRecord> mTriggerRecords, *mTriggerRecordsPtr = &mTriggerRecords;
   std::vector<o2::MCCompLabel> mLabels, *mLabelsPtr = &mLabels;
 };
 
 /// create a processor spec
 /// read TRD tracklets from a root file
-framework::DataProcessorSpec getTRDTrackletReaderSpec(bool useMC);
+framework::DataProcessorSpec getTRDTrackletReaderSpec(bool useMC, bool useCalibratedTracklets);
 
 } // namespace trd
 } // namespace o2

Detectors/TRD/workflow/src/TRDGlobalTrackingSpec.cxx

@@ -20,6 +20,7 @@
 #include "DetectorsBase/Propagator.h"
 #include "ReconstructionDataFormats/TrackTPCITS.h"
 #include "DataFormatsTRD/Tracklet64.h"
+#include "DataFormatsTRD/CalibratedTracklet.h"
 #include "DataFormatsTRD/TriggerRecord.h"
 
 // GPU header
@@ -43,6 +44,7 @@ namespace trd
 
 void TRDGlobalTracking::init(InitContext& ic)
 {
+
   //-------- init geometry and field --------//
   o2::base::GeometryManager::loadGeometry();
   o2::base::Propagator::initFieldFromGRP(o2::base::NameConf::getGRPFileName());
@@ -65,6 +67,7 @@ void TRDGlobalTracking::init(InitContext& ic)
   mTracker->SetNCandidates(mRec->GetProcessingSettings().trdNCandidates); // must be set before initialization
   mTracker->SetProcessPerTimeFrame();
   mTracker->SetNMaxCollisions(mRec->GetProcessingSettings().trdNMaxCollisions);
+  mTracker->SetTrkltTransformNeeded(!mUseTrackletTransform);
 
   mRec->RegisterGPUProcessor(mTracker, false);
   mChainTracking->SetTRDGeometry(std::move(mFlatGeo));
@@ -90,6 +93,21 @@ void TRDGlobalTracking::run(ProcessingContext& pc)
   int nTracklets = trackletsTRD.size();
   LOGF(INFO, "There are %i tracklets in total from %i trigger records", nTracklets, nCollisions);
 
+  const gsl::span<const CalibratedTracklet>* cTrkltsPtr = nullptr;
+  using cTrkltType = std::decay_t<decltype(pc.inputs().get<gsl::span<CalibratedTracklet>>(""))>;
+  std::optional<cTrkltType> cTrklts;
+  int nTrackletsCal = 0;
+
+  if (mUseTrackletTransform) {
+    cTrklts.emplace(pc.inputs().get<gsl::span<CalibratedTracklet>>("trdctracklets")); // MC labels associated to the input digits
+    cTrkltsPtr = &cTrklts.value();
+    nTrackletsCal = cTrkltsPtr->size();
+    LOGF(INFO, "Got %i calibrated tracklets as input", nTrackletsCal);
+    if (nTracklets != nTrackletsCal) {
+      LOGF(ERROR, "Number of calibrated tracklets (%i) differs from the number of uncalibrated tracklets (%i)", nTrackletsCal, nTracklets);
+    }
+  }
+
   std::vector<float> trdTriggerTimes;
   std::vector<int> trdTriggerIndices;
 
@@ -141,10 +159,13 @@ void TRDGlobalTracking::run(ProcessingContext& pc)
   for (int iTrklt = 0; iTrklt < nTracklets; ++iTrklt) {
     auto trklt = trackletsTRD[iTrklt];
     GPUTRDTrackletWord trkltLoad(trklt.getTrackletWord());
-    trkltLoad.SetId(iTrklt);
     if (mTracker->LoadTracklet(trkltLoad) > 0) {
       LOG(WARNING) << "Could not load tracklet " << iTrklt;
     }
+    if (mUseTrackletTransform) {
+      const CalibratedTracklet cTrklt = (cTrkltsPtr->data())[iTrklt];
+      mTracker->SetInternalSpacePoint(iTrklt, cTrklt.getX(), cTrklt.getY(), cTrklt.getZ(), cTrklt.getDy());
+    }
   }
   mTracker->SetTriggerRecordTimes(&(trdTriggerTimes[0]));
   mTracker->SetTriggerRecordIndices(&(trdTriggerIndices[0]));
@@ -166,11 +187,14 @@ void TRDGlobalTracking::endOfStream(EndOfStreamContext& ec)
        mTimer.CpuTime(), mTimer.RealTime(), mTimer.Counter() - 1);
 }
 
-DataProcessorSpec getTRDGlobalTrackingSpec(bool useMC)
+DataProcessorSpec getTRDGlobalTrackingSpec(bool useMC, bool useTrkltTransf)
 {
   std::vector<InputSpec> inputs;
   std::vector<OutputSpec> outputs;
   inputs.emplace_back("tpcitstrack", "GLO", "TPCITS", 0, Lifetime::Timeframe);
+  if (useTrkltTransf) {
+    inputs.emplace_back("trdctracklets", o2::header::gDataOriginTRD, "CTRACKLETS", 0, Lifetime::Timeframe);
+  }
   inputs.emplace_back("trdtracklets", o2::header::gDataOriginTRD, "TRACKLETS", 0, Lifetime::Timeframe);
   inputs.emplace_back("trdtriggerrec", o2::header::gDataOriginTRD, "TRKTRGRD", 0, Lifetime::Timeframe);
 
@@ -184,7 +208,7 @@ DataProcessorSpec getTRDGlobalTrackingSpec(bool useMC)
     "trd-globaltracking",
     inputs,
     outputs,
-    AlgorithmSpec{adaptFromTask<TRDGlobalTracking>(useMC)},
+    AlgorithmSpec{adaptFromTask<TRDGlobalTracking>(useMC, useTrkltTransf)},
     Options{}};
 }
 

Detectors/TRD/workflow/src/TRDTrackingWorkflow.cxx

@@ -23,16 +23,16 @@ namespace o2
 namespace trd
 {
 
-framework::WorkflowSpec getTRDTrackingWorkflow(bool disableRootInp, bool disableRootOut)
+framework::WorkflowSpec getTRDTrackingWorkflow(bool disableRootInp, bool disableRootOut, bool useTrkltTransf)
 {
   framework::WorkflowSpec specs;
   bool useMC = false;
   if (!disableRootInp) {
     specs.emplace_back(o2::globaltracking::getTrackTPCITSReaderSpec(useMC));
-    specs.emplace_back(o2::trd::getTRDTrackletReaderSpec(useMC));
+    specs.emplace_back(o2::trd::getTRDTrackletReaderSpec(useMC, useTrkltTransf));
   }
 
-  specs.emplace_back(o2::trd::getTRDGlobalTrackingSpec(useMC));
+  specs.emplace_back(o2::trd::getTRDGlobalTrackingSpec(useMC, useTrkltTransf));
 
   if (!disableRootOut) {
     specs.emplace_back(o2::trd::getTRDTrackWriterSpec(useMC));

Detectors/TRD/workflow/src/TRDTrackletReaderSpec.cxx

@@ -28,49 +28,72 @@ void TRDTrackletReader::init(InitContext& ic)
 {
   // get the option from the init context
   LOG(INFO) << "Init TRD tracklet reader!";
-  mInFileName = ic.options().get<std::string>("trd-tracklet-infile");
-  mInTreeName = ic.options().get<std::string>("treename");
-  connectTree(mInFileName);
+  mInFileNameTrklt = ic.options().get<std::string>("trd-tracklet-infile");
+  mInTreeNameTrklt = ic.options().get<std::string>("treename");
+  connectTree();
+  if (mUseTrackletTransform) {
+    connectTreeCTracklet();
+  }
+}
+
+void TRDTrackletReader::connectTreeCTracklet()
+{
+  mTreeCTrklt.reset(nullptr); // in case it was already loaded
+  mFileCTrklt.reset(TFile::Open("trdcalibratedtracklets.root"));
+  assert(mFileCTrklt && !mFileCTrklt->IsZombie());
+  mTreeCTrklt.reset((TTree*)mFileCTrklt->Get("ctracklets"));
+  assert(mTreeCTrklt);
+  mTreeCTrklt->SetBranchAddress("CTracklets", &mTrackletsCalPtr);
+  LOG(INFO) << "Loaded tree from trdcalibratedtracklets.root with " << mTreeCTrklt->GetEntries() << " entries";
 }
 
-void TRDTrackletReader::connectTree(const std::string& filename)
+void TRDTrackletReader::connectTree()
 {
-  mTree.reset(nullptr); // in case it was already loaded
-  mFile.reset(TFile::Open(filename.c_str()));
-  assert(mFile && !mFile->IsZombie());
-  mTree.reset((TTree*)mFile->Get(mInTreeName.c_str()));
-  assert(mTree);
-  mTree->SetBranchAddress("Tracklet", &mTrackletsPtr);
-  mTree->SetBranchAddress("TrackTrg", &mTriggerRecordsPtr);
+  mTreeTrklt.reset(nullptr); // in case it was already loaded
+  mFileTrklt.reset(TFile::Open(mInFileNameTrklt.c_str()));
+  assert(mFileTrklt && !mFileTrklt->IsZombie());
+  mTreeTrklt.reset((TTree*)mFileTrklt->Get(mInTreeNameTrklt.c_str()));
+  assert(mTreeTrklt);
+  mTreeTrklt->SetBranchAddress("Tracklet", &mTrackletsPtr);
+  mTreeTrklt->SetBranchAddress("TrackTrg", &mTriggerRecordsPtr);
   if (mUseMC) {
     LOG(FATAL) << "MC information not yet included for TRD tracklets";
   }
-  LOG(INFO) << "Loaded tree from " << filename << " with " << mTree->GetEntries() << " entries";
+  LOG(INFO) << "Loaded tree from " << mInFileNameTrklt << " with " << mTreeTrklt->GetEntries() << " entries";
 }
 
 void TRDTrackletReader::run(ProcessingContext& pc)
 {
-  auto currEntry = mTree->GetReadEntry() + 1;
-  assert(currEntry < mTree->GetEntries()); // this should not happen
-  mTree->GetEntry(currEntry);
+  auto currEntry = mTreeTrklt->GetReadEntry() + 1;
+  assert(currEntry < mTreeTrklt->GetEntries()); // this should not happen
+  mTreeTrklt->GetEntry(currEntry);
   LOG(INFO) << "Pushing " << mTriggerRecords.size() << " TRD trigger records at entry " << currEntry;
-  LOG(INFO) << "Pushing " << mTracklets.size() << " TRD tracklets for these trigger records";
-
+  LOG(INFO) << "Pushing " << mTracklets.size() << " uncalibrated TRD tracklets for these trigger records";
   pc.outputs().snapshot(Output{o2::header::gDataOriginTRD, "TRACKLETS", 0, Lifetime::Timeframe}, mTracklets);
+  if (mUseTrackletTransform) {
+    assert(mTreeTrklt->GetEntries() == mTreeCTrklt->GetEntries());
+    mTreeCTrklt->GetEntry(currEntry);
+    LOG(INFO) << "Pushing " << mTrackletsCal.size() << " calibrated TRD tracklets for these trigger records";
+    pc.outputs().snapshot(Output{o2::header::gDataOriginTRD, "CTRACKLETS", 0, Lifetime::Timeframe}, mTrackletsCal);
+  }
+
   pc.outputs().snapshot(Output{o2::header::gDataOriginTRD, "TRKTRGRD", 0, Lifetime::Timeframe}, mTriggerRecords);
   if (mUseMC) {
     LOG(FATAL) << "MC information not yet included for TRD tracklets";
   }
 
-  if (mTree->GetReadEntry() + 1 >= mTree->GetEntries()) {
+  if (mTreeTrklt->GetReadEntry() + 1 >= mTreeTrklt->GetEntries()) {
     pc.services().get<ControlService>().endOfStream();
     pc.services().get<ControlService>().readyToQuit(QuitRequest::Me);
   }
 }
 
-DataProcessorSpec getTRDTrackletReaderSpec(bool useMC)
+DataProcessorSpec getTRDTrackletReaderSpec(bool useMC, bool useCalibratedTracklets)
 {
   std::vector<OutputSpec> outputs;
+  if (useCalibratedTracklets) {
+    outputs.emplace_back(o2::header::gDataOriginTRD, "CTRACKLETS", 0, Lifetime::Timeframe);
+  }
   outputs.emplace_back(o2::header::gDataOriginTRD, "TRACKLETS", 0, Lifetime::Timeframe);
   outputs.emplace_back(o2::header::gDataOriginTRD, "TRKTRGRD", 0, Lifetime::Timeframe);
   if (useMC) {
@@ -81,7 +104,7 @@ DataProcessorSpec getTRDTrackletReaderSpec(bool useMC)
     "TRDTrackletReader",
     Inputs{},
     outputs,
-    AlgorithmSpec{adaptFromTask<TRDTrackletReader>(useMC)},
+    AlgorithmSpec{adaptFromTask<TRDTrackletReader>(useMC, useCalibratedTracklets)},
     Options{
       {"trd-tracklet-infile", VariantType::String, "trdtracklets.root", {"Name of the input file"}},
       {"treename", VariantType::String, "o2sim", {"Name of top-level TTree"}},

Detectors/TRD/workflow/src/TRDTrackletTransformerWorkflow.cxx

@@ -35,7 +35,7 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
   WorkflowSpec spec;
 
   if (rootIn) {
-    spec.emplace_back(o2::trd::getTRDTrackletReaderSpec(0));
+    spec.emplace_back(o2::trd::getTRDTrackletReaderSpec(false, false));
   }
 
   spec.emplace_back(o2::trd::getTRDTrackletTransformerSpec());

Detectors/TRD/workflow/src/trd-tracking-workflow.cxx

@@ -20,6 +20,8 @@ using namespace o2::framework;
 void customize(std::vector<o2::framework::ConfigParamSpec>& workflowOptions)
 {
   // option allowing to set parameters
+  workflowOptions.push_back(ConfigParamSpec{"disable-mc", o2::framework::VariantType::Bool, false, {"Disable MC labels"}});
+  workflowOptions.push_back(ConfigParamSpec{"use-tracklet-transformer", VariantType::Bool, false, {"Use calibrated tracklets instead raw Tracklet64"}});
   workflowOptions.push_back(ConfigParamSpec{
     "disable-root-input", o2::framework::VariantType::Bool, false, {"disable root-files input readers"}});
   workflowOptions.push_back(ConfigParamSpec{
@@ -40,5 +42,6 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
   o2::conf::ConfigurableParam::writeINI("o2trdtracking-workflow_configuration.ini");
   auto disableRootInp = configcontext.options().get<bool>("disable-root-input");
   auto disableRootOut = configcontext.options().get<bool>("disable-root-output");
-  return std::move(o2::trd::getTRDTrackingWorkflow(disableRootInp, disableRootOut));
+  auto useTrackletTransformer = configcontext.options().get<bool>("use-tracklet-transformer");
+  return std::move(o2::trd::getTRDTrackingWorkflow(disableRootInp, disableRootOut, useTrackletTransformer));
 }