GPU QA: Add option to dump clusters and tracks to ROOT file

davidrohr · davidrohr · commit 797ab95caf03 · 2023-01-20T11:01:42.000+01:00
diff --git a/GPU/GPUTracking/Benchmark/standalone.cxx b/GPU/GPUTracking/Benchmark/standalone.cxx
@@ -221,6 +221,9 @@ int ReadConfiguration(int argc, char** argv)
   if (configStandalone.QA.inputHistogramsOnly) {
     configStandalone.rundEdx = false;
   }
+  if (configStandalone.QA.dumpToROOT) {
+    configStandalone.proc.outputSharedClusterMap = true;
+  }
   if (configStandalone.eventDisplay) {
     configStandalone.noprompt = 1;
   }
diff --git a/GPU/GPUTracking/Definitions/GPUSettingsList.h b/GPU/GPUTracking/Definitions/GPUSettingsList.h
@@ -317,6 +317,7 @@ AddOption(filterCharge, int, 0, "", 0, "Filter for positive (+1) or negative (-1
 AddOption(filterPID, int, -1, "", 0, "Filter for Particle Type (0 Electron, 1 Muon, 2 Pion, 3 Kaon, 4 Proton)")
 AddOption(nativeFitResolutions, bool, false, "", 0, "Create resolution histograms in the native fit units (sin(phi), tan(lambda), Q/Pt)")
 AddOption(enableLocalOutput, bool, true, "", 0, "Enable normal output to local PDF files / console")
+AddOption(dumpToROOT, int, 0, "", 0, "Dump all clusters and tracks to a ROOT file, 1 = combined TNTUple dump, 2 = also individual cluster / track branch dump")
 AddOption(writeMCLabels, bool, false, "", 0, "Store mc labels to file for later matching")
 AddOption(writeRootFiles, bool, false, "", 0, "Create ROOT canvas files")
 AddOptionVec(matchMCLabels, std::string, "", 0, "Read labels from files and match them, only process tracks where labels differ")
diff --git a/GPU/GPUTracking/qa/GPUQA.cxx b/GPU/GPUTracking/qa/GPUQA.cxx
@@ -50,6 +50,7 @@
 #include "GPUTPCConvertImpl.h"
 #include "TPCFastTransform.h"
 #include "CorrectionMapsHelper.h"
+#include "GPUROOTDump.h"
 #ifdef GPUCA_HAVE_O2HEADERS
 #include "SimulationDataFormat/ConstMCTruthContainer.h"
 #include "SimulationDataFormat/MCCompLabel.h"
@@ -1708,6 +1709,52 @@ void GPUQA::RunQA(bool matchOnly, const std::vector<o2::tpc::TrackTPC>* tracksEx
     GPUInfo("QA Time: Cluster Counts:\t%6.0f us", timer.GetCurrentElapsedTime(true) * 1e6);
   }
 
+  if (mConfig.dumpToROOT) {
+    if (!clNative || !mTracking || !mTracking->mIOPtrs.mergedTrackHitAttachment || !mTracking->mIOPtrs.mergedTracks) {
+      throw std::runtime_error("Cannot dump non o2::tpc::clusterNative clusters, need also hit attachmend and GPU tracks");
+    }
+    static auto cldump = GPUROOTDump<o2::tpc::ClusterNative, GPUTPCGMMergedTrack, GPUTPCGMMergedTrackHit, unsigned int, unsigned int, float, float, float, unsigned int, unsigned int, unsigned int>::getNew("cluster", "track", "trackHit", "attach", "extState", "x", "y", "z", "sector", "row", "nEv", "clusterTree");
+    unsigned int clid = 0;
+    for (unsigned int i = 0; i < GPUChainTracking::NSLICES; i++) {
+      for (unsigned int j = 0; j < GPUCA_ROW_COUNT; j++) {
+        for (unsigned int k = 0; k < mClNative->nClusters[i][j]; k++) {
+          const auto& cl = mClNative->clusters[i][j][k];
+          unsigned int attach = mTracking->mIOPtrs.mergedTrackHitAttachment[clid];
+          GPUTPCGMMergedTrack trk;
+          GPUTPCGMMergedTrackHit trkHit;
+          memset((void*)&trk, 0, sizeof(trk));
+          memset((void*)&trkHit, 0, sizeof(trkHit));
+          float x = 0, y = 0, z = 0;
+          if (attach & gputpcgmmergertypes::attachFlagMask) {
+            unsigned int track = attach & gputpcgmmergertypes::attachTrackMask;
+            trk = mTracking->mIOPtrs.mergedTracks[track];
+            for (unsigned int l = 0; l < trk.NClusters(); l++) {
+              const auto& tmp = mTracking->mIOPtrs.mergedTrackHits[trk.FirstClusterRef() + l];
+              if (tmp.num == clid) {
+                trkHit = tmp;
+                break;
+              }
+            }
+            mTracking->GetTPCTransformHelper()->Transform(i, j, cl.getPad(), cl.getTime(), x, y, z, trk.GetParam().GetTZOffset());
+            mTracking->GetParam().Slice2Global(i, x, y, z, &x, &y, &z);
+          }
+
+          unsigned int extState = mTracking->mIOPtrs.mergedTrackHitStates ? mTracking->mIOPtrs.mergedTrackHitStates[clid] : 0;
+
+          cldump.Fill(cl, trk, trkHit, attach, extState, x, y, z, i, j, mNEvents - 1);
+          clid++;
+        }
+      }
+    }
+  }
+
+  static auto trkdump = GPUROOTDump<unsigned int, GPUTPCGMMergedTrack>::getNew("nEv", "tracks");
+  for (unsigned int i = 0; i < mTracking->mIOPtrs.nMergedTracks; i++) {
+    if (mTracking->mIOPtrs.mergedTracks[i].OK()) {
+      trkdump.Fill(mNEvents - 1, mTracking->mIOPtrs.mergedTracks[i]);
+    }
+  }
+
   mTrackingScratchBuffer.clear();
 }
 

Original file line number	Diff line number	Diff line change
`@@ -221,6 +221,9 @@ int ReadConfiguration(int argc, char** argv)`
`221`	`221`	`if (configStandalone.QA.inputHistogramsOnly) {`
`222`	`222`	`configStandalone.rundEdx = false;`
`223`	`223`	`}`
	`224`	`+ if (configStandalone.QA.dumpToROOT) {`
	`225`	`+ configStandalone.proc.outputSharedClusterMap = true;`
	`226`	`+ }`
`224`	`227`	`if (configStandalone.eventDisplay) {`
`225`	`228`	`configStandalone.noprompt = 1;`
`226`	`229`	`}`