From 3f548f72ab6ae675a368f6f7020e44cd4fea0965 Mon Sep 17 00:00:00 2001
From: Felix Schlepper <felix.schlepper@cern.ch>
Date: Sun, 13 Oct 2024 14:58:38 +0200
Subject: [PATCH] ITS: TrackExtensionStudy more plots

---
 .../studies/macros/PostTrackExtension.C       | 220 ++++++++++++++++--
 .../studies/src/TrackExtension.cxx            |  41 +++-
 2 files changed, 243 insertions(+), 18 deletions(-)

diff --git a/Detectors/ITSMFT/ITS/postprocessing/studies/macros/PostTrackExtension.C b/Detectors/ITSMFT/ITS/postprocessing/studies/macros/PostTrackExtension.C
index 773b4a2066a0d..78331c8dd5dc3 100644
--- a/Detectors/ITSMFT/ITS/postprocessing/studies/macros/PostTrackExtension.C
+++ b/Detectors/ITSMFT/ITS/postprocessing/studies/macros/PostTrackExtension.C
@@ -16,9 +16,12 @@
 #include "TColor.h"
 #include "TCanvas.h"
 #include "TH1F.h"
+#include "TF1.h"
 #include "TEfficiency.h"
+#include "TMarker.h"
 #include "TLegend.h"
 #include "TTree.h"
+#include "TLatex.h"
 
 #include <memory>
 #include <array>
@@ -282,21 +285,208 @@ void PostTrackExtension(const char* fileName = "TrackExtensionStudy.root")
     auto t = fIn->Get<TTree>("tree");
     auto c = new TCanvas("cKG", "", 800, 600);
     c->Divide(3, 2);
-    auto p = c->cd(1);
-    p->SetGrid();
-    auto h = p->DrawFrame(-.5, 0., .5, 30.);
-    h->GetXaxis()->SetTitle("#it{p}_{T,TRK}-#it{p}_{T,MC}");
-    h->GetYaxis()->SetTitle("n. counts");
-    t->Draw("trk.getPt()-mcTrk.getPt()>>hPtNo(100,-.5,.5)", "isGood&&!isExtended", "HIST;SAME");
-    auto htemp = (TH1F*)p->GetPrimitive("hPtNo");
-    htemp->Scale(1.0 / htemp->Integral("width"));
-    htemp->SetLineColor(kRed);
-    t->Draw("trk.getPt()-mcTrk.getPt()>>hPtYes(100,-.5,.5)", "isGood&&isExtended", "HIST;SAME");
-    htemp = (TH1F*)p->GetPrimitive("hPtYes");
-    htemp->Scale(1.0 / htemp->Integral("width"));
-    htemp->SetLineColor(kBlue);
-    p->Modified();
-    p->Update();
+    {
+      auto p = c->cd(1);
+      p->SetGrid();
+      auto h = p->DrawFrame(-.6, 0., .6, 9.);
+      h->GetXaxis()->SetTitle("#frac{Q^{2}}{p_{T,TRK}}-#frac{Q^{2}}{p_{T,MC}}");
+      h->GetYaxis()->SetTitle("n. counts");
+      t->Draw("trk.getQ2Pt()-mcTrk.getQ2Pt()>>hPtNo(100,-.6,.6)", "isGood&&!isExtended", "HIST;SAME");
+      auto hNo = (TH1F*)p->GetPrimitive("hPtNo");
+      hNo->Scale(1.0 / hNo->Integral("width"));
+      hNo->SetLineColor(kRed);
+      auto fitNo = new TF1("fitNo", "gaus", -0.04, 0.04);
+      hNo->Fit(fitNo, "QR");
+      fitNo->SetLineColor(kRed);
+      fitNo->Draw("SAME");
+      auto textNo = new TLatex(-0.55, 8.2, Form("#mu = %.3f, #sigma = %.3f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textNo->SetTextColor(kRed);
+      textNo->SetNDC(false);
+      textNo->SetTextSize(0.05);
+      textNo->Draw();
+
+      t->Draw("trk.getQ2Pt()-mcTrk.getQ2Pt()>>hPtYes(100,-.6,.6)", "isGood&&isExtended", "HIST;SAME");
+      auto hYes = (TH1F*)p->GetPrimitive("hPtYes");
+      hYes->Scale(1.0 / hYes->Integral("width"));
+      hYes->SetLineColor(kBlue);
+      auto fitYes = new TF1("fitYes", "gaus", -0.04, 0.04);
+      hYes->Fit(fitYes, "QR");
+      fitYes->SetLineColor(kBlue);
+      fitYes->Draw("SAME");
+      auto textYes = new TLatex(-0.55, 7, Form("#mu = %.4f, #sigma = %.4f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textYes->SetTextColor(kBlue);
+      textYes->SetNDC(false);
+      textYes->SetTextSize(0.05);
+      textYes->Draw();
+
+      p->Modified();
+      p->Update();
+    }
+    {
+      auto p = c->cd(2);
+      p->SetGrid();
+      auto h = p->DrawFrame(-3, 0., 3, 2.);
+      h->GetXaxis()->SetTitle("Y_{TRK}-Y_{MC}");
+      h->GetYaxis()->SetTitle("n. counts");
+      t->Draw("trk.getY()-mcTrk.getY()>>hYNo(100,-3,3)", "isGood&&!isExtended", "HIST;SAME");
+      auto hNo = (TH1F*)p->GetPrimitive("hYNo");
+      hNo->Scale(1.0 / hNo->Integral("width"));
+      hNo->SetLineColor(kRed);
+      auto fitNo = new TF1("fitNo", "gaus", -0.5, 0.5);
+      hNo->Fit(fitNo, "QR");
+      fitNo->SetLineColor(kRed);
+      fitNo->Draw("SAME");
+      auto textNo = new TLatex(-2, 1.7, Form("#mu = %.3f, #sigma = %.3f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textNo->SetTextColor(kRed);
+      textNo->SetNDC(false);
+      textNo->SetTextSize(0.05);
+      textNo->Draw();
+
+      t->Draw("trk.getY()-mcTrk.getY()>>hYYes(100,-3,3)", "isGood&&isExtended", "HIST;SAME");
+      auto hYes = (TH1F*)p->GetPrimitive("hYYes");
+      hYes->Scale(1.0 / hYes->Integral("width"));
+      hYes->SetLineColor(kBlue);
+      auto fitYes = new TF1("fitYes", "gaus", -0.5, 0.5);
+      hYes->Fit(fitYes, "QR");
+      fitYes->SetLineColor(kBlue);
+      fitYes->Draw("SAME");
+      auto textYes = new TLatex(-2, 1.5, Form("#mu = %.4f, #sigma = %.4f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textYes->SetTextColor(kBlue);
+      textYes->SetNDC(false);
+      textYes->SetTextSize(0.05);
+      textYes->Draw();
+
+      p->Modified();
+      p->Update();
+    }
+    {
+      auto p = c->cd(3);
+      p->SetGrid();
+      auto h = p->DrawFrame(-2, 0., 2, 4.2);
+      h->GetXaxis()->SetTitle("Z_{TRK}-Z_{MC}");
+      h->GetYaxis()->SetTitle("n. counts");
+      t->Draw("trk.getZ()-mcTrk.getZ()>>hZNo(100,-2,2)", "isGood&&!isExtended", "HIST;SAME");
+      auto hNo = (TH1F*)p->GetPrimitive("hZNo");
+      hNo->Scale(1.0 / hNo->Integral("width"));
+      hNo->SetLineColor(kRed);
+      auto fitNo = new TF1("fitNo", "gaus", -0.2, 0.2);
+      hNo->Fit(fitNo, "QR");
+      fitNo->SetLineColor(kRed);
+      fitNo->Draw("SAME");
+      auto textNo = new TLatex(-1.7, 3.8, Form("#mu = %.3f, #sigma = %.3f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textNo->SetTextColor(kRed);
+      textNo->SetNDC(false);
+      textNo->SetTextSize(0.05);
+      textNo->Draw();
+
+      t->Draw("trk.getZ()-mcTrk.getZ()>>hZYes(100,-2,2)", "isGood&&isExtended", "HIST;SAME");
+      auto hYes = (TH1F*)p->GetPrimitive("hZYes");
+      hYes->Scale(1.0 / hYes->Integral("width"));
+      hYes->SetLineColor(kBlue);
+      auto fitYes = new TF1("fitYes", "gaus", -0.2, 0.2);
+      hYes->Fit(fitYes, "QR");
+      fitYes->SetLineColor(kBlue);
+      fitYes->Draw("SAME");
+      auto textYes = new TLatex(-1.7, 3.5, Form("#mu = %.4f, #sigma = %.4f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textYes->SetTextColor(kBlue);
+      textYes->SetNDC(false);
+      textYes->SetTextSize(0.05);
+      textYes->Draw();
+
+      p->Modified();
+      p->Update();
+    }
+    {
+      auto p = c->cd(4);
+      p->SetGrid();
+      auto h = p->DrawFrame(-0.02, 0., 0.02, 370.);
+      h->GetXaxis()->SetTitle("TGL_{TRK}-TGL_{MC}");
+      h->GetYaxis()->SetTitle("n. counts");
+      t->Draw("trk.getTgl()-mcTrk.getTgl()>>hTglNo(100,-0.02,0.02)", "isGood&&!isExtended", "HIST;SAME");
+      auto hNo = (TH1F*)p->GetPrimitive("hTglNo");
+      hNo->Scale(1.0 / hNo->Integral("width"));
+      hNo->SetLineColor(kRed);
+      auto fitNo = new TF1("fitNo", "gaus", -0.003, 0.003);
+      hNo->Fit(fitNo, "QR");
+      fitNo->SetLineColor(kRed);
+      fitNo->Draw("SAME");
+      auto textNo = new TLatex(-0.018, 330, Form("#mu = %.3f, #sigma = %.3f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textNo->SetTextColor(kRed);
+      textNo->SetNDC(false);
+      textNo->SetTextSize(0.05);
+      textNo->Draw();
+
+      t->Draw("trk.getTgl()-mcTrk.getTgl()>>hTglYes(100,-0.02,0.02)", "isGood&&isExtended", "HIST;SAME");
+      auto hYes = (TH1F*)p->GetPrimitive("hTglYes");
+      hYes->Scale(1.0 / hYes->Integral("width"));
+      hYes->SetLineColor(kBlue);
+      auto fitYes = new TF1("fitYes", "gaus", -0.003, 0.003);
+      hYes->Fit(fitYes, "QR");
+      fitYes->SetLineColor(kBlue);
+      fitYes->Draw("SAME");
+      auto textYes = new TLatex(-0.018, 310, Form("#mu = %.6f, #sigma = %.6f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textYes->SetTextColor(kBlue);
+      textYes->SetNDC(false);
+      textYes->SetTextSize(0.05);
+      textYes->Draw();
+
+      p->Modified();
+      p->Update();
+    }
+    {
+      auto p = c->cd(5);
+      p->SetGrid();
+      auto h = p->DrawFrame(-0.08, 0., 0.08, 80.);
+      h->GetXaxis()->SetTitle("SNP_{TRK}-SNP_{MC}");
+      h->GetYaxis()->SetTitle("n. counts");
+      t->Draw("trk.getSnp()-mcTrk.getSnp()>>hSnpNo(100,-0.08,0.08)", "isGood&&!isExtended", "HIST;SAME");
+      auto hNo = (TH1F*)p->GetPrimitive("hSnpNo");
+      hNo->Scale(1.0 / hNo->Integral("width"));
+      hNo->SetLineColor(kRed);
+      auto fitNo = new TF1("fitNo", "gaus", -0.03, 0.03);
+      hNo->Fit(fitNo, "QR");
+      fitNo->SetLineColor(kRed);
+      fitNo->Draw("SAME");
+      auto textNo = new TLatex(-0.07, 72, Form("#mu = %.3f, #sigma = %.3f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textNo->SetTextColor(kRed);
+      textNo->SetNDC(false);
+      textNo->SetTextSize(0.05);
+      textNo->Draw();
+
+      t->Draw("trk.getSnp()-mcTrk.getSnp()>>hSnpYes(100,-0.08,0.08)", "isGood&&isExtended", "HIST;SAME");
+      auto hYes = (TH1F*)p->GetPrimitive("hSnpYes");
+      hYes->Scale(1.0 / hYes->Integral("width"));
+      hYes->SetLineColor(kBlue);
+      auto fitYes = new TF1("fitYes", "gaus", -0.03, 0.03);
+      hYes->Fit(fitYes, "QR");
+      fitYes->SetLineColor(kBlue);
+      fitYes->Draw("SAME");
+      auto textYes = new TLatex(-0.07, 66, Form("#mu = %.6f, #sigma = %.6f", fitNo->GetParameter(1), fitNo->GetParameter(2)));
+      textYes->SetTextColor(kBlue);
+      textYes->SetNDC(false);
+      textYes->SetTextSize(0.05);
+      textYes->Draw();
+
+      p->Modified();
+      p->Update();
+    }
+    {
+      auto p = c->cd(6);
+      auto legend = new TLegend(0.2, 0.2, 0.8, 0.8);
+      legend->SetTextSize(0.06);
+      legend->SetLineWidth(3);
+      legend->SetHeader("GOOD tracks", "C");
+      auto mBlue = new TMarker();
+      mBlue->SetMarkerColor(kBlue);
+      mBlue->SetMarkerSize(4);
+      legend->AddEntry(mBlue, "extended", "p");
+      auto mRed = new TMarker();
+      mRed->SetMarkerColor(kRed);
+      mRed->SetMarkerSize(4);
+      legend->AddEntry(mRed, "normal", "p");
+      legend->SetLineColor(kRed);
+      legend->Draw();
+    }
     c->SaveAs("trkExt_kinematics.pdf");
   }
 }
diff --git a/Detectors/ITSMFT/ITS/postprocessing/studies/src/TrackExtension.cxx b/Detectors/ITSMFT/ITS/postprocessing/studies/src/TrackExtension.cxx
index 8c01cd803c903..4cc5f9d9df97f 100644
--- a/Detectors/ITSMFT/ITS/postprocessing/studies/src/TrackExtension.cxx
+++ b/Detectors/ITSMFT/ITS/postprocessing/studies/src/TrackExtension.cxx
@@ -20,8 +20,11 @@
 #include "ITSBase/GeometryTGeo.h"
 #include "ITSStudies/Helpers.h"
 #include "ITSStudies/TrackExtension.h"
+#include "SimulationDataFormat/MCEventHeader.h"
 #include "SimulationDataFormat/MCTrack.h"
 #include "Steer/MCKinematicsReader.h"
+#include "ReconstructionDataFormats/Vertex.h"
+#include "ReconstructionDataFormats/DCA.h"
 
 #include <bitset>
 
@@ -40,7 +43,7 @@ using o2::steer::MCKinematicsReader;
 class TrackExtensionStudy : public Task
 {
   struct ParticleInfo {
-    int event;
+    dataformats::MCEventHeader event;
     int pdg;
     float pt;
     float eta;
@@ -117,6 +120,9 @@ class TrackExtensionStudy : public Task
   std::unique_ptr<TH1D> mEExtensionMixNum, mEExtensionMixPurityNum, mEExtensionMixFakeNum;
   std::array<std::unique_ptr<TH1D>, mBitPatternsBefore.size()> mEExtensionPatternGoodNum, mEExtensionPatternFakeNum;
   std::array<std::array<std::unique_ptr<TH1D>, mBitPatternsAfter.size()>, mBitPatternsBefore.size()> mEExtensionPatternIndGoodNum, mEExtensionPatternIndFakeNum;
+  // DCA
+  std::unique_ptr<TH2D> mDCAxyVsPtPionsNormal, mDCAxyVsPtPionsExtended;
+  std::unique_ptr<TH2D> mDCAzVsPtPionsNormal, mDCAzVsPtPionsExtended;
 
   template <class T, typename... C, typename... F>
   std::unique_ptr<T> createHistogram(C... n, F... b)
@@ -227,6 +233,12 @@ void TrackExtensionStudy::init(InitContext& ic)
     }
   }
 
+  /// DCA
+  mDCAxyVsPtPionsNormal = createHistogram<TH2D>("hDCAxyVsPtResNormal", "DCA_{#it{xy}} NORMAL Pions;#it{p}_{T} (GeV/#it{c});#sigma(DCA_{#it{xy}}) (#mum)", effHistBins, xbins.data(), 1000, -500, 500);
+  mDCAxyVsPtPionsExtended = createHistogram<TH2D>("hDCAxyVsPtResExtended", "DCA_{#it{xy}} EXTENDED Pions;#it{p}_{T} (GeV/#it{c});#sigma(DCA_{#it{xy}}) (#mum)", effHistBins, xbins.data(), 1000, -500, 500);
+  mDCAzVsPtPionsNormal = createHistogram<TH2D>("hDCAzVsPtResNormal", "DCA_{#it{z}} NORMAL Pions;#it{p}_{T} (GeV/#it{c});#sigma(DCA_{#it{z}}) (#mum)", effHistBins, xbins.data(), 1000, -500, 500);
+  mDCAzVsPtPionsExtended = createHistogram<TH2D>("hDCAzVsPtResExtended", "DCA_{#it{z}} EXTENDED Pions;#it{p}_{T} (GeV/#it{c});#sigma(DCA_{#it{z}}) (#mum)", effHistBins, xbins.data(), 1000, -500, 500);
+
   mStream = std::make_unique<utils::TreeStreamRedirector>(mOutFileName.c_str(), "RECREATE");
 }
 
@@ -257,10 +269,11 @@ void TrackExtensionStudy::process()
   for (int iSource{0}; iSource < mKineReader->getNSources(); ++iSource) {
     mParticleInfo[iSource].resize(mKineReader->getNEvents(iSource)); // events
     for (int iEvent{0}; iEvent < mKineReader->getNEvents(iSource); ++iEvent) {
+      const auto& mcEvent = mKineReader->getMCEventHeader(iSource, iEvent);
       mParticleInfo[iSource][iEvent].resize(mKineReader->getTracks(iSource, iEvent).size()); // tracks
       for (auto iPart{0}; iPart < mKineReader->getTracks(iEvent).size(); ++iPart) {
-        auto& part = mKineReader->getTracks(iSource, iEvent)[iPart];
-        mParticleInfo[iSource][iEvent][iPart].event = iEvent;
+        const auto& part = mKineReader->getTracks(iSource, iEvent)[iPart];
+        mParticleInfo[iSource][iEvent][iPart].event = mcEvent;
         mParticleInfo[iSource][iEvent][iPart].pdg = part.GetPdgCode();
         mParticleInfo[iSource][iEvent][iPart].pt = part.GetPt();
         mParticleInfo[iSource][iEvent][iPart].phi = part.GetPhi();
@@ -336,6 +349,8 @@ void TrackExtensionStudy::process()
   LOGP(info, "\t- Total number of good: {} ({:.2f} %)", good, good * 100. / mTracks.size());
   LOGP(info, "\t- Total number of extensions: {} ({:.2f} %)", extended, extended * 100. / mTracks.size());
 
+  o2::dataformats::VertexBase collision;
+  o2::dataformats::DCA impactParameter;
   LOGP(info, "** Filling histograms ... ");
   for (auto iTrack{0}; iTrack < mTracks.size(); ++iTrack) {
     auto& lab = mTracksMCLabels[iTrack];
@@ -389,6 +404,26 @@ void TrackExtensionStudy::process()
       break;
     }
 
+    // impact parameter
+    while (isGood && std::abs(part.pdg) == 211) {
+      auto trkC = part.track;
+      collision.setXYZ(part.vx, part.vy, part.vz);
+      if (!o2::base::Propagator::Instance()->propagateToDCA(collision, trkC, o2::base::Propagator::Instance()->getNominalBz(), 2.0, o2::base::Propagator::MatCorrType::USEMatCorrTGeo, &impactParameter)) {
+        break;
+      }
+
+      auto dcaXY = impactParameter.getY() * 1e4;
+      auto dcaZ = impactParameter.getZ() * 1e4;
+      if (!extPattern.any()) {
+        mDCAxyVsPtPionsNormal->Fill(part.pt, dcaXY);
+        mDCAzVsPtPionsNormal->Fill(part.pt, dcaZ);
+      } else {
+        mDCAxyVsPtPionsExtended->Fill(part.pt, dcaXY);
+        mDCAzVsPtPionsExtended->Fill(part.pt, dcaZ);
+      }
+      break;
+    }
+
     mEExtensionDen->Fill(trk.getPt());
 
     if (!extPattern.any()) {