Add MC reco analysis (AliceO2Group#814)

giovannimalfattore · web-flow · commit be39bdd4460a · 2022-06-03T12:17:31.000+02:00
* Add MC reco analysis

* Minor formatting fixes
diff --git a/PWGLF/DataModel/LFNucleiTables.h b/PWGLF/DataModel/LFNucleiTables.h
@@ -72,13 +72,21 @@ DECLARE_SOA_COLUMN(DCAz, dcaz, float);
 DECLARE_SOA_COLUMN(TPCInnerParam, tpcInnerParam, float);
 DECLARE_SOA_COLUMN(TPCSignal, tpcSignal, float);
 DECLARE_SOA_COLUMN(Beta, beta, float);
-//TPC and ITS QA
+// TPC and ITS QA
 DECLARE_SOA_COLUMN(NcrTPC, ncrTPC, int);
 DECLARE_SOA_COLUMN(RTPC, rTPC, float);
 DECLARE_SOA_COLUMN(Chi2TPC, chi2TPC, float);
 DECLARE_SOA_COLUMN(Chi2ITS, chi2ITS, float);
 } // namespace full
 
+/*
+namespace fullMC
+{
+DECLARE_SOA_INDEX_COLUMN(LfCandNucleusFullEvent, lfCandNucleusFullEvent);
+DECLARE_SOA_COLUMN(PdgCode, pdgCode, int);
+}
+*/
+
 DECLARE_SOA_TABLE(LfCandNucleusFull, "AOD", "LFNUCL",
                   o2::soa::Index<>,
                   full::LfCandNucleusFullEventId,
@@ -104,6 +112,8 @@ DECLARE_SOA_TABLE(LfCandNucleusFull, "AOD", "LFNUCL",
                   full::RTPC,
                   full::Chi2TPC,
                   full::Chi2ITS);
+DECLARE_SOA_TABLE(LfCandNucleusMC, "AOD", "LFNUCLMC",
+                  mcparticle::PdgCode);
 
 } // namespace o2::aod
 #endif
