[PWGHF] treeCreatorDstarSpinAlignMixing: Adding new task for local mix candidate analysis in Dstar spin alignment measurement. (AliceO2Group#16721)

Co-authored-by: ALICE Action Bot <alibuild@cern.ch>

Author: Mingze129

PWGHF/D2H/TableProducer/CMakeLists.txt

@@ -106,3 +106,10 @@ o2physics_add_dpl_workflow(converter-reduced-3-prongs-ml
                      SOURCES converterReducedHadronDausPid.cxx
                      PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
                      COMPONENT_NAME Analysis)
+
+# Tree creator
+
+o2physics_add_dpl_workflow(tree-creator-dstar-spin-align-mixing
+                    SOURCES treeCreatorDstarSpinAlignMixing.cxx
+                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
+                    COMPONENT_NAME Analysis)

PWGHF/D2H/TableProducer/treeCreatorDstarSpinAlignMixing.cxx

@@ -0,0 +1,233 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file treeCreatorDstarSpinAlignMixing.cxx
+/// \brief Writer of D*+ → D0 ( → π+ K-) π+ candidates in the form of flat tables to be stored in TTrees.
+///        Intended for Mix-candidate analysis in spin alignment measurement.
+///        Serving as a correction for detector acceptance and reconstruction efficiency.
+///
+/// \author Mingze li <mingze.li@cern.ch>, CCNU/UniTo
+
+#include "PWGHF/Core/CentralityEstimation.h"
+#include "PWGHF/Core/DecayChannels.h"
+#include "PWGHF/D2H/Utils/utilsFlow.h"
+#include "PWGHF/DataModel/CandidateReconstructionTables.h"
+#include "PWGHF/DataModel/CandidateSelectionTables.h"
+
+#include "Common/Core/RecoDecay.h"
+
+#include <CommonConstants/PhysicsConstants.h>
+#include <Framework/ASoA.h>
+#include <Framework/AnalysisDataModel.h>
+#include <Framework/AnalysisHelpers.h>
+#include <Framework/AnalysisTask.h>
+#include <Framework/Configurable.h>
+#include <Framework/InitContext.h>
+#include <Framework/runDataProcessing.h>
+
+#include <cstdint>
+
+using namespace o2;
+using namespace o2::framework;
+using namespace o2::framework::expressions;
+using namespace o2::hf_centrality;
+using namespace o2::hf_occupancy;
+using namespace o2::analysis::hf_flow_utils;
+
+namespace o2::aod
+{
+namespace mixing_dstar
+{
+DECLARE_SOA_INDEX_COLUMN(Collision, collision);
+// D0 related variables
+DECLARE_SOA_COLUMN(MD0, mD0, float);
+// DECLARE_SOA_COLUMN(PtD0, ptD0, float);
+// DECLARE_SOA_COLUMN(PD0, pD0, float);
+// DECLARE_SOA_COLUMN(EtaD0, etaD0, float);
+// DECLARE_SOA_COLUMN(PhiD0, phiD0, float);
+// DECLARE_SOA_COLUMN(YD0, yD0, float);
+// soft pion related variables
+DECLARE_SOA_COLUMN(PtSoftPi, ptSoftPi, float);
+// DECLARE_SOA_COLUMN(PSoftPi, pSoftPi, float);
+DECLARE_SOA_COLUMN(EtaSoftPi, etaSoftPi, float);
+DECLARE_SOA_COLUMN(PhiSoftPi, phiSoftPi, float);
+// DECLARE_SOA_COLUMN(YSoftPi, ySoftPi, float);
+// Dstar related variables
+DECLARE_SOA_COLUMN(M, m, float);
+DECLARE_SOA_COLUMN(Pt, pt, float);
+// DECLARE_SOA_COLUMN(P, p, float);
+DECLARE_SOA_COLUMN(Eta, eta, float);
+DECLARE_SOA_COLUMN(Phi, phi, float);
+DECLARE_SOA_COLUMN(Y, y, float);
+DECLARE_SOA_COLUMN(MlProbDstarToD0PiBkg, mlProbDstarToD0PiBkg, float);
+DECLARE_SOA_COLUMN(MlProbDstarToD0PiPrompt, mlProbDstarToD0PiPrompt, float);
+DECLARE_SOA_COLUMN(MlProbDstarToD0PiNonPrompt, mlProbDstarToD0PiNonPrompt, float);
+// Events
+DECLARE_SOA_COLUMN(ZVtx, zVtx, float);
+DECLARE_SOA_COLUMN(Centrality, centrality, float);
+DECLARE_SOA_COLUMN(NPvContrib, nPvContrib, uint16_t);
+DECLARE_SOA_COLUMN(Occupancy, occupancy, int);
+DECLARE_SOA_COLUMN(XQVec, xQVec, float);
+DECLARE_SOA_COLUMN(YQVec, yQVec, float);
+DECLARE_SOA_COLUMN(GIndexCol, gIndexCol, int);
+DECLARE_SOA_COLUMN(TimeStamp, timeStamp, int64_t);
+// Tracks
+DECLARE_SOA_COLUMN(MinAbsEtaTrack, minAbsEtaTrack, float);
+DECLARE_SOA_COLUMN(MinNumItsCls, minNumItsCls, uint8_t);
+DECLARE_SOA_COLUMN(MinNumTpcCls, minNumTpcCls, uint8_t);
+} // namespace mixing_dstar
+
+DECLARE_SOA_TABLE(HfCandDstMix, "AOD", "HFCANDDSTMIX",
+                  mixing_dstar::MD0,
+                  mixing_dstar::PtSoftPi,
+                  mixing_dstar::EtaSoftPi,
+                  mixing_dstar::PhiSoftPi,
+                  mixing_dstar::M,
+                  mixing_dstar::Pt,
+                  mixing_dstar::Eta,
+                  mixing_dstar::Phi,
+                  mixing_dstar::Y,
+                  mixing_dstar::MlProbDstarToD0PiBkg,
+                  mixing_dstar::MlProbDstarToD0PiPrompt,
+                  mixing_dstar::MlProbDstarToD0PiNonPrompt,
+                  mixing_dstar::ZVtx,
+                  mixing_dstar::Centrality,
+                  mixing_dstar::NPvContrib,
+                  mixing_dstar::Occupancy,
+                  mixing_dstar::XQVec,
+                  mixing_dstar::YQVec,
+                  mixing_dstar::MinAbsEtaTrack,
+                  mixing_dstar::MinNumItsCls,
+                  mixing_dstar::MinNumTpcCls,
+                  mixing_dstar::GIndexCol,
+                  mixing_dstar::TimeStamp);
+} // namespace o2::aod
+
+/// Writes the full information in an output TTree
+struct HfTreeCreatorDstarSpinAlignMixing {
+  Produces<o2::aod::HfCandDstMix> rowCandidateMix;
+
+  Configurable<int> qVecDetector{"qVecDetector", 2, "Detector for Q vector estimation (FV0A: 0, FT0M: 1, FT0C: 2)"};
+  Configurable<int> centEstimator{"centEstimator", 2, "Centrality estimator ((None: 0, FT0C: 2, FT0M: 3))"};
+  Configurable<int> occEstimator{"occEstimator", 2, "If enabled, replace number of PV contributors with occupancy estimation (0: don't use, 1: ITS, 2: FT0C)"};
+
+  using CollsWithQVecs = soa::Join<aod::Collisions, aod::EvSels, aod::QvectorFT0Cs, aod::QvectorFT0As, aod::QvectorFT0Ms, aod::QvectorFV0As, aod::QvectorBPoss, aod::QvectorBNegs, aod::QvectorBTots, aod::CentFT0Ms, aod::CentFT0Cs>;
+  using TracksWithExtra = soa::Join<aod::Tracks, aod::TracksExtra>;
+  using CandDstarWSelFlag = soa::Join<aod::HfCandDstars, aod::HfSelDstarToD0Pi>;
+  using FilteredCandDstarWSelFlagAndMl = soa::Filtered<soa::Join<CandDstarWSelFlag, aod::HfMlDstarToD0Pi>>;
+
+  Filter filterSelectDstarCandidates = aod::hf_sel_candidate_dstar::isSelDstarToD0Pi == true;
+
+  Preslice<FilteredCandDstarWSelFlagAndMl> dstarWithMlPerCollision = aod::hf_cand::collisionId;
+
+  void init(InitContext const&)
+  {
+  }
+
+  /// prongTracks is the vector of daughter tracks
+  /// etaMin is the minimum eta
+  /// nItsClsMin is the minumum number of clusters in ITS
+  /// nTpcClsMin is the minumum number of clusters in TPC
+  template <typename Trk>
+  void getTrackingInfos(std::vector<Trk> const& prongTracks, float& etaMin, int& nItsClsMin, int& nTpcClsMin)
+  {
+    etaMin = 10.f;
+    nItsClsMin = 10;
+    nTpcClsMin = 1000;
+
+    for (const auto& track : prongTracks) {
+      if (std::abs(track.eta()) < etaMin) {
+        etaMin = std::abs(track.eta());
+      }
+      if (track.itsNCls() < nItsClsMin) {
+        nItsClsMin = track.itsNCls();
+      }
+      if (track.tpcNClsCrossedRows() < nTpcClsMin) {
+        nTpcClsMin = track.tpcNClsCrossedRows();
+      }
+    }
+  }
+
+  template <typename CollType, typename T, typename Trk, typename BcType>
+  void fillCandidateTable(CollType const& collision, const T& candidate, Trk const& /*tracks*/, BcType const& /*bcs*/)
+  {
+    const auto bc = collision.template bc_as<BcType>();
+    const int64_t timeStamp = bc.timestamp();
+
+    float massD0{-1.f};
+    float massDStar{-1.f};
+    float etaSoftPi{-1.f};
+    float phiSoftPi{-1.f};
+    if (candidate.signSoftPi() > 0) {
+      massD0 = candidate.invMassD0();
+      massDStar = candidate.invMassDstar();
+    } else {
+      massD0 = candidate.invMassD0Bar();
+      massDStar = candidate.invMassAntiDstar();
+    }
+    etaSoftPi = RecoDecay::eta(std::array{candidate.pxSoftPi(), candidate.pySoftPi(), candidate.pzSoftPi()});
+    phiSoftPi = RecoDecay::phi(std::array{candidate.pxSoftPi(), candidate.pySoftPi(), candidate.pzSoftPi()});
+
+    float absEtaTrackMin{-1.f};
+    int numItsClsMin{-1}, numTpcClsMin{-1};
+
+    auto trackProng0 = candidate.template prong0_as<Trk>();
+    auto trackProng1 = candidate.template prong1_as<Trk>();
+    auto trackProng2 = candidate.template prongPi_as<Trk>();
+    getTrackingInfos(std::vector{trackProng0, trackProng1, trackProng2}, absEtaTrackMin, numItsClsMin, numTpcClsMin);
+    std::array<float, 3> const Qvector = getQvec(collision, qVecDetector.value);
+
+    rowCandidateMix(
+      massD0,
+      candidate.ptSoftPi(),
+      etaSoftPi,
+      phiSoftPi,
+      massDStar,
+      candidate.pt(),
+      candidate.eta(),
+      candidate.phi(),
+      candidate.y(constants::physics::MassDStar),
+      candidate.mlProbDstarToD0Pi()[0],
+      candidate.mlProbDstarToD0Pi()[1],
+      candidate.mlProbDstarToD0Pi()[2],
+      collision.posZ(),
+      getCentralityColl(collision, centEstimator),
+      collision.numContrib(),
+      getOccupancyColl(collision, occEstimator),
+      Qvector[0],
+      Qvector[1],
+      absEtaTrackMin,
+      numItsClsMin,
+      numTpcClsMin,
+      collision.globalIndex(),
+      timeStamp);
+  }
+
+  void processData(CollsWithQVecs const& collisions,
+                   FilteredCandDstarWSelFlagAndMl const& dstarCandidates,
+                   TracksWithExtra const& tracks,
+                   aod::BCsWithTimestamps const& bcWithTimeStamps)
+  {
+    for (const auto& collision : collisions) {
+      auto thisCollId = collision.globalIndex();
+      auto groupedDstarCandidates = dstarCandidates.sliceBy(dstarWithMlPerCollision, thisCollId);
+      for (const auto& dstarCandidate : groupedDstarCandidates) {
+        fillCandidateTable(collision, dstarCandidate, tracks, bcWithTimeStamps);
+      }
+    }
+  }
+  PROCESS_SWITCH(HfTreeCreatorDstarSpinAlignMixing, processData, "Process data", true);
+};
+
+WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
+{
+  return WorkflowSpec{adaptAnalysisTask<HfTreeCreatorDstarSpinAlignMixing>(cfgc)};
+}