diff --git a/PWGLF/TableProducer/LFTreeCreatorNuclei.cxx b/PWGLF/TableProducer/LFTreeCreatorNuclei.cxx
@@ -48,82 +48,118 @@ using namespace o2::framework::expressions;
 struct LfTreeCreatorNuclei {
   Produces<o2::aod::LfCandNucleusFull> rowCandidateFull;
   Produces<o2::aod::LfCandNucleusFullEvents> rowCandidateFullEvents;
+  Produces<o2::aod::LfCandNucleusMC> rowCandidateMC;
 
   void init(o2::framework::InitContext&)
   {
   }
 
-  //track
+  // track
   Configurable<float> yMin{"yMin", -0.5, "Maximum rapidity"};
   Configurable<float> yMax{"yMax", 0.5, "Minimum rapidity"};
   Configurable<float> cfgCutDCAxy{"cfgCutDCAxy", 2.0f, "DCAxy range for tracks"};
   Configurable<float> cfgCutDCAz{"cfgCutDCAz", 2.0f, "DCAz range for tracks"};
   Configurable<float> cfgCutEta{"cfgCutEta", 0.8f, "Eta range for tracks"};
   Configurable<float> nsigmacutLow{"nsigmacutLow", -8.0, "Value of the Nsigma cut"};
   Configurable<float> nsigmacutHigh{"nsigmacutHigh", +8.0, "Value of the Nsigma cut"};
-  //events
+  // events
   Configurable<float> cfgCutVertex{"cfgCutVertex", 10.0f, "Accepted z-vertex range"};
   Configurable<bool> useEvsel{"useEvsel", true, "Use sel8 for run3 Event Selection"};
 
   Filter collisionFilter = nabs(aod::collision::posZ) < cfgCutVertex;
   Filter trackFilter = (nabs(aod::track::eta) < cfgCutEta) && (requireGlobalTrackInFilter());
   Filter DCAcutFilter = (nabs(aod::track::dcaXY) < cfgCutDCAxy) && (nabs(aod::track::dcaZ) < cfgCutDCAz);
-  using TrackCandidates = soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra, aod::TracksExtended, aod::TrackSelection,
-                                                  aod::pidTOFbeta, aod::TOFSignal,
-                                                  aod::pidTPCFullPi, aod::pidTOFFullPi,
-                                                  aod::pidTPCFullKa, aod::pidTOFFullKa,
-                                                  aod::pidTPCFullPr, aod::pidTOFFullPr,
-                                                  aod::pidTPCFullDe, aod::pidTOFFullDe,
-                                                  aod::pidTPCFullHe, aod::pidTOFFullHe>>;
-
-  void process(soa::Filtered<soa::Join<aod::Collisions,
-                                       aod::EvSels, aod::Mults>>::iterator const& collision,
-               TrackCandidates const& tracks,
-               aod::BCs const&)
+  using TrackCandidates = soa::Join<aod::Tracks, aod::TracksExtra, aod::TracksExtended, aod::TrackSelection,
+                                    aod::pidTOFbeta, aod::TOFSignal,
+                                    aod::pidTPCFullPi, aod::pidTOFFullPi,
+                                    aod::pidTPCFullKa, aod::pidTOFFullKa,
+                                    aod::pidTPCFullPr, aod::pidTOFFullPr,
+                                    aod::pidTPCFullDe, aod::pidTOFFullDe,
+                                    aod::pidTPCFullHe, aod::pidTOFFullHe>;
+  int nevs = 0;
+  void process(soa::Filtered<soa::Join<aod::Collisions, aod::EvSels, aod::Mults>> const& collisions,
+               soa::Filtered<TrackCandidates> const& tracks, aod::BCs const&)
   {
-    if (useEvsel && !collision.sel8()) {
-      return;
-    }
-    // Filling event properties
-    rowCandidateFullEvents(
-      collision.bcId(),
-      collision.numContrib(),
-      collision.posX(),
-      collision.posY(),
-      collision.posZ(),
-      collision.multFV0M(),
-      collision.sel8(),
-      collision.bc().runNumber());
+    for (const auto& collision : collisions) {
+      LOG(INFO) << "nevs=" << nevs++;
+      if (useEvsel && !collision.sel8()) {
+        return;
+      }
+      // std::cout<<"mc Z-vertex ====>"<<mcColl.posZ()<<std::endl;
+      /*for(auto& mcCollItr : mcCollisions)
+      {}*/
+      // Filling event properties
+      rowCandidateFullEvents(
+        collision.bcId(),
+        collision.numContrib(),
+        collision.posX(),
+        collision.posY(),
+        collision.posZ(),
+        collision.multFV0M(),
+        collision.sel8(),
+        collision.bc().runNumber());
 
-    // Filling candidate properties
-    rowCandidateFull.reserve(tracks.size());
-    for (auto& track : tracks) {
-      rowCandidateFull(
-        rowCandidateFullEvents.lastIndex(),
-        track.dcaXY(),
-        track.dcaZ(),
-        track.tpcNSigmaPi(), track.tpcNSigmaKa(), track.tpcNSigmaPr(),
-        track.tpcNSigmaDe(), track.tpcNSigmaHe(),
-        track.tofNSigmaPi(), track.tofNSigmaKa(), track.tofNSigmaPr(),
-        track.tofNSigmaDe(), track.tofNSigmaHe(),
-        track.hasTOF(),
-        track.tpcInnerParam(),
-        track.tpcSignal(),
-        track.beta(),
-        track.px(),
-        track.py(),
-        track.pz(),
-        track.pt(),
-        track.p(),
-        track.eta(),
-        track.phi(),
-        track.sign(),
-        track.tpcNClsCrossedRows(),
-        track.tpcCrossedRowsOverFindableCls(),
-        track.tpcChi2NCl(),
-        track.itsChi2NCl());
+      // Filling candidate properties
+      const auto& tracksInCollision = tracks.sliceBy(aod::track::collisionId, collision.globalIndex());
+      rowCandidateFull.reserve(tracksInCollision.size());
+      for (auto& track : tracksInCollision) {
+        // auto const& mcParticle = track.mcParticle();
+        // std::cout<<"mc pdg code ====>"<<mcParticle.pdgCode()<<std::endl;
+        // std::cout<<"mc physical primary ====>"<<mcParticle.isPhysicalPrimary()<<std::endl;
+        // std::cout<<"eta-gen ====>"<<mcParticle.eta()<<"  eta-reco"<<track.eta()<<std::endl;
+        rowCandidateFull(
+          rowCandidateFullEvents.lastIndex(),
+          track.dcaXY(),
+          track.dcaZ(),
+          track.tpcNSigmaPi(), track.tpcNSigmaKa(), track.tpcNSigmaPr(),
+          track.tpcNSigmaDe(), track.tpcNSigmaHe(),
+          track.tofNSigmaPi(), track.tofNSigmaKa(), track.tofNSigmaPr(),
+          track.tofNSigmaDe(), track.tofNSigmaHe(),
+          track.hasTOF(),
+          track.tpcInnerParam(),
+          track.tpcSignal(),
+          track.beta(),
+          track.px(),
+          track.py(),
+          track.pz(),
+          track.pt(),
+          track.p(),
+          track.eta(),
+          track.phi(),
+          track.sign(),
+          track.tpcNClsCrossedRows(),
+          track.tpcCrossedRowsOverFindableCls(),
+          track.tpcChi2NCl(),
+          track.itsChi2NCl());
+      }
+    }
+  }
+  int nevsmc = 0;
+  void processMC(soa::Filtered<soa::Join<aod::Collisions, aod::EvSels, aod::Mults, aod::McCollisionLabels>> const& collisions,
+                 soa::Filtered<soa::Join<TrackCandidates, aod::McTrackLabels>> const& tracks,
+                 aod::BCs const&, aod::McCollisions const& mcCollisions, aod::McParticles const& mcParticles)
+  {
+    for (const auto& collision : collisions) {
+      LOG(INFO) << "nevsmc=" << nevsmc++;
+      if (useEvsel && !collision.sel8()) {
+        return;
+      }
+      const auto& tracksInCollision = tracks.sliceBy(aod::track::collisionId, collision.globalIndex());
+      rowCandidateMC.reserve(tracksInCollision.size());
+      LOG(info) << tracks.size() << " " << tracksInCollision.size();
+      for (const auto& track : tracksInCollision)
+      // for(const auto& track: tracks)
+      {
+        if (track.has_mcParticle()) {
+          const auto& particle = track.mcParticle();
+          rowCandidateMC(particle.pdgCode());
+          continue;
+        }
+        rowCandidateMC(0);
+      }
     }
   }
+  PROCESS_SWITCH(LfTreeCreatorNuclei, processMC, "process MC", true);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
diff --git a/PWGLF/Tasks/LFNucleiBATask.cxx b/PWGLF/Tasks/LFNucleiBATask.cxx
@@ -35,7 +35,7 @@ using namespace o2;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
 
-struct LFNucleiDeuteronTask {
+struct LFNucleiBATask {
   HistogramRegistry histos{"Histos", {}, OutputObjHandlingPolicy::AnalysisObject};
   HistogramRegistry spectraGen{"spectraGen", {}, OutputObjHandlingPolicy::AnalysisObject, false, true};
 
@@ -90,17 +90,18 @@ struct LFNucleiDeuteronTask {
 
     // MC histograms
     AxisSpec ptAxis = {2000, 0.f, 20.f, "#it{p}_{T} (GeV/#it{c})"};
-    spectraGen.add("histGenVetxZ", "PosZ generated events", HistType::kTH1F, {{2000, -20.f, 20.f, "Vertex Z (cm)"}});
-    spectraGen.add("histGenPtPion", "generated particles", HistType::kTH1F, {ptAxis});
-    spectraGen.add("histGenPtKaon", "generated particles", HistType::kTH1F, {ptAxis});
-    spectraGen.add("histGenPtProton", "generated particles", HistType::kTH1F, {ptAxis});
-    spectraGen.add("histGenPtD", "generated particles", HistType::kTH1F, {ptAxis});
-    spectraGen.add("histGenPtantiD", "generated particles", HistType::kTH1F, {ptAxis});
-    spectraGen.add("histGenPtHe", "generated particles", HistType::kTH1F, {ptAxis});
-    spectraGen.add("histGenPtantiHe", "generated particles", HistType::kTH1F, {ptAxis});
+    histos.add("spectraGen/histGenVetxZ", "PosZ generated events", HistType::kTH1F, {{2000, -20.f, 20.f, "Vertex Z (cm)"}});
+    histos.add("spectraGen/histGenPtPion", "generated particles", HistType::kTH1F, {ptAxis});
+    histos.add("spectraGen/histGenPtKaon", "generated particles", HistType::kTH1F, {ptAxis});
+    histos.add("spectraGen/histGenPtProton", "generated particles", HistType::kTH1F, {ptAxis});
+    histos.add("spectraGen/histGenPtD", "generated particles", HistType::kTH1F, {ptAxis});
+    histos.add("spectraGen/histGenPtantiD", "generated particles", HistType::kTH1F, {ptAxis});
+    histos.add("spectraGen/histGenPtHe", "generated particles", HistType::kTH1F, {ptAxis});
+    histos.add("spectraGen/histGenPtantiHe", "generated particles", HistType::kTH1F, {ptAxis});
   }
 
-  void processData(o2::aod::LfCandNucleusFullEvents::iterator const& event, o2::aod::LfCandNucleusFull const& tracks)
+  template <bool IsMC, typename CollisionType, typename TracksType>
+  void fillHistograms(const CollisionType& event, const TracksType& tracks)
   {
     constexpr float fMassProton = 0.938272088f;
     constexpr float fMassDeuteron = 1.87561f;
@@ -173,55 +174,74 @@ struct LFNucleiDeuteronTask {
         if (std::abs(track.nsigTPC3He()) < nsigmaTPCcut)
           histos.fill(HIST("tracks/helium/h2TOFmass2HeliumVsPt"), massTOF * massTOF - fMassHelium * fMassHelium, track.pt());
       }
+      if constexpr (IsMC) {
+        track.pdgCode();
+      }
     }
-    // LOG(info)<<"Vertex Z ==="<<coll.posZ();
+  }
+
+  void processData(o2::aod::LfCandNucleusFullEvents::iterator const& event,
+                   o2::aod::LfCandNucleusFull const& tracks)
+  {
+    fillHistograms<false>(event, tracks);
   } // CLOSING PROCESS DATA
+  PROCESS_SWITCH(LFNucleiBATask, processData, "process data", true);
 
-  PROCESS_SWITCH(LFNucleiDeuteronTask, processData, "process data", true);
+  void processMCReco(o2::aod::LfCandNucleusFullEvents::iterator const& event,
+                     soa::Join<o2::aod::LfCandNucleusFull, o2::aod::LfCandNucleusMC> const& tracks)
+  {
+    fillHistograms<true>(event, tracks);
+  } // CLOSING PROCESS MC RECO
+  PROCESS_SWITCH(LFNucleiBATask, processMCReco, "process mc reco", false);
   Int_t nCount = 0;
 
   void processMCGen(aod::McCollision const& mcCollision, aod::McParticles_001& mcParticles)
   {
+    constexpr int PDGPion = 211;
+    constexpr int PDGKaon = 321;
+    constexpr int PDGProton = 2212;
+    constexpr int PDGDeuteron = 1000010020;
+    constexpr int PDGHelium = 1000020030;
     nCount++;
-    spectraGen.fill(HIST("histGenVetxZ"), mcCollision.posZ());
+    histos.fill(HIST("spectraGen/histGenVetxZ"), mcCollision.posZ());
     for (auto& mcParticleGen : mcParticles) {
       if (abs(mcParticleGen.y()) > std::abs(etaCut)) {
         continue;
       }
 
-      if (std::abs(mcParticleGen.pdgCode()) == 211) {
-        spectraGen.fill(HIST("histGenPtPion"), mcParticleGen.pt());
+      if (std::abs(mcParticleGen.pdgCode()) == PDGPion) {
+        histos.fill(HIST("spectraGen/histGenPtPion"), mcParticleGen.pt());
       }
 
-      if (std::abs(mcParticleGen.pdgCode()) == 321) {
-        spectraGen.fill(HIST("histGenPtKaon"), mcParticleGen.pt());
+      if (std::abs(mcParticleGen.pdgCode()) == PDGKaon) {
+        histos.fill(HIST("spectraGen/histGenPtKaon"), mcParticleGen.pt());
       }
 
-      if (std::abs(mcParticleGen.pdgCode()) == 2212) {
-        spectraGen.fill(HIST("histGenPtProton"), mcParticleGen.pt());
+      if (std::abs(mcParticleGen.pdgCode()) == PDGProton) {
+        histos.fill(HIST("spectraGen/histGenPtProton"), mcParticleGen.pt());
       }
 
-      if (mcParticleGen.pdgCode() == 1000010020) {
-        spectraGen.fill(HIST("histGenPtD"), mcParticleGen.pt());
+      if (mcParticleGen.pdgCode() == PDGDeuteron) {
+        histos.fill(HIST("spectraGen/histGenPtD"), mcParticleGen.pt());
       }
 
-      if (mcParticleGen.pdgCode() == -1000010020) {
-        spectraGen.fill(HIST("histGenPtantiD"), mcParticleGen.pt());
+      if (mcParticleGen.pdgCode() == -PDGDeuteron) {
+        histos.fill(HIST("spectraGen/histGenPtantiD"), mcParticleGen.pt());
       }
 
-      if (mcParticleGen.pdgCode() == 1000020030) {
-        spectraGen.fill(HIST("histGenPtHe"), mcParticleGen.pt());
+      if (mcParticleGen.pdgCode() == PDGHelium) {
+        histos.fill(HIST("spectraGen/histGenPtHe"), mcParticleGen.pt());
       }
 
-      if (mcParticleGen.pdgCode() == -1000020030) {
-        spectraGen.fill(HIST("histGenPtantiHe"), mcParticleGen.pt());
+      if (mcParticleGen.pdgCode() == -PDGHelium) {
+        histos.fill(HIST("spectraGen/histGenPtantiHe"), mcParticleGen.pt());
       }
     }
-  }
-  PROCESS_SWITCH(LFNucleiDeuteronTask, processMCGen, "process MC Generated", true);
+  } // Close processMCGen
+  PROCESS_SWITCH(LFNucleiBATask, processMCGen, "process MC Generated", true);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
 {
-  return WorkflowSpec{adaptAnalysisTask<LFNucleiDeuteronTask>(cfgc)};
+  return WorkflowSpec{adaptAnalysisTask<LFNucleiBATask>(cfgc)};
 }